diff --git a/docs/source/models/supported-models.md b/docs/source/models/supported-models.md index ba99d619a957..b9588323b1af 100644 --- a/docs/source/models/supported-models.md +++ b/docs/source/models/supported-models.md @@ -91,7 +91,7 @@ Note: Support for other models may vary. Features marked "N/A" are not applicabl | `Gemma4ForConditionalGeneration` | Untested | Yes | Untested | No | Yes | MTP | Yes | Untested | No | Yes | Untested | Yes | | `Gemma4UnifiedForConditionalGeneration` | Untested | Untested | Untested | No | Yes | No | Yes | Untested | No | Yes | Untested | Yes | | `Step3p7ForConditionalGeneration`| Yes | Yes | Yes | Untested | Untested | MTP | Yes | Untested | Untested | Yes | Untested | Yes | -| `MiniMaxM3SparseForConditionalGeneration` [^12] | Yes | Yes | Yes | Untested | Untested | No | Yes | Untested | No | N/A | Untested | Yes | +| `MiniMaxM3SparseForConditionalGeneration` [^12] | Yes | Yes | Yes | Untested | Untested | EAGLE-3 (Linear) | Yes | Untested | No | N/A | Untested | Yes | [^1]: Chunked Prefill for MLA can only be enabled on SM90/SM100/SM103/SM120. [^2]: KV cache reuse for MLA can only be enabled on SM90/SM100/SM103/SM120/SM121 and in BF16/FP8 KV cache dtype. diff --git a/tensorrt_llm/_torch/attention/backends/fmha/flashinfer_trtllm_gen.py b/tensorrt_llm/_torch/attention/backends/fmha/flashinfer_trtllm_gen.py index d4648893dd68..70a824cd2187 100644 --- a/tensorrt_llm/_torch/attention/backends/fmha/flashinfer_trtllm_gen.py +++ b/tensorrt_llm/_torch/attention/backends/fmha/flashinfer_trtllm_gen.py @@ -760,8 +760,12 @@ def _is_supported_with_reason( return False, f"non-positive tokens_per_block ({tokens_per_block})." if tokens_per_block & (tokens_per_block - 1) != 0: return False, f"tokens_per_block ({tokens_per_block}) that is not a power of 2." - if tokens_per_block not in self.SUPPORTED_TOKENS_PER_BLOCK: - supported = sorted(self.SUPPORTED_TOKENS_PER_BLOCK) + # A KV cache manager may allow extra page sizes, e.g. MiniMax-M3 adds 128. + supported_tokens_per_block = self.SUPPORTED_TOKENS_PER_BLOCK | set( + getattr(meta.kv_cache_manager, "trtllm_gen_extra_tokens_per_block", ()) + ) + if tokens_per_block not in supported_tokens_per_block: + supported = sorted(supported_tokens_per_block) return False, f"tokens_per_block ({tokens_per_block}). Supported: {supported}." return True, "" diff --git a/tensorrt_llm/_torch/attention/backends/sparse/minimax_m3/cache_manager.py b/tensorrt_llm/_torch/attention/backends/sparse/minimax_m3/cache_manager.py index e65b168c5d14..4c52f9e677c6 100644 --- a/tensorrt_llm/_torch/attention/backends/sparse/minimax_m3/cache_manager.py +++ b/tensorrt_llm/_torch/attention/backends/sparse/minimax_m3/cache_manager.py @@ -20,7 +20,8 @@ construction). * :class:`MiniMaxM3KVCacheManagerV2` — :class:`KVCacheManagerV2` subclass that registers a per-sparse-layer ``Role.INDEX_KEY`` paged - buffer alongside the standard K/V buffers. + buffer alongside the standard K/V buffers; shared Eagle3 draft layers get + their own virtual attention-op pools. """ from __future__ import annotations @@ -39,6 +40,10 @@ ) from tensorrt_llm.bindings import DataType from tensorrt_llm.bindings.internal.batch_manager import CacheType as CacheTypeCpp +from tensorrt_llm.bindings.internal.batch_manager.kv_cache_manager_v2_utils import ( + copy_batch_block_offsets_to_device, +) +from tensorrt_llm.logger import logger from tensorrt_llm.runtime.kv_cache_manager_v2 import BufferConfig, PageIndexMode from tensorrt_llm.runtime.kv_cache_manager_v2._common import BAD_PAGE_INDEX from tensorrt_llm.runtime.kv_cache_manager_v2._config import DataRole @@ -133,6 +138,87 @@ def set_index_v(self, layer_idx: int, out_cache_loc: torch.Tensor, idx_v: torch. buf.index_copy_(0, out_cache_loc.to(torch.long), idx_v.to(buf.dtype)) +def shared_draft_layer_count(spec_config, layer_mask) -> int: + """How many one-model draft layers the base manager appends to this one. + + Same rule as ``get_pp_layers``: only with a speculative config and no + ``layer_mask`` (a separate draft manager setup always passes a mask). + """ + if spec_config is None or layer_mask is not None: + return 0 + from tensorrt_llm._torch.speculative.utils import get_num_spec_layers + + return int(get_num_spec_layers(spec_config)) + + +def derive_shared_draft_layout( + num_layers: Optional[int], + num_kv_heads, + num_draft: int, +) -> Tuple[List[int], Optional[int]]: + """Locate the draft layers the base manager appends after the target's. + + ``num_layers`` is the target count; a per-layer ``num_kv_heads`` list (a + drafter with a different head count) already includes the draft tail and + pins the total. Returns ``(draft_layer_ids, num_target_layers)``, or + ``([], None)`` when nothing pins the range. + """ + num_draft = max(0, int(num_draft)) + if isinstance(num_kv_heads, (list, tuple)): + total = len(num_kv_heads) + if num_layers is not None: + total = max(total, int(num_layers)) + elif num_layers is not None: + total = int(num_layers) + num_draft + else: + return [], None + num_target = total - num_draft + return list(range(num_target, total)), num_target + + +def extend_attention_op_pools_for_shared_draft_layers( + pool_pointers: torch.Tensor, + pool_mapping: torch.Tensor, + num_pools: int, + draft_layers: Sequence[Tuple[int, int, int]], +) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, List[Tuple[int, int]]]: + """Give each shared draft layer its own attention-op pool rooted at its K page. + + ``draft_layers`` holds ``(local_layer_idx, key_base_addr, sub_pages_per_slot)``. + With ``index_scale = sub_pages_per_slot`` and ``kv_offset = 1``, slot ``s`` + maps to page ``s * scale`` (K) and ``s * scale + 1`` (V). Returns + ``(pool_pointers, pool_mapping, index_scales, kv_offsets, op_pools)`` with + ``op_pools = [(attention_op_pool_id, source_storage_pool_id)]``. + """ + if pool_pointers.dim() == 3: + # NVFP4 pools carry [data, scale] pointer pairs; the draft layer's + # block-scale pages would need their own root and stride. + raise NotImplementedError( + "MiniMax-M3 shared Eagle3 draft layers do not support an NVFP4 KV cache." + ) + pointer_rows = pool_pointers.tolist() + mapping_rows = pool_mapping.tolist() + index_scales: List[int] = [] + kv_offsets: List[int] = [] + op_pools: List[Tuple[int, int]] = [] + for i, (local_layer_idx, key_base_addr, sub_pages_per_slot) in enumerate(draft_layers): + op_pool_id = num_pools + i + source_pool_id = int(mapping_rows[local_layer_idx][0]) + pointer_rows.append([key_base_addr, 0]) + mapping_rows[local_layer_idx] = [op_pool_id, 0] + index_scales.append(int(sub_pages_per_slot)) + kv_offsets.append(1) + op_pools.append((op_pool_id, source_pool_id)) + pinned = prefer_pinned() + return ( + torch.tensor(pointer_rows, dtype=pool_pointers.dtype, device="cpu", pin_memory=pinned), + torch.tensor(mapping_rows, dtype=pool_mapping.dtype, device="cpu", pin_memory=pinned), + torch.tensor(index_scales, dtype=torch.int32, device="cpu", pin_memory=pinned), + torch.tensor(kv_offsets, dtype=torch.int32, device="cpu", pin_memory=pinned), + op_pools, + ) + + class MiniMaxM3KVCacheManagerV2(KVCacheManagerV2): """KVCacheManagerV2 subclass with a V2-managed paged index-K cache per sparse layer. @@ -152,9 +238,29 @@ class MiniMaxM3KVCacheManagerV2(KVCacheManagerV2): * ``disable_index_value_layer_ids`` — subset whose index-V is omitted. * ``sparse_index_dim`` — width of the index-K/V vectors. + + Shared Eagle3 draft layers: the base manager appends them after the target + layers. They run the generic TRTLLM attention op, which needs a uniform + per-layer stride inside a pool. M3's pool is not uniform: here an index-K + page is as large as a K or V page, so V2 packs K, V and index-K of all + layers into one pool (3 sub-pages per sparse layer, 2 per dense or draft + layer). M3's own kernels don't care (they use :meth:`get_buffers`), but the + draft layer would read the wrong pages. So each draft layer gets its own + virtual attention-op pool rooted at its K page, like + ``DeepseekV4CacheManager`` does for SWA layers; see + :meth:`_prepare_page_table_tensor`. Can go once V2 keeps index-K in its + own pool. """ _main_kv_mapper_kind = MapperKind.NHD + # Virtual attention-op pools for shared draft layers: + # (attention_op_pool_id, source_storage_pool_id) plus their copy parameters. + _draft_op_pools: Tuple[Tuple[int, int], ...] = () + _draft_index_scales: Optional[torch.Tensor] = None + _draft_kv_offsets: Optional[torch.Tensor] = None + # Extra page sizes trtllm-gen may use with this manager (see + # FlashInferTrtllmGenFmha); set with the virtual pools. + trtllm_gen_extra_tokens_per_block: frozenset = frozenset() def __init__( self, @@ -170,6 +276,7 @@ def __init__( # disable_index_value=True, sparse_index_dim=128). Honoring the # executor keyword also makes non-default sparse_index_dim values # authoritative for the cache layout instead of falling back to 128. + # Peeked (not popped) so the base __init__ still receives them. sparse_attention_config = kwargs.get("sparse_attention_config") num_layers = kwargs.get("num_layers") implementation = getattr(sparse_attention_config, "implementation", "triton") @@ -184,9 +291,16 @@ def __init__( raise ValueError( f"MiniMax M3 sparse_index_dim must be greater than 0, got {sparse_index_dim}." ) + # Shared draft layers sit above the target layers and have no index-K + # cache, so the sparse-layer default stops at the target range. + self._shared_draft_layer_ids, num_target_layers = derive_shared_draft_layout( + num_layers, + kwargs.get("num_kv_heads"), + shared_draft_layer_count(kwargs.get("spec_config"), kwargs.get("layer_mask")), + ) if sparse_layer_ids is None: - if num_layers is not None: - sparse_layer_ids = list(range(3, int(num_layers))) + if num_target_layers is not None: + sparse_layer_ids = list(range(3, num_target_layers)) else: sparse_layer_ids = [] if disable_index_value_layer_ids is None: @@ -232,6 +346,87 @@ def __init__( device=device, ) + def _prepare_page_table_tensor(self, index_mapper_capacity: int) -> None: + """Base pool tables plus one virtual attention-op pool per shared draft layer. + + See the class docstring for why. The virtual pool reuses the source + pool's slot ids; :meth:`copy_batch_block_offsets` scales them. + """ + super()._prepare_page_table_tensor(index_mapper_capacity) + draft_layers = [ + layer_idx + for layer_idx in self._shared_draft_layer_ids + if layer_idx in self.layer_offsets + ] + if self.is_draft or not draft_layers: + return + if self.enable_swa_scratch_reuse: + raise NotImplementedError( + "MiniMax-M3 shared Eagle3 draft layers do not support SWA scratch reuse." + ) + # Draft layers run at the target's 128-token pages. trtllm-gen has P128 + # kernels for their dense-GQA shapes but not for every shape, so opt in + # here rather than in the global allowlist. + if self.tokens_per_block == 128: + self.trtllm_gen_extra_tokens_per_block = frozenset({128}) + if self._use_per_layer_page_tables: + # Per-layer page tables already give every layer its own pool. + return + geometry = [] + for layer_idx in draft_layers: + key_base_addr, _dtype, _num_slots, sub_pages_per_slot, _shape = self._kv_slot_geometry( + layer_idx + ) + geometry.append((self.layer_offsets[layer_idx], key_base_addr, sub_pages_per_slot)) + ( + self.kv_cache_pool_pointers, + self.kv_cache_pool_mapping, + self._draft_index_scales, + self._draft_kv_offsets, + op_pools, + ) = extend_attention_op_pools_for_shared_draft_layers( + self.kv_cache_pool_pointers, self.kv_cache_pool_mapping, self.num_pools, geometry + ) + self._draft_op_pools = tuple(op_pools) + self.num_attention_op_pools = self.num_pools + len(op_pools) + logger.info( + f"[unified-kv] draft layers {draft_layers} share the target KV cache manager; " + f"attention-op pools {[pool for pool, _ in op_pools]} address their pages." + ) + + def copy_batch_block_offsets( + self, + dst_tensor: torch.Tensor, + request_ids: List[int], + beam_width: int, + num_contexts: int, + num_seqs: int, + max_blocks: Optional[int] = None, + ) -> None: + super().copy_batch_block_offsets( + dst_tensor, request_ids, beam_width, num_contexts, num_seqs, max_blocks=max_blocks + ) + if not self._draft_op_pools: + return + # Fill each virtual pool from its source pool's slot ids with the draft + # layer's scale. + copy_idx = self.index_mapper.get_copy_index(request_ids, num_contexts, beam_width) + for i, (op_pool_id, source_pool_id) in enumerate(self._draft_op_pools): + copy_batch_block_offsets_to_device( + self.host_kv_cache_block_offsets[source_pool_id : source_pool_id + 1], + dst_tensor[op_pool_id : op_pool_id + 1], + copy_idx, + self._draft_index_scales[i : i + 1], + self._draft_kv_offsets[i : i + 1], + self._stream.cuda_stream, + ) + + def update_resources(self, scheduled_batch, attn_metadata=None, kv_cache_dtype_byte_size=None): + # Only tree acceptance relocates draft KV, which M3's pool layout cannot do. + if any(r.py_num_accepted_draft_tokens_indices for r in scheduled_batch.generation_requests): + raise NotImplementedError("MiniMax-M3 does not relocate accepted draft tokens.") + super().update_resources(scheduled_batch, attn_metadata, kv_cache_dtype_byte_size) + def _extra_buffers_per_layer(self, *, tokens_per_block): """Register a per-sparse-layer ``Role.INDEX_KEY`` :class:`BufferConfig`. @@ -311,13 +506,12 @@ def has_index_value(self, layer_idx: int) -> bool: def _kv_slot_geometry( self, layer_idx: int, kv_layout: Optional[str] = None ) -> Tuple[int, torch.dtype, int, int, List[int]]: - """Resolve one layer's position in the coalesced K/V pool. + """Where a layer's K/V live in the coalesced pool. - Returns (addr_key, torch_dtype, num_slots, scale, page_shape), where - scale is the number of equal-sized sub-pages a slot packs and - page_shape is one sub-page's shape in kv_layout. This layer's K is - sub-page 0 and its V sub-page 1, counting from addr_key. - When omitted, ``kv_layout`` follows the selected sparse backend. + Returns ``(addr_key, torch_dtype, num_slots, scale, page_shape)``: + ``scale`` sub-pages per slot, this layer's K at sub-page 0 and V at + sub-page 1 from ``addr_key``. Used by :meth:`get_buffers` and the draft + layers' virtual pools. ``kv_layout`` defaults to the backend's layout. """ if kv_layout is None: kv_layout = self._main_kv_layout_name() @@ -534,5 +728,8 @@ def get_minimax_m3_kv_cache_manager_cls(): __all__ = [ "MiniMaxM3KVCacheManagerV2", "MiniMaxM3SparseIndexCache", + "derive_shared_draft_layout", + "extend_attention_op_pools_for_shared_draft_layers", "get_minimax_m3_kv_cache_manager_cls", + "shared_draft_layer_count", ] diff --git a/tensorrt_llm/_torch/attention/backends/sparse/minimax_m3/msa_backend.py b/tensorrt_llm/_torch/attention/backends/sparse/minimax_m3/msa_backend.py index bc55aa6cf51f..5a3f652fbe8c 100644 --- a/tensorrt_llm/_torch/attention/backends/sparse/minimax_m3/msa_backend.py +++ b/tensorrt_llm/_torch/attention/backends/sparse/minimax_m3/msa_backend.py @@ -17,6 +17,11 @@ into them. The standard CUDAGraphRunner clones one metadata per graph batch size (create_cuda_graph_metadata), so no per-batch-size cache is needed here. + * With Eagle3 a decode row has 1 + draft_len query tokens. Slots and + valid-block counts are per token, and the decode kernels take that uniform + query length from msa_decode_span. on_update_kv_lens re-derives the + per-request lengths, the slots and the counts after the overlap scheduler + corrects kv_lens on device. The classes subclass TrtllmAttention and TrtllmAttentionMetadata, imported at module scope. That is cycle-free only because the dependency runs one way: the @@ -134,6 +139,14 @@ class MiniMaxM3MsaSparseAttentionMetadata(TrtllmAttentionMetadata): graph-stable storage. The plans themselves (msa_prefill_*_plan) cover the context rows alone, and a step carrying those is never captured, so they need none either. + + With Eagle3 a decode row has 1 + draft_len query tokens: the span's + query_len is that count, slots and valid-block counts are per token, and + the overlap scheduler corrects kv_lens on device after prepare(). The + decode kernels read their lengths from msa_seq_lens_cuda, so + on_update_kv_lens patches that buffer, the slots and the counts. The + fmha_sm100 plans cover context rows only, whose lengths the correction + never touches, so they need no patch. """ # Graph-stable buffers; consumers slice to the live count at the call @@ -154,6 +167,13 @@ class MiniMaxM3MsaSparseAttentionMetadata(TrtllmAttentionMetadata): # factor, or 0 where the pool has no single one; see msa_subpage_rows. msa_subpage_block_table: Optional[torch.Tensor] = None _msa_subpages_per_slot: int = 0 + # Inputs for on_update_kv_lens: each query token's request row and offset + # within the request, and the kv_lens prepare() staged (the upper bound the + # correction is clamped to). The write slots are re-derived from + # msa_block_table, so no per-token slot table is kept. + msa_q_batch_row: Optional[torch.Tensor] = None + msa_q_intra: Optional[torch.Tensor] = None + msa_kv_lens_staged: Optional[torch.Tensor] = None # _msa_buffers_ready gates the once-only device buffers; # _msa_fields_ready marks that the current step's buffers are populated. @@ -186,6 +206,9 @@ class MiniMaxM3MsaSparseAttentionMetadata(TrtllmAttentionMetadata): _msa_max_kv_len: int = 0 # See msa_worst_case_max_k_tiles. _msa_worst_case_max_k_tiles: int = 0 + # True only with speculative decoding; otherwise on_update_kv_lens and its + # staging are skipped and non-speculative steps run as before. + _msa_kv_lens_dynamic: bool = False def __post_init__(self) -> None: super().__post_init__() @@ -224,6 +247,14 @@ def as_pinned_int32(lens: torch.Tensor) -> torch.Tensor: seq_lens = self.seq_lens kv_lens = getattr(self, "kv_lens", None) + # The host kv_lens carries num_extra_kv_tokens (speculative draft + # slots); MSA needs the attended length, which is what kv_lens_cuda + # holds. Every consumer below, msa_seq_lens_cuda included, inherits + # the correction from here. + params = self.kv_cache_params + extra = int(params.num_extra_kv_tokens) if params is not None else 0 + if kv_lens is not None and extra: + kv_lens = kv_lens - extra self._msa_qo_lens_cpu = None if seq_lens is None else as_pinned_int32(seq_lens) self._msa_kv_lens_cpu = ( None if seq_lens is None or kv_lens is None else as_pinned_int32(kv_lens) @@ -332,16 +363,21 @@ def _validate_decode_kernel_support(self) -> None: # No MSA geometry to check, as for a structural test's metadata. return page_size = int(kv_cache_manager.tokens_per_block) + # 1 + draft_len query tokens per generation request under speculative + # decoding, one otherwise; see msa_decode_span. The scorer tiles + # num_index_heads * query tokens into one Q block, which bounds the + # draft length it can verify. + decode_query_len = self._msa_max_decode_query_len() if not self._cutedsl_indexer_supported( num_index_heads=params.num_index_heads, page_size=page_size, - # One query token per generation request; see msa_decode_span. - decode_query_len=1, + decode_query_len=decode_query_len, ): raise RuntimeError( "The MiniMax-M3 CuTe DSL indexer scorer does not support this " f"configuration: {params.num_index_heads} index heads, page size " - f"{page_size}, index dtype {self._msa_index_kv_dtype()}." + f"{page_size}, index dtype {self._msa_index_kv_dtype()}, up to " + f"{decode_query_len} query tokens per generation request." ) dense_unsupported = dense_decode_unsupported_reason(kv_cache_manager, MSA_REQUIRED_HEAD_DIM) if dense_unsupported is not None: @@ -408,6 +444,28 @@ def _create_msa_buffers(self) -> None: dtype=torch.int32, capture_graph=capture_graph, ) + # Inputs for on_update_kv_lens. + self.msa_q_batch_row = self.get_empty( + buffers, + (max_num_tokens,), + cache_name="msa_q_batch_row", + dtype=torch.int32, + capture_graph=capture_graph, + ) + self.msa_q_intra = self.get_empty( + buffers, + (max_num_tokens,), + cache_name="msa_q_intra", + dtype=torch.int32, + capture_graph=capture_graph, + ) + self.msa_kv_lens_staged = self.get_empty( + buffers, + (max_num_sequences,), + cache_name="msa_kv_lens_staged", + dtype=torch.int32, + capture_graph=capture_graph, + ) # The proxy scratch needs the fmha_sm100 plan geometry. This metadata # exists only for the MSA backend, whose selection already required the # kernels, so a failed import here is a hard error rather than a reason @@ -424,37 +482,63 @@ def _create_msa_buffers(self) -> None: self._msa_worst_case_max_k_tiles = int(max_k_tiles) self._alloc_msa_proxy_scratch( num_index_heads=params.num_index_heads, - max_batch=max_num_sequences, + max_tokens=self._msa_max_decode_tokens(), max_k_tiles=max_k_tiles, capture_graph=capture_graph, ) self._msa_buffers_ready = True + def _msa_max_decode_query_len(self) -> int: + """Worst-case query tokens per generation request: 1 + draft_len. + + The KV cache manager knows the speculative config when this metadata is + built (__post_init__ runs before update_spec_dec_param), so it is the + one source for both the up-front kernel validation and the scratch + sizing. A run without speculative decoding reports 1. + """ + draft_len = int(getattr(self.kv_cache_manager, "max_total_draft_tokens", 0) or 0) + return 1 + max(0, draft_len) + + def _msa_max_decode_tokens(self) -> int: + """Worst-case query tokens in one decode step: 1 + draft_len per row. + + Sizes the proxy scratch the CuTe DSL scorer writes one column per query + token into, and the per-token valid-block buffer beside it. Capped at + 16384, which also keeps a whole-batch fmha_sm100 proxy plan under its + 65536 total_q * num_qo_heads limit with 4 index heads. + """ + max_seqs = int(self.max_num_sequences) + tokens = max_seqs * self._msa_max_decode_query_len() + max_toks = int(self.max_num_tokens or 0) + if max_toks > 0: + tokens = min(tokens, max_toks) + return max(max_seqs, min(tokens, 16384)) + def _alloc_msa_proxy_scratch( self, *, num_index_heads: int, - max_batch: int, + max_tokens: int, max_k_tiles: int, capture_graph: bool, ) -> None: """Allocate the flat proxy max-score store and the valid-block scratch. - The store is sized for the worst-case max_k_tiles so one allocation - serves every decode step. msa_proxy_max_score_view slices the per-step - shape out of it. + Sized for the worst-case max_k_tiles and query-token count (more than + the batch size under speculative verify), so one allocation serves every + decode step. msa_proxy_max_score_view slices the per-step shape. """ buffers = self.cuda_graph_buffers self.msa_max_score = self.get_empty( buffers, - (num_index_heads * max_k_tiles * max_batch,), + (num_index_heads * max_k_tiles * max_tokens,), cache_name="msa_max_score", dtype=torch.float32, capture_graph=capture_graph, ) self.msa_n_valid_blocks = self.get_empty( buffers, - (max_batch,), + (max_tokens,), cache_name="msa_n_valid_blocks", dtype=torch.int32, capture_graph=capture_graph, @@ -469,14 +553,15 @@ def _ensure_msa_decode_scratch_buffers( required_max_k_tiles: int, ) -> None: """Ensure proxy scratch buffers exist and cover the current plan.""" - required_numel = num_index_heads * required_max_k_tiles * max_batch + max_tokens = max(int(max_batch), self._msa_max_decode_tokens()) + required_numel = num_index_heads * required_max_k_tiles * max_tokens if self.msa_max_score is not None: if self.msa_max_score.numel() < required_numel: raise ValueError( f"msa_max_score backing store ({self.msa_max_score.numel()} " f"elements) is smaller than the decode plan needs " f"({required_numel} = {num_index_heads} heads * " - f"{required_max_k_tiles} k-tiles * {max_batch} batch)." + f"{required_max_k_tiles} k-tiles * {max_tokens} tokens)." ) return @@ -499,7 +584,7 @@ def _ensure_msa_decode_scratch_buffers( self._msa_worst_case_max_k_tiles = int(max_k_tiles) self._alloc_msa_proxy_scratch( num_index_heads=num_index_heads, - max_batch=max_batch, + max_tokens=max_tokens, max_k_tiles=max_k_tiles, capture_graph=capture_graph, ) @@ -556,7 +641,10 @@ def _set_decode_span(self) -> None: The one property of the rows themselves that has to hold is a single positive query length across them, which the kernels derive the request - id from, so a batch without it is rejected rather than served. + id from, so a batch without it is rejected rather than served. Under + speculative decoding that length is the verify window, 1 + draft_len, + and may not exceed what _validate_decode_kernel_support settled the + scorer for. """ self._msa_decode_span = None self._msa_max_kv_len = 0 @@ -572,12 +660,12 @@ def _set_decode_span(self) -> None: # Host-side tensors, so these reads do not sync the device. gen_qo_lens = qo_lens_cpu[row_first:] qo_min, qo_max = int(gen_qo_lens.min()), int(gen_qo_lens.max()) - if qo_max > 1: - raise NotImplementedError( - "MiniMax-M3 MSA attention does not support speculative decoding " - "(multiple query tokens per decode step): generation rows " - f"[{row_first}, {row_last}) carry up to {qo_max} query tokens. " - "Disable speculative decoding or use the non-MSA MiniMax-M3 backend." + max_query_len = self._msa_max_decode_query_len() + if qo_max > max_query_len: + raise RuntimeError( + "MiniMax-M3 MSA attention validated its decode kernels for at most " + f"{max_query_len} query tokens per generation request, but rows " + f"[{row_first}, {row_last}) carry up to {qo_max}." ) if qo_min != qo_max or qo_max <= 0: raise RuntimeError( @@ -670,6 +758,85 @@ def _cutedsl_indexer_supported( ) ) + def _msa_kv_lens_may_change(self) -> bool: + """Whether kv_lens can change after prepare(): only with speculative decoding. + + max_total_draft_tokens is set by the engine's update_spec_dec_param; the + other checks mirror TrtllmAttentionMetadata's spec_active. + """ + params = self.kv_cache_params + runtime_features = self.runtime_features + return bool( + self.max_total_draft_tokens + or self.draft_kv_cache_manager is not None + or self.is_spec_decoding_enabled + or (params is not None and params.num_extra_kv_tokens) + or (runtime_features is not None and runtime_features.has_speculative_draft_tokens) + ) + + def on_update_kv_lens(self) -> None: + """Re-derive lengths, slots and valid-block counts from the corrected kv_lens_cuda. + + The overlap scheduler shortens kv_lens on device after prepare() staged + full-acceptance values (CUDA-graph warmup restores them the same way + between forwards). Shrinking keeps the staged page table valid, so only + the per-row lengths and what derives from them are patched; the clamp + to msa_kv_lens_staged enforces that. Device-only, capture-safe and + idempotent; skipped without speculative decoding. + + Three buffers carry the correction to the kernels: msa_seq_lens_cuda, + which the CuTe DSL scorer, the Triton sparse decode and the trtllm-gen + dense decode all read their lengths from; msa_out_cache_loc, the K/V + and index-K write slots; and the per-token valid-block count the top-k + selection is bounded by, on whichever buffer this step staged it. The + fmha_sm100 plans need no patch: they cover context rows only, whose + lengths the correction never touches. msa_max_kv_len is a host upper + bound and stays valid as lengths shrink. + """ + super().on_update_kv_lens() + if not self._msa_fields_ready or not self._msa_kv_lens_dynamic: + return + batch = int(self.num_seqs) + total_q = int(self.num_tokens) + if batch <= 0 or total_q <= 0: + return + # kv_lens_cuda is the attended length (no num_extra_kv_tokens), the same + # domain as the staged bound and msa_seq_lens_cuda. + kv_true = torch.minimum(self.kv_lens_cuda[:batch], self.msa_kv_lens_staged[:batch]) + self.msa_seq_lens_cuda[:batch].copy_(kv_true) + + qbr = self.msa_q_batch_row[:total_q].to(torch.long) + qo_dev = self.seq_lens_cuda[:batch] + # Position each query token attends up to. + pos = kv_true[qbr] - qo_dev[qbr] + self.msa_q_intra[:total_q] + + # KV/idx-K write slots, by the formula build_paged_kv_slot_mapping used + # for the staged ones: block_table[request, pos // page] * page + + # pos % page. Positions only move down, so the page is one prepare() + # staged for that request; a padding row (pos -1) keeps its own first + # slot, as on the host. + page = int(self.kv_cache_manager.tokens_per_block) + table = self.msa_block_table + pos_slot = pos.clamp_min(0) + block_col = torch.div(pos_slot, page, rounding_mode="floor").clamp( + max=int(table.shape[1]) - 1 + ) + slots = table[qbr, block_col.to(torch.long)] * page + torch.remainder(pos_slot, page) + self.msa_out_cache_loc[:total_q].copy_(slots) + + # Per-token valid-block counts for top-k, as per_token_valid_blocks + # derives them on the host: ceil((pos + 1) / page), 0 for a row that + # attends nothing (a CUDA-graph padding row), which the selector and + # the decode kernel both accept. + n_valid = torch.div((pos + 1).clamp_min(0) + (page - 1), page, rounding_mode="floor") + n_valid_buf = ( + self._msa_prefill_n_valid_blocks + if self._msa_prefill_n_valid_blocks is not None + else self.msa_n_valid_blocks + ) + if n_valid_buf is not None: + n_valid_buf[:total_q].copy_(n_valid.to(torch.int32)) + def _build_step_plans(self) -> None: """Build the layer-invariant fmha_sm100 plans this step still needs. @@ -714,10 +881,13 @@ def _build_step_plans(self) -> None: # No proxy plan, so the worst case is the only bound available. required_max_k_tiles=self._msa_worst_case_max_k_tiles, ) + # One entry per query token: batch * (1 + draft_len) under + # speculative decoding, batch otherwise. n_valid = per_token_valid_blocks( qo_lens_cpu, kv_lens_cpu, qo_offset_cpu, causal=True, block_size=page_size ) - self.msa_n_valid_blocks[:batch].copy_(n_valid.to(torch.int32), non_blocking=True) + total_q = int(n_valid.shape[0]) + self.msa_n_valid_blocks[:total_q].copy_(n_valid.to(torch.int32), non_blocking=True) return fmha_sm100 = require_msa_module() @@ -788,7 +958,8 @@ def _build_msa_fields(self) -> None: The page table and per-new-token cache slots are derived via the build_paged_kv_slot_mapping helper, then copied into the persistent - buffers. The transient builder tensors are discarded. + buffers. The transient builder tensors are discarded. With speculative + decoding the inputs for on_update_kv_lens are staged as well. """ self._msa_fields_ready = False if not self._msa_buffers_ready: @@ -810,7 +981,7 @@ def _build_msa_fields(self) -> None: # Built in prepare() (outside capture), so these transients are # fine: forwards read only the persistent buffers filled below. # qo_offset is the prefix length, so one build covers prefill - # (num_cached) and decode (kv_len - 1 with qo_len 1). + # (num_cached) and decode (kv_len - qo_len). mapping = build_paged_kv_slot_mapping( kv_cache_manager=kv_cache_manager, request_ids=request_ids, @@ -870,6 +1041,32 @@ def _build_msa_fields(self) -> None: self._msa_subpages_per_slot, self.msa_subpage_block_table[:batch_size], ) + + self._msa_kv_lens_dynamic = self._msa_kv_lens_may_change() + if not self._msa_kv_lens_dynamic: + self._msa_fields_ready = True + return + + # Inputs for on_update_kv_lens: each query token's (request row, offset + # in request). Pinned and non-blocking so the copies do not synchronize + # the stream. The slots themselves come from msa_block_table above. + qo_long = qo_lens_cpu.to(torch.long) + batch_row_cpu = torch.repeat_interleave( + torch.arange(batch_size, dtype=torch.int32), qo_long + ) + starts = torch.cumsum(qo_long, 0) - qo_long + intra_cpu = ( + torch.arange(total_new_tokens, dtype=torch.int64) + - torch.repeat_interleave(starts, qo_long) + ).to(torch.int32) + self.msa_q_batch_row[:total_new_tokens].copy_( + maybe_pin_memory(batch_row_cpu), non_blocking=True + ) + self.msa_q_intra[:total_new_tokens].copy_(maybe_pin_memory(intra_cpu), non_blocking=True) + # The staged lens are the upper bound on_update_kv_lens clamps to: the + # same attended lengths msa_seq_lens_cuda was just filled from, kept + # apart so the clamp stays fixed however often the hook runs. + self.msa_kv_lens_staged[:batch_size].copy_(kv_lens_cpu, non_blocking=True) self._msa_fields_ready = True def msa_idx_k_cache(self, layer_idx: int) -> torch.Tensor: diff --git a/tensorrt_llm/_torch/attention/backends/sparse/minimax_m3/triton_metadata.py b/tensorrt_llm/_torch/attention/backends/sparse/minimax_m3/triton_metadata.py index 696f0f68055c..a1bea4f942fc 100644 --- a/tensorrt_llm/_torch/attention/backends/sparse/minimax_m3/triton_metadata.py +++ b/tensorrt_llm/_torch/attention/backends/sparse/minimax_m3/triton_metadata.py @@ -20,7 +20,7 @@ import torch -from ...interface import AttentionMetadata +from ...trtllm import TrtllmAttentionMetadata from .common import build_paged_kv_slot_mapping @@ -165,7 +165,10 @@ def prepare(self) -> None: if batch_size == 0: self.max_seqlen_k = 1 else: - self.max_seqlen_k = int(self.seq_lens_cpu[:batch_size].max().item()) + max_k = int(self.seq_lens_cpu[:batch_size].max().item()) + # With the overlap scheduler, optimistic lengths can run past the page + # table; SDPA uses max_seqlen_k as the mask width, so clamp it. + self.max_seqlen_k = min(max_k, int(self.req_to_token.shape[1])) def ensure_metadata_on_device( @@ -369,6 +372,40 @@ def _build_runtime_metadata_fresh( return meta, out_cache_loc +def derive_q_positions_and_cache_slots( + req_to_token: torch.Tensor, + prefix_lens: torch.Tensor, + cu_seqlens_q: torch.Tensor, + q_batch_row: torch.Tensor, +) -> Tuple[torch.Tensor, torch.Tensor]: + """Per-query-token K positions and KV slot ids, computed on device. + + Shared by the metadata builder and on_update_kv_lens so they agree. + """ + total_q = int(q_batch_row.shape[0]) + qbr = q_batch_row.to(torch.long) + tok = torch.arange(total_q, dtype=torch.int32, device=q_batch_row.device) + q_positions = prefix_lens[qbr] + (tok - cu_seqlens_q[qbr]) + # Optimistic prefix_lens may point past the last allocated page. Those + # slots are placeholders that on_update_kv_lens fixes before use, but the + # gather must stay in bounds. Not in-place: .to() may alias the input. + idx = q_positions.to(torch.long).clamp(min=0, max=req_to_token.shape[1] - 1) + flat = qbr * req_to_token.shape[1] + idx + return q_positions, req_to_token.reshape(-1).index_select(0, flat) + + +def derive_decode_cache_slots(req_to_token: torch.Tensor, seq_lens: torch.Tensor) -> torch.Tensor: + """Decode-row KV slot ids (the new token sits at ``seq_lens[b] - 1``). + + Same clamp as derive_q_positions_and_cache_slots; ``min=0`` also covers + empty dummy rows. + """ + rows = torch.arange(seq_lens.shape[0], device=seq_lens.device, dtype=torch.long) + idx = (seq_lens.to(torch.long) - 1).clamp_(min=0, max=req_to_token.shape[1] - 1) + flat = rows * req_to_token.shape[1] + idx + return req_to_token.reshape(-1).index_select(0, flat) + + def build_runtime_metadata_from_kv_manager( *, kv_cache_manager, @@ -539,21 +576,13 @@ def build_runtime_metadata_from_kv_manager( req_to_token = req_to_token_fresh slot_ids = torch.arange(batch, device=device, dtype=torch.int32) - # Compute out_cache_loc: per-new-token slot ids, in flattened order - # matching the q-token order the model layer projects. The Python - # loops below run on CPU lists derived from the CPU-resident - # ``seq_lens_cpu`` / ``prefix_lens`` / ``extend_seq_lens_cpu``, so - # no GPU sync is needed at this point. The resulting - # ``out_cache_loc`` tensor is constructed directly on ``device``. - # The ``int(...item())`` reads against ``req_to_token`` are a CPU - # sync but only ever run from ``prepare()`` (outside any CUDA-graph - # capture window) — they are not in the forward path. + # out_cache_loc must be flattened in the q-token order the model layer + # projects, or K/V lands in the wrong requests' slots. if is_prefill: if extend_seq_lens_cpu is None: raise ValueError("prefill metadata requires extend_seq_lens_cpu") if prefix_lens is None: raise ValueError("prefill metadata requires prefix_lens") - prefix_lens_cpu = prefix_lens.to("cpu").tolist() if static_buffers is not None: prefix_buf = static_buffers["prefix_lens"] prefix_src = prefix_lens.to(device=device, dtype=torch.int32, non_blocking=True) @@ -563,17 +592,18 @@ def build_runtime_metadata_from_kv_manager( prefix_lens_dev = ( prefix_lens.to(device) if prefix_lens.device != device else prefix_lens ) - out_cache_loc_list: List[int] = [] cu_q: List[int] = [0] - req_to_token_cpu = req_to_token_fresh.to("cpu") - for b in range(batch): - pref = int(prefix_lens_cpu[b]) - ext = int(extend_seq_lens_cpu[b]) - for offset in range(ext): - slot = int(req_to_token_cpu[b, pref + offset].item()) - out_cache_loc_list.append(slot) - cu_q.append(cu_q[-1] + ext) + for ext in extend_seq_lens_cpu: + cu_q.append(cu_q[-1] + int(ext)) total_q = cu_q[-1] + cu_seqlens_q_src = torch.tensor(cu_q, dtype=torch.int32, device=device) + q_batch_row_src = torch.repeat_interleave( + torch.arange(batch, device=device, dtype=torch.int32), + torch.tensor(extend_seq_lens_cpu, dtype=torch.int64, device=device), + ) + q_positions_src, out_cache_loc_src = derive_q_positions_and_cache_slots( + req_to_token, prefix_lens_dev, cu_seqlens_q_src, q_batch_row_src + ) if static_buffers is not None: if total_q > static_buffers["max_num_tokens"]: raise ValueError( @@ -581,33 +611,22 @@ def build_runtime_metadata_from_kv_manager( f"is smaller than current total_q={total_q}" ) out_cache_loc_buf = static_buffers["out_cache_loc"] - out_cache_loc_src = torch.tensor(out_cache_loc_list, dtype=torch.int32, device=device) out_cache_loc_buf[:total_q].copy_(out_cache_loc_src, non_blocking=True) out_cache_loc = out_cache_loc_buf[:total_q] cu_seqlens_q_buf = static_buffers["cu_seqlens_q"] - cu_seqlens_q_src = torch.tensor(cu_q, dtype=torch.int32, device=device) cu_seqlens_q_buf[: batch + 1].copy_(cu_seqlens_q_src, non_blocking=True) cu_seqlens_q = cu_seqlens_q_buf[: batch + 1] - # Populate persistent q_batch_row / q_positions in-place so - # the inner metadata's prepare() can leave them alone. q_batch_row_buf = static_buffers["q_batch_row"] q_positions_buf = static_buffers["q_positions"] - for b in range(batch): - start, end = cu_q[b], cu_q[b + 1] - if end > start: - q_batch_row_buf[start:end] = b - pref = int(prefix_lens_cpu[b]) - offsets = ( - torch.arange(start, end, device=device, dtype=torch.int32) - start + pref - ) - q_positions_buf[start:end].copy_(offsets, non_blocking=True) + q_batch_row_buf[:total_q].copy_(q_batch_row_src, non_blocking=True) + q_positions_buf[:total_q].copy_(q_positions_src, non_blocking=True) q_batch_row = q_batch_row_buf[:total_q] q_positions = q_positions_buf[:total_q] else: - out_cache_loc = torch.tensor(out_cache_loc_list, dtype=torch.int32, device=device) - cu_seqlens_q = torch.tensor(cu_q, dtype=torch.int32, device=device) - q_batch_row = None - q_positions = None + out_cache_loc = out_cache_loc_src + cu_seqlens_q = cu_seqlens_q_src + q_batch_row = q_batch_row_src + q_positions = q_positions_src meta = MiniMaxM3TritonSparseAttentionMetadata( is_prefill=True, req_to_token=req_to_token, @@ -622,12 +641,7 @@ def build_runtime_metadata_from_kv_manager( ) else: # Decode: the new token sits at position seq_lens[b] - 1. - seq_lens_cpu_list = seq_lens_cpu.to("cpu").tolist() - out_cache_loc_list = [] - req_to_token_cpu = req_to_token_fresh.to("cpu") - for b in range(batch): - pos = int(seq_lens_cpu_list[b]) - 1 - out_cache_loc_list.append(int(req_to_token_cpu[b, pos].item())) + out_cache_loc_src = derive_decode_cache_slots(req_to_token, seq_lens_dev) if static_buffers is not None: if batch > static_buffers["max_num_tokens"]: raise ValueError( @@ -635,11 +649,10 @@ def build_runtime_metadata_from_kv_manager( f"is smaller than current batch={batch}" ) out_cache_loc_buf = static_buffers["out_cache_loc"] - out_cache_loc_src = torch.tensor(out_cache_loc_list, dtype=torch.int32, device=device) out_cache_loc_buf[:batch].copy_(out_cache_loc_src, non_blocking=True) out_cache_loc = out_cache_loc_buf[:batch] else: - out_cache_loc = torch.tensor(out_cache_loc_list, dtype=torch.int32, device=device) + out_cache_loc = out_cache_loc_src meta = MiniMaxM3TritonSparseAttentionMetadata( is_prefill=False, req_to_token=req_to_token, @@ -651,8 +664,12 @@ def build_runtime_metadata_from_kv_manager( return meta, out_cache_loc -class MiniMaxM3AttentionMetadata(AttentionMetadata): - """:class:`AttentionMetadata` that pre-builds MiniMax-M3 metadata. +class MiniMaxM3AttentionMetadata(TrtllmAttentionMetadata): + """:class:`TrtllmAttentionMetadata` that pre-builds MiniMax-M3 metadata. + + Subclasses :class:`TrtllmAttentionMetadata` (like + ``DSAtrtllmAttentionMetadata``) so one-model Eagle3 draft layers can run + :class:`TrtllmAttention` on this metadata. Overrides :meth:`prepare` so the M3-sparse :class:`MiniMaxM3TritonSparseAttentionMetadata` and the per-new-token @@ -808,23 +825,9 @@ def prepare(self) -> None: static_buffers = self._maybe_get_m3_static_buffers(cache_device, kv_cache_manager) - # Any batch containing a context (prefill or chunked extend) - # request takes the extend path. For prefill rows - # ``num_cached_per_seq`` is ``prefix_lens`` and the full new - # chunk is ``extend_seq_len``; for decode rows - # ``num_cached`` is ``kv_len - 1`` and ``extend_seq_len`` is - # 1, so the same builder produces the correct one-slot - # entry. Pure-decode batches (``num_contexts == 0``) still - # take the decode optimization for CUDA-graph warmup - # geometry. - # - # Mixed prefill+decode batches always take the extend path: - # the prefill kernel handles decode rows as 1-slot extends. - # The decode branch below is a pure-decode-only perf - # specialization. (iter-131 regression: previously a wrong - # predicate routed mixed batches into the decode branch and - # crashed in index_copy_.) - is_extend = num_contexts > 0 + # Multi-token generation rows (spec verify) also take the extend path; + # plain decode stays one token per row. + is_extend = num_contexts > 0 or int(seq_lens_cpu[:batch_size].max().item()) > 1 if is_extend: prefix_lens_list = [int(num_cached_per_seq[b]) for b in range(batch_size)] extend_seq_lens_cpu = [ @@ -862,11 +865,50 @@ def prepare(self) -> None: "out_cache_loc": out_cache_loc, } + def on_update_kv_lens(self) -> None: + """Re-derive the M3 attachment from the corrected ``kv_lens_cuda``. + + With the overlap scheduler and speculative decoding, prepare() runs + with optimistic lengths and the engine corrects ``kv_lens_cuda`` on + device before calling this (as DSAtrtllmAttentionMetadata does). + Device-only and idempotent. ``seq_lens_cpu`` / ``max_seqlen_k`` keep + the optimistic values; they only bound widths the kernels mask by + ``seq_lens``. + """ + super().on_update_kv_lens() + attachment = self.minimax_m3 + if not attachment: + return + meta = attachment["metadata"] + out_cache_loc = attachment["out_cache_loc"] + batch = int(meta.slot_ids.shape[0]) + kv_lens = self.kv_lens_cuda[:batch] + meta.seq_lens[:batch].copy_(kv_lens) + if meta.is_prefill: + # Only the K-side prefix moves with rejections; the Q-side layout + # (cu_seqlens_q, q_batch_row) is fixed per step. + total_q = int(meta.q_positions.shape[0]) + cu = meta.cu_seqlens_q + meta.prefix_lens[:batch].copy_(kv_lens - (cu[1 : batch + 1] - cu[:batch])) + q_positions, cache_slots = derive_q_positions_and_cache_slots( + meta.req_to_token, + meta.prefix_lens[:batch], + cu, + meta.q_batch_row[:total_q], + ) + meta.q_positions[:total_q].copy_(q_positions) + out_cache_loc[:total_q].copy_(cache_slots) + else: + # No-op today (decode rows are not corrected); kept for symmetry. + out_cache_loc[:batch].copy_(derive_decode_cache_slots(meta.req_to_token, kv_lens)) + __all__ = [ "MiniMaxM3AttentionMetadata", "MiniMaxM3TritonSparseAttentionMetadata", "allocate_minimax_m3_static_buffers", "build_runtime_metadata_from_kv_manager", + "derive_decode_cache_slots", + "derive_q_positions_and_cache_slots", "ensure_metadata_on_device", ] diff --git a/tensorrt_llm/_torch/models/modeling_minimaxm3.py b/tensorrt_llm/_torch/models/modeling_minimaxm3.py index 1176262d960e..2495e8dd446b 100644 --- a/tensorrt_llm/_torch/models/modeling_minimaxm3.py +++ b/tensorrt_llm/_torch/models/modeling_minimaxm3.py @@ -64,16 +64,12 @@ from ..modules.rms_norm import RMSNorm from ..moe.fused_moe import MiniMaxM3MoeRoutingMethod, SwigluBiasActivation, create_moe from ..pyexecutor.breakable_cuda_graph import eager_on_graph, is_in_breakable_cuda_graph +from ..speculative import SpecMetadata from ..utils import AuxStreamType, EventType, get_model_extra_attrs, is_torch_compiling from .checkpoints.base_weight_mapper import BaseWeightMapper from .checkpoints.hf.minimaxm3_weight_mapper import MINIMAX_M3_PARAMS_MAP, MiniMaxM3HfWeightMapper -from .modeling_utils import ( - DecoderModel, - DecoderModelForCausalLM, - ModelConfig, - filter_weights, - register_auto_model, -) +from .modeling_speculative import SpecDecOneEngineForCausalLM +from .modeling_utils import DecoderModel, ModelConfig, filter_weights, register_auto_model # Dense layers use SDPA with non-contiguous Q/K/V and a bool attn_mask. # Limit backends to memory-efficient and math; cuDNN SDPA fails for this layout, @@ -1211,7 +1207,14 @@ def _sdpa_dense_attention_core( # 7. Gather padded K/V for every batch row and run dense GQA. batch = int(m3_meta.slot_ids.shape[0]) - max_k = int(m3_meta.max_seqlen_k) + # Under CUDA graphs bake a fixed gather/mask width: min(page-table width, + # max_seq_len). The raw table width is inflated by the KV-estimation pass + # and would OOM the gather; the seq_lens mask hides the slack. + if attn_metadata.is_cuda_graph: + capacity = int(m3_meta.req_to_token.shape[1]) + max_k = min(capacity, int(attn_metadata.max_seq_len or capacity)) + else: + max_k = int(m3_meta.max_seqlen_k) if max_k <= 0: max_k = 1 # ``_gather_paged_batched`` decomposes the flat slot id into @@ -1238,15 +1241,11 @@ def _sdpa_dense_attention_core( f"by num_key_value_heads ({self.num_key_value_heads})" ) group = self.num_heads // max(self.num_key_value_heads, 1) - if group > 1: - k_padded = k_padded.repeat_interleave(group, dim=2) - v_padded = v_padded.repeat_interleave(group, dim=2) # Build the per-query attention mask that masks out padded KV # positions beyond each sequence's true ``seq_lens`` and (for # prefill) preserves causality. ``q_positions`` from the prefill - # metadata names each Q token's K-side position; for decode - # there is one Q token per request at position ``seq_lens - 1``. + # metadata names each Q token's K-side position. # The metadata tensors are produced by # :meth:`MiniMaxM3AttentionMetadata.prepare` on the cache # device, so ``.to(dtype=torch.long)`` is a same-device dtype @@ -1255,6 +1254,9 @@ def _sdpa_dense_attention_core( kv_positions = torch.arange(max_k, device=q.device).unsqueeze(0) # [1, max_k] if m3_meta.is_prefill: + if group > 1: + k_padded = k_padded.repeat_interleave(group, dim=2) + v_padded = v_padded.repeat_interleave(group, dim=2) # Prefill: build [total_q, max_k] mask using q_positions / q_batch_row. # Prefill never runs inside the CUDA-graph capture window # (capture is decode-only), so the per-batch Python loop and @@ -1295,34 +1297,31 @@ def _sdpa_dense_attention_core( ) # [1, H, q, d] output_view[start:end].copy_(out_b.squeeze(0).transpose(0, 1)) else: - # Decode: one Q token per request at position seq_lens - 1. - # Every input tensor here is already on q.device (set up by - # prepare()), so SDPA captures cleanly. + # Decode: one query token per request at position seq_lens - 1. valid = kv_positions < seq_lens_dev.unsqueeze(-1) # [batch, max_k] - q_b = q_view.unsqueeze(1).transpose(1, 2) # [batch, H, 1, d] - k_b = k_padded.transpose(1, 2) # [batch, H, k, d] - v_b = v_padded.transpose(1, 2) # [batch, H, k, d] + q_b = q_view.view(batch, 1, self.num_heads, self.head_dim).transpose( + 1, 2 + ) # [batch, H, 1, d] mask_b = valid.unsqueeze(1).unsqueeze(1) # [batch, 1, 1, k] + # Expand K/V one KV head at a time: all heads at once needs an + # O(batch * max_k * num_heads) temporary that overflows the CUDA-graph + # pool under attention DP. + out_b = q.new_empty(batch, self.num_heads, 1, self.head_dim) with sdpa_kernel(_DENSE_SDPA_BACKENDS): - out_b = torch.nn.functional.scaled_dot_product_attention( - q_b.to(q.dtype), - k_b.to(q.dtype), - v_b.to(q.dtype), - attn_mask=mask_b, - dropout_p=0.0, - is_causal=False, - ) # [batch, H, 1, d] - # Drop the singleton Q-length axis and write the resulting - # ``[batch, num_heads, head_dim]`` tensor into the final buffer. - # The prior ``.transpose(1, 2).reshape(batch, H, d)`` pattern - # was wrong: with ``H != head_dim`` (M3 TP=8 has H=8, d=128) - # the non-contiguous transpose forces ``reshape`` to copy the - # data in C-order under its current ``[batch, d, H]`` shape, - # then reinterpret as ``[batch, H, d]`` — which scrambles - # ``(head, head_dim)`` ordering and feeds permuted activations - # into ``o_proj``. Prefill is unaffected because its - # ``transpose(0, 1)`` runs between q-len and num_heads axes - # which the per-batch loop already laid out correctly. + for h in range(max(self.num_key_value_heads, 1)): + qh = slice(h * group, (h + 1) * group) + k_h = k_padded[:, :, h : h + 1].repeat_interleave(group, dim=2) + v_h = v_padded[:, :, h : h + 1].repeat_interleave(group, dim=2) + out_b[:, qh] = torch.nn.functional.scaled_dot_product_attention( + q_b[:, qh].to(q.dtype), + k_h.transpose(1, 2).to(q.dtype), + v_h.transpose(1, 2).to(q.dtype), + attn_mask=mask_b, + dropout_p=0.0, + is_causal=False, + ) # [batch, group, 1, d] + # Drop the singleton query axis; a transpose(1, 2).reshape would + # scramble (head, head_dim) when H != head_dim. output.view(batch, self.num_heads, self.head_dim).copy_(out_b.squeeze(2)) return output @@ -1747,6 +1746,7 @@ def forward( hidden_states: torch.Tensor, attn_metadata: AttentionMetadata, residual: Optional[torch.Tensor], + spec_metadata: Optional[SpecMetadata] = None, **kwargs, ) -> torch.Tensor: # Layer-0 prologue only. For every subsequent layer the input_layernorm @@ -1772,10 +1772,19 @@ def forward( **kwargs, ) + if spec_metadata is not None and spec_metadata.is_layer_capture(self.layer_idx): + # Eagle3 captures this layer's raw output, so keep the feed-forward + # AllReduce in the module and the boundary norm unfused (as DeepSeek-V3 + # does on capture layers). + self.post_feed_forward_fusion = False if self.block_sparse_moe is not None: - hidden_states, residual = self.forward_MoE(hidden_states, attn_metadata, residual) + hidden_states, residual = self.forward_MoE( + hidden_states, attn_metadata, residual, spec_metadata=spec_metadata + ) else: - hidden_states, residual = self.forward_mlp(hidden_states, residual) + hidden_states, residual = self.forward_mlp( + hidden_states, residual, spec_metadata=spec_metadata + ) return hidden_states, residual @@ -1849,6 +1858,7 @@ def forward_MoE( hidden_states: torch.Tensor, attn_metadata: AttentionMetadata, residual: torch.Tensor, + spec_metadata: Optional[SpecMetadata] = None, ) -> Tuple[torch.Tensor, torch.Tensor]: hidden_states, residual = self._apply_pre_feed_forward_norm(hidden_states, residual) @@ -1858,6 +1868,8 @@ def forward_MoE( final_all_reduce_params=self._feed_forward_all_reduce_params(), ) + if spec_metadata is not None and spec_metadata.is_layer_capture(self.layer_idx): + spec_metadata.maybe_capture_hidden_states(self.layer_idx, hidden_states, residual) hidden_states, residual = self._apply_next_layer_layernorm(hidden_states, residual) return hidden_states, residual @@ -1865,6 +1877,7 @@ def forward_mlp( self, hidden_states: torch.Tensor, residual: torch.Tensor, + spec_metadata: Optional[SpecMetadata] = None, ) -> Tuple[torch.Tensor, torch.Tensor]: hidden_states, residual = self._apply_pre_feed_forward_norm(hidden_states, residual) @@ -1873,6 +1886,8 @@ def forward_mlp( final_all_reduce_params=self._feed_forward_all_reduce_params(), ) + if spec_metadata is not None and spec_metadata.is_layer_capture(self.layer_idx): + spec_metadata.maybe_capture_hidden_states(self.layer_idx, hidden_states, residual) hidden_states, residual = self._apply_next_layer_layernorm(hidden_states, residual) return hidden_states, residual @@ -1936,6 +1951,7 @@ def forward( input_ids: Optional[torch.IntTensor] = None, position_ids: Optional[torch.IntTensor] = None, inputs_embeds: Optional[torch.FloatTensor] = None, + spec_metadata: Optional[SpecMetadata] = None, **kwargs, ) -> torch.Tensor: if (input_ids is None) ^ (inputs_embeds is not None): @@ -1952,6 +1968,7 @@ def forward( hidden_states=hidden_states, attn_metadata=attn_metadata, residual=residual, + spec_metadata=spec_metadata, ) # When setup_aliases has chained the final norm into the last decoder @@ -2031,7 +2048,7 @@ def _fold_gemma_boundary_norm_weights(weights): @register_auto_model("MiniMaxM3SparseForCausalLM") -class MiniMaxM3ForCausalLM(DecoderModelForCausalLM[MiniMaxM3Model, PretrainedConfig]): +class MiniMaxM3ForCausalLM(SpecDecOneEngineForCausalLM[MiniMaxM3Model, PretrainedConfig]): """Text-only M3 model.""" @classmethod @@ -2059,12 +2076,7 @@ def __init__(self, model_config: "ModelConfig[PretrainedConfig]"): raw_pretrained = model_config.pretrained_config if is_minimax_m3_vl_config(raw_pretrained): model_config = get_text_model_config(model_config) - super().__init__( - MiniMaxM3Model(model_config), - config=model_config, - hidden_size=model_config.pretrained_config.hidden_size, - vocab_size=model_config.pretrained_config.vocab_size, - ) + super().__init__(MiniMaxM3Model(model_config), model_config) def load_weights( self, diff --git a/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py b/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py index a3b8a1f4be9d..517eb5072d31 100644 --- a/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py +++ b/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py @@ -502,8 +502,24 @@ def _create_py_executor_impl( # drafters, which it stranded on their private max_seq_len-dense arena. is_standalone_drafter = (spec_config.spec_dec_mode.is_dflash() or spec_config.spec_dec_mode.is_dspark()) - if cache_transceiver_config is not None and not is_standalone_drafter: + # MiniMax-M3 supports one-model Eagle3 only; its drafter shares the + # target KV cache. + is_m3_eagle3 = (is_minimax_m3(m3_sparse_config) + and spec_config.spec_dec_mode.is_eagle3_one_model()) + if ((cache_transceiver_config is not None and not is_standalone_drafter) + or is_m3_eagle3): spec_config._allow_separate_draft_kv_cache = False + # The triton reference backend runs multi-token verify through its + # prefill builder, which cannot be CUDA-graph captured. + if (is_m3_eagle3 and m3_sparse_config.implementation != "msa" + and llm_args.cuda_graph_config is not None): + raise ValueError( + "MiniMax-M3 Eagle3 on the triton reference backend does not " + "support CUDA graphs; use implementation='msa' or set " + "cuda_graph_config=None.") + if is_m3_eagle3 and not spec_config.is_linear_tree: + raise ValueError( + "MiniMax-M3 Eagle3 supports the linear draft chain only.") # chunk_unit_size may be changed to 64 when using flash mla attn_runtime_features = AttentionRuntimeFeatures( diff --git a/tests/integration/defs/accuracy/references/gsm8k.yaml b/tests/integration/defs/accuracy/references/gsm8k.yaml index dffc36089764..7f4783acde49 100644 --- a/tests/integration/defs/accuracy/references/gsm8k.yaml +++ b/tests/integration/defs/accuracy/references/gsm8k.yaml @@ -478,6 +478,10 @@ nvidia/MiniMax-M3-NVFP4: - quant_algo: MIXED_PRECISION kv_cache_quant_algo: FP8 accuracy: 86 + - quant_algo: MIXED_PRECISION + kv_cache_quant_algo: FP8 + spec_dec_algo: Eagle3 + accuracy: 86 nvidia/NVIDIA-Nemotron-Nano-9B-v2: - accuracy: 85.027 - quant_algo: FP8 diff --git a/tests/integration/defs/accuracy/references/mmlu.yaml b/tests/integration/defs/accuracy/references/mmlu.yaml index da6cb0b2fabb..1d2a1defa4cc 100644 --- a/tests/integration/defs/accuracy/references/mmlu.yaml +++ b/tests/integration/defs/accuracy/references/mmlu.yaml @@ -151,6 +151,10 @@ nvidia/MiniMax-M3-NVFP4: - quant_algo: MIXED_PRECISION kv_cache_quant_algo: FP8 accuracy: 81 + - quant_algo: MIXED_PRECISION + kv_cache_quant_algo: FP8 + spec_dec_algo: Eagle3 + accuracy: 81 moonshotai/Kimi-K2-Instruct: - quant_algo: FP8_BLOCK_SCALES accuracy: 87.65 diff --git a/tests/integration/defs/accuracy/test_llm_api_pytorch.py b/tests/integration/defs/accuracy/test_llm_api_pytorch.py index ff009c92a449..b3362c103214 100644 --- a/tests/integration/defs/accuracy/test_llm_api_pytorch.py +++ b/tests/integration/defs/accuracy/test_llm_api_pytorch.py @@ -7946,3 +7946,104 @@ def test_nvfp4(self, use_msa): task.evaluate(llm) task = GSM8K(model_name) task.evaluate(llm) + + @pytest.mark.skip_less_device(4) + @pytest.mark.skip_less_device_memory(140000) + @parametrize_with_ids("overlap_scheduler", [False, True]) + @parametrize_with_ids("attention_dp", [False, True]) + @parametrize_with_ids("tp_size,ep_size", [(4, 4)]) + def test_nvfp4_eagle3(self, tp_size, ep_size, attention_dp, + overlap_scheduler): + # One-model Eagle3 on the MSA backend with an FP8 KV cache and CUDA + # graphs; the GQA drafter shares the target KV cache. MMLU + GSM8K plus + # a chat-GSM8K acceptance probe, since accuracy alone does not notice a + # corrupted drafter KV. + from tensorrt_llm._torch.attention.backends.sparse.minimax_m3.kernels.msa_utils import \ + msa_package_available + if not msa_package_available(): + pytest.skip("MSA kernels (fmha_sm100) not available") + model_name = "nvidia/MiniMax-M3-NVFP4" + model_path = f"{llm_models_root()}/MiniMax-M3-NVFP4" + max_draft_len = 3 + spec_config = Eagle3DecodingConfig( + max_draft_len=max_draft_len, + speculative_model=f"{llm_models_root()}/MiniMax-M3-EAGLE3-GQA", + ) + # The MSA path runs an FP8 KV cache, as in test_nvfp4. + kv_cache_config = KvCacheConfig(free_gpu_memory_fraction=0.6, + enable_block_reuse=False, + dtype="fp8") + with LLM( + model_path, + tensor_parallel_size=tp_size, + moe_expert_parallel_size=ep_size, + kv_cache_config=kv_cache_config, + sparse_attention_config=MiniMaxM3SparseAttentionConfig( + implementation="msa", indexer_kv_dtype="fp8"), + moe_config=MoeConfig(backend="CUTLASS"), + max_seq_len=4096, + # fmha_sm100 caps total_q x heads at 65536; with 4 verify + # tokens per row that is 512 (256 with unsharded heads). + max_batch_size=256 if attention_dp else 512, + speculative_config=spec_config, + cuda_graph_config=CudaGraphConfig( + enable_padding=True, + max_batch_size=64 if attention_dp else 128, + ), + disable_overlap_scheduler=not overlap_scheduler, + enable_attention_dp=attention_dp, + enable_iter_perf_stats=True, + # Keep the whole acceptance probe: the default buffer holds + # only the latest 1000 iterations. + max_stats_len=5000, + trust_remote_code=True) as llm: + assert llm.args.quant_config.quant_algo == QuantAlgo.MIXED_PRECISION + + def drain_spec_stats(llm): + drafted = accepted = steps = 0 + for s in llm.get_stats(timeout=2): + s = json.loads(s) if isinstance(s, str) else s + sd = s.get("specDecodingStats") or {} + drafted += sd.get("numDraftTokens", 0) + accepted += sd.get("numAcceptedTokens", 0) + steps += sd.get("numRequestsWithDraftTokens", 0) + return drafted, accepted, steps + + task = MMLU(model_name) + task.evaluate(llm) + task = GSM8K(model_name) + task.evaluate(llm) + + # Acceptance probe: 200 chat-format GSM8K questions, greedy, 512 tokens. + questions = [ + r["question"] + for r in load_dataset(GSM8K.DATASET_DIR, "main", split="test") + ][:200] + chat_prompts = [ + llm.tokenizer.apply_chat_template([{ + "role": "user", + "content": q + }], + tokenize=False, + add_generation_prompt=True) + for q in questions + ] + drain_spec_stats(llm) + llm.generate(chat_prompts, + SamplingParams(max_tokens=512, temperature=0)) + drafted, accepted, steps = drain_spec_stats(llm) + assert steps > 0, "no speculative iterations recorded" + chat_rate = accepted / drafted + chat_length = 1 + accepted / steps + # Reference: the MHA drafter card (Inferact/MiniMax-M3-EAGLE3) + # reports 0.839 / 3.518; the GQA head measures 0.838 / 3.515 here. + print(f"MiniMax-M3 Eagle3 chat-GSM8K acceptance: rate=" + f"{chat_rate:.3f}, mean acceptance length=" + f"{chat_length:.3f} ({steps} spec iterations)") + assert chat_rate > 0.80, \ + f"Eagle3 chat-GSM8K acceptance rate too low: {chat_rate:.3f} " \ + f"(threshold 0.80, reference 0.839 from the drafter card)" + assert chat_length > 3.4, \ + f"Eagle3 chat-GSM8K acceptance length too low: " \ + f"{chat_length:.3f} (threshold 3.4, reference 3.518 from " \ + f"the drafter card)" diff --git a/tests/integration/test_lists/qa/llm_function_core.txt b/tests/integration/test_lists/qa/llm_function_core.txt index 2490f94f8613..45c482876549 100644 --- a/tests/integration/test_lists/qa/llm_function_core.txt +++ b/tests/integration/test_lists/qa/llm_function_core.txt @@ -491,6 +491,7 @@ accuracy/test_llm_api_pytorch.py::TestMiniMaxM3::test_mxfp8_piecewise_cuda_graph accuracy/test_llm_api_pytorch.py::TestMiniMaxM3::test_nvfp4[use_msa=False] TIMEOUT (180) accuracy/test_llm_api_pytorch.py::TestMiniMaxM3::test_nvfp4[use_msa=True] accuracy/test_llm_api_pytorch.py::TestMiniMaxM3::test_nvfp4[use_msa=True] TIMEOUT (60) +accuracy/test_llm_api_pytorch.py::TestMiniMaxM3::test_nvfp4_eagle3[tp_size=4-ep_size=4-attention_dp=False-overlap_scheduler=True] accuracy/test_llm_api_pytorch.py::TestMinistral8BInstruct::test_auto_dtype accuracy/test_llm_api_pytorch.py::TestMinistral8BInstruct::test_fp8 accuracy/test_llm_api_pytorch.py::TestMistralLarge3_675B::test_fp8[latency_moe_deepgemm] diff --git a/tests/integration/test_lists/test-db/l0_dgx_b200.yml b/tests/integration/test_lists/test-db/l0_dgx_b200.yml index 381a82ce99de..982755de8e34 100644 --- a/tests/integration/test_lists/test-db/l0_dgx_b200.yml +++ b/tests/integration/test_lists/test-db/l0_dgx_b200.yml @@ -53,6 +53,7 @@ l0_dgx_b200: - disaggregated/test_disaggregated.py::test_disaggregated_gpt_oss_120b_harmony[gpt_oss/gpt-oss-120b] - accuracy/test_llm_api_pytorch.py::TestDeepSeekR1::test_nvfp4_multi_gpus[latency_adp_lmtp_tp4] - accuracy/test_llm_api_pytorch.py::TestMiniMaxM3::test_nvfp4[use_msa=True] TIMEOUT (60) + - accuracy/test_llm_api_pytorch.py::TestMiniMaxM3::test_nvfp4_eagle3[tp_size=4-ep_size=4-attention_dp=False-overlap_scheduler=True] - accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_bfloat16_4gpus[pp4-mtp_nextn=0-attention_dp=False-cuda_graph=False-overlap_scheduler=False-torch_compile=False] TIMEOUT (60) - unittest/_torch/modeling/test_modeling_deepseekv4.py - accuracy/test_llm_api_pytorch.py::TestDeepSeekV4Flash::test_auto_dtype TIMEOUT (60) diff --git a/tests/unittest/_torch/attention/sparse/msa/test_minimax_m3_shared_draft_layers.py b/tests/unittest/_torch/attention/sparse/msa/test_minimax_m3_shared_draft_layers.py new file mode 100644 index 000000000000..ceda71f86b41 --- /dev/null +++ b/tests/unittest/_torch/attention/sparse/msa/test_minimax_m3_shared_draft_layers.py @@ -0,0 +1,175 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Pure-logic tests for MiniMax-M3's shared Eagle3 draft layers: each gets its own +virtual attention-op pool rooted at its K page inside the mega-slot. +""" + +from types import SimpleNamespace + +import pytest +import torch + +from tensorrt_llm._torch.attention.backends.sparse.minimax_m3 import ( + cache_manager as m3_cache_manager, +) +from tensorrt_llm._torch.attention.backends.sparse.minimax_m3.cache_manager import ( + MiniMaxM3KVCacheManagerV2, + derive_shared_draft_layout, + extend_attention_op_pools_for_shared_draft_layers, + shared_draft_layer_count, +) +from tensorrt_llm._torch.pyexecutor.kv_cache.kv_cache_manager_v2 import KVCacheManagerV2 + +DRAFT_LOCAL_LAYER = 60 +SCALE = 179 # sub-pages per M3 mega-slot: 3 dense x 2 + 57 sparse x 3 + draft x 2 +DRAFT_K_ADDR = 0x7000_0000 + + +def test_virtual_pool_is_rooted_at_the_draft_k_page(): + pool_pointers = torch.tensor([[0x6000_0000, 0]], dtype=torch.int64) + pool_mapping = torch.tensor([[0, i] for i in range(61)], dtype=torch.int32) + + pointers, mapping, index_scales, kv_offsets, op_pools = ( + extend_attention_op_pools_for_shared_draft_layers( + pool_pointers, pool_mapping, 1, [(DRAFT_LOCAL_LAYER, DRAFT_K_ADDR, SCALE)] + ) + ) + + assert pointers.tolist() == [[0x6000_0000, 0], [DRAFT_K_ADDR, 0]] + # The draft layer moves to the new pool at offset 0; target rows are untouched. + assert mapping[DRAFT_LOCAL_LAYER].tolist() == [1, 0] + assert mapping[:DRAFT_LOCAL_LAYER].tolist() == pool_mapping[:DRAFT_LOCAL_LAYER].tolist() + # Slot s -> page s * SCALE for K and s * SCALE + 1 for V. + assert index_scales.tolist() == [SCALE] + assert kv_offsets.tolist() == [1] + assert op_pools == [(1, 0)] + + +def test_block_offset_copy_fills_the_virtual_pool_from_the_source_pool(monkeypatch): + """The base copy fills the storage pools; the override then fills the virtual pool + from the source pool's slot ids, and does nothing extra without draft layers. + """ + calls = [] + + def fake_base_copy( + self, dst_tensor, request_ids, beam_width, num_contexts, num_seqs, max_blocks=None + ): + calls.append(("base", request_ids, num_seqs, max_blocks)) + + def fake_device_copy(host, dst, copy_idx, index_scales, kv_offsets, stream): + calls.append( + ("virtual", host, dst, copy_idx, index_scales.tolist(), kv_offsets.tolist(), stream) + ) + + monkeypatch.setattr(KVCacheManagerV2, "copy_batch_block_offsets", fake_base_copy) + monkeypatch.setattr(m3_cache_manager, "copy_batch_block_offsets_to_device", fake_device_copy) + + manager = MiniMaxM3KVCacheManagerV2.__new__(MiniMaxM3KVCacheManagerV2) + manager._draft_op_pools = ((1, 0),) + manager._draft_index_scales = torch.tensor([SCALE], dtype=torch.int32) + manager._draft_kv_offsets = torch.tensor([1], dtype=torch.int32) + manager.host_kv_cache_block_offsets = torch.zeros((1, 4, 2, 8), dtype=torch.int32) + copy_idx = torch.tensor([2, 0], dtype=torch.int32) + manager.index_mapper = SimpleNamespace(get_copy_index=lambda ids, nc, bw: copy_idx) + manager._stream = SimpleNamespace(cuda_stream=1234) + dst = torch.zeros((2, 4, 2, 8), dtype=torch.int32) + + manager.copy_batch_block_offsets(dst, [7, 9], 1, 0, 2, max_blocks=5) + + assert calls[0] == ("base", [7, 9], 2, 5) + kind, host, dst_slice, idx, scales, offsets, stream = calls[1] + assert kind == "virtual" + assert ( + host.shape == (1, 4, 2, 8) + and host.data_ptr() == manager.host_kv_cache_block_offsets.data_ptr() + ) + assert dst_slice.shape == (1, 4, 2, 8) and dst_slice.data_ptr() == dst[1].data_ptr() + assert idx is copy_idx + assert scales == [SCALE] and offsets == [1] and stream == 1234 + + # Non-speculative MiniMax-M3: the base copy is all that runs. + calls.clear() + plain = MiniMaxM3KVCacheManagerV2.__new__(MiniMaxM3KVCacheManagerV2) + plain.copy_batch_block_offsets(dst, [7], 1, 0, 1) + assert calls == [("base", [7], 1, None)] + + +def test_per_layer_page_tables_get_no_virtual_pools(monkeypatch): + """With per-layer page tables every layer already has its own pool.""" + pointers = torch.tensor([[0x6000_0000 + i, 0] for i in range(61)]) + mapping = torch.tensor([[i, 0] for i in range(61)], dtype=torch.int32) + + def fake_base_prepare(self, index_mapper_capacity): + self._use_per_layer_page_tables = True + self.kv_cache_pool_pointers = pointers.clone() + self.kv_cache_pool_mapping = mapping.clone() + self.num_attention_op_pools = 61 + + monkeypatch.setattr(KVCacheManagerV2, "_prepare_page_table_tensor", fake_base_prepare) + manager = MiniMaxM3KVCacheManagerV2.__new__(MiniMaxM3KVCacheManagerV2) + manager._shared_draft_layer_ids = [DRAFT_LOCAL_LAYER] + manager.layer_offsets = {i: i for i in range(61)} + manager.is_draft = False + manager.enable_swa_scratch_reuse = False + manager.tokens_per_block = 128 + + manager._prepare_page_table_tensor(8) + + assert torch.equal(manager.kv_cache_pool_pointers, pointers) + assert torch.equal(manager.kv_cache_pool_mapping, mapping) + assert manager.num_attention_op_pools == 61 + assert manager._draft_op_pools == () + assert manager.trtllm_gen_extra_tokens_per_block == frozenset({128}) + + +def test_update_resources_refuses_tree_relocation(monkeypatch): + """Linear acceptance rewinds through the base; tree acceptance is refused.""" + calls = [] + monkeypatch.setattr(KVCacheManagerV2, "update_resources", lambda self, *a: calls.append(a)) + manager = MiniMaxM3KVCacheManagerV2.__new__(MiniMaxM3KVCacheManagerV2) + linear = SimpleNamespace( + py_num_accepted_draft_tokens=2, py_num_accepted_draft_tokens_indices=[] + ) + tree = SimpleNamespace( + py_num_accepted_draft_tokens=2, py_num_accepted_draft_tokens_indices=[0, 2] + ) + + manager.update_resources(SimpleNamespace(generation_requests=[linear]), None, 2) + assert len(calls) == 1 + + with pytest.raises(NotImplementedError, match="relocate"): + manager.update_resources(SimpleNamespace(generation_requests=[linear, tree]), None, 2) + + +def test_draft_layout_locates_the_appended_tail(): + # num_layers is the target count; a per-layer head list already includes the + # draft tail. + assert derive_shared_draft_layout(60, 4, 1) == ([60], 60) + assert derive_shared_draft_layout(60, [4] * 60 + [64], 1) == ([60], 60) + assert derive_shared_draft_layout(60, 4, 0) == ([], 60) + assert derive_shared_draft_layout(None, 4, 1) == ([], None) + + +def test_shared_draft_layer_count_follows_the_base_manager(): + # Same rule as get_pp_layers: a spec config and no layer_mask. + mode = SimpleNamespace( + is_mtp_eagle_one_model=lambda: False, + is_mtp_vanilla=lambda: False, + is_eagle3_one_model=lambda: True, + ) + eagle3 = SimpleNamespace(spec_dec_mode=mode, _num_draft_hidden_layers=None) + assert shared_draft_layer_count(eagle3, None) == 1 + assert shared_draft_layer_count(eagle3, [True] * 60) == 0 + assert shared_draft_layer_count(None, None) == 0 diff --git a/tests/unittest/_torch/attention/sparse/msa/test_msa_backend.py b/tests/unittest/_torch/attention/sparse/msa/test_msa_backend.py index bba4c7918d58..210db3093c41 100644 --- a/tests/unittest/_torch/attention/sparse/msa/test_msa_backend.py +++ b/tests/unittest/_torch/attention/sparse/msa/test_msa_backend.py @@ -192,9 +192,12 @@ def test_msa_metadata_rejects_undersized_max_score_buffer(): metadata_cls = MiniMaxM3MsaSparseAttention.Metadata metadata = metadata_cls.__new__(metadata_cls) # Flat backing store sized for 4 heads * 8 k-tiles * 2 batch = 64 elements, - # too small for the plan's required 4 * 16 * 2 = 128. + # too small for the plan's required 4 * 16 * 2 = 128. One query token per + # decode row, so the token bound equals the batch here. metadata.msa_max_score = torch.zeros(4 * 8 * 2) metadata.kv_cache_manager = None + metadata.max_num_sequences = 2 + metadata.max_num_tokens = 2 with pytest.raises(ValueError, match=r"msa_max_score backing store"): metadata._ensure_msa_decode_scratch_buffers( @@ -417,6 +420,31 @@ def test_msa_paged_hnd_input_materializes_unaligned_outer_stride() -> None: torch.testing.assert_close(prepared, view) +def test_per_token_valid_blocks_multi_token_decode(): + """Spec-verify decode rows get one entry per query token, following the causal + ladder inside the verify window. + """ + from tensorrt_llm._torch.attention.backends.sparse.minimax_m3.kernels.msa_utils import ( + per_token_valid_blocks, + ) + + # One request verifying 4 tokens against kv_len 10 (offset 6): token t + # attends 7 + t positions; with 2-token blocks that is ceil((7+t)/2). + qo = torch.tensor([4], dtype=torch.int32) + kv = torch.tensor([10], dtype=torch.int32) + off = torch.tensor([6], dtype=torch.int32) + n_valid = per_token_valid_blocks(qo, kv, off, causal=True, block_size=2) + assert n_valid.tolist() == [4, 4, 5, 5] + + # Mixed batch: an ordinary decode row (qo=1) alongside a verify row. + qo = torch.tensor([1, 3], dtype=torch.int32) + kv = torch.tensor([9, 6], dtype=torch.int32) + off = kv - qo + n_valid = per_token_valid_blocks(qo, kv, off, causal=True, block_size=4) + # Row 0: 9 positions -> 3 blocks. Row 1 tokens attend 4, 5, 6 -> 1, 2, 2. + assert n_valid.tolist() == [3, 1, 2, 2] + + def test_msa_index_k_uses_hnd_cache_view_and_writer(): metadata_cls = MiniMaxM3MsaSparseAttention.Metadata metadata = metadata_cls.__new__(metadata_cls) @@ -1380,3 +1408,171 @@ def run_gqa(output, *, token_first, token_last, row_first, num_rows): torch.testing.assert_close( split[num_ctx_tokens:], reference[num_ctx_tokens:], rtol=6e-2, atol=6e-2 ) + + +def test_msa_scratch_sizing_covers_spec_verify_tokens(): + """With Eagle3 a decode step has 1 + draft_len query tokens per request, so the + proxy scratch must be sized by tokens, not by batch. The draft length comes + from the KV cache manager, which knows the speculative config at build time. + """ + metadata_cls = MiniMaxM3MsaSparseAttention.Metadata + metadata = metadata_cls.__new__(metadata_cls) + # 2 sequences, 4 tokens each (draft_len=3): 8 decode tokens per step. + metadata.kv_cache_manager = SimpleNamespace(max_total_draft_tokens=3) + metadata.max_num_sequences = 2 + metadata.max_num_tokens = 8 + # Store sized for batch-only sizing (4 heads * 16 k-tiles * 2), which is + # too small once tokens are accounted for (4 * 16 * 8). + metadata.msa_max_score = torch.zeros(4 * 16 * 2) + + with pytest.raises(ValueError, match=r"msa_max_score backing store"): + metadata._ensure_msa_decode_scratch_buffers( + num_index_heads=4, + max_batch=2, + capture_graph=False, + required_max_k_tiles=16, + ) + + +def _kv_lens_update_metadata( + monkeypatch, *, qo_lens, kv_staged, kv_corrected, page_size, prefill_counts=False +): + """Metadata with just enough staged state for on_update_kv_lens: requests with + several query tokens each, optimistic kv_lens from prepare(), then the + corrected kv_lens. Everything lives on the host. `prefill_counts` stages the + per-step valid-block buffer a mixed step uses in place of the graph-stable + one. + """ + from tensorrt_llm._torch.attention.backends.sparse.minimax_m3 import msa_backend + + # The base hook only touches MLA caches this metadata never builds. + monkeypatch.setattr(msa_backend.TrtllmAttentionMetadata, "on_update_kv_lens", lambda self: None) + metadata_cls = MiniMaxM3MsaSparseAttention.Metadata + metadata = metadata_cls.__new__(metadata_cls) + qo = torch.tensor(qo_lens, dtype=torch.int32) + batch = int(qo.shape[0]) + total_q = int(qo.sum()) + metadata._seq_lens = qo + metadata._seq_lens_cuda = qo.clone() + metadata._num_tokens = total_q + metadata.kv_lens_cuda = torch.tensor(kv_corrected, dtype=torch.int32) + metadata.msa_kv_lens_staged = torch.tensor(kv_staged, dtype=torch.int32) + # What prepare() handed the decode kernels, before the correction. + metadata.msa_seq_lens_cuda = torch.tensor(kv_staged, dtype=torch.int32) + metadata.kv_cache_manager = SimpleNamespace(tokens_per_block=page_size) + # Page table: request b's page p is block 100 // page_size * b + p, so its + # position pos maps to slot 100 * b + pos (16 positions per request). + assert 100 % page_size == 0 + metadata.msa_block_table = (100 // page_size) * torch.arange( + batch, dtype=torch.int32 + ).unsqueeze(1) + torch.arange(16 // page_size, dtype=torch.int32).unsqueeze(0) + qo_long = qo.to(torch.long) + metadata.msa_q_batch_row = torch.repeat_interleave( + torch.arange(batch, dtype=torch.int32), qo_long + ) + starts = torch.cumsum(qo_long, 0) - qo_long + metadata.msa_q_intra = ( + torch.arange(total_q, dtype=torch.int64) - torch.repeat_interleave(starts, qo_long) + ).to(torch.int32) + metadata.msa_out_cache_loc = torch.full((total_q + 3,), -1, dtype=torch.int32) + metadata.msa_n_valid_blocks = torch.zeros(total_q + 3, dtype=torch.int32) + metadata._msa_prefill_n_valid_blocks = ( + torch.zeros(total_q, dtype=torch.int32) if prefill_counts else None + ) + metadata._msa_fields_ready = True + metadata._msa_kv_lens_dynamic = True + return metadata + + +def test_on_update_kv_lens_rederives_lengths_slots_and_counts(monkeypatch): + """Request 0 loses one rejected draft token (staged 9 -> corrected 8), request 1 + is unchanged. The per-request length the decode kernels read, the write + slots and the per-token valid-block counts must all follow. + """ + metadata = _kv_lens_update_metadata( + monkeypatch, qo_lens=(2, 3), kv_staged=(9, 12), kv_corrected=(8, 12), page_size=4 + ) + + metadata.on_update_kv_lens() + + # One length per request, in the buffer every decode kernel reads. + assert metadata.msa_seq_lens_cuda.tolist() == [8, 12] + # Token positions: request 0 attends 8 tokens with 2 queries -> 6, 7; + # request 1 attends 12 with 3 queries -> 9, 10, 11. + assert metadata.msa_out_cache_loc[:5].tolist() == [6, 7, 109, 110, 111] + assert metadata.msa_out_cache_loc[5:].tolist() == [-1, -1, -1] + # ceil((pos + 1) / 4) per token, into the graph-stable buffer on a + # pure-decode step; the tail past the step's tokens is left alone. + assert metadata.msa_n_valid_blocks[:5].tolist() == [2, 2, 3, 3, 3] + assert metadata.msa_n_valid_blocks[5:].tolist() == [0, 0, 0] + + +def test_on_update_kv_lens_clamps_to_the_staged_lens(monkeypatch): + """A correction can only shrink lengths; anything above the staged value is + clamped, so the staged page table stays valid. + """ + metadata = _kv_lens_update_metadata( + monkeypatch, qo_lens=(1, 1), kv_staged=(5, 7), kv_corrected=(9, 7), page_size=4 + ) + + metadata.on_update_kv_lens() + + assert metadata.msa_seq_lens_cuda.tolist() == [5, 7] + assert metadata.msa_out_cache_loc[:2].tolist() == [4, 106] + assert metadata.msa_n_valid_blocks[:2].tolist() == [2, 2] + + +def test_on_update_kv_lens_patches_a_mixed_steps_per_step_counts(monkeypatch): + """A mixed step stages its valid-block counts in the per-step prefill buffer + rather than the graph-stable one, and the correction must land there. The + context row (row 0) is untouched by the correction; the generation row + (row 1) loses one token. + """ + metadata = _kv_lens_update_metadata( + monkeypatch, + qo_lens=(5, 3), + kv_staged=(5, 12), + kv_corrected=(5, 11), + page_size=4, + prefill_counts=True, + ) + + metadata.on_update_kv_lens() + + assert metadata.msa_seq_lens_cuda.tolist() == [5, 11] + # Context row: positions 0..4 -> slots 0..4. Generation row attends 11 + # tokens with 3 queries: positions 8, 9, 10. + assert metadata.msa_out_cache_loc[:8].tolist() == [0, 1, 2, 3, 4, 108, 109, 110] + # ceil((pos + 1) / 4): context [1, 1, 1, 1, 2], generation [3, 3, 3]. + assert metadata._msa_prefill_n_valid_blocks.tolist() == [1, 1, 1, 1, 2, 3, 3, 3] + assert metadata.msa_n_valid_blocks.tolist() == [0] * 11 + + +def test_on_update_kv_lens_is_a_noop_without_speculative_decoding(monkeypatch): + """Without speculative decoding the patch and its staging are skipped. The gate + also sees max_total_draft_tokens, set by update_spec_dec_param, which + covers draft_len=1 under trtllm-gen (no extra KV tokens, linear-tree spec + decoding reported disabled). + """ + metadata = _kv_lens_update_metadata( + monkeypatch, qo_lens=(1, 1), kv_staged=(4, 6), kv_corrected=(3, 5), page_size=4 + ) + metadata.max_total_draft_tokens = None + metadata.draft_kv_cache_manager = None + metadata.is_spec_decoding_enabled = False + metadata.kv_cache_params = SimpleNamespace(num_extra_kv_tokens=0) + metadata.runtime_features = None + assert metadata._msa_kv_lens_may_change() is False + metadata.max_total_draft_tokens = 1 + assert metadata._msa_kv_lens_may_change() is True + metadata.max_total_draft_tokens = None + metadata.kv_cache_params = SimpleNamespace(num_extra_kv_tokens=2) + assert metadata._msa_kv_lens_may_change() is True + + metadata._msa_kv_lens_dynamic = False + + metadata.on_update_kv_lens() + + assert metadata.msa_seq_lens_cuda.tolist() == [4, 6] + assert metadata.msa_out_cache_loc.tolist() == [-1] * 5 + assert metadata.msa_n_valid_blocks.tolist() == [0] * 5 diff --git a/tests/unittest/_torch/models/checkpoints/hf/test_minimaxm3_weight_mapper.py b/tests/unittest/_torch/models/checkpoints/hf/test_minimaxm3_weight_mapper.py index 896539beaa71..04e479f05fd8 100644 --- a/tests/unittest/_torch/models/checkpoints/hf/test_minimaxm3_weight_mapper.py +++ b/tests/unittest/_torch/models/checkpoints/hf/test_minimaxm3_weight_mapper.py @@ -184,6 +184,8 @@ def test_load_weights_accepts_base_mapper_without_params_map() -> None: model = object.__new__(MiniMaxM3ForCausalLM) torch.nn.Module.__init__(model) model.model_config = model_config + # Set by SpecDecOneEngineForCausalLM.__init__, which this test bypasses. + model.spec_config = None mapper = HfWeightMapper() source_name = "model.layers.3.block_sparse_moe.e_score_correction_bias" target_name = "model.layers.3.block_sparse_moe.gate.e_score_correction_bias"