diff --git a/tensorrt_llm/_torch/attention_backend/fmha/msa_sparse_gqa.py b/tensorrt_llm/_torch/attention_backend/fmha/msa_sparse_gqa.py index d6dcf26fd834..d71c083c76ca 100644 --- a/tensorrt_llm/_torch/attention_backend/fmha/msa_sparse_gqa.py +++ b/tensorrt_llm/_torch/attention_backend/fmha/msa_sparse_gqa.py @@ -124,9 +124,20 @@ def run_msa_paged_gqa( with the sparse plan) and the dense layers (kv_block_indexes None, with the dense plan, attending the full page table). fmha_sm100 reads the paged cache directly, so the new-token K/V must be resident before the run. + + This is also where a step splits by request phase. TensorRT-LLM orders a + batch context-first, so the generation requests are its token suffix: a + ported decode kernel takes that suffix and fmha_sm100 keeps the context + prefix, running under the plan prepare() built over those rows alone (see + _msa_fmha_plan_rows). The prefix is empty on a pure-decode step, so one code + path covers both. The split lives here rather than in a PhasedFmha subclass + because MiniMaxM3MsaSparseAttention.forward_prepopulated_kv also calls this + helper directly, bypassing TrtllmAttention.forward. """ from tensorrt_llm._torch.attention_backend.sparse.minimax_m3.msa_utils import ( + msa_decode_span_bounds, msa_paged_kv, + msa_ported_decode_active, write_msa_main_kv, ) @@ -155,29 +166,140 @@ def run_msa_paged_gqa( k_paged, v_paged = msa_paged_kv(kv_cache_manager, layer_idx) sm_scale = (head_dim**-0.5) / float(attn.q_scaling) + # Query tokens and batch rows the ported kernels own; (num_tokens, batch) + # of them on a pure-decode step, the trailing generation slice on a mixed + # one. gen_tok0 is 0 whenever nothing is ported. + gen_tok0, gen_row0, gen_row1, decode_query_len = msa_decode_span_bounds(metadata, num_tokens) + # Leading query tokens fmha_sm100 must still run: the whole batch until a + # ported kernel takes the generation slice, then the context prefix alone. + fmha_tokens = num_tokens + ported = msa_ported_decode_active(metadata) + + if kv_block_indexes is not None and ported: + from tensorrt_llm._torch.attention_backend.sparse.minimax_m3.triton_sparse_decode import ( + minimax_m3_sparse_attn_decode, + ) + + # The Triton kernel dequantizes the FP8 cache into q's dtype, so q stays + # wide here. An already-E4M3 q from the fused producer widens exactly, + # leaving the same values the fmha_sm100 path would have used. + gen_q = q_view[gen_tok0:] + if gen_q.dtype != out_view.dtype: + gen_q = gen_q.to(out_view.dtype) + minimax_m3_sparse_attn_decode( + gen_q, + k_paged, + v_paged, + # [total_q, num_kv_heads, topk] -> head-major, contiguous when the + # indexer emitted a head-major table and the slice is the whole + # batch (see msa_ported_decode_active, which both sites read). The + # kernel takes every stride, so a mixed step's strided suffix is fine. + kv_block_indexes[gen_tok0:].permute(1, 0, 2), + metadata.msa_block_table[gen_row0:gen_row1], + metadata.msa_seq_lens_cuda[gen_row0:gen_row1], + sm_scale=sm_scale, + output=out_view[gen_tok0:], + decode_query_len=decode_query_len, + ) + fmha_tokens = gen_tok0 + + elif kv_block_indexes is None and ported: + from tensorrt_llm._torch.attention_backend.sparse.minimax_m3.trtllm_gen_dense_decode import ( + dense_decode_unsupported_reason, + minimax_m3_trtllm_gen_dense_decode, + ) + + unsupported = dense_decode_unsupported_reason(kv_cache_manager, head_dim) + if unsupported is not None: + raise RuntimeError( + "MiniMax-M3 resolved a generation span for this step's dense " + f"layers and skipped the fmha_sm100 dense plan, but {unsupported} " + "The two must agree, and there is no plan left to run the span; " + "see _resolve_decode_kernels." + ) + # The sub-page block table prepare() staged, if it could; the kernel + # expands its own when the factor does not match this layer's. + staged_subpage_rows = getattr(metadata, "msa_subpage_rows", None) + staged_table, staged_factor = ( + staged_subpage_rows(gen_row0, gen_row1) + if staged_subpage_rows is not None + else (None, 0) + ) + minimax_m3_trtllm_gen_dense_decode( + q_view[gen_tok0:], + kv_cache_manager, + layer_idx, + metadata.msa_block_table[gen_row0:gen_row1], + metadata.msa_seq_lens_cuda[gen_row0:gen_row1], + sm_scale=sm_scale, + output=out_view[gen_tok0:], + decode_query_len=decode_query_len, + # Bounded by the span's own rows, so a long context request + # cannot inflate the kernel's scheduling hint. + max_seq_len=int(metadata.msa_max_kv_len), + max_num_requests=int(metadata.max_num_requests), + staged_subpage_table=staged_table, + staged_subpages_per_slot=staged_factor, + ) + fmha_tokens = gen_tok0 + + if fmha_tokens == 0: + return + + if fmha_tokens == num_tokens: + # Reaching fmha_sm100 for the whole batch when a span was resolved means + # neither branch above took it, after prepare had already skipped this + # layer's plan and, on a pure-decode step, the flattened page table the + # call below reads. Fail loudly: running on with a stale msa_kv_indices + # would silently attend the wrong pages. + if ported: + raise RuntimeError( + "MiniMax-M3 paged GQA reached fmha_sm100 with no plan for a " + f"{'sparse' if kv_block_indexes is not None else 'dense'} layer. " + "The step resolved a generation span, which the ported decode " + "kernels own; see _resolve_decode_kernels." + ) + fmha_rows = None + else: + # The context prefix, matching the rows `plan` was built over. The + # flattened page table needs no slice: context pages are its prefix and + # the plan implies how many of them to read. + fmha_rows = gen_row0 + # The fmha_sm100 variant is chosen from q.dtype and shares one dtype across # q/k/v, so q must be FP8 to match an FP8 paged K/V. MiniMax-M3 has no # KV-cache scales, so the scale is 1.0 and this is a plain E4M3 cast. When the # model's fused QK-norm+RoPE already emitted FP8 q/k/v (the FP8-KV fast path), # this .to() is a no-op; it stays as a safety net for callers that pass bf16 q. + fmha_q = q_view[:fmha_tokens] use_fp8 = k_paged.dtype == torch.float8_e4m3fn - if use_fp8 and q_view.dtype != torch.float8_e4m3fn: - q_view = q_view.to(torch.float8_e4m3fn) + if use_fp8 and fmha_q.dtype != torch.float8_e4m3fn: + fmha_q = fmha_q.to(torch.float8_e4m3fn) + + def rows_of(lens: Optional[torch.Tensor]) -> Optional[torch.Tensor]: + """Narrow a per-request host length tensor to the rows fmha_sm100 runs. + + Slicing a pinned tensor keeps the pinned backing, so the inline planner + still stages these with non-blocking copies. + """ + if lens is None or fmha_rows is None: + return lens + return lens[:fmha_rows] run_msa_sparse_gqa( - q_view, + fmha_q, k_paged, v_paged, - kv_block_indexes, + None if kv_block_indexes is None else kv_block_indexes[:fmha_tokens], kv_indices=metadata.msa_kv_indices, sm_scale=sm_scale, - qo_lens_cpu=metadata.msa_qo_lens_cpu, - kv_lens_cpu=metadata.msa_kv_lens_cpu, - qo_offset_cpu=metadata.msa_qo_offset_cpu, + qo_lens_cpu=rows_of(metadata.msa_qo_lens_cpu), + kv_lens_cpu=rows_of(metadata.msa_kv_lens_cpu), + qo_offset_cpu=rows_of(metadata.msa_qo_offset_cpu), causal=True, head_dim=head_dim, plan=plan, - out=out_view, + out=out_view[:fmha_tokens], use_fp8=use_fp8, ) @@ -190,11 +312,13 @@ class MsaSparseGqaFmha(Fmha): and attend those blocks; dense layers leave the indices None and attend the full page table. - Inherits Fmha rather than PhasedFmha: fmha_sm100 takes a single plan and - the selected block indices span the whole batch, so it handles a mixed - context and generation batch in one call and there is no - context/generation split from PhasedFmha to reuse. Requires head_dim 128 - and 4-D HND paged K/V. + Inherits Fmha rather than PhasedFmha even though a mixed batch is split + by phase, because that split has to happen in run_msa_paged_gqa rather + than in forward: MiniMaxM3MsaSparseAttention.forward_prepopulated_kv + calls that helper directly, so a split placed in PhasedFmha.forward + would miss it. PhasedFmha also cannot reach the third ported kernel, the + indexer scorer, which runs before forward from run_indexer. Requires + head_dim 128 and 4-D HND paged K/V. """ @classmethod diff --git a/tensorrt_llm/_torch/attention_backend/sparse/minimax_m3/cache_manager.py b/tensorrt_llm/_torch/attention_backend/sparse/minimax_m3/cache_manager.py index 63bf145c7bf0..d79372833d6c 100644 --- a/tensorrt_llm/_torch/attention_backend/sparse/minimax_m3/cache_manager.py +++ b/tensorrt_llm/_torch/attention_backend/sparse/minimax_m3/cache_manager.py @@ -25,7 +25,7 @@ from __future__ import annotations -from typing import List, Optional, Sequence +from typing import List, Optional, Sequence, Tuple import torch @@ -324,25 +324,17 @@ def get_index_v_buffer(self, layer_idx: int) -> Optional[torch.Tensor]: def has_index_value(self, layer_idx: int) -> bool: return layer_idx in self._index_v_buffers - def get_buffers( - self, layer_idx: int, kv_layout: Optional[str] = None - ) -> Optional[torch.Tensor]: - """Return a paged K+V view with strides spanning the coalesced pool. - - The base :meth:`KVCacheManagerV2.get_buffers` produces a - ``[num_pages, kv_factor, ...]`` view with contiguous strides - that assume the slot holds exactly one layer's K+V. In M3's - pool the slot packs K+V for *all* layers of the group - (``scale >= 2 * num_layers_in_group``), so the base view's - dim-0 stride does not reach the next slot's K for this layer. - (When INDEX_KEY's per-block size coincides with K/V's, it is - coalesced into the same pool and contributes to ``scale`` too.) - - The override builds a ``[num_slots, scale, ...]`` view rooted - at K's base, then slices ``[:, :2]`` to extract K+V. The slice - preserves the dim-0 stride (``scale * page_stride``), so - ``view[s, 0/1, ...]`` lands on this layer's K/V at slot ``s``. - When omitted, ``kv_layout`` follows the selected sparse backend. + def _kv_slot_geometry( + self, layer_idx: int, kv_layout: Optional[str] + ) -> Tuple[int, torch.dtype, int, int, List[int]]: + """Resolve one layer's position in the coalesced K/V 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 ``kv_layout`` is None it follows the selected + sparse backend. """ if kv_layout is None: kv_layout = self._main_kv_layout @@ -350,7 +342,7 @@ def get_buffers( raise ValueError(f"Unsupported kv_layout: {kv_layout}") if self.kv_cache_type == CacheTypeCpp.SELFKONLY: raise NotImplementedError( - "MiniMaxM3KVCacheManagerV2.get_buffers does not support SELFKONLY cache type" + "MiniMaxM3KVCacheManagerV2 does not support the SELFKONLY cache type" ) layer_offset = self.layer_offsets[layer_idx] @@ -361,13 +353,13 @@ def get_buffers( # V2 always lays V immediately after K within the per-layer # contribution to a slot. The slice ``[:, :2]`` depends on this. assert addr_key + page_stride_value == addr_value, ( - f"MiniMaxM3 get_buffers requires addr_K + page_stride " + f"MiniMaxM3 requires addr_K + page_stride " f"== addr_V (V immediately after K in slot); got " f"addr_K={addr_key} page_stride_V={page_stride_value} " f"addr_V={addr_value} for layer {layer_idx}." ) assert page_stride_key == page_stride_value, ( - f"MiniMaxM3 get_buffers requires equal K and V page " + f"MiniMaxM3 requires equal K and V page " f"strides; got K={page_stride_key} V=" f"{page_stride_value}." ) @@ -394,27 +386,67 @@ def get_buffers( layer_head_dim = self.head_dim_per_layer[layer_offset] num_kv_heads = self.num_kv_heads_per_layer[layer_offset] + containers = layer_head_dim // element_per_container if kv_layout == "NHD": - full_slot_shape = [ - num_slots, - scale, - self.tokens_per_block, - num_kv_heads, - layer_head_dim // element_per_container, - ] + page_shape = [self.tokens_per_block, num_kv_heads, containers] else: - full_slot_shape = [ - num_slots, - scale, - num_kv_heads, - self.tokens_per_block, - layer_head_dim // element_per_container, - ] + page_shape = [num_kv_heads, self.tokens_per_block, containers] + return addr_key, torch_dtype, num_slots, scale, page_shape + + def get_buffers( + self, layer_idx: int, kv_layout: Optional[str] = None + ) -> Optional[torch.Tensor]: + """Return a paged K+V view with strides spanning the coalesced pool. + + The base :meth:`KVCacheManagerV2.get_buffers` produces a + ``[num_pages, kv_factor, ...]`` view with contiguous strides + that assume the slot holds exactly one layer's K+V. In M3's + pool the slot packs K+V for *all* layers of the group + (``scale >= 2 * num_layers_in_group``), so the base view's + dim-0 stride does not reach the next slot's K for this layer. + (When INDEX_KEY's per-block size coincides with K/V's, it is + coalesced into the same pool and contributes to ``scale`` too.) + The override builds a ``[num_slots, scale, ...]`` view rooted + at K's base, then slices ``[:, :2]`` to extract K+V. The slice + preserves the dim-0 stride (``scale * page_stride``), so + ``view[s, 0/1, ...]`` lands on this layer's K/V at slot ``s``. + When omitted, ``kv_layout`` follows the selected sparse backend. + """ + addr_key, torch_dtype, num_slots, scale, page_shape = self._kv_slot_geometry( + layer_idx, kv_layout + ) + full_slot_shape = [num_slots, scale, *page_shape] full_view = convert_to_torch_tensor(TensorWrapper(addr_key, torch_dtype, full_slot_shape)) return full_view[:, :2] + def get_kv_subpage_pool( + self, layer_idx: int, kv_layout: str = "HND" + ) -> Tuple[torch.Tensor, int]: + """Return ``(flat_pool, subpages_per_slot)`` for flat-block consumers. + + trtllm-gen addresses K and V pages independently, through a + ``[batch, 2, max_blocks]`` block table into one flat + ``[num_subpages, *page_shape]`` pool. That is expressible here even + though the per-layer stride is not uniform: a slot packs ``scale`` + equal-sized sub-pages, of which this layer owns two adjacent ones, so + rooting the flat pool at this layer's K puts slot ``s``'s K at + ``s * scale`` and its V at ``s * scale + 1``. + + The view stops two sub-pages past the last slot's K rather than + spanning ``num_slots * scale``, which would run off the pool by + whatever this layer's K offset is inside a slot. + """ + addr_key, torch_dtype, num_slots, scale, page_shape = self._kv_slot_geometry( + layer_idx, kv_layout + ) + num_subpages = (num_slots - 1) * scale + 2 + flat = convert_to_torch_tensor( + TensorWrapper(addr_key, torch_dtype, [num_subpages, *page_shape]) + ) + return flat, scale + def _kv_pool_mapping_offset(self, layer_id, layer_group_id, key_base_addr) -> int: """Pool-mapping offset from the layer's physical position in its pool. diff --git a/tensorrt_llm/_torch/attention_backend/sparse/minimax_m3/msa_backend.py b/tensorrt_llm/_torch/attention_backend/sparse/minimax_m3/msa_backend.py index 1d3d456c2bc5..4d0f336a5db8 100644 --- a/tensorrt_llm/_torch/attention_backend/sparse/minimax_m3/msa_backend.py +++ b/tensorrt_llm/_torch/attention_backend/sparse/minimax_m3/msa_backend.py @@ -40,14 +40,21 @@ build_paged_kv_slot_mapping, write_kv_slots, ) -from .msa_indexer import MsaIndexer +from .msa_indexer import MsaIndexer, cutedsl_score_runner from .msa_utils import ( MSA_REQUIRED_HEAD_DIM, MSA_REQUIRED_TOPK, build_kv_page_indices, + msa_decode_span_bounds, + msa_ported_decode_active, per_token_valid_blocks, require_msa_module, ) +from .trtllm_gen_dense_decode import ( + dense_decode_unsupported_reason, + uniform_subpages_per_slot, + write_subpage_block_table, +) def _cache_device(meta) -> torch.device: @@ -135,6 +142,44 @@ def _msa_plan_length_keys(sub_plan: dict) -> tuple: return _MSA_SPARSE_LENGTH_KEYS if sub_plan.get("MM-SA-Nv") else _MSA_DENSE_LENGTH_KEYS +@dataclass(frozen=True) +class _MsaDecodeSpan: + """The generation rows and query tokens the ported decode kernels own. + + TensorRT-LLM orders a batch context-first, so the generation requests are + always the row suffix [row_first, row_last) and their query tokens the + token suffix [token_first, token_last). A pure-decode step has + row_first == token_first == 0, which is why one span covers both it and a + mixed step: the ported kernels always take the suffix and fmha_sm100 always + takes the prefix, which is empty on pure decode. + + query_len is the per-request query length, uniform across the generation + rows. The ported kernels address query tokens as request * query_len + + intra, so a span is only resolved where that holds. + """ + + row_first: int + row_last: int + token_first: int + token_last: int + query_len: int + + @property + def batch(self) -> int: + """Number of generation requests in the span.""" + return self.row_last - self.row_first + + @property + def num_tokens(self) -> int: + """Number of generation query tokens in the span.""" + return self.token_last - self.token_first + + @property + def is_mixed(self) -> bool: + """Whether context requests precede the span, so fmha_sm100 also runs.""" + return self.token_first > 0 + + class _MsaGraphSafePlan: """CUDA-graph-stable mirror of one fmha_sm100 decode plan. @@ -228,12 +273,13 @@ class MiniMaxM3MsaSparseAttentionMetadata(TrtllmAttentionMetadata): """TrtllmAttentionMetadata for MiniMax-M3 MSA sparse layers. Tensors read inside the captured forward are CUDA-graph-stable: the - cache slots (msa_out_cache_loc), page table (msa_kv_indices), and proxy - scratch (msa_max_score, msa_n_valid_blocks) are allocated once from the - manager's worst-case geometry. msa_out_cache_loc, msa_kv_indices, and - msa_n_valid_blocks are refreshed via copy_, while the fmha_sm100 proxy pass - writes msa_max_score directly (see msa_proxy_max_score_view). Decode-plan - worklists live on _MsaGraphSafePlan owners, surfaced via msa_decode_*_plan. + cache slots (msa_out_cache_loc), page tables (msa_kv_indices, + msa_block_table), lengths (msa_seq_lens_cuda) and proxy scratch + (msa_max_score, msa_n_valid_blocks) are allocated once from the manager's + worst-case geometry. All of those except msa_max_score are refreshed via + copy_; the fmha_sm100 proxy pass writes msa_max_score directly (see + msa_proxy_max_score_view). Decode-plan worklists live on _MsaGraphSafePlan + owners, surfaced via msa_decode_*_plan. Length inputs to fmha_sm100_plan (msa_qo_lens_cpu, msa_kv_lens_cpu, msa_qo_offset_cpu) are host properties of the base seq_lens/kv_lens, @@ -249,6 +295,18 @@ class MiniMaxM3MsaSparseAttentionMetadata(TrtllmAttentionMetadata): msa_kv_indices: Optional[torch.Tensor] = None msa_max_score: Optional[torch.Tensor] = None msa_n_valid_blocks: Optional[torch.Tensor] = None + # The same page table and lengths as msa_kv_indices / msa_kv_lens, in the + # per-request 2-D form the ported decode kernels index directly + # (block_table[request, block] and seq_lens[request]). fmha_sm100 instead + # takes the flattened msa_kv_indices with the page count implied by its + # plan, so both forms are kept rather than one being derived at call time. + msa_block_table: Optional[torch.Tensor] = None + msa_seq_lens_cuda: Optional[torch.Tensor] = None + # msa_block_table with each slot expanded into the K and V sub-pages the + # trtllm-gen dense kernel indexes. _msa_subpages_per_slot is the expansion + # 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 # Per-request kv_lens as staged by prepare(), before the overlap scheduler # corrects them. on_update_kv_lens clamps against this; see there. msa_kv_lens_staged: Optional[torch.Tensor] = None @@ -282,6 +340,23 @@ class MiniMaxM3MsaSparseAttentionMetadata(TrtllmAttentionMetadata): # persistent backing store for the view. _msa_eager_n_valid_buf: Optional[torch.Tensor] = None _msa_eager_n_valid_blocks: Optional[torch.Tensor] = None + # The generation rows and tokens the ported decode kernels own this step, + # resolved once by _resolve_decode_kernels() before any preparation work + # and the whole of this step's kernel resolution: every site reads it + # through msa_ported_decode_active, so none of them can disagree about + # which kernel ran or therefore about which preparation was needed. None + # only when the step has no generation row at all. See _MsaDecodeSpan. + _msa_decode_span: Optional[_MsaDecodeSpan] = None + # Staged max per-request KV length over the span's generation rows, a + # scheduling upper bound for the ported decode kernels. + _msa_max_kv_len: int = 0 + # max_k_tiles of a proxy plan at the manager's worst case, kept from the + # one-off buffer sizing so the indexer can shape its max_score view + # without a per-step proxy plan to read it from. + _msa_worst_case_max_k_tiles: int = 0 + # First resolution seen by a CUDA-graph metadata, which every later step + # replaying that graph must match. See _check_capture_stable_resolution. + _msa_captured_resolution: Optional[tuple] = None def __post_init__(self) -> None: super().__post_init__() @@ -314,7 +389,7 @@ def msa_qo_lens_cpu(self) -> Optional[torch.Tensor]: def msa_kv_lens_cpu(self) -> Optional[torch.Tensor]: """Per-request KV length, cached plus new tokens (host int32). - The base ``kv_lens`` includes ``num_extra_kv_tokens`` (speculative + The base kv_lens includes num_extra_kv_tokens (speculative draft-loop slots consumed by the C++ kernels); the MSA plans, ladder slots and page counts need the true attended length, so it is excluded here. @@ -358,6 +433,9 @@ def msa_decode_dense_plan(self) -> Optional[tuple]: plan = self._msa_dense_plan return plan.plan if plan is not None else None + # The eager plans cover the rows fmha_sm100 still owns at their site, which + # on a mixed step whose generation span went to a ported kernel is the + # context prefix rather than the whole batch; see _msa_fmha_plan_rows. @property def msa_eager_proxy_plan(self) -> Optional[tuple]: """Prebuilt indexer proxy plan for the eager (prefill/mixed) path.""" @@ -373,12 +451,53 @@ def msa_eager_dense_plan(self) -> Optional[tuple]: """Prebuilt dense GQA plan for the eager (prefill/mixed) path.""" return self._msa_eager_dense_plan + @property + def msa_decode_span(self) -> Optional[_MsaDecodeSpan]: + """Generation rows and tokens the ported decode kernels own, else None.""" + return self._msa_decode_span + + @property + def msa_decode_query_len(self) -> Optional[int]: + """Uniform per-request query length over this step's generation rows. + + None when no span was resolved, which is what msa_ported_decode_active + tests to keep a step on fmha_sm100. + """ + span = self._msa_decode_span + return span.query_len if span is not None else None + + @property + def msa_max_kv_len(self) -> int: + """Staged max KV length over this step's generation rows.""" + return self._msa_max_kv_len + + @property + def msa_worst_case_max_k_tiles(self) -> int: + """max_k_tiles of a proxy plan at the manager's worst-case KV length. + + The bound the proxy scratch was allocated against, so it is valid for + any step and lets a step that skipped the proxy plan still shape its + max_score view. + """ + return self._msa_worst_case_max_k_tiles + @property def msa_eager_n_valid_blocks(self) -> Optional[torch.Tensor]: """Device int32 valid-block count for the eager path, or None if no eager step was prepared (a decode step or a structural test).""" return self._msa_eager_n_valid_blocks + def msa_subpage_rows(self, row_first: int, row_last: int) -> Tuple[Optional[torch.Tensor], int]: + """Staged sub-page block table for the given rows, with its factor. + + (None, 0) when the pool has no single sub-pages-per-slot factor, which + leaves the caller to expand its own layer's table. + """ + table = self.msa_subpage_block_table + if table is None: + return None, 0 + return table[row_first:row_last], self._msa_subpages_per_slot + def _msa_main_kv_is_fp8(self) -> bool: """Whether the main paged K/V cache is stored as FP8 E4M3. @@ -434,6 +553,31 @@ def _create_msa_buffers(self) -> None: dtype=torch.int32, capture_graph=capture_graph, ) + self.msa_block_table = self.get_empty( + buffers, + (max_num_sequences, max_blocks_per_seq), + cache_name="msa_block_table", + dtype=torch.int32, + capture_graph=capture_graph, + ) + self.msa_seq_lens_cuda = self.get_empty( + buffers, + (max_num_sequences,), + cache_name="msa_seq_lens_cuda", + dtype=torch.int32, + capture_graph=capture_graph, + ) + # Resolved once here rather than per step: the factor is fixed by the + # pool's layout for the life of the manager. + self._msa_subpages_per_slot = uniform_subpages_per_slot(kv_cache_manager) + if self._msa_subpages_per_slot > 0: + self.msa_subpage_block_table = self.get_empty( + buffers, + (max_num_sequences, 2, max_blocks_per_seq), + cache_name="msa_subpage_block_table", + dtype=torch.int32, + capture_graph=capture_graph, + ) # Staging for on_update_kv_lens: re-derives slots/bounds on device # from the corrected kv_lens_cuda, sync-free. tokens_per_block = int(kv_cache_manager.tokens_per_block) @@ -486,6 +630,9 @@ def _create_msa_buffers(self) -> None: kv_cache_manager=kv_cache_manager, max_batch=max_num_sequences, ) + # Kept so a step that skips the proxy plan can still shape the + # max_score view; see msa_proxy_max_score_view. + self._msa_worst_case_max_k_tiles = int(max_k_tiles) self._alloc_msa_proxy_scratch( num_index_heads=num_index_heads, max_tokens=self._msa_max_decode_tokens(), @@ -497,7 +644,7 @@ def _create_msa_buffers(self) -> None: def _msa_max_decode_tokens(self) -> int: """Worst-case decode-step query tokens (spec verify emits 1 + draft_len per row), bounded by max_num_tokens. getattr fallbacks cover metadata - built via ``__new__`` in structural tests. + built via __new__ in structural tests. """ max_seqs = int(getattr(self, "max_num_sequences", 0) or 0) max_toks = int(getattr(self, "max_num_tokens", 0) or 0) @@ -576,6 +723,7 @@ def _ensure_msa_decode_scratch_buffers( f"Worst-case max_k_tiles ({max_k_tiles}) is less than the " f"decode plan ({required_max_k_tiles})." ) + self._msa_worst_case_max_k_tiles = int(max_k_tiles) self._alloc_msa_proxy_scratch( num_index_heads=num_index_heads, max_tokens=max_tokens, @@ -599,31 +747,251 @@ def _ensure_eager_n_valid_buffer(self, total_q: int, device: torch.device) -> to def prepare(self) -> None: super().prepare() + # Resolved first: both _build_msa_fields and _build_step_plans skip the + # fmha_sm100 preparation the chosen kernels replace. + self._resolve_decode_kernels() self._build_msa_fields() + if not self._msa_fields_ready: + # Nothing was staged, so nothing may claim a ported decode kernel. + self._clear_decode_kernel_resolution() + # Checked here rather than inside the resolver so it sees the final + # answer, including a resolution the field build withdrew. + self._check_capture_stable_resolution() self._build_step_plans() - def _msa_live_plans(self) -> tuple: - """The fmha_sm100 plans in play for this step. + def _clear_decode_kernel_resolution(self) -> None: + """Drop any resolved span, leaving this step's rows to fmha_sm100.""" + self._msa_decode_span = None + self._msa_max_kv_len = 0 + + def _resolve_decode_kernels(self) -> None: + """Resolve this step's generation span once, before any preparation work. + + Every generation row of every step runs on the ported decode kernels: + the Triton sparse kernel, the trtllm-gen dense kernel and the CuTe DSL + indexer scorer. There is no second decode implementation to choose + between and nothing to fall back to, so this either commits a span or + raises. Falling back would not be a small cost: fmha_sm100 schedules a + generation row like a context row, and on a mixed step cannot even + split the two apart (see _mixed_batch_split in fmha_sm100/api.py), so + every decode row would ride the prefill schedule. + + The generation requests are the row suffix of the batch, so they form a + span the kernels own on a mixed step just as on a pure decode step, + leaving the context prefix to fmha_sm100; see _MsaDecodeSpan. Only + query-length uniformity within the span matters, not whether context + requests precede it. + + prepare() and the per-layer call sites both read the span, so they + cannot disagree about which kernel runs, and therefore cannot disagree + about which preparation was needed. Every input is a host-side + fact that is either fixed for the whole run (static kernel support, + cache geometry) or fixed for the whole step (batch composition, + query-length uniformity); nothing here may depend on a per-layer + tensor, which a call site could see differently. + + The geometry checks mirror the ones the call sites still make. They are + exact rather than optimistic: the MSA backend rejects any head_dim, + sparse_index_dim or topk other than the MSA_REQUIRED_* values at + construction (see MiniMaxM3MsaSparseAttention.__init__), and the index + Q/K dtype is the cache's, so prepare can evaluate them without a live + tensor. + """ + self._clear_decode_kernel_resolution() + params = self._msa_params + kv_cache_manager = self.kv_cache_manager + qo_lens_cpu = self.msa_qo_lens_cpu + kv_lens_cpu = self.msa_kv_lens_cpu + if params is None or kv_cache_manager is None: + return + if qo_lens_cpu is None or kv_lens_cpu is None or int(qo_lens_cpu.shape[0]) == 0: + return + row_first = int(self.num_contexts or 0) + row_last = int(qo_lens_cpu.shape[0]) + if row_first >= row_last: + # Pure prefill: no generation row to place. + return + # The ported kernels address the page table and lengths directly, so + # they need those buffers allocated even though the plans are not. They + # are, from __post_init__, for every manager that carries an index-K + # cache, which is every real MiniMax-M3 run; see _create_msa_buffers. + if self.msa_block_table is None or self.msa_seq_lens_cuda is None: + return + # 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_min != qo_max or qo_max <= 0: + raise NotImplementedError( + "MiniMax-M3 decode requires one query length across a step's " + f"generation rows, got lengths {gen_qo_lens.tolist()} for rows " + f"[{row_first}, {row_last}). The ported decode kernels derive " + "the request id as token // query_len, which needs a uniform " + "positive length." + ) + num_index_heads = params.sharded_index_head_count(self.mapping) + page_size = int(kv_cache_manager.tokens_per_block) + if not self._cutedsl_indexer_supported( + num_index_heads=num_index_heads, + page_size=page_size, + decode_query_len=qo_max, + ): + raise NotImplementedError( + "MiniMax-M3 decode requires the CuTe DSL indexer scorer, which " + f"does not support this geometry: {num_index_heads} index heads, " + f"page size {page_size}, query length {qo_max}, index dtype " + f"{self._msa_index_kv_dtype()}. See is_supported on the runner in " + "cute_dsl_custom_ops." + ) + dense_unsupported = dense_decode_unsupported_reason(kv_cache_manager, MSA_REQUIRED_HEAD_DIM) + if dense_unsupported is not None: + raise NotImplementedError( + "MiniMax-M3 decode requires the trtllm-gen dense kernel for its " + f"dense layers, but {dense_unsupported}" + ) + # Staged, i.e. before the overlap scheduler's correction, which only + # shrinks lengths. That keeps it a valid upper bound for the ported + # kernels' scheduling hints even when it is baked into a CUDA graph. + # Taken over the span alone, so a long context request cannot inflate + # the generation kernels' scheduling bound. + self._msa_max_kv_len = int(kv_lens_cpu[row_first:].max()) + # token_first is derived from the same host lengths the plans are built + # from rather than from num_ctx_tokens, so the token boundary cannot + # drift from the plan the context prefix runs under. + token_first = int(qo_lens_cpu[:row_first].sum()) + self._msa_decode_span = _MsaDecodeSpan( + row_first=row_first, + row_last=row_last, + token_first=token_first, + token_last=token_first + qo_max * (row_last - row_first), + query_len=qo_max, + ) - Decode steps populate the graph-safe owners, prefill and mixed steps the - plain eager tuples. _build_step_plans clears whichever set does not - apply, so only one set is ever live. + def _check_capture_stable_resolution(self) -> None: + """Fail if the resolved span moved after a CUDA graph captured it. + + A replay reruns prepare() to restage the graph's input buffers, but the + kernels inside the graph are fixed at capture. A span that changed + afterwards would stage inputs for one kernel while the graph ran + another, so it has to be caught rather than tolerated. Every input to + the resolution is stable across a graph's replays (the static support + checks and the bucket's own decode_query_len), which is what makes this + a check and not a re-capture. + + Whether a span exists decides whether any fmha_sm100 work runs at all, + and whether it is mixed decides how much of the batch the plans that + survive were built over, so those two are what must hold. The span's row + and token bounds feed them but are otherwise the graph bucket's + business. A captured batch is pure decode, so is_mixed is False at both + capture and replay; comparing it makes that an assertion rather than an + assumption. """ - plans = [ - owner.plan - for owner in (self._msa_proxy_plan, self._msa_gqa_plan, self._msa_dense_plan) - if owner is not None and owner.plan is not None - ] - plans.extend( - plan - for plan in ( - self._msa_eager_proxy_plan, - self._msa_eager_gqa_plan, - self._msa_eager_dense_plan, + if not self.is_cuda_graph: + return + span = self._msa_decode_span + resolution = (span is not None, span is not None and span.is_mixed) + captured = self._msa_captured_resolution + if captured is None: + self._msa_captured_resolution = resolution + elif captured != resolution: + raise RuntimeError( + "MiniMax-M3 decode span changed under a captured CUDA graph: " + f"(resolved, mixed) was {captured} at capture and is " + f"{resolution} now. It must hold for every replay; see " + "_resolve_decode_kernels." + ) + + def _msa_runs_no_fmha(self) -> bool: + """Whether nothing at all this step reaches fmha_sm100. + + When True the whole of its per-step preparation is dead: the three + plans, the graph-safe mirrors of their worklists, the length mirrors + on_update_kv_lens patches into them, and the flattened msa_kv_indices + page table. + + That is every pure-decode step, since the ported kernels own all of its + rows. A mixed step never qualifies: they own only the generation span, + so fmha_sm100 still runs the context prefix and needs its plans and page + table. + """ + span = self._msa_decode_span + return span is not None and not span.is_mixed + + def _msa_fmha_plan_rows(self) -> Optional[Tuple[int, int]]: + """Batch rows this step's fmha_sm100 plans must cover. + + One answer for all three plans, because one span decides all three + sites. The outcomes are + + * no span: the whole batch, as fmha_sm100 runs every row; + * a mixed span: the context prefix only, since the span takes the + generation suffix; + * a pure-decode span: None, no rows left to plan. + + The range always starts at batch row 0, since fmha_sm100 keeps the batch + prefix and the ported kernels take the suffix. It is returned as a range + rather than a bare "context only" flag so that on_update_kv_lens can + rebase each plan's row indices onto the batch rows without assuming + that. + """ + span = self._msa_decode_span + if span is None: + return (0, self._msa_live_batch) + return (0, span.row_first) if span.is_mixed else None + + def _msa_index_kv_dtype(self) -> torch.dtype: + """dtype of the paged index-K cache, which index Q is cast to. + + run_indexer casts index Q to the cache dtype, and the CuTe DSL scorer + requires the two to match, so the cache is the authority on the dtype + the scorer will actually see. + """ + indexer_kv_dtype = str(getattr(self.kv_cache_manager, "indexer_kv_dtype", "bf16")) + return torch.float8_e4m3fn if indexer_kv_dtype == "fp8" else torch.bfloat16 + + def _cutedsl_indexer_supported( + self, *, num_index_heads: int, page_size: int, decode_query_len: int + ) -> bool: + """Whether the CuTe DSL scorer accepts this step's geometry.""" + runner = cutedsl_score_runner() + if runner is None: + return False + return bool( + runner.is_supported( + q_dtype=self._msa_index_kv_dtype(), + num_heads=int(num_index_heads), + # Pinned to MSA_REQUIRED_HEAD_DIM by the backend's constructor. + head_dim=MSA_REQUIRED_HEAD_DIM, + page_size=int(page_size), + max_decode_query_len=int(decode_query_len), ) - if plan is not None ) - return tuple(plans) + + def _msa_live_plans(self) -> tuple: + """The fmha_sm100 plans in play this step, with the rows each covers. + + Yields (plan, row_first, row_last). Decode steps populate the + graph-safe owners, prefill and mixed steps the plain eager tuples; + _build_step_plans clears whichever set does not apply, so only one set + is ever live per site. The row range is narrower than the batch when the + ported kernels took the generation span (see _msa_fmha_plan_rows). + """ + rows = self._msa_fmha_plan_rows() + if rows is None: + # The ported kernels took the whole step, so any plan still held is + # stale and must not be patched. + return () + live = [] + for owner, eager in ( + (self._msa_proxy_plan, self._msa_eager_proxy_plan), + (self._msa_gqa_plan, self._msa_eager_gqa_plan), + (self._msa_dense_plan, self._msa_eager_dense_plan), + ): + plan = owner.plan if owner is not None else None + if plan is None: + plan = eager + if plan is not None: + live.append((plan, rows[0], rows[1])) + return tuple(live) def on_update_kv_lens(self) -> None: """Re-derive length-dependent MSA state from the corrected kv_lens_cuda. @@ -661,23 +1029,37 @@ def on_update_kv_lens(self) -> None: # block, which would NaN the fully-masked GQA row. page = self._msa_page_size n_valid = torch.div((pos + 1).clamp_min(1) + (page - 1), page, rounding_mode="floor") + # Keyed on which buffer the step staged, not on the proxy plan: the + # CuTe DSL scorer reads msa_n_valid_blocks too, so a decode step that + # skipped the plan still needs its counts corrected here. n_valid_buf = ( - self.msa_n_valid_blocks - if self.msa_decode_proxy_plan is not None - else self._msa_eager_n_valid_blocks + self._msa_eager_n_valid_blocks + if self._msa_eager_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)) + # Per-request attended length for the ported decode kernels, patched + # from the same kv_true the plan mirrors below use so the two paths + # can never disagree on how far to walk. msa_block_table needs no + # patch: the correction only shrinks lengths, so the pages already + # listed stay valid and only the walk bound moves. + if self.msa_seq_lens_cuda is not None: + self.msa_seq_lens_cuda[:batch].copy_(kv_true) + # Plan length mirrors. A plan is (has_mixed, split, batch, decode_sub, - # prefill_sub), whose last two entries cover batch rows [0, split) and - # [split, batch). Which of the two holds the prefill requests depends on - # the batch order, so those names are positional only and nothing here - # may key off them; each sub-plan mirrors just its own range. Within a - # sub-plan holds either one row per request or one per query token, - # since the planner row-expands dense plans over query tokens, so the row - # count selects the source. qo_offset must stay non-negative: negative - # values hit the kernel's packed-length sentinel fallback. + # prefill_sub), whose last two entries cover the plan's own rows + # [0, split) and [split, batch). Which of the two holds the prefill + # requests depends on the batch order, so those names are positional + # only and nothing here may key off them; each sub-plan mirrors just its + # own range. Those row indices are relative to the plan, so plan_first + # rebases them onto the batch rows the corrected lengths are indexed by + # (see _msa_fmha_plan_rows). A sub-plan holds either one row per request + # or one per query token, since the planner row-expands dense plans over + # query tokens, so the row count selects the source. qo_offset must stay + # non-negative: negative values hit the kernel's packed-length sentinel + # fallback. per_request = { "kv_segment_lens": kv_true, "qo_offset": (kv_true - qo_dev).clamp_min(0), @@ -689,11 +1071,13 @@ def on_update_kv_lens(self) -> None: "seqused_k": (pos + 1).clamp_min(0), } starts = self._msa_q_token_starts - for has_mixed, split, _, decode_sub, prefill_sub in self._msa_live_plans(): + for plan, plan_first, plan_last in self._msa_live_plans(): + has_mixed, split, _, decode_sub, prefill_sub = plan + plan_split = plan_first + split subs = ( - ((decode_sub, 0, split), (prefill_sub, split, batch)) + ((decode_sub, plan_first, plan_split), (prefill_sub, plan_split, plan_last)) if has_mixed - else ((decode_sub, 0, batch),) + else ((decode_sub, plan_first, plan_last),) ) for sub, first, last in subs: if sub is None: @@ -720,7 +1104,7 @@ def on_update_kv_lens(self) -> None: dst.copy_(src) def _build_step_plans(self) -> None: - """Build the three layer-invariant fmha_sm100 plans once per step. + """Build the layer-invariant fmha_sm100 plans this step still needs. Runs in prepare(), outside CUDA graph capture. The proxy, GQA, and dense plans depend only on the per-step sparse geometry (qo/kv lengths, @@ -734,6 +1118,15 @@ def _build_step_plans(self) -> None: * Prefill, chunked-prefill, and mixed batches run eagerly (never captured), so the plans are stored as plain tuples (msa_eager_*_plan) that every sparse and dense layer reuses. + + Each plan covers only the rows fmha_sm100 still owns (see + _msa_fmha_plan_rows). A pure-decode step is not planned at all, since + the ported kernels took every row: planning is host work on the critical + path, and the plan tuple, its graph-safe mirror and the per-step length + patching in on_update_kv_lens all fall away with it. A mixed step is + planned over the context prefix, which is also the half fmha_sm100 would + have planned into its own sub-plan (see _mixed_batch_split in + fmha_sm100/api.py). """ # Drop any plan tuples from the previous step; the msa_decode_*_plan and # msa_eager_*_plan properties then report None until rebuilt below. @@ -773,42 +1166,46 @@ def _build_step_plans(self) -> None: # bf16 index-K cache, so it never needs the flag. use_fp8 = self._msa_main_kv_is_fp8() + plan_rows = self._msa_fmha_plan_rows() + + def plan_for(**plan_kwargs) -> Optional[tuple]: + """Plan one site over the rows fmha_sm100 still owns, or None. + + Slicing a pinned length tensor keeps the pinned backing, so a + context-only plan stages just as cheaply as a whole-batch one. + """ + if plan_rows is None: + return None + first, last = plan_rows + whole = (first, last) == (0, int(qo_lens_cpu.shape[0])) + return fmha_sm100.fmha_sm100_plan( + qo_lens_cpu if whole else qo_lens_cpu[first:last], + kv_lens_cpu if whole else kv_lens_cpu[first:last], + qo_offset=qo_offset_cpu if whole else qo_offset_cpu[first:last], + page_size=page_size, + num_kv_splits=1, + causal=True, + **plan_kwargs, + ) + # Proxy plan: MQA (num_kv_heads=1) max-score pass over the index # branch; output_maxscore feeds the indexer's top-k block selection. - proxy_plan = fmha_sm100.fmha_sm100_plan( - qo_lens_cpu, - kv_lens_cpu, - num_index_heads, + proxy_plan = plan_for( + num_qo_heads=num_index_heads, num_kv_heads=1, - qo_offset=qo_offset_cpu, - page_size=page_size, output_maxscore=True, - num_kv_splits=1, - causal=True, ) # Sparse-layer plan: kv_block_num=topk limits attention to top-k blocks. - gqa_plan = fmha_sm100.fmha_sm100_plan( - qo_lens_cpu, - kv_lens_cpu, - num_q_heads, + gqa_plan = plan_for( + num_qo_heads=num_q_heads, num_kv_heads=num_kv_heads, - qo_offset=qo_offset_cpu, - page_size=page_size, kv_block_num=topk, - num_kv_splits=1, - causal=True, use_fp8_kvcache=use_fp8, ) # Dense-layer plan: no kv_block_num, so it attends the full page table. - dense_plan = fmha_sm100.fmha_sm100_plan( - qo_lens_cpu, - kv_lens_cpu, - num_q_heads, + dense_plan = plan_for( + num_qo_heads=num_q_heads, num_kv_heads=num_kv_heads, - qo_offset=qo_offset_cpu, - page_size=page_size, - num_kv_splits=1, - causal=True, use_fp8_kvcache=use_fp8, ) @@ -834,12 +1231,19 @@ def _build_step_plans(self) -> None: self._msa_eager_n_valid_blocks = dev_buf[:total_q] return - required_max_k_tiles = int(proxy_plan[3]["max_k_tiles"]) + # The decode query length and max KV length were resolved before any + # preparation ran; see _resolve_decode_kernels. self._ensure_msa_decode_scratch_buffers( num_index_heads=num_index_heads, max_batch=max_batch, capture_graph=capture_graph, - required_max_k_tiles=required_max_k_tiles, + # Without a proxy plan the worst case is the only bound available, + # and it is the one msa_proxy_max_score_view will shape against. + required_max_k_tiles=( + self._msa_worst_case_max_k_tiles + if proxy_plan is None + else int(proxy_plan[3]["max_k_tiles"]) + ), ) # Allocate the graph-safe plan owners once per metadata; later steps @@ -873,10 +1277,16 @@ def _build_step_plans(self) -> None: ) # refresh() stores each plan tuple on its owner, surfaced by the - # msa_decode_*_plan properties. - self._msa_proxy_plan.refresh(proxy_plan) - self._msa_gqa_plan.refresh(gqa_plan) - self._msa_dense_plan.refresh(dense_plan) + # msa_decode_*_plan properties. A skipped plan leaves its owner reset, + # so both that property and _msa_live_plans keep reporting None and the + # graph-safe mirror copies never run. + for owner, plan in ( + (self._msa_proxy_plan, proxy_plan), + (self._msa_gqa_plan, gqa_plan), + (self._msa_dense_plan, dense_plan), + ): + if plan is not None: + owner.refresh(plan) n_valid = per_token_valid_blocks( qo_lens_cpu, kv_lens_cpu, qo_offset_cpu, causal=True, block_size=page_size @@ -891,6 +1301,10 @@ 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. + + Two of those buffers exist only for fmha_sm100 and are skipped when + _resolve_decode_kernels left it with nothing to run this step; see + _msa_runs_no_fmha. """ self._msa_fields_ready = False # Drop any prewritten marker a failed prior step left unconsumed, so @@ -925,25 +1339,59 @@ def _build_msa_fields(self) -> None: ) req_to_token = mapping.req_to_token out_cache_loc = mapping.out_cache_loc - # The page table comes from the same host block ids the mapping was - # built from, so it costs no device work. - kv_indices = build_kv_page_indices(mapping.block_ids_cpu, kv_lens_cpu, page_size) + # Only fmha_sm100 reads the flattened page table (the ported kernels + # index msa_block_table directly), so a step with no fmha_sm100 work + # left skips building and staging it. + needs_flat_page_table = not self._msa_runs_no_fmha() + kv_indices = ( + # Comes from the same host block ids the mapping was built from, + # so it costs no device work. + build_kv_page_indices(mapping.block_ids_cpu, kv_lens_cpu, page_size) + if needs_flat_page_table + else None + ) total_new_tokens = int(out_cache_loc.shape[0]) - total_pages = int(kv_indices.shape[0]) if total_new_tokens > self.msa_out_cache_loc.shape[0]: raise ValueError( f"MSA out_cache_loc buffer ({self.msa_out_cache_loc.shape[0]}) is " f"smaller than the step's new-token count ({total_new_tokens})." ) - if total_pages > self.msa_kv_indices.shape[0]: + if kv_indices is not None and int(kv_indices.shape[0]) > self.msa_kv_indices.shape[0]: raise ValueError( f"MSA kv_indices buffer ({self.msa_kv_indices.shape[0]}) is " - f"smaller than the step's page count ({total_pages})." + f"smaller than the step's page count ({int(kv_indices.shape[0])})." + ) + block_ids_cpu = mapping.block_ids_cpu + block_table_cols = int(block_ids_cpu.shape[1]) + if block_table_cols > self.msa_block_table.shape[1]: + raise ValueError( + f"MSA block_table buffer ({self.msa_block_table.shape[1]} columns) is " + f"smaller than the step's per-request page count ({block_table_cols})." ) self.msa_out_cache_loc[:total_new_tokens].copy_(out_cache_loc, non_blocking=True) - self.msa_kv_indices[:total_pages].copy_(kv_indices, non_blocking=True) + if kv_indices is not None: + self.msa_kv_indices[: int(kv_indices.shape[0])].copy_(kv_indices, non_blocking=True) + + # 2-D page table and per-request length for the ported decode kernels, + # from the same host block ids the flat page table was built from. + # Columns past a request's page count are left stale rather than + # cleared: every consumer bounds its walk by seq_lens. + self.msa_block_table[:batch_size, :block_table_cols].copy_( + maybe_pin_memory(block_ids_cpu.to(torch.int32)), non_blocking=True + ) + self.msa_seq_lens_cuda[:batch_size].copy_(kv_lens_cpu, non_blocking=True) + # Sub-page expansion for the trtllm-gen dense layers, staged once here + # instead of once per layer. It runs outside capture and writes a + # graph-stable buffer, so a replay reads what this step staged, exactly + # as it does for the slot table above. + if self.msa_subpage_block_table is not None: + write_subpage_block_table( + self.msa_block_table[:batch_size], + self._msa_subpages_per_slot, + self.msa_subpage_block_table[:batch_size], + ) # Staging for on_update_kv_lens. step_width = int(req_to_token.shape[1]) @@ -971,8 +1419,11 @@ def _build_msa_fields(self) -> None: self.msa_kv_lens_staged[:batch_size].copy_(kv_lens_cuda[:batch_size], non_blocking=True) # Token offset of each request, plus total_new_tokens as the tail. Host # side, so on_update_kv_lens can slice a sub-plan's token range without - # a device read. - self._msa_q_token_starts = (0, *torch.cumsum(qo_long, 0).tolist()) + # a device read. Only the plan-mirror patching reads it, so a step with + # no plan to mirror does not pay for the transfer off the device. + self._msa_q_token_starts = ( + (0,) if self._msa_runs_no_fmha() else (0, *torch.cumsum(qo_long, 0).tolist()) + ) self._msa_live_batch = batch_size self._msa_live_total_q = total_new_tokens self._msa_page_size = page_size @@ -1054,6 +1505,13 @@ def msa_proxy_max_score_view( decode plan at the worst-case max_k_tiles, so replays only shrink it. """ store = self.msa_max_score + if plan_max_k_tiles <= 0: + raise ValueError( + "The proxy max-score view has no block extent (max_k_tiles=" + f"{plan_max_k_tiles}). Both the fmha_sm100 proxy and the CuTe " + "DSL scorer address it by block id, so a zero extent would put " + "their writes past the end of the view." + ) numel = num_index_heads * plan_max_k_tiles * num_tokens if numel > store.numel(): raise ValueError( @@ -1149,9 +1607,13 @@ def run_indexer( config = self.m3_config idx_sm_scale = idx_sm_scale if idx_sm_scale is not None else config.sparse_index_dim**-0.5 num_tokens = int(idx_q.shape[0]) + # Whether the ported kernels own this step's generation span, which + # decides both the layout the top-k table is produced in and how the + # scoring below is split. + ported = msa_ported_decode_active(metadata) head_major_output = ( int(metadata.num_contexts or 0) > 0 and int(metadata.num_generations or 0) == 0 - ) + ) or ported # idx_q and idx_k may be strided column-views of a fused buffer, so # reshape to keep them zero-copy. The proxy fmha_sm100 and the index-K # scatter below both honor the source strides. @@ -1160,8 +1622,8 @@ def run_indexer( # Index-K may already be in the cache by two routes: the fused per-layer # write (msa_write_layer_caches, idx_k_prewritten=True) stored a live # bf16 idx_k, or the FP8 producer inserted FP8 index-K and passed - # idx_k=None. Write here only when neither owns it — i.e. a live idx_k - # that was not pre-written. + # idx_k=None. Write here only when neither owns it: a live idx_k that + # was not pre-written. if idx_k is not None and not idx_k_prewritten: idx_k_view = idx_k.reshape(num_tokens, 1, config.sparse_index_dim) metadata.msa_write_idx_k(self.layer_idx, idx_k_view) @@ -1185,29 +1647,71 @@ def run_indexer( if idx_q_view.dtype != torch.float8_e4m3fn: idx_q_view = idx_q_view.to(torch.float8_e4m3fn) - # One selection path. Decode passes the graph-safe proxy plan plus the - # proxy scratch shaped to the live query count. Prefill and mixed batches - # pass the eager proxy plan and the device-staged valid-block count. When - # neither is present (a standalone test that skips prepare) select_blocks - # plans inline and computes the valid-block count itself. + # Alternate inputs for the CuTe DSL scorer, which owns this step's + # generation span. Left as None on a step that resolved no span (a pure + # prefill) so the scorer declines and the proxy plan runs the whole + # batch instead. + block_table = None + seq_lens_cuda = None + decode_query_len = None + # gen_first is the span's first query token: the scorer takes + # [gen_first, num_tokens) over rows [ctx_rows, row_last) and the proxy + # the context prefix ahead of both, which is empty on a pure-decode step. + gen_first = 0 + ctx_rows = 0 + if ported: + gen_first, ctx_rows, row_last, decode_query_len = msa_decode_span_bounds( + metadata, num_tokens + ) + block_table = metadata.msa_block_table[ctx_rows:row_last] + seq_lens_cuda = metadata.msa_seq_lens_cuda[ctx_rows:row_last] + # One selection path, and one thing decides which scratch it reads: + # only a pure-decode step is CUDA-graph captured, so only there did + # prepare() mirror the plans into the graph-safe buffers and refresh the + # decode valid-block scratch. Everything else runs eagerly off the + # per-step eager plan and count. When neither is present (a standalone + # test that skips prepare) select_blocks plans inline and computes the + # valid-block count itself. + pure_decode = int(metadata.num_contexts or 0) == 0 proxy_plan = metadata.msa_decode_proxy_plan - if proxy_plan is not None: + if pure_decode and (proxy_plan is not None or ported): # proxy_plan is (has_mixed, split, batch, decode_dict, prefill); - # decode_dict carries max_k_tiles for the contiguous score view. - plan_max_k_tiles = int(proxy_plan[3]["max_k_tiles"]) + # decode_dict carries max_k_tiles for the contiguous score view. A + # step that resolved to the CuTe DSL scorer has no proxy plan to + # read it from and shapes the view against the worst case instead, + # which the scorer accepts: it takes every score stride at runtime. + plan_max_k_tiles = ( + metadata.msa_worst_case_max_k_tiles + if proxy_plan is None + else int(proxy_plan[3]["max_k_tiles"]) + ) max_score = metadata.msa_proxy_max_score_view( config.num_index_heads, plan_max_k_tiles, num_tokens ) n_valid_blocks = metadata.msa_n_valid_blocks[:num_tokens] else: proxy_plan = metadata.msa_eager_proxy_plan - max_score = None # No host-side empty check: the staged counts are clamped to at # least one block, and the kernel masks each query to its own # valid-block extent. n_valid_blocks = metadata.msa_eager_n_valid_blocks if n_valid_blocks is not None: n_valid_blocks = n_valid_blocks[:num_tokens] + # The scorer never allocates: it fills the buffer it is handed. + # Shaped to the span's tokens alone, because the proxy writes its + # own half as a contiguous [heads, k_tiles, tokens] block (see + # msa_proxy_max_score_view) and so cannot take a slice of this one. + # The span's tokens are at most a decode step's worth, which is what + # the store was sized for. + max_score = ( + metadata.msa_proxy_max_score_view( + config.num_index_heads, + metadata.msa_worst_case_max_k_tiles, + num_tokens - gen_first, + ) + if ported + else None + ) return self.indexer.select_blocks( idx_q_view, idx_k_cache, @@ -1219,7 +1723,13 @@ def run_indexer( proxy_plan=proxy_plan, max_score=max_score, n_valid_blocks=n_valid_blocks, + require_cutedsl=ported, head_major_output=head_major_output, + block_table=block_table, + seq_lens_cuda=seq_lens_cuda, + decode_query_len=decode_query_len, + gen_token_first=gen_first, + ctx_rows=ctx_rows, ) def sparse_attn_predict( @@ -1251,10 +1761,10 @@ def forward_prepopulated_kv( ) -> None: """Run MSA after the eager-prefill producer inserted main K/V. - ``TrtllmAttention.forward`` interprets ``k=None`` as a fused QKV - buffer, so it cannot represent compact Q with prewritten paged K/V. - Dispatch the same MSA paged-GQA helper directly; the #16755 prewritten - marker is consumed there exactly as on its general scatter path. + TrtllmAttention.forward reads k=None as a fused QKV buffer, so it + cannot express compact Q with prewritten paged K/V. Dispatch the same + MSA paged-GQA helper directly; it consumes the prewritten-layer marker + exactly as it does on its general scatter path. """ output = forward_args.output if output is None: diff --git a/tensorrt_llm/_torch/attention_backend/sparse/minimax_m3/msa_indexer.py b/tensorrt_llm/_torch/attention_backend/sparse/minimax_m3/msa_indexer.py index 482ddda18d38..457f9c540f76 100644 --- a/tensorrt_llm/_torch/attention_backend/sparse/minimax_m3/msa_indexer.py +++ b/tensorrt_llm/_torch/attention_backend/sparse/minimax_m3/msa_indexer.py @@ -4,15 +4,20 @@ Mirrors the DSA indexer pattern: a submodule owned by the sparse backend that runs the predictor pass and returns the per-query selected KV block -indices the main attention consumes. It calls fmha_sm100 directly in -output_maxscore mode, reduces the per-index-head max score to KV-head -granularity, and selects the top-k blocks per query. +indices the main attention consumes. It scores the index branch, reduces the +per-index-head max score to KV-head granularity, and selects the top-k blocks +per query. + +Scoring has two implementations, and a step can use both: the dedicated CuTe +DSL kernel takes the generation span, and fmha_sm100 in output_maxscore mode +takes whatever is left, which on a mixed batch is the context prefix. Results are [total_q, num_kv_heads, topk] int32, ascending with -1 padding. """ from __future__ import annotations +import functools from typing import TYPE_CHECKING, Optional import torch @@ -28,6 +33,78 @@ from .common import MiniMaxM3SparseConfig +@functools.lru_cache(maxsize=1) +def cutedsl_score_runner(): + """Return the CuTe DSL indexer scoring runner, or None if unavailable. + + The CuTe DSL ops are registered only when the nvidia-cutlass-dsl package is + importable, so this stays a soft dependency. prepare() consults the same + runner to decide whether to skip the fmha_sm100 proxy plan, so the answer + here and there must come from one place. + + Resolved once for the process: package availability cannot change under a + running model, and every sparse layer of every step scores through here. + """ + try: + from tensorrt_llm._torch.custom_ops import cute_dsl_custom_ops + except ImportError: + return None + return getattr(cute_dsl_custom_ops, "CuteDSLMiniMaxM3IndexDecodeScoreRunner", None) + + +def _cutedsl_score( + idx_q: torch.Tensor, + idx_k_paged: torch.Tensor, + max_score: torch.Tensor, + *, + block_table: torch.Tensor, + seq_lens_cuda: torch.Tensor, + decode_query_len: int, +) -> bool: + """Try to fill `max_score` with the CuTe DSL scorer; report whether it ran. + + `max_score` is the [num_index_heads, max_k_tiles, total_q] buffer the block + selector consumes. The kernel writes [head, token, block], so it is handed + the transposed view: same backing store, no copy, and the stores end up + coalesced across tokens rather than strided by max_k_tiles. + + The buffer is deliberately not pre-filled with -inf. The kernel writes + blocks [0, ceil(seq_len / page_size)) for every token of a request, and the + selector reads only [0, n_valid_blocks[token])), which is bounded by that + same count for every token including the shorter ones in a multi-token + speculative step. So every entry the selector reads has just been written. + """ + runner = cutedsl_score_runner() + if runner is None: + return False + + total_q, num_index_heads, head_dim = idx_q.shape + page_size = int(idx_k_paged.shape[2]) + if not runner.is_supported( + q_dtype=idx_q.dtype, + num_heads=num_index_heads, + head_dim=head_dim, + page_size=page_size, + max_decode_query_len=decode_query_len, + ): + return False + if idx_k_paged.dtype != idx_q.dtype or max_score.shape[2] != total_q: + return False + + # The kernel wants MQA index-K as [num_pages, page_size, head_dim]; the + # squeeze is zero-copy and keeps the pool's real per-page stride, which the + # TMA descriptor reads at runtime. + torch.ops.trtllm.cute_dsl_minimax_m3_index_decode_score( + idx_q, + idx_k_paged.squeeze(1), + block_table, + seq_lens_cuda, + max_score.transpose(1, 2), + decode_query_len, + ) + return True + + def _proxy_max_score( idx_q: torch.Tensor, idx_k_paged: torch.Tensor, @@ -84,6 +161,32 @@ def _proxy_max_score( return max_score +def _combined_topk_table( + ctx_table: torch.Tensor, + gen_table: torch.Tensor, + *, + head_major: bool, +) -> torch.Tensor: + """Concatenate the context and generation top-k tables along the token axis. + + Both halves are [tokens, num_kv_heads, topk]. `head_major` reproduces the + head-major backing select_blocks_from_maxscore gives its own output, so the + combined table permutes to a contiguous [num_kv_heads, total_q, topk] + exactly as an unsplit one would and the Triton decode kernel sees the + layout it expects. + """ + ctx_tokens = int(ctx_table.shape[0]) + total_q = ctx_tokens + int(gen_table.shape[0]) + num_kv_heads, topk = int(ctx_table.shape[1]), int(ctx_table.shape[2]) + shape = (num_kv_heads, total_q, topk) if head_major else (total_q, num_kv_heads, topk) + out = torch.empty(shape, dtype=ctx_table.dtype, device=ctx_table.device) + if head_major: + out = out.transpose(0, 1) + out[:ctx_tokens].copy_(ctx_table) + out[ctx_tokens:].copy_(gen_table) + return out + + def _group_max_reduce( max_score: torch.Tensor, config: "MiniMaxM3SparseConfig", @@ -132,43 +235,92 @@ def select_blocks( max_score: Optional[torch.Tensor] = None, n_valid_blocks: Optional[torch.Tensor] = None, head_major_output: bool = False, + block_table: Optional[torch.Tensor] = None, + seq_lens_cuda: Optional[torch.Tensor] = None, + decode_query_len: Optional[int] = None, + require_cutedsl: bool = False, + gen_token_first: int = 0, + ctx_rows: int = 0, ) -> torch.Tensor: """Return [total_q, num_kv_heads, topk] selected block indices. Plan/run split, mirroring the sparse GQA. Both production paths pass a - prebuilt `proxy_plan` and a precomputed device `n_valid_blocks` (decode - from the graph-safe scratch, eager from the step-level device buffer); - decode additionally runs into the preallocated `max_score` buffer inside - the captured region. + precomputed device `n_valid_blocks` (decode from the graph-safe + scratch, eager from the step-level device buffer) and, unless the step + resolved to the CuTe DSL scorer, a prebuilt `proxy_plan`; decode + additionally runs into the preallocated `max_score` buffer inside the + captured region. + + A step with generation rows additionally passes `block_table`, + `seq_lens_cuda` and `decode_query_len`, which is what puts the dedicated + CuTe DSL scorer on that span in place of the fmha_sm100 proxy pass. Every + such step does, from prepare(); a caller that leaves them unset (the + standalone kernel tests) gets the proxy over the whole batch. + + `gen_token_first` and `ctx_rows` are where the span starts: the scorer + takes query tokens [gen_token_first, total_q) and rows + [ctx_rows, batch), and the proxy keeps the context prefix ahead of both + under the plan prepare() built over exactly those rows. Both are 0 on a + pure-decode step, where the scorer owns everything. The two halves are + scored into separate buffers rather than one, because fmha_sm100 writes + a contiguous [heads, k_tiles, tokens] block and so cannot fill a slice + of the scorer's; they are selected separately and the tables joined. + + `require_cutedsl` says prepare() committed to the scorer and narrowed + the proxy plan to the context prefix, leaving the span nothing to + decline to, so a decline raises. """ config = self.config + total_q = int(idx_q.shape[0]) + page_size = int(idx_k_paged.shape[2]) + gen_first = int(gen_token_first) - if proxy_plan is None: - max_score = _proxy_max_score( + scored = False + if ( + max_score is not None + and block_table is not None + and seq_lens_cuda is not None + and decode_query_len is not None + ): + # Like the fmha_sm100 proxy, whose max_score is read off the MMA + # accumulator before the softmax scale, the CuTe DSL scorer emits + # raw Q.K rather than idx_sm_scale * Q.K. Block ranking, and the + # +inf forcing of the init/local blocks in + # select_blocks_from_maxscore, are both invariant under a positive + # scale, so neither depends on the omission. + scored = _cutedsl_score( + idx_q[gen_first:], + idx_k_paged, + max_score, + block_table=block_table, + seq_lens_cuda=seq_lens_cuda, + decode_query_len=decode_query_len, + ) + + if require_cutedsl and not scored: + raise RuntimeError( + "MiniMax-M3 prepare() resolved a generation span for this step " + "and narrowed the fmha_sm100 proxy plan to the context prefix, " + "but the CuTe DSL indexer scorer declined the span. The two must " + "agree, and there is no proxy pass left to score it; see " + "_resolve_decode_kernels." + ) + + # Nothing was ported, so the proxy runs every token and there is no + # split: the plan it was handed covers the whole batch. + if not scored: + gen_first = 0 + max_score = self._proxy_scores( idx_q, idx_k_paged, + proxy_plan=proxy_plan, + max_score=max_score, qo_lens_cpu=qo_lens_cpu, kv_lens_cpu=kv_lens_cpu, qo_offset_cpu=qo_offset_cpu, kv_indices=kv_indices, - sm_scale=idx_sm_scale, - causal=True, + idx_sm_scale=idx_sm_scale, ) - else: - fmha_sm100 = require_msa_module() - _, max_score = fmha_sm100.fmha_sm100( - idx_q, - idx_k_paged, - idx_k_paged, - proxy_plan, - kv_indices=kv_indices, - output_o=False, - output_maxscore=True, - max_score=max_score, - sm_scale=idx_sm_scale, - ) - - max_score_kv = _group_max_reduce(max_score, config) if n_valid_blocks is None: n_valid_blocks = per_token_valid_blocks( @@ -176,15 +328,15 @@ def select_blocks( kv_lens_cpu, qo_offset_cpu, causal=True, - block_size=int(idx_k_paged.shape[2]), + block_size=page_size, ) # Empty-selection guard. n_valid_blocks is a host tensor on # this path, so the .item() read does not sync the device. if n_valid_blocks.numel() == 0 or int(n_valid_blocks.max().item()) <= 0: output_shape = ( - (config.num_kv_heads, idx_q.shape[0], MSA_REQUIRED_TOPK) + (config.num_kv_heads, total_q, MSA_REQUIRED_TOPK) if head_major_output - else (idx_q.shape[0], config.num_kv_heads, MSA_REQUIRED_TOPK) + else (total_q, config.num_kv_heads, MSA_REQUIRED_TOPK) ) output = torch.full( output_shape, @@ -193,14 +345,93 @@ def select_blocks( device=idx_q.device, ) return output.permute(1, 0, 2) if head_major_output else output + + gen_table = self._select( + max_score, + n_valid_blocks[gen_first:], + head_major_output=head_major_output, + ) + if gen_first == 0: + return gen_table + # The context prefix, whose scores the proxy produces into its own + # buffer under the plan built over rows [0, ctx_rows). Context pages are + # the prefix of the flattened page table, so kv_indices needs no slice. + ctx_table = self._select( + self._proxy_scores( + idx_q[:gen_first], + idx_k_paged, + proxy_plan=proxy_plan, + max_score=None, + qo_lens_cpu=None if qo_lens_cpu is None else qo_lens_cpu[:ctx_rows], + kv_lens_cpu=None if kv_lens_cpu is None else kv_lens_cpu[:ctx_rows], + qo_offset_cpu=None if qo_offset_cpu is None else qo_offset_cpu[:ctx_rows], + kv_indices=kv_indices, + idx_sm_scale=idx_sm_scale, + ), + n_valid_blocks[:gen_first], + head_major_output=head_major_output, + ) + return _combined_topk_table(ctx_table, gen_table, head_major=head_major_output) + + def _proxy_scores( + self, + idx_q: torch.Tensor, + idx_k_paged: torch.Tensor, + *, + proxy_plan: Optional[tuple], + max_score: Optional[torch.Tensor], + qo_lens_cpu: Optional[torch.Tensor], + kv_lens_cpu: Optional[torch.Tensor], + qo_offset_cpu: Optional[torch.Tensor], + kv_indices: torch.Tensor, + idx_sm_scale: float, + ) -> torch.Tensor: + """Run the fmha_sm100 proxy pass over `idx_q` and return its max score. + + Uses the prebuilt plan when prepare() supplied one, and plans inline + from the host lengths otherwise (standalone callers that skip prepare). + """ + if proxy_plan is None: + return _proxy_max_score( + idx_q, + idx_k_paged, + qo_lens_cpu=qo_lens_cpu, + kv_lens_cpu=kv_lens_cpu, + qo_offset_cpu=qo_offset_cpu, + kv_indices=kv_indices, + sm_scale=idx_sm_scale, + causal=True, + ) + fmha_sm100 = require_msa_module() + _, scores = fmha_sm100.fmha_sm100( + idx_q, + idx_k_paged, + idx_k_paged, + proxy_plan, + kv_indices=kv_indices, + output_o=False, + output_maxscore=True, + max_score=max_score, + sm_scale=idx_sm_scale, + ) + return scores + + def _select( + self, + max_score: torch.Tensor, + n_valid_blocks: torch.Tensor, + *, + head_major_output: bool, + ) -> torch.Tensor: + """Reduce scores to KV-head granularity and take the top-k blocks.""" return select_blocks_from_maxscore( - max_score_kv, + _group_max_reduce(max_score, self.config), topk=MSA_REQUIRED_TOPK, n_valid_blocks=n_valid_blocks, - init_blocks=config.init_blocks, - local_blocks=config.local_blocks, + init_blocks=self.config.init_blocks, + local_blocks=self.config.local_blocks, head_major_output=head_major_output, ) -__all__ = ["MsaIndexer"] +__all__ = ["MsaIndexer", "cutedsl_score_runner"] diff --git a/tensorrt_llm/_torch/attention_backend/sparse/minimax_m3/msa_scatter.py b/tensorrt_llm/_torch/attention_backend/sparse/minimax_m3/msa_scatter.py index 94e333bcf653..c42c6ee4a568 100644 --- a/tensorrt_llm/_torch/attention_backend/sparse/minimax_m3/msa_scatter.py +++ b/tensorrt_llm/_torch/attention_backend/sparse/minimax_m3/msa_scatter.py @@ -4,11 +4,11 @@ One Triton launch writes a layer's new-token main K, main V, and (sparse layers) index-K into their paged HND caches at the step's write slots. -The legacy path costs three aten advanced-indexing writes per layer plus -their index preprocessing; at 60 layers per forward step, all captured -into decode CUDA graphs, the launch count dominates the cost. The kernel -derives each token's (page, within-page) split from ``out_cache_loc`` -in-register, so it needs no precomputed index tensors at all. +Writing them separately costs three aten advanced-indexing writes per +layer plus their index preprocessing; at 60 layers per forward step, all +captured into decode CUDA graphs, the launch count dominates the cost. +The kernel derives each token's (page, within-page) split from +out_cache_loc in-register, so it needs no precomputed index tensors. Sources may be strided row views (slices of the fused QKV projection); only the innermost [num_heads * head_dim] extent must be contiguous. diff --git a/tensorrt_llm/_torch/attention_backend/sparse/minimax_m3/msa_utils.py b/tensorrt_llm/_torch/attention_backend/sparse/minimax_m3/msa_utils.py index ded4775f63d5..04b24eb78e37 100644 --- a/tensorrt_llm/_torch/attention_backend/sparse/minimax_m3/msa_utils.py +++ b/tensorrt_llm/_torch/attention_backend/sparse/minimax_m3/msa_utils.py @@ -27,6 +27,56 @@ _MSA_PYTHON_RELPATH = Path("3rdparty") / "MSA" / "python" +def msa_ported_decode_active(metadata) -> bool: + """Whether this step's generation rows run on the ported decode kernels. + + They always do, when the step has generation rows at all: the Triton sparse + kernel, the trtllm-gen dense kernel and the CuTe DSL indexer scorer own the + generation span together (the whole of a pure-decode batch, or the row and + token suffix of a mixed one) and fmha_sm100 is left with the context prefix. + There is no per-kernel switch and no fallback for a generation row: a + geometry the span's kernels cannot serve raises in prepare() instead of + being routed back to fmha_sm100 at several times the decode cost. See + _resolve_decode_kernels. + + So False means only that this step has no span: a pure-prefill step, or one + that opted out by leaving the query length or the buffers the ported kernels + address unset, which the standalone kernel tests do since their metadata + never runs prepare(). + + One predicate serves every site that has to agree on the span: both + attention branches and the indexer, which additionally orients its top-k + table for whichever of them consumes it. + """ + return ( + getattr(metadata, "msa_decode_query_len", None) is not None + and getattr(metadata, "msa_block_table", None) is not None + and getattr(metadata, "msa_seq_lens_cuda", None) is not None + ) + + +def msa_decode_span_bounds(metadata, num_tokens: int) -> Tuple[int, int, int, int]: + """Bounds of the generation span, as (token_first, row_first, row_last, query_len). + + The span is what the ported decode kernels own this step; see + _MsaDecodeSpan in msa_backend. Reading it through getattr keeps this module + free of an import back into the backend, and covers the standalone kernel + tests, whose metadata never ran prepare() and so carries no span: there the + whole batch is the span, derived from the query length alone. + + Returns zeros when no query length is resolved either, in which case every + caller is on the fmha_sm100 path and ignores these bounds. + """ + span = getattr(metadata, "msa_decode_span", None) + if span is not None: + return span.token_first, span.row_first, span.row_last, span.query_len + query_len = getattr(metadata, "msa_decode_query_len", None) + if query_len is None: + return 0, 0, 0, 0 + query_len = int(query_len) + return 0, 0, num_tokens // query_len, query_len + + @functools.lru_cache(maxsize=1) def _find_msa_python_dir() -> Optional[Path]: """Locate the fmha_sm100 package dir by walking up from this file. @@ -233,9 +283,8 @@ def select_blocks_from_maxscore( Applies init and local forced blocks and per-query valid-block masking on the amax-reduced scores [num_kv_heads, n_blocks, total_q]. Returns [total_q, num_kv_heads, topk] int32 ascending block ids with -1 tail - padding. When ``head_major_output`` is set, the logical result uses a - head-major backing so ``result.permute(1, 0, 2)`` is contiguous without a - copy. + padding. head_major_output backs that result head-major instead, so + result.permute(1, 0, 2) is contiguous without a copy. """ nvb = n_valid_blocks.to(device=max_score_kv.device, dtype=torch.int32).contiguous() return torch.ops.trtllm.minimax_m3_select_blocks( @@ -252,8 +301,10 @@ def select_blocks_from_maxscore( "MSA_REQUIRED_HEAD_DIM", "MSA_REQUIRED_TOPK", "build_kv_page_indices", + "msa_decode_span_bounds", "msa_package_available", "msa_paged_kv", + "msa_ported_decode_active", "per_token_valid_blocks", "require_msa_module", "select_blocks_from_maxscore", diff --git a/tensorrt_llm/_torch/attention_backend/sparse/minimax_m3/triton_sparse_decode.py b/tensorrt_llm/_torch/attention_backend/sparse/minimax_m3/triton_sparse_decode.py new file mode 100644 index 000000000000..01d4200cc3fa --- /dev/null +++ b/tensorrt_llm/_torch/attention_backend/sparse/minimax_m3/triton_sparse_decode.py @@ -0,0 +1,410 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-License-Identifier: Apache-2.0 +# Vendored from vLLM (Apache-2.0): +# https://github.com/vllm-project/vllm/blob/6f91edf96d3f3272945809c04702380053bff4de/vllm/models/minimax_m3/common/ops/sparse_attn.py +"""Triton block-sparse GQA decode attention for MiniMax-M3. + +Flash-decoding over the blocks the indexer selected: one CTA per (query token, +top-k chunk, KV head) accumulates a partial output plus its log-sum-exp, and a +second kernel merges the chunks by LSE weight. That beats running the +context-schedule FMHA kernel at decode, where a single query token leaves most +of a 128-row Q tile idle. + +Vendored from the vLLM source linked in the file header (v0.26.1rc0-77-g6f91edf96). +Differences from upstream: + +* K and V are separate HND paged views ([num_pages, num_kv_heads, page_size, + head_dim]) with independent strides rather than one cache fused along the + last dim, because that is how the M3 KV pool is laid out here. +* Only the scalar KV-scale mode is kept; M3 stores unscaled E4M3 K/V. +* Partial-output and LSE scratch come from the persistent buffer arena so their + addresses survive CUDA graph replay. +* The merge kernel zeroes rows whose chunks are all empty instead of letting + them produce NaN, since CUDA-graph padding rows flow on into the rest of the + network here. +""" + +from __future__ import annotations + +import os +from typing import Optional + +import torch +import triton +import triton.language as tl + +from tensorrt_llm._torch.memory_buffer_utils import get_memory_buffers + +# One sparse block is exactly one KV page. +SPARSE_BLOCK_SIZE = 128 + +_FP8_DTYPES = (torch.float8_e4m3fn, torch.float8_e5m2) + +# Total CTAs the split-K partitioning aims for before it stops splitting. +_TARGET_GRID = 256 + + +def _pdl_enabled() -> bool: + return os.environ.get("TRTLLM_ENABLE_PDL", "1") == "1" + + +@triton.heuristics( + { + "BLOCK_SIZE_H": lambda args: max(16, triton.next_power_of_2(args["gqa_group_size"])), + "BLOCK_SIZE_D": lambda args: triton.next_power_of_2(args["head_dim"]), + } +) +@triton.jit(do_not_specialize=["decode_query_len"]) +def _gqa_sparse_decode_kernel( + q_ptr, # [total_q, num_heads, head_dim] + k_ptr, # paged K: [num_pages, num_kv_heads, page_size, head_dim] + v_ptr, # paged V: same layout as K + kv_scale_ptr, # scalar dequant scale, or a dummy when USE_SCALE is False + t_ptr, # topk_idx: [num_kv_heads, total_q, topk] + o_ptr, # partial out: [NUM_TOPK_CHUNKS, total_q, num_heads, head_dim] + lse_ptr, # partial lse (log2): [NUM_TOPK_CHUNKS, total_q, num_heads] + block_table_ptr, # [num_reqs, max_blocks] + seq_lens, # [num_reqs] + total_q, + gqa_group_size, + head_dim, + max_topk, + sm_scale, + decode_query_len, + stride_qn, + stride_qh, + stride_qd, + stride_k_blk, + stride_k_h, + stride_k_pos, + stride_k_d, + stride_v_blk, + stride_v_h, + stride_v_pos, + stride_v_d, + stride_th, + stride_tn, + stride_tk, + stride_o_c, + stride_o_b, + stride_o_h, + stride_o_d, + stride_l_c, + stride_l_b, + stride_l_h, + stride_bt_b, + BLOCK_SIZE_K: tl.constexpr, # == SPARSE_BLOCK_SIZE (128) + NUM_TOPK_CHUNKS: tl.constexpr, + BLOCK_SIZE_H: tl.constexpr, + BLOCK_SIZE_D: tl.constexpr, + USE_SCALE: tl.constexpr, # apply the scalar dequant scale to an fp8 cache + USE_PDL: tl.constexpr, +): + sm_scale_log2e = sm_scale * 1.4426950409 + # Split-K over the topk dimension: pid(0) folds (query token, chunk). + pid_bc, pid_kh = tl.program_id(0), tl.program_id(1) + pid_b = pid_bc % total_q + pid_c = pid_bc // total_q + req_id = pid_b // decode_query_len + q_offset = pid_b - req_id * decode_query_len + pid_h = pid_kh * gqa_group_size + chunk_size_topk = (max_topk + NUM_TOPK_CHUNKS - 1) // NUM_TOPK_CHUNKS + chunk_start = pid_c * chunk_size_topk + + if USE_PDL: + tl.extra.cuda.gdc_wait() + + seq_len = tl.load(seq_lens + req_id) + query_pos = seq_len - decode_query_len + q_offset + # Full-CG padding uses zero-length request rows. Clamp to an empty + # attention range instead of letting padded rows produce negative lengths. + kv_len = tl.maximum(query_pos + 1, 0) + + # Bound the walk by the token's own valid block count rather than a + # sentinel: the selector emits ascending ids and only pads with -1 past + # min(topk, cdiv(kv_len, blk)) entries, so those are never dereferenced. + num_blocks = (kv_len + BLOCK_SIZE_K - 1) // BLOCK_SIZE_K + real_topk = tl.minimum(max_topk, num_blocks) + chunk_end = tl.minimum(chunk_start + chunk_size_topk, real_topk) + + off_n = tl.arange(0, BLOCK_SIZE_K) + off_d = tl.arange(0, BLOCK_SIZE_D) + off_h = tl.arange(0, BLOCK_SIZE_H) + d_mask = off_d < head_dim + h_mask = off_h < gqa_group_size + hd_mask = h_mask[:, None] & d_mask[None, :] + bt_row = block_table_ptr + req_id * stride_bt_b + + m_i = tl.full((BLOCK_SIZE_H,), float("-inf"), dtype=tl.float32) + lse_i = tl.full((BLOCK_SIZE_H,), float("-inf"), dtype=tl.float32) + acc_o = tl.zeros((BLOCK_SIZE_H, BLOCK_SIZE_D), dtype=tl.float32) + q = tl.load( + q_ptr + + pid_b * stride_qn + + (pid_h + off_h[:, None]) * stride_qh + + off_d[None, :] * stride_qd, + mask=hd_mask, + other=0.0, + ) + kv_scale = tl.load(kv_scale_ptr) if USE_SCALE else 1.0 + + cur_idx_ptr = t_ptr + pid_kh * stride_th + pid_b * stride_tn + chunk_start * stride_tk + for _ in tl.range(chunk_start, chunk_end): + blk = tl.load(cur_idx_ptr) + cur_idx_ptr = cur_idx_ptr + stride_tk + # int64 page offsets: a large cache overflows int32 well before the + # per-page block offsets do. + page = tl.load(bt_row + blk).to(tl.int64) + pos_mask = blk * BLOCK_SIZE_K + off_n < kv_len + k = tl.load( + k_ptr + + page * stride_k_blk + + pid_kh * stride_k_h + + off_n[None, :] * stride_k_pos + + off_d[:, None] * stride_k_d, + mask=d_mask[:, None] & pos_mask[None, :], + other=0.0, + ).to(q.dtype) + if USE_SCALE: + k = (k * kv_scale).to(q.dtype) + qk = tl.where(pos_mask[None, :], 0.0, float("-inf")) + qk += tl.dot(q, k) * sm_scale_log2e + m_ij = tl.maximum(m_i, tl.max(qk, axis=1)) + p = tl.exp2(qk - m_ij[:, None]) + l_ij = tl.sum(p, axis=1) + acc_o = acc_o * tl.exp2(m_i - m_ij)[:, None] + v = tl.load( + v_ptr + + page * stride_v_blk + + pid_kh * stride_v_h + + off_n[:, None] * stride_v_pos + + off_d[None, :] * stride_v_d, + mask=pos_mask[:, None] & d_mask[None, :], + other=0.0, + ).to(q.dtype) + if USE_SCALE: + v = (v * kv_scale).to(q.dtype) + acc_o += tl.dot(p.to(v.dtype), v) + m_i = m_ij + lse_i = m_ij + tl.log2(tl.exp2(lse_i - m_ij) + l_ij) + + # An empty chunk of an active row must store zero, or the merge hits 0 * NaN. + scale = tl.where(lse_i > float("-inf"), tl.exp2(m_i - lse_i), 0.0) + acc_o = acc_o * scale[:, None] + o_base = o_ptr + pid_c * stride_o_c + pid_b * stride_o_b + pid_h * stride_o_h + tl.store( + o_base + off_h[:, None] * stride_o_h + off_d[None, :] * stride_o_d, + acc_o.to(o_ptr.dtype.element_ty), + mask=hd_mask, + ) + lse_base = lse_ptr + pid_c * stride_l_c + pid_b * stride_l_b + pid_h * stride_l_h + tl.store(lse_base + off_h * stride_l_h, lse_i, mask=h_mask) + + # After the stores, never before: the merge grid's gdc_wait() releases on + # this trigger, and it reads exactly the partials written above. + if USE_PDL: + tl.extra.cuda.gdc_launch_dependents() + + +@triton.heuristics({"BLOCK_SIZE_D": lambda args: triton.next_power_of_2(args["head_dim"])}) +@triton.jit +def _merge_topk_attn_out_kernel( + o_ptr, # partials: [NUM_TOPK_CHUNKS, total_q, num_heads, head_dim] + lse_ptr, # partials (log2): [NUM_TOPK_CHUNKS, total_q, num_heads] + out_ptr, # merged out: [total_q, num_heads, head_dim] + head_dim, + stride_o_c, + stride_o_b, + stride_o_h, + stride_o_d, + stride_l_c, + stride_l_b, + stride_l_h, + stride_out_n, + stride_out_h, + stride_out_d, + NUM_TOPK_CHUNKS: tl.constexpr, + BLOCK_SIZE_D: tl.constexpr, + USE_PDL: tl.constexpr, +): + pid_b, pid_h = tl.program_id(0), tl.program_id(1) + + if USE_PDL: + tl.extra.cuda.gdc_wait() + + off_c = tl.arange(0, NUM_TOPK_CHUNKS) + off_d = tl.arange(0, BLOCK_SIZE_D) + o = tl.load( + o_ptr + + pid_b * stride_o_b + + pid_h * stride_o_h + + off_c[:, None] * stride_o_c + + off_d[None, :] * stride_o_d, + mask=off_d[None, :] < head_dim, + other=0.0, + ).to(tl.float32) + # Empty chunks contribute -inf, hence weight 0. + lse = tl.load(lse_ptr + pid_b * stride_l_b + pid_h * stride_l_h + off_c * stride_l_c) + lse_max = tl.max(lse, axis=0) + # A row whose every chunk is empty (a CUDA-graph padding row) would give + # -inf - -inf here. Zero it instead: the padded output is discarded, but a + # NaN would survive into the residual stream and the all-reduce. + lse_max = tl.where(lse_max == float("-inf"), 0.0, lse_max) + weights = tl.exp2(lse - lse_max) + denom = tl.sum(weights, axis=0) + weights = weights / tl.where(denom > 0, denom, 1.0) + o_merged = tl.sum(o * weights[:, None], axis=0) + out_ptrs = out_ptr + pid_b * stride_out_n + pid_h * stride_out_h + off_d * stride_out_d + tl.store(out_ptrs, o_merged.to(out_ptr.dtype.element_ty), mask=off_d < head_dim) + + if USE_PDL: + tl.extra.cuda.gdc_launch_dependents() + + +def resolve_num_topk_chunks(total_q: int, num_kv_heads: int, max_topk: int) -> int: + """Split-K factor over the top-k blocks, as a power of two. + + Depends only on shapes that are fixed for a captured batch size, so the + launch geometry is frozen inside a CUDA graph. + """ + target = max(1, min(max_topk, _TARGET_GRID // max(1, total_q * num_kv_heads))) + return 1 << (target.bit_length() - 1) + + +@torch.no_grad() +def minimax_m3_sparse_attn_decode( + q: torch.Tensor, # [total_q, num_heads, head_dim] + k_paged: torch.Tensor, # [num_pages, num_kv_heads, page_size, head_dim] + v_paged: torch.Tensor, # same layout as k_paged + topk_idx: torch.Tensor, # [num_kv_heads, total_q, topk] + block_table: torch.Tensor, # [num_reqs, max_blocks] + seq_lens: torch.Tensor, # [num_reqs] int32 + *, + sm_scale: float, + output: torch.Tensor, # [total_q, num_heads, head_dim] + decode_query_len: int, + kv_scale: Optional[torch.Tensor] = None, + num_topk_chunks: Optional[int] = None, +) -> None: + """Block-sparse GQA decode attention, written into output in place. + + kv_scale is an optional scalar dequantization factor for an FP8 cache; + MiniMax-M3 stores unscaled E4M3, so it is normally None. num_topk_chunks + overrides the split-K factor and exists for tests: the merged result must + not depend on it. + """ + total_q, num_heads, head_dim = q.shape + num_kv_heads = int(k_paged.shape[1]) + if total_q != int(seq_lens.shape[0]) * decode_query_len: + raise ValueError( + f"total_q ({total_q}) must be batch ({int(seq_lens.shape[0])}) * " + f"decode_query_len ({decode_query_len})." + ) + if int(k_paged.shape[2]) != SPARSE_BLOCK_SIZE: + raise ValueError( + f"MiniMax-M3 sparse decode requires page_size={SPARSE_BLOCK_SIZE}; " + f"got {int(k_paged.shape[2])}." + ) + max_topk = int(topk_idx.shape[-1]) + gqa_group_size = num_heads // num_kv_heads + use_scale = k_paged.dtype in _FP8_DTYPES and kv_scale is not None + # Triton needs a real pointer even for the unused argument. + scale_arg = kv_scale if use_scale else output + + if num_topk_chunks is None: + num_topk_chunks = resolve_num_topk_chunks(total_q, num_kv_heads, max_topk) + elif num_topk_chunks & (num_topk_chunks - 1): + raise ValueError(f"num_topk_chunks must be a power of two; got {num_topk_chunks}.") + + # Persistent arena rather than torch.empty, so the partials keep one address + # across CUDA graph replays. + reserve = torch.cuda.is_current_stream_capturing() + # fp32 partials: at decode these are a couple of MB, and keeping them wide + # means the split-K factor cannot perturb the merged result. + o_partial = get_memory_buffers().get_buffer( + [num_topk_chunks, total_q, num_heads, head_dim], + torch.float32, + buffer_name="m3_sparse_decode_o_partial", + reserve_buffer=reserve, + ) + lse_partial = get_memory_buffers().get_buffer( + [num_topk_chunks, total_q, num_heads], + torch.float32, + buffer_name="m3_sparse_decode_lse_partial", + reserve_buffer=reserve, + ) + + use_pdl = _pdl_enabled() + pdl_launch = {"launch_pdl": True} if use_pdl else {} + + _gqa_sparse_decode_kernel[(total_q * num_topk_chunks, num_kv_heads)]( + q, + k_paged, + v_paged, + scale_arg, + topk_idx, + o_partial, + lse_partial, + block_table, + seq_lens, + total_q, + gqa_group_size, + head_dim, + max_topk, + sm_scale, + decode_query_len, + q.stride(0), + q.stride(1), + q.stride(2), + k_paged.stride(0), + k_paged.stride(1), + k_paged.stride(2), + k_paged.stride(3), + v_paged.stride(0), + v_paged.stride(1), + v_paged.stride(2), + v_paged.stride(3), + topk_idx.stride(0), + topk_idx.stride(1), + topk_idx.stride(2), + o_partial.stride(0), + o_partial.stride(1), + o_partial.stride(2), + o_partial.stride(3), + lse_partial.stride(0), + lse_partial.stride(1), + lse_partial.stride(2), + block_table.stride(0), + BLOCK_SIZE_K=SPARSE_BLOCK_SIZE, + NUM_TOPK_CHUNKS=num_topk_chunks, + USE_SCALE=use_scale, + USE_PDL=use_pdl, + **pdl_launch, + ) + _merge_topk_attn_out_kernel[(total_q, num_heads)]( + o_partial, + lse_partial, + output, + head_dim, + o_partial.stride(0), + o_partial.stride(1), + o_partial.stride(2), + o_partial.stride(3), + lse_partial.stride(0), + lse_partial.stride(1), + lse_partial.stride(2), + output.stride(0), + output.stride(1), + output.stride(2), + NUM_TOPK_CHUNKS=num_topk_chunks, + USE_PDL=use_pdl, + **pdl_launch, + ) + + +__all__ = [ + "SPARSE_BLOCK_SIZE", + "minimax_m3_sparse_attn_decode", + "resolve_num_topk_chunks", +] diff --git a/tensorrt_llm/_torch/attention_backend/sparse/minimax_m3/trtllm_gen_dense_decode.py b/tensorrt_llm/_torch/attention_backend/sparse/minimax_m3/trtllm_gen_dense_decode.py new file mode 100644 index 000000000000..23609a9fce7b --- /dev/null +++ b/tensorrt_llm/_torch/attention_backend/sparse/minimax_m3/trtllm_gen_dense_decode.py @@ -0,0 +1,241 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""trtllm-gen decode attention for MiniMax-M3's dense layers (0-2). + +Those layers attend the whole page table, so nothing about them needs MSA; +they only run there because MsaSparseGqaFmha claims every M3 layer. MSA's +kernel uses the context schedule, spending a 128-row Q tile on one decode +token, while trtllm-gen has a generation tile scheduler for exactly this shape. + +FlashInferTrtllmGenFmha cannot be reused as-is. It reaches the pool through +build_trtllm_gen_kv_cache_metadata, which assumes each layer contributes +exactly K+V to a pool slot. M3 packs K+V for every layer of a group into one +slot, and sparse layers add an index-K sub-page on top, so there is no uniform +per-layer stride and _kv_pool_mapping_offset is only a ranking, not an +addressable offset. This module goes around that: it builds the same flat +sub-page pool and [batch, 2, max_blocks] block table the kernel expects +directly out of M3's own slot geometry, then calls the same flashinfer entry +point the generic path calls. +""" + +from __future__ import annotations + +import functools +from typing import Optional + +import torch + +from tensorrt_llm._torch.memory_buffer_utils import get_memory_buffers + + +@functools.lru_cache(maxsize=None) +def _counter_size(num_heads: int, max_num_requests: int, device_index: int) -> int: + """Byte size of the multi-CTA KV counter block. + + Cached like the workspace size below: it depends only on the head count, + the request bound and the device's SM count, none of which move during a + run, while computing it reaches C++ through a deferred import and a + device-properties query on every dense layer of every step. + """ + from tensorrt_llm._torch.attention_backend.fmha.flashinfer_trtllm_gen import ( + _get_multi_ctas_kv_counter_size, + ) + + multi_processor_count = torch.cuda.get_device_properties(device_index).multi_processor_count + return int(_get_multi_ctas_kv_counter_size(num_heads, max_num_requests, multi_processor_count)) + + +def _counter_buffer( + device: torch.device, num_heads: int, max_num_requests: int, reserve: bool +) -> torch.Tensor: + """Zeroed multi-CTA KV counters for one call. + + Taken from the shared arena, like the workspace beside it, so the block + joins the graph memory pool and clear_memory_buffers(). The arena hands + back uninitialized memory, so the zeroing is per call rather than per + allocation; a few KB of memset costs nothing next to the kernel it feeds. + """ + device_index = device.index if device.index is not None else torch.cuda.current_device() + counters = get_memory_buffers().get_buffer( + [_counter_size(num_heads, max_num_requests, device_index)], + torch.uint8, + buffer_name="m3_trtllm_gen_kv_counters", + reserve_buffer=reserve, + ) + counters.zero_() + return counters + + +@functools.lru_cache(maxsize=None) +def _workspace(q_dtype: torch.dtype, num_heads: int, head_dim: int, num_kv_heads: int) -> int: + """Byte size of the trtllm-gen scratch slab. + + It is a fixed slab (kTrtllmGenWorkspaceSize), independent of the batch, but + the size is read from the C++ layout rather than hardcoded. Cached so that + read, and the layout dict it builds, happen once instead of once per dense + layer per step. + """ + from tensorrt_llm._torch.attention_backend.fmha.flashinfer_trtllm_gen import ( + _get_generation_workspace_layout, + ) + + layout = _get_generation_workspace_layout(q_dtype, 1, 1, num_heads, head_dim, num_kv_heads, 0) + return int(layout["trtllm_gen_workspace_size"]) + + +def subpage_block_table( + block_table: torch.Tensor, subpages_per_slot: int, reserve: bool = False +) -> torch.Tensor: + """Expand a slot table into trtllm-gen's separate K and V page rows. + + uses_shared_paged_kv_idx is False for TensorRT-LLM, so the kernel takes + [batch, 2, max_blocks] and indexes K and V independently. Rooting the pool + at this layer's K (see get_kv_subpage_pool) puts slot s's K at + s * subpages_per_slot and its V one sub-page later. + + The result is a function of the slot table and that factor alone, so every + dense layer of a step would compute the same one. prepare() therefore + stages it once into a graph-stable buffer, and this runs only where it + could not: a manager whose layers disagree on the factor, or a caller that + skipped prepare(). + """ + batch, max_blocks = block_table.shape + out = get_memory_buffers().get_buffer( + [batch, 2, max_blocks], + torch.int32, + buffer_name="m3_trtllm_gen_subpage_block_table", + reserve_buffer=reserve, + ) + write_subpage_block_table(block_table, subpages_per_slot, out) + return out + + +def write_subpage_block_table( + block_table: torch.Tensor, subpages_per_slot: int, out: torch.Tensor +) -> None: + """Write the K and V sub-page rows of block_table into out.""" + torch.mul(block_table, subpages_per_slot, out=out[:, 0]) + torch.add(out[:, 0], 1, out=out[:, 1]) + + +def uniform_subpages_per_slot(kv_cache_manager) -> int: + """Sub-pages per slot when every layer of the pool agrees, else 0. + + The factor is a property of a layer group, so a single-group model has one + for the whole pool and prepare() can expand the block table without naming + a layer (see subpage_block_table). A manager with no sub-page pool, or one + whose groups disagree, reports 0 rather than a guess. + """ + get_pool = getattr(kv_cache_manager, "get_kv_subpage_pool", None) + layer_offsets = getattr(kv_cache_manager, "layer_offsets", None) + if get_pool is None or not layer_offsets: + return 0 + factors = {int(get_pool(layer_idx, "HND")[1]) for layer_idx in layer_offsets} + return factors.pop() if len(factors) == 1 else 0 + + +def minimax_m3_trtllm_gen_dense_decode( + q: torch.Tensor, # [total_q, num_heads, head_dim] + kv_cache_manager, + layer_idx: int, + block_table: torch.Tensor, # [batch, max_blocks] slot ids + seq_lens: torch.Tensor, # [batch] int32 + *, + sm_scale: float, + output: torch.Tensor, # [total_q, num_heads, head_dim] + decode_query_len: int, + max_seq_len: int, + max_num_requests: int, + staged_subpage_table: Optional[torch.Tensor] = None, + staged_subpages_per_slot: int = 0, + enable_pdl: bool = True, +) -> None: + """Full-context decode attention through trtllm-gen, in place into output. + + staged_subpage_table is the expansion of block_table prepare() already + staged, used when staged_subpages_per_slot matches this layer's factor and + expanded here otherwise; see subpage_block_table. + """ + from tensorrt_llm._torch.attention_backend.fmha.flashinfer_trtllm_gen import ( + _trtllm_gen_batch_decode_with_kv_cache, + ) + + kv_pool, subpages_per_slot = kv_cache_manager.get_kv_subpage_pool(layer_idx, "HND") + num_heads = int(q.shape[1]) + + # The kernel variant is picked from the Q dtype and shares one dtype across + # q/k/v, so an FP8 pool needs FP8 Q. M3 stores unscaled E4M3, so this is a + # plain cast and a no-op when the fused producer already emitted FP8. + if kv_pool.dtype == torch.float8_e4m3fn and q.dtype != torch.float8_e4m3fn: + q = q.to(torch.float8_e4m3fn) + + reserve = torch.cuda.is_current_stream_capturing() + workspace = get_memory_buffers().get_buffer( + [_workspace(q.dtype, num_heads, int(q.shape[2]), int(kv_pool.shape[1]))], + torch.uint8, + buffer_name="m3_trtllm_gen_workspace", + reserve_buffer=reserve, + ) + if staged_subpage_table is None or staged_subpages_per_slot != subpages_per_slot: + staged_subpage_table = subpage_block_table(block_table, subpages_per_slot, reserve) + + _trtllm_gen_batch_decode_with_kv_cache( + q, # query + kv_pool, # kv_pool + workspace, # workspace_buffer + _counter_buffer( + q.device, num_heads, max_num_requests, reserve + ), # multi_ctas_kv_counter_buffer + staged_subpage_table, # block_tables + seq_lens, # seq_lens + max_seq_len, # max_seq_len + sm_scale, # bmm1_scale + 1.0, # bmm2_scale + -1, # window_left: M3 dense layers are fully causal + output, # out + None, # sinks + enable_pdl, # enable_pdl + decode_query_len, # q_len_per_req + None, # max_q_len + None, # cum_seq_lens_q + None, # kv_scale_pool: M3 stores unscaled E4M3 + False, # uses_shared_paged_kv_idx + ) + + +@functools.lru_cache(maxsize=1) +def _flashinfer_available() -> bool: + """Whether flashinfer can be imported, resolved once for the process. + + The verdict below is consulted once per step by prepare() and again per + dense layer, so the import statement would otherwise be on the hot path. + """ + try: + import flashinfer # noqa: F401 + except ImportError: + return False + return True + + +def dense_decode_unsupported_reason(kv_cache_manager, head_dim: int) -> Optional[str]: + """Return None when the geometry is supported, else why it is not. + + Takes head_dim rather than a query tensor so prepare() can reach the same + verdict as the call site without one, and so the two cannot drift. + """ + if not hasattr(kv_cache_manager, "get_kv_subpage_pool"): + return "the KV cache manager does not expose a flat sub-page pool." + if int(head_dim) != 128: + return f"head_dim {int(head_dim)}; only 128 has trtllm-gen H128 cubins." + if not _flashinfer_available(): + return "flashinfer is not installed." + return None + + +__all__ = [ + "dense_decode_unsupported_reason", + "minimax_m3_trtllm_gen_dense_decode", + "subpage_block_table", + "uniform_subpages_per_slot", + "write_subpage_block_table", +] diff --git a/tensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.py b/tensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.py index 0db351140735..b7009f5b4b37 100644 --- a/tensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.py +++ b/tensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.py @@ -7670,6 +7670,211 @@ def _( dtype=output_dtype, device=q.device) + # ------------------------------------------------------------------ # + # CuTe DSL MiniMax-M3 index decode scoring (Blackwell SM100) # + # ------------------------------------------------------------------ # + from ..cute_dsl_kernels.blackwell.cute_ptx_utils import \ + TORCH_TO_CUTE_DTYPE as _M3_TORCH_TO_CUTE_DTYPE + from ..cute_dsl_kernels.blackwell.minimax_m3_index_decode_score import \ + IndexDecodeScoreKernel + + class CuteDSLMiniMaxM3IndexDecodeScoreRunner: + """Runner for the MiniMax-M3 indexer decode block-scoring kernel. + + Caches compiled kernels keyed on the static params + (dtype, num_heads, max_decode_query_len, head_dim); batch, query-token + count, page count and block-table width stay symbolic, so one compile + covers every decode step of a given model shape. + """ + + kernel_cache = dict() + # One CTA per (request, split); each walks blocks split, split + 256, + # ... so a request longer than 256 pages just loops. Matches upstream. + SPLIT_K = 256 + # The only geometry the kernel has been validated on. + SUPPORTED_HEAD_DIM = 128 + SUPPORTED_PAGE_SIZE = IndexDecodeScoreKernel.BLOCK_K + # BLOCK_Q must fit one warp's worth of epilogue lanes. + MAX_BLOCK_Q = 32 + + @classmethod + def is_supported( + cls, + *, + q_dtype: torch.dtype, + num_heads: int, + head_dim: int, + page_size: int, + max_decode_query_len: int, + ) -> bool: + """Whether this kernel can serve the given decode geometry. + + Callers use this to pick between the CuTe DSL scorer and the + fallback rather than catching an exception on the hot path. + """ + return (is_sm_100f() and q_dtype in _M3_TORCH_TO_CUTE_DTYPE + and head_dim == cls.SUPPORTED_HEAD_DIM + and page_size == cls.SUPPORTED_PAGE_SIZE + and num_heads * max_decode_query_len <= cls.MAX_BLOCK_Q + and max_decode_query_len >= 1) + + @classmethod + def _compile(cls, q_dtype: torch.dtype, num_heads: int, + max_decode_query_len: int, head_dim: int): + key = (q_dtype, num_heads, max_decode_query_len, head_dim) + if key in cls.kernel_cache: + return + + cute_dtype = _M3_TORCH_TO_CUTE_DTYPE[q_dtype] + page_size = cls.SUPPORTED_PAGE_SIZE + + sym_total_tokens = cute.sym_int() + sym_batch = cute.sym_int() + + # 16-element divisibility on every non-innermost stride is what the + # TMA descriptors for Q and K assume. + def _sym_stride(): + return cute.sym_int64(divisibility=16) + + q_fake = cute.runtime.make_fake_tensor( + cute_dtype, (sym_total_tokens, num_heads, head_dim), + stride=(_sym_stride(), _sym_stride(), 1)) + + # The index-K pool may be coalesced with the main K/V cache, in + # which case the per-page stride exceeds page_size * head_dim, so + # dim 0 is read at runtime. + k_fake = cute.runtime.make_fake_tensor( + cute_dtype, (cute.sym_int(), page_size, head_dim), + stride=(_sym_stride(), _sym_stride(), 1)) + + bt_fake = cute.runtime.make_fake_compact_tensor( + cutlass.Int32, (sym_batch, cute.sym_int()), stride_order=(1, 0)) + + # Every stride is symbolic because production passes a transposed + # view of the [heads, blocks, tokens] selector buffer, whose + # innermost logical stride is the token count rather than 1. + score_fake = cute.runtime.make_fake_tensor( + cutlass.Float32, (num_heads, sym_total_tokens, cute.sym_int()), + stride=(cute.sym_int64(), cute.sym_int64(), cute.sym_int64())) + + sl_fake = cute.runtime.make_fake_compact_tensor(cutlass.Int32, + (sym_batch, ), + stride_order=(0, )) + + fake_stream = cute.runtime.make_fake_stream( + use_tvm_ffi_env_stream=True) + + kernel = IndexDecodeScoreKernel( + cute_dtype, + num_heads, + max_decode_query_len, + cls.SPLIT_K, + head_dim, + ) + cls.kernel_cache[key] = cute.compile( + kernel, + q_fake, + k_fake, + bt_fake, + score_fake, + sl_fake, + fake_stream, + options="--enable-tvm-ffi", + ) + logger.debug( + f"[compile cute_dsl minimax_m3_index_decode_score] {key}") + + @classmethod + def forward( + cls, + idx_q: torch.Tensor, + index_k_cache: torch.Tensor, + block_table: torch.Tensor, + seq_lens: torch.Tensor, + score: torch.Tensor, + max_decode_query_len: int, + ) -> None: + """Score one decode step, compiling the kernel on first use. + + The compile allocates, so it must not land inside a CUDA graph + capture. It does not: CUDAGraphRunner.capture runs eager warmup + forwards first, and those cover every geometry the graphs it then + captures will replay. + """ + _, num_heads, head_dim = idx_q.shape + key = (idx_q.dtype, num_heads, max_decode_query_len, head_dim) + if key not in cls.kernel_cache: + cls._compile(idx_q.dtype, num_heads, max_decode_query_len, + head_dim) + cls.kernel_cache[key](idx_q, index_k_cache, block_table, score, + seq_lens) + + @torch.library.custom_op("trtllm::cute_dsl_minimax_m3_index_decode_score", + mutates_args=("score", ), + device_types="cuda") + def cute_dsl_minimax_m3_index_decode_score( + idx_q: torch.Tensor, + index_k_cache: torch.Tensor, + block_table: torch.Tensor, + seq_lens: torch.Tensor, + score: torch.Tensor, + max_decode_query_len: int, + ) -> None: + """Write per-block max index scores for one decode step, in place. + + Args: + idx_q: [total_q, num_index_heads, head_dim], BF16 or FP8 E4M3. + total_q must be batch * decode_query_len with a uniform + decode_query_len, which the kernel infers. + index_k_cache: [num_pages, page_size, head_dim], same dtype as + idx_q. May be a strided view of a coalesced pool. + block_table: [batch, max_blocks_per_seq] int32 page table. + seq_lens: [batch] int32 attended KV length per request. + score: [num_index_heads, total_q, max_blocks] float32, mutated in + place. Only blocks below ceil(seq_len / page_size) are written, + which is exactly the range the block selector reads. Arbitrary + strides are accepted so a transposed selector buffer can be + passed without a copy. + max_decode_query_len: compile-time bound on decode_query_len; + num_index_heads * max_decode_query_len must not exceed 32. + """ + if not is_sm_100f(): + raise ValueError( + f"CuteDSL: SM version {get_sm_version()} is not supported. " + f"CuteDSL MiniMax-M3 index decode score only supports SM 100 " + f"family.") + logger.info_once( + f"cute_dsl_minimax_m3_index_decode_score inputs: " + f"idx_q dtype={idx_q.dtype} shape={tuple(idx_q.shape)} stride={idx_q.stride()}; " + f"index_k_cache dtype={index_k_cache.dtype} shape={tuple(index_k_cache.shape)} " + f"stride={index_k_cache.stride()}; " + f"block_table shape={tuple(block_table.shape)} stride={block_table.stride()}; " + f"seq_lens shape={tuple(seq_lens.shape)}; " + f"score dtype={score.dtype} shape={tuple(score.shape)} stride={score.stride()}; " + f"max_decode_query_len={max_decode_query_len}", + key="cute_dsl_minimax_m3_index_decode_score_inputs", + ) + CuteDSLMiniMaxM3IndexDecodeScoreRunner.forward( + idx_q, + index_k_cache, + block_table, + seq_lens, + score, + max_decode_query_len, + ) + + @torch.library.register_fake( + "trtllm::cute_dsl_minimax_m3_index_decode_score") + def _( + idx_q: torch.Tensor, + index_k_cache: torch.Tensor, + block_table: torch.Tensor, + seq_lens: torch.Tensor, + score: torch.Tensor, + max_decode_query_len: int, + ) -> None: + return None + # ====================================================================== # BF16 Dense Persistent BMM (CuTe DSL) for Blackwell # ====================================================================== diff --git a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/cute_ptx_utils.py b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/cute_ptx_utils.py new file mode 100644 index 000000000000..127604893b83 --- /dev/null +++ b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/cute_ptx_utils.py @@ -0,0 +1,161 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-License-Identifier: Apache-2.0 +# Vendored from vLLM (Apache-2.0): +# https://github.com/vllm-project/vllm/blob/6f91edf96d3f3272945809c04702380053bff4de/vllm/cute_utils/__init__.py +# https://github.com/vllm-project/vllm/blob/6f91edf96d3f3272945809c04702380053bff4de/vllm/cute_utils/cvt.py +"""Warp-level PTX intrinsics for CuTe DSL kernels that use ``mma.sync``. + +Most Blackwell kernels in this tree drive the tensor cores through +``cute.gemm`` / tcgen05, which the CuTe DSL exposes directly. A kernel whose +GEMM has a tiny N dimension does better with warp-level ``mma.sync`` and high +CTA occupancy than with a deep single-CTA tcgen05 pipeline, and the DSL has no +wrapper for that instruction, so it is spelled out as inline PTX here. + +Vendored from the vLLM sources linked in the file header +(v0.26.1rc0-77-g6f91edf96), reduced to the symbols +:mod:`minimax_m3_index_decode_score` needs. +""" + +import torch +from cutlass import BFloat16, Float8E4M3FN, Float16, Float32, Int32, Int64, Uint32, cute +from cutlass._mlir import ir +from cutlass._mlir.dialects import llvm, vector +from cutlass.cute.nvgpu import cpasync +from cutlass.cutlass_dsl import T, dsl_user_op + +__all__ = [ + "EVICT_FIRST", + "TORCH_TO_CUTE_DTYPE", + "fp8x4_to_fp16x4", + "mma_sync", + "simple_tma_copy", +] + +TORCH_TO_CUTE_DTYPE = { + torch.bfloat16: BFloat16, + torch.float8_e4m3fn: Float8E4M3FN, +} + +_CUTE_TO_PTX_DTYPE = { + BFloat16: "bf16", + Float16: "f16", + Float8E4M3FN: "e4m3", + Float32: "f32", +} + +# L2 cache-eviction policy descriptor; see CUTLASS +# include/cute/arch/copy_sm90_desc.hpp (v4.3.2, L193-L197). +EVICT_FIRST = Int64(0x12F0000000000000) + + +def simple_tma_copy(atom, src, dst, mbar=None, cache_policy=None): + """Wrap ``group_modes()`` + ``tma_partition()`` for a whole-tile TMA copy. + + Call this WITHOUT ``cute.elect_one()``: ``tma_partition`` already reduces + the copy to a single issuing lane. + """ + if isinstance(atom.op, cpasync.CopyBulkTensorTileG2SOp): + gmem = src + smem = dst + elif isinstance(atom.op, cpasync.CopyBulkTensorTileS2GOp): + smem = src + gmem = dst + else: + raise ValueError(f"simple_tma_copy expects a bulk-tensor TMA atom, got {atom.op!r}.") + + s_part, g_part = cpasync.tma_partition( + atom, + 0, + cute.make_layout(1), + cute.group_modes(smem, 0), + cute.group_modes(gmem, 0), + ) + + if isinstance(atom.op, cpasync.CopyBulkTensorTileG2SOp): + cute.copy(atom, g_part, s_part, tma_bar_ptr=mbar, cache_policy=cache_policy) + else: + cute.copy(atom, s_part, g_part, cache_policy=cache_policy) + + +@dsl_user_op +def mma_sync(a, b, c: cute.Tensor, *, loc=None, ip=None): + """Warp-level ``mma.sync.aligned.m16n8kK`` accumulating into ``c``. + + ``K`` follows from the operand width (32B of A per lane), so this covers + m16n8k16 for 16-bit operands and m16n8k32 for 8-bit ones. + """ + a_ty = _CUTE_TO_PTX_DTYPE[a.element_type] + b_ty = _CUTE_TO_PTX_DTYPE[b.element_type] + c_ty = _CUTE_TO_PTX_DTYPE[c.element_type] + mlir_ty = c.element_type.mlir_type + K = 256 // a.element_type.width # 32B + + # recast_tensor needs tensor-backed fragments, so materialize SSA values + # here and let callers pass converted FP8 fragments straight through. + if isinstance(a, cute.TensorSSA): + a_ = cute.make_rmem_tensor_like(a) + a_.store(a, loc=loc, ip=ip) + a = a_ + if isinstance(b, cute.TensorSSA): + b_ = cute.make_rmem_tensor_like(b) + b_.store(b, loc=loc, ip=ip) + b = b_ + + a = cute.recast_tensor(a, Int32, loc=loc, ip=ip) + b = cute.recast_tensor(b, Int32, loc=loc, ip=ip) + out = llvm.inline_asm( + llvm.StructType.get_literal([mlir_ty] * 4), + [a[i].ir_value(loc=loc, ip=ip) for i in range(4)] + + [b[i].ir_value(loc=loc, ip=ip) for i in range(2)] + + [c[i].ir_value(loc=loc, ip=ip) for i in range(4)], + f"mma.sync.aligned.m16n8k{K}.row.col.{c_ty}.{a_ty}.{b_ty}.{c_ty} " + "{$0, $1, $2, $3}, " + "{$4, $5, $6, $7}, " + "{$8, $9}, " + "{$10, $11, $12, $13};", + "=f,=f,=f,=f,r,r,r,r,r,r,f,f,f,f", + has_side_effects=False, + is_align_stack=False, + loc=loc, + ip=ip, + ) + vec = vector.from_elements( + ir.VectorType.get([4], mlir_ty, loc=loc), + [llvm.extractvalue(mlir_ty, out, [i], loc=loc, ip=ip) for i in range(4)], + loc=loc, + ip=ip, + ) + return cute.TensorSSA(vec, 4, c.element_type) + + +@dsl_user_op +def fp8x4_to_fp16x4(x: Uint32, *, loc=None, ip=None) -> cute.TensorSSA: + """Convert four packed E4M3 values to four FP16 values, as two ``Uint32``. + + SM100 has no native ``mma.sync.f8``; ptxas lowers it to F2FP.F16.E4M3 plus + HMMA anyway, so converting explicitly gives better codegen and keeps the + two FP16 k-fragments visible to the caller. + """ + out = llvm.inline_asm( + llvm.StructType.get_literal([T.i32()] * 2), + [x.ir_value(loc=loc, ip=ip)], + "{\n\t" + ".reg .b16 lo, hi;\n\t" + "mov.b32 {lo, hi}, $2;\n\t" + "cvt.rn.f16x2.e4m3x2 $0, lo;\n\t" + "cvt.rn.f16x2.e4m3x2 $1, hi;\n\t" + "}\n", + "=r,=r,r", + has_side_effects=False, + is_align_stack=False, + loc=loc, + ip=ip, + ) + vec = vector.from_elements( + ir.VectorType.get([2], T.i32(), loc=loc), + [llvm.extractvalue(T.i32(), out, [i], loc=loc, ip=ip) for i in range(2)], + loc=loc, + ip=ip, + ) + return cute.TensorSSA(vec, 2, Uint32) diff --git a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/minimax_m3_index_decode_score.py b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/minimax_m3_index_decode_score.py new file mode 100644 index 000000000000..d5a2a7ceb3ea --- /dev/null +++ b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/minimax_m3_index_decode_score.py @@ -0,0 +1,429 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-License-Identifier: Apache-2.0 +# Vendored from vLLM (Apache-2.0): +# https://github.com/vllm-project/vllm/blob/6f91edf96d3f3272945809c04702380053bff4de/vllm/models/minimax_m3/nvidia/ops/index_decode_score.py +"""CuTe DSL MiniMax-M3 index decode block-scoring kernel (Blackwell SM100). + +Computes, for every (index head, decode query token, KV block), the maximum +causally-valid Q . K dot product over the 128 index-K positions of that block. +Those per-block maxima are what minimax_m3_select_blocks ranks to pick the +top-k blocks the sparse attention then attends. + +Uses TMA plus warp-level mma.sync rather than tcgen05: the score GEMM's N +dimension is one decode token times a handful of index heads, so CTA occupancy +matters far more than a deep single-CTA pipeline. + +Vendored from the vLLM source linked in the file header (v0.26.1rc0-77-g6f91edf96). +Differences from upstream: + +* cpasync.make_tiled_tma_atom returns (atom, tensor) in the CuTe DSL version + pinned here rather than a TmaInfo, so the shared-memory layouts are rebuilt + from the same compile-time constants instead of being read back off the + descriptor. +* PDL follows TRTLLM_ENABLE_PDL and grid dependency control comes from + blackwell.utils rather than cute.arch. +* Compilation and caching live in the trtllm::cute_dsl_minimax_m3_index_decode_score + runner, matching the other CuTe DSL ops in this tree. +""" + +import cutlass +from cuda.bindings.driver import CUstream +from cutlass import Float8E4M3FN, Float16, Float32, Int64, Uint32, cute +from cutlass.cute.nvgpu import cpasync, warp +from cutlass.utils.smem_allocator import SmemAllocator + +from .cute_ptx_utils import EVICT_FIRST, fp8x4_to_fp16x4, mma_sync, simple_tma_copy +from .utils import TRTLLM_ENABLE_PDL, griddepcontrol_launch_dependents, griddepcontrol_wait + +__all__ = ["IndexDecodeScoreKernel"] + + +@cute.jit +def _fp8_to_f16_mma_fragments(src: cute.Tensor): + """Split an FP8 ldmatrix fragment into the two FP16 k-fragments MMA wants.""" + src_elems = cute.size(src) + src_u32 = cute.recast_tensor(src, Uint32) + src_f16 = cute.make_rmem_tensor(src_elems, Float16) + src_f16_u32 = cute.recast_tensor(src_f16, Uint32) + # Packed conversion; faster and fewer SASS instructions than + # src.load().to(Float16). + for i in cutlass.range_constexpr(src_elems // 4): + converted = fp8x4_to_fp16x4(src_u32[i]) + src_f16_u32[i * 2] = converted[0] + src_f16_u32[i * 2 + 1] = converted[1] + lower = cute.make_rmem_tensor(src_elems // 2, Float16) + upper = cute.make_rmem_tensor(src_elems // 2, Float16) + + # FP8 ldmatrix yields four consecutive values along K. Split each group + # into the lower two and upper two for the two FP16 MMA k-fragments. + for i in cutlass.range_constexpr(src_elems // 2): + lower[i] = src_f16[(i // 2) * 4 + i % 2] + upper[i] = src_f16[(i // 2) * 4 + 2 + i % 2] + return lower, upper + + +class IndexDecodeScoreKernel: + """Per-block max index score for one decode step. + + Grid is (batch, split_k): CTA (b, s) walks KV blocks + s, s + split_k, s + 2 * split_k, ... of request b and writes each block's + score directly, so no cross-CTA reduction is needed. CTAs whose split_id + exceeds the request's block count exit immediately. + """ + + BLOCK_K = 128 + BAR_MMA = 1 + num_stages = 2 + + def __init__( + self, + dtype: type[cutlass.Numeric], + num_heads: int, + max_decode_query_len: int, + split_k: int, + head_dim: int = 128, + ): + self.dtype = dtype + self.num_heads = num_heads + self.max_decode_query_len = max_decode_query_len + self.split_k = split_k + self.head_dim = head_dim + + def _swizzle_elems(self) -> int: + """Elements spanned by one 128-byte swizzle atom.""" + return 128 * 8 // self.dtype.width + + def _sq_layout(self): + """Composed SMEM layout for the Q tile, shared by the descriptor and the kernel.""" + elems = self._swizzle_elems() + head_dim = self.head_dim + block_q = self.num_heads * self.max_decode_query_len + layout = cute.make_layout( + (self.max_decode_query_len, self.num_heads, (elems, head_dim // elems)), + stride=(elems, self.max_decode_query_len * elems, (1, block_q * elems)), + ) + return cute.make_composed_layout(cute.make_swizzle(3, 4, 3), 0, layout) + + def _sk_layout(self): + """Composed SMEM layout for the pipelined K tiles.""" + elems = self._swizzle_elems() + head_dim = self.head_dim + block_k = self.BLOCK_K + layout = cute.make_layout( + (1, block_k, (elems, head_dim // elems), self.num_stages), + stride=(0, elems, (1, block_k * elems), block_k * head_dim), + ) + return cute.make_composed_layout(cute.make_swizzle(3, 4, 3), 0, layout) + + @cute.jit + def __call__( + self, + gQ: cute.Tensor, # [bs * runtime_decode_query_len, num_heads, head_dim] + gK_cache: cute.Tensor, # [num_pages, page_size, head_dim] + block_table: cute.Tensor, # [bs, max_pages] + score: cute.Tensor, # [num_heads, bs * runtime_decode_query_len, max_pages] + seq_lens: cute.Tensor, # [bs] + stream: CUstream, + ): + num_heads = self.num_heads + head_dim = self.head_dim + MAX_DQL = self.max_decode_query_len + assert num_heads * MAX_DQL <= 32 + + batch = seq_lens.shape[0] + decode_query_len = gQ.shape[0] // batch + grid = (batch, self.split_k, 1) + block = (32 * 5, 1, 1) + + tma_g2s = cpasync.CopyBulkTensorTileG2SOp() + elems = self._swizzle_elems() + + q_tma_atom, q_tma_tensor = cpasync.make_tiled_tma_atom( + tma_g2s, + cute.logical_divide(gQ, (None, None, elems)), + self._sq_layout(), + (MAX_DQL, num_heads, head_dim), + ) + k_tma_atom, k_tma_tensor = cpasync.make_tiled_tma_atom( + tma_g2s, + cute.logical_divide(gK_cache, (None, None, elems)), + self._sk_layout(), + (1, self.BLOCK_K, head_dim), + ) + + self.kernel( + q_tma_atom, + q_tma_tensor, + k_tma_atom, + k_tma_tensor, + block_table, + score, + seq_lens, + decode_query_len, + ).launch(grid=grid, block=block, stream=stream, use_pdl=TRTLLM_ENABLE_PDL) + + @cute.kernel + def kernel( + self, + q_tma_atom: cute.CopyAtom, + q_tma_tensor: cute.Tensor, + k_tma_atom: cute.CopyAtom, + k_tma_tensor: cute.Tensor, + block_table: cute.Tensor, + score: cute.Tensor, + seq_lens: cute.Tensor, + decode_query_len, + ): + batch_id, split_id, _ = cute.arch.block_idx() + _, split_k, _ = cute.arch.grid_dim() + warp_id = cute.arch.make_warp_uniform(cute.arch.warp_idx()) + lane_id = cute.arch.lane_idx() + + NUM_HEADS = self.num_heads + MAX_DQL = self.max_decode_query_len + BLOCK_Q = NUM_HEADS * MAX_DQL + BLOCK_K = self.BLOCK_K + head_dim = self.head_dim + dtype = self.dtype + MMA_N = 8 + num_stages = self.num_stages + Q_TILES = cute.ceil_div(BLOCK_Q, MMA_N) + EPI_Q = Q_TILES * MMA_N + + sq_layout = self._sq_layout() + sk_layout = self._sk_layout() + + smem = SmemAllocator() + sK = smem.allocate_tensor( + dtype, + sk_layout.outer, + byte_alignment=128, + swizzle=sk_layout.inner, + )[0, None, None, None] + # sQ aliases the first K stage: Q is consumed into registers before the + # first K tile is needed, so the two never overlap in time. + sQ_tma = cute.make_tensor(sK[None, None, 0].iterator, layout=sq_layout.outer) + # TMA sees Q as (query, head, dim) while ldmatrix consumes a flattened Q + # column mode. The target profile keeps the rank-2 view even for + # degenerate shapes such as MAX_DQL == 1. + q_tma_elems = self._swizzle_elems() + sQ = cute.coalesce( + cute.group_modes(sQ_tma, 0, 2), + target_profile=(BLOCK_Q, (q_tma_elems, head_dim // q_tma_elems)), + ) + epi_buffer = smem.allocate_tensor(Float32, cute.make_layout((EPI_Q, 4))) + + tma_full_mbar = smem.allocate_array(Int64, num_stages) + tma_empty_mbar = smem.allocate_array(Int64, num_stages) + + # TODO: this load precedes the griddepcontrol_wait() below, so under PDL + # it can observe seq_lens as the predecessor grid left it. num_blocks + # bounds both the block_table load and the score store, neither of which + # is otherwise masked, so a stale length here is an out-of-bounds write. + # Safe only while no PDL predecessor writes seq_lens (nothing between + # on_update_kv_lens and this kernel does today). Either move the load + # past the wait, as triton_sparse_decode.py orders its gdc_wait against + # the same tensor, or record the constraint here. + seqlen = seq_lens[batch_id] + num_blocks = cute.ceil_div(seqlen, BLOCK_K) + + if split_id < num_blocks: + if warp_id == 0: + with cute.arch.elect_one(): + for i in cutlass.range_constexpr(num_stages): + cute.arch.mbarrier_init(tma_full_mbar + i, 1) + cute.arch.mbarrier_init(tma_empty_mbar + i, 128) + cute.arch.mbarrier_init_fence() + elif warp_id == 1: + cpasync.prefetch_descriptor(q_tma_atom) + cpasync.prefetch_descriptor(k_tma_atom) + cute.arch.sync_threads() + + griddepcontrol_wait() + # TODO: releasing dependents here, rather than after the epilogue + # stores, is safe only while no PDL-launched successor reads score. + # Confirm that and either move the release past the stores or record + # the constraint here. + griddepcontrol_launch_dependents() + + if warp_id == 4: + # TMA warp + tma_stage = 0 + tma_parity = 1 + + gQ_tile = cute.local_tile( + cute.domain_offset((batch_id * decode_query_len, 0, 0), q_tma_tensor), + tiler=(MAX_DQL, NUM_HEADS, head_dim), + coord=(0, 0, 0), + ) + cute.arch.mbarrier_wait(tma_empty_mbar, tma_parity) + with cute.arch.elect_one(): + Q_size = BLOCK_Q * head_dim * (dtype.width // 8) + cute.arch.mbarrier_arrive_and_expect_tx(tma_full_mbar, Q_size) + # TMA bounds-checks rows when the runtime decode_query_len is + # below MAX_DQL; padded Q columns are masked before the stores. + simple_tma_copy(q_tma_atom, gQ_tile, sQ_tma, tma_full_mbar) + + tma_stage = (tma_stage + 1) % num_stages + if tma_stage == 0: + tma_parity ^= 1 + + for block_id in range(split_id, num_blocks, split_k): + page_id = block_table[batch_id, block_id] + gK_tile = k_tma_tensor[page_id, None, None] + k_mbar = tma_full_mbar + tma_stage + + cute.arch.mbarrier_wait(tma_empty_mbar + tma_stage, tma_parity) + with cute.arch.elect_one(): + K_size = BLOCK_K * head_dim * (dtype.width // 8) + cute.arch.mbarrier_arrive_and_expect_tx(k_mbar, K_size) + simple_tma_copy( + k_tma_atom, + gK_tile, + sK[None, None, tma_stage], + k_mbar, + cache_policy=EVICT_FIRST, + ) + + tma_stage = (tma_stage + 1) % num_stages + if tma_stage == 0: + tma_parity ^= 1 + + else: + # MMA warps; each handles K[32, head_dim] @ Q[BLOCK_Q, head_dim].T + sK_warp = cute.local_tile(sK, (32, head_dim, num_stages), (warp_id, 0, 0)) + q_start = seqlen - decode_query_len + + elems = 128 // dtype.width # 16B + MMA_K = 32 * 8 // dtype.width # 32B + + # Pre-compute ldmatrix addresses. + # sK loads a [16 x 16B] tile: + # ((16, (16B, 2), 1), (32 / 16, head_dim / 32B, num_stages)) + # sQ loads an [8 x 32B] tile: + # ((8, (16B, 4)), (BLOCK_Q / MMA_N, head_dim / 64B)) + sK_ldsm = cute.zipped_divide(sK_warp, (16, cute.make_layout((elems, 2)), 1)) + sQ_ldsm = cute.zipped_divide(sQ, (MMA_N, cute.make_layout((elems, 4)))) + + # sK: (16B, (32 / 16, head_dim / 32B, num_stages)) + # sQ: (16B, (BLOCK_Q / MMA_N, head_dim / 64B)) + sK_ldsm = sK_ldsm[(lane_id % 16, (None, lane_id // 16), 0), None] + sQ_ldsm = sQ_ldsm[(lane_id % MMA_N, (None, lane_id // 8)), None] + + ldsm_op = warp.LdMatrix8x8x16bOp(num_matrices=4) + ldsm_atom = cute.make_copy_atom(ldsm_op, dtype) + + rQ = cute.make_rmem_tensor( + ((elems // 2, 2), head_dim // (MMA_K * 2), Q_TILES), dtype + ) + rK = cute.make_rmem_tensor((elems, 2, head_dim // MMA_K), dtype) + rC = cute.make_rmem_tensor((4, 2, Q_TILES), Float32) + + if warp_id == 0: + cute.arch.mbarrier_wait(tma_full_mbar, 0) + cute.arch.barrier(barrier_id=self.BAR_MMA, number_of_threads=128) + for q in cutlass.range_constexpr(Q_TILES): + cute.copy(ldsm_atom, sQ_ldsm[None, (q, None)], rQ[None, None, q]) + cute.arch.mbarrier_arrive(tma_empty_mbar) + + tma_stage = 1 % self.num_stages + tma_parity = 0 + if tma_stage == 0: + tma_parity ^= 1 + + if cutlass.const_expr(dtype is Float8E4M3FN): + rQ_f16 = cute.make_rmem_tensor((4, head_dim // MMA_K, Q_TILES, 2), Float16) + q_lower, q_upper = _fp8_to_f16_mma_fragments(rQ) + rQ_f16[None, None, None, 0].store(q_lower.load()) + rQ_f16[None, None, None, 1].store(q_upper.load()) + + for block_id in range(split_id, num_blocks, split_k): + rC.fill(0.0) + + if warp_id == 0: + cute.arch.mbarrier_wait(tma_full_mbar + tma_stage, tma_parity) + cute.arch.barrier(barrier_id=self.BAR_MMA, number_of_threads=128) + + for k in cutlass.range_constexpr(head_dim // MMA_K): + cute.copy( + ldsm_atom, + sK_ldsm[None, (None, k, tma_stage)], + rK[None, None, k], + ) + for m in cutlass.range_constexpr(2): + if cutlass.const_expr(dtype is Float8E4M3FN): + rK_lower, rK_upper = _fp8_to_f16_mma_fragments(rK[None, m, k]) + for n in cutlass.range_constexpr(Q_TILES): + rC[None, m, n] = mma_sync( + rK_lower, + rQ_f16[None, k, n, 0], + rC[None, m, n], + ) + rC[None, m, n] = mma_sync( + rK_upper, + rQ_f16[None, k, n, 1], + rC[None, m, n], + ) + else: + for n in cutlass.range_constexpr(Q_TILES): + rC[None, m, n] = mma_sync( + rK[None, m, k], + rQ[(None, k % 2), k // 2, n], + rC[None, m, n], + ) + + cute.arch.mbarrier_arrive(tma_empty_mbar + tma_stage) + + k_start = block_id * BLOCK_K + warp_id * 32 + + for q in cutlass.range_constexpr(Q_TILES): + for i in cutlass.range_constexpr(4): + for j in cutlass.range_constexpr(2): + col = q * 8 + (lane_id % 4) * 2 + j + q_local_pos = col % MAX_DQL + q_pos = q_start + q_local_pos + k_pos = k_start + i * 8 + lane_id // 4 + rC[q * 8 + i * 2 + j] = ( + rC[q * 8 + i * 2 + j] if q_pos >= k_pos else float("-inf") + ) + + for q in cutlass.range_constexpr(Q_TILES): + # Thread-local reduction along the BLOCK_K dim. + rScore = cute.make_rmem_tensor(2, Float32) + rScore.fill(float("-inf")) + for i in cutlass.range_constexpr(4): + rScore[0] = cute.arch.fmax(rScore[0], rC[i * 2 + 0 + q * 8]) + rScore[1] = cute.arch.fmax(rScore[1], rC[i * 2 + 1 + q * 8]) + + # Warp reduction among lanes 0, 4, 8, 12, ... + for i in cutlass.range_constexpr(3): + offset = 4 << i + other0 = cute.arch.shuffle_sync_bfly( + rScore[0], offset=offset, mask=-1, mask_and_clamp=31 + ) + other1 = cute.arch.shuffle_sync_bfly( + rScore[1], offset=offset, mask=-1, mask_and_clamp=31 + ) + rScore[0] = cute.arch.fmax(rScore[0], other0) + rScore[1] = cute.arch.fmax(rScore[1], other1) + + if lane_id * 2 < MMA_N: + epi_buffer[q * MMA_N + lane_id * 2 + 0, warp_id] = rScore[0] + epi_buffer[q * MMA_N + lane_id * 2 + 1, warp_id] = rScore[1] + cute.arch.barrier(barrier_id=self.BAR_MMA, number_of_threads=128) + + head_id = lane_id // MAX_DQL + q_local_pos = lane_id - head_id * MAX_DQL + valid_q = head_id < NUM_HEADS and q_local_pos < decode_query_len + if lane_id < BLOCK_Q and valid_q: + final_score = epi_buffer[lane_id, 0] + for i in cutlass.range_constexpr(1, 4): + final_score = cute.arch.fmax(final_score, epi_buffer[lane_id, i]) + + t = batch_id * decode_query_len + q_local_pos + score[head_id, t, block_id] = final_score + + tma_stage = (tma_stage + 1) % self.num_stages + if tma_stage == 0: + tma_parity ^= 1 diff --git a/tests/integration/test_lists/test-db/l0_b200_m3.yml b/tests/integration/test_lists/test-db/l0_b200_m3.yml index a147e4e36411..1dddcb53df42 100644 --- a/tests/integration/test_lists/test-db/l0_b200_m3.yml +++ b/tests/integration/test_lists/test-db/l0_b200_m3.yml @@ -13,7 +13,10 @@ l0_b200_m3: stage: pre_merge backend: pytorch tests: + - unittest/_torch/attention/sparse/test_minimax_m3_dense_decode.py + - unittest/_torch/attention/sparse/test_minimax_m3_index_decode_score.py - unittest/_torch/attention/sparse/test_minimax_m3_msa_backend.py - unittest/_torch/attention/sparse/test_minimax_m3_msa_selector.py + - unittest/_torch/attention/sparse/test_minimax_m3_sparse_attn_decode.py - unittest/_torch/models/test_minimax_m3.py - unittest/_torch/models/checkpoints/hf/test_minimaxm3_weight_mapper.py diff --git a/tests/microbenchmarks/minimax_m3_index_decode_score.py b/tests/microbenchmarks/minimax_m3_index_decode_score.py new file mode 100644 index 000000000000..520c8f9b8ae4 --- /dev/null +++ b/tests/microbenchmarks/minimax_m3_index_decode_score.py @@ -0,0 +1,101 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Latency of the MiniMax-M3 indexer decode scorer, CuTe DSL vs the MSA proxy. + +Both produce the per-block index scores the selector ranks: the CuTe DSL +kernel on a resolved decode span, the fmha_sm100 output_maxscore pass on +everything else. Run this on one SM100 GPU to see the per-call cost of each at +a given batch size and context length; no model weights are needed. + + python tests/microbenchmarks/minimax_m3_index_decode_score.py \ + --batch 1 --seq-len 8192 --num-heads 1 +""" + +import argparse + +import torch + +import tensorrt_llm._torch.custom_ops # noqa: F401 +from tensorrt_llm._torch.attention_backend.sparse.minimax_m3.msa_indexer import _proxy_max_score +from tensorrt_llm._torch.attention_backend.sparse.minimax_m3.msa_utils import ( + build_kv_page_indices, + msa_package_available, +) + +PAGE_SIZE = 128 +HEAD_DIM = 128 + + +def _time_us(fn, warmup: int = 20, iters: int = 100) -> float: + for _ in range(warmup): + fn() + torch.cuda.synchronize() + start, end = torch.cuda.Event(True), torch.cuda.Event(True) + start.record() + for _ in range(iters): + fn() + end.record() + torch.cuda.synchronize() + return start.elapsed_time(end) * 1000.0 / iters + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--batch", type=int, default=1) + parser.add_argument("--seq-len", type=int, default=8192) + parser.add_argument("--num-heads", type=int, default=1, help="Sharded index heads.") + parser.add_argument("--decode-query-len", type=int, default=1) + parser.add_argument("--dtype", choices=("bfloat16", "fp8_e4m3"), default="fp8_e4m3") + args = parser.parse_args() + + dtype = torch.bfloat16 if args.dtype == "bfloat16" else torch.float8_e4m3fn + batch, seq_len, num_heads = args.batch, args.seq_len, args.num_heads + dql = args.decode_query_len + total_q = batch * dql + num_blocks = (seq_len + PAGE_SIZE - 1) // PAGE_SIZE + score_width = ((num_blocks + 15) // 16) * 16 + num_pages = batch * num_blocks + + block_table = ( + torch.randperm(num_pages, device="cuda").to(torch.int32).reshape(batch, num_blocks) + ) + seq_lens = torch.full((batch,), seq_len, device="cuda", dtype=torch.int32) + idx_q = torch.randn(total_q, num_heads, HEAD_DIM, device="cuda").to(dtype) + k_cache = torch.randn(num_pages, PAGE_SIZE, HEAD_DIM, device="cuda").to(dtype) + backing = torch.full((num_heads, score_width, total_q), -float("inf"), device="cuda") + score = backing.transpose(1, 2) + + def run_cutedsl(): + torch.ops.trtllm.cute_dsl_minimax_m3_index_decode_score( + idx_q, k_cache, block_table, seq_lens, score, dql + ) + + print(f"batch={batch} seq_len={seq_len} heads={num_heads} dql={dql} dtype={args.dtype}") + print(f" cutedsl : {_time_us(run_cutedsl):8.2f} us/call") + + if not msa_package_available(): + print(" msa : skipped (fmha_sm100 submodule not available)") + return + + qo_lens_cpu = torch.full((batch,), dql, dtype=torch.int32) + kv_lens_cpu = torch.full((batch,), seq_len, dtype=torch.int32) + kv_indices = build_kv_page_indices(block_table.cpu(), kv_lens_cpu, PAGE_SIZE).cuda() + k_paged = k_cache.unsqueeze(1) + + def run_msa(): + _proxy_max_score( + idx_q, + k_paged, + qo_lens_cpu=qo_lens_cpu, + kv_lens_cpu=kv_lens_cpu, + qo_offset_cpu=kv_lens_cpu - qo_lens_cpu, + kv_indices=kv_indices, + sm_scale=HEAD_DIM**-0.5, + causal=True, + ) + + print(f" msa : {_time_us(run_msa):8.2f} us/call") + + +if __name__ == "__main__": + main() diff --git a/tests/unittest/_torch/attention/sparse/test_minimax_m3_dense_decode.py b/tests/unittest/_torch/attention/sparse/test_minimax_m3_dense_decode.py new file mode 100644 index 000000000000..98959e6e2187 --- /dev/null +++ b/tests/unittest/_torch/attention/sparse/test_minimax_m3_dense_decode.py @@ -0,0 +1,413 @@ +# 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. +"""trtllm-gen decode for MiniMax-M3's dense attention layers. + +Two things are under test and they fail differently. The pool geometry +(get_kv_subpage_pool) is pure addressing against a real cache manager: if the +flat sub-page view disagrees with get_buffers, the kernel silently reads +another layer's cache. The kernel call itself is checked against a PyTorch +oracle through a stub manager, so the flashinfer argument conventions are +exercised without standing up a model. +""" + +from __future__ import annotations + +from typing import List + +import pytest +import torch + +from tensorrt_llm._torch.attention_backend.sparse.minimax_m3.trtllm_gen_dense_decode import ( + dense_decode_unsupported_reason, + minimax_m3_trtllm_gen_dense_decode, + subpage_block_table, + uniform_subpages_per_slot, + write_subpage_block_table, +) + +pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA") + +PAGE_SIZE = 32 +HEAD_DIM = 128 +# The bmm1 scale, spelled as run_msa_paged_gqa spells it at q_scaling 1. +SM_SCALE = HEAD_DIM**-0.5 + + +def _is_sm100f() -> bool: + major, minor = torch.cuda.get_device_capability() + return major == 10 and minor in (0, 3) + + +# -------------------------------------------------------------------------- +# Pool geometry against a real MiniMaxM3KVCacheManagerV2 +# -------------------------------------------------------------------------- + + +def _create_manager(tp_size: int, sparse_layers: List[int], num_layers: int = 4): + from tensorrt_llm import Mapping + from tensorrt_llm._torch.attention_backend.sparse.minimax_m3 import MiniMaxM3KVCacheManagerV2 + from tensorrt_llm.bindings import DataType + from tensorrt_llm.bindings.internal.batch_manager import CacheType as CacheTypeCpp + from tensorrt_llm.llmapi.llm_args import KvCacheConfig + + max_num_tokens = 2048 + return MiniMaxM3KVCacheManagerV2( + kv_cache_config=KvCacheConfig( + enable_block_reuse=False, + max_tokens=max_num_tokens, + event_buffer_max_size=0, + dtype="auto", + ), + kv_cache_type=CacheTypeCpp.SELF, + num_layers=num_layers, + num_kv_heads=2, + head_dim=HEAD_DIM, + tokens_per_block=PAGE_SIZE, + max_seq_len=512, + max_batch_size=4, + mapping=Mapping(world_size=tp_size, rank=0, tp_size=tp_size, pp_size=1), + dtype=DataType.BF16, + vocab_size=1024, + max_num_tokens=max_num_tokens, + sparse_layer_ids=list(sparse_layers), + disable_index_value_layer_ids=list(sparse_layers), + sparse_index_dim=HEAD_DIM, + ) + + +@pytest.mark.parametrize( + "tp_size,sparse_layers", + [ + # TP=2 makes K == V == INDEX_KEY bytes per block, so V2 coalesces + # index-K into the K/V pool and the per-layer stride goes non-uniform. + # That is the layout build_trtllm_gen_kv_cache_metadata cannot express. + pytest.param(2, [1, 3], id="coalesced-index-k"), + pytest.param(1, [1, 3], id="separate-index-pool"), + pytest.param(2, [], id="all-dense"), + ], +) +@pytest.mark.parametrize("kv_layout", ["HND", "NHD"]) +def test_subpage_pool_addresses_match_get_buffers(tp_size, sparse_layers, kv_layout): + """flat[s * scale + {0,1}] must be exactly this layer's K and V at slot s.""" + manager = _create_manager(tp_size, sparse_layers) + try: + for layer_idx in range(4): + kv = manager.get_buffers(layer_idx, kv_layout=kv_layout) + flat, scale = manager.get_kv_subpage_pool(layer_idx, kv_layout) + + assert flat.is_contiguous() + assert list(flat.shape[1:]) == list(kv.shape[2:]) + assert flat.dtype == kv.dtype + num_slots = int(kv.shape[0]) + assert int(flat.shape[0]) == (num_slots - 1) * scale + 2 + + for slot in (0, 1, num_slots // 2, num_slots - 1): + assert flat[slot * scale].data_ptr() == kv[slot, 0].data_ptr() + assert flat[slot * scale + 1].data_ptr() == kv[slot, 1].data_ptr() + finally: + manager.shutdown() + + +def test_subpage_pool_stops_at_the_last_slots_v(): + """The tail bound matters: every layer but the first starts mid-slot, so a + naive num_slots * scale view would run off the end of the pool.""" + manager = _create_manager(2, [1, 3]) + try: + for layer_idx in range(4): + flat, scale = manager.get_kv_subpage_pool(layer_idx, "HND") + kv = manager.get_buffers(layer_idx, kv_layout="HND") + last_v = kv[int(kv.shape[0]) - 1, 1] + flat_end = flat.data_ptr() + flat.numel() * flat.element_size() + v_end = last_v.data_ptr() + last_v.numel() * last_v.element_size() + assert flat_end == v_end + finally: + manager.shutdown() + + +# -------------------------------------------------------------------------- +# Block-table expansion +# -------------------------------------------------------------------------- + + +def test_subpage_block_table_splits_k_and_v_rows(): + slots = torch.tensor([[0, 3, 7], [2, 5, 11]], device="cuda", dtype=torch.int32) + table = subpage_block_table(slots, subpages_per_slot=9) + + assert table.shape == (2, 2, 3) + assert table.dtype == torch.int32 + assert table[:, 0].tolist() == [[0, 27, 63], [18, 45, 99]] + assert table[:, 1].tolist() == [[1, 28, 64], [19, 46, 100]] + + +def test_subpage_block_table_reuses_one_buffer(): + """All dense layers share the arena block, so a later call must not alias + a live earlier one within a step; they are written before every use.""" + slots_a = torch.zeros((2, 4), device="cuda", dtype=torch.int32) + slots_b = torch.full((2, 4), 5, device="cuda", dtype=torch.int32) + first = subpage_block_table(slots_a, 4) + second = subpage_block_table(slots_b, 4) + + assert first.data_ptr() == second.data_ptr() + assert second[:, 0].tolist() == [[20] * 4] * 2 + + +def test_write_subpage_block_table_fills_a_caller_owned_buffer(): + """prepare() stages the expansion into its own graph-stable buffer, so the + in-place writer has to agree with the arena-backed helper.""" + slots = torch.tensor([[0, 3, 7], [2, 5, 11]], device="cuda", dtype=torch.int32) + out = torch.empty((2, 2, 3), device="cuda", dtype=torch.int32) + + write_subpage_block_table(slots, 9, out) + + assert torch.equal(out, subpage_block_table(slots, 9)) + + +@pytest.mark.parametrize("sparse_layers", [[1, 3], []], ids=["mixed", "all-dense"]) +def test_uniform_subpages_per_slot_matches_every_layer(sparse_layers): + """prepare() expands the block table without naming a layer, so the factor + it uses has to be one every layer of the real pool agrees on.""" + manager = _create_manager(2, sparse_layers) + try: + per_layer = {manager.get_kv_subpage_pool(i, "HND")[1] for i in range(4)} + factor = uniform_subpages_per_slot(manager) + assert factor == (per_layer.pop() if len(per_layer) == 1 else 0) + finally: + manager.shutdown() + + +def test_uniform_subpages_per_slot_reports_zero_without_a_pool(): + """A manager with no sub-page pool leaves each dense layer to expand its + own table, rather than being staged against a guessed factor.""" + + class _NoPool: + layer_offsets = {0: 0} + + assert uniform_subpages_per_slot(_NoPool()) == 0 + + +# -------------------------------------------------------------------------- +# Kernel parity against a PyTorch oracle +# -------------------------------------------------------------------------- + + +class _StubManager: + """Presents a flat sub-page pool the way MiniMaxM3KVCacheManagerV2 does.""" + + def __init__(self, pool: torch.Tensor, subpages_per_slot: int): + self._pool = pool + self._scale = subpages_per_slot + + def get_kv_subpage_pool(self, layer_idx: int, kv_layout: str = "HND"): + assert kv_layout == "HND" + return self._pool, self._scale + + +def _reference_dense_decode( + q: torch.Tensor, # [total_q, num_heads, head_dim] + pool: torch.Tensor, # [num_subpages, num_kv_heads, page, head_dim] + subpages_per_slot: int, + block_table: torch.Tensor, # [batch, max_blocks] + seq_lens: torch.Tensor, # [batch] + decode_query_len: int, + sm_scale: float, +) -> torch.Tensor: + pool_f32 = pool.float() + num_heads = q.shape[1] + num_kv_heads = pool.shape[1] + group = num_heads // num_kv_heads + out = torch.zeros_like(q, dtype=torch.float32) + positions = torch.arange(PAGE_SIZE, device=q.device) + + for req in range(block_table.shape[0]): + kv_len_full = int(seq_lens[req]) + num_pages = (kv_len_full + PAGE_SIZE - 1) // PAGE_SIZE + slots = [int(block_table[req, p]) for p in range(num_pages)] + for intra in range(decode_query_len): + token = req * decode_query_len + intra + kv_len = kv_len_full - decode_query_len + intra + 1 + for kv_head in range(num_kv_heads): + keys = torch.cat([pool_f32[s * subpages_per_slot, kv_head] for s in slots]) + values = torch.cat([pool_f32[s * subpages_per_slot + 1, kv_head] for s in slots]) + valid = torch.cat([p * PAGE_SIZE + positions < kv_len for p in range(num_pages)]) + q_rows = q[token, kv_head * group : (kv_head + 1) * group].float() + logits = (q_rows @ keys.T) * sm_scale + logits = logits.masked_fill(~valid[None, :], -float("inf")) + probs = torch.softmax(logits, dim=-1) + out[token, kv_head * group : (kv_head + 1) * group] = probs @ values + return out + + +def _make_pool_inputs( + seq_lens: List[int], + num_heads: int, + num_kv_heads: int, + decode_query_len: int, + subpages_per_slot: int, + dtype: torch.dtype, + seed: int = 0, +): + generator = torch.Generator(device="cuda").manual_seed(seed) + batch = len(seq_lens) + max_blocks = max((s + PAGE_SIZE - 1) // PAGE_SIZE for s in seq_lens) + num_slots = batch * max_blocks + 3 + + pool = torch.randn( + (num_slots * subpages_per_slot, num_kv_heads, PAGE_SIZE, HEAD_DIM), + device="cuda", + generator=generator, + dtype=torch.float32, + ).to(dtype) + # Slots are handed out non-contiguously, exactly as the block manager does. + perm = torch.randperm(num_slots, generator=generator, device="cuda")[: batch * max_blocks] + block_table = perm.reshape(batch, max_blocks).to(torch.int32) + + total_q = batch * decode_query_len + q = torch.randn( + (total_q, num_heads, HEAD_DIM), device="cuda", generator=generator, dtype=torch.float32 + ).to(torch.bfloat16) + return q, pool, block_table, torch.tensor(seq_lens, device="cuda", dtype=torch.int32) + + +def _run_dense(q, pool, subpages_per_slot, block_table, seq_lens, decode_query_len, **kwargs): + out = torch.zeros_like(q, dtype=torch.bfloat16) + minimax_m3_trtllm_gen_dense_decode( + q, + _StubManager(pool, subpages_per_slot), + 0, + block_table, + seq_lens, + sm_scale=SM_SCALE, + output=out, + decode_query_len=decode_query_len, + max_seq_len=int(seq_lens.max()), + max_num_requests=int(seq_lens.shape[0]), + **kwargs, + ) + return out + + +@pytest.mark.skipif(not _is_sm100f(), reason="trtllm-gen decode kernels are SM100/SM103 only") +@pytest.mark.parametrize("kv_dtype", [torch.bfloat16, torch.float8_e4m3fn]) +@pytest.mark.parametrize("num_heads,num_kv_heads", [(8, 1), (16, 2)], ids=["gqa8x1", "gqa16x2"]) +@pytest.mark.parametrize("decode_query_len", [1, 2]) +@pytest.mark.parametrize("subpages_per_slot", [2, 9]) +def test_matches_reference(kv_dtype, num_heads, num_kv_heads, decode_query_len, subpages_per_slot): + """subpages_per_slot=9 is the M3 coalesced case; 2 is the uniform one.""" + seq_lens = [PAGE_SIZE * 3, PAGE_SIZE + 5, PAGE_SIZE * 2 - 1] + q, pool, block_table, seq_lens_t = _make_pool_inputs( + seq_lens, num_heads, num_kv_heads, decode_query_len, subpages_per_slot, kv_dtype + ) + out = _run_dense(q, pool, subpages_per_slot, block_table, seq_lens_t, decode_query_len) + # The kernel runs Q in the KV dtype, so the oracle gets the same rounding. + q_ref = q.to(kv_dtype).to(torch.float32) if kv_dtype == torch.float8_e4m3fn else q + ref = _reference_dense_decode( + q_ref, pool, subpages_per_slot, block_table, seq_lens_t, decode_query_len, SM_SCALE + ) + tol = 6e-2 if kv_dtype == torch.float8_e4m3fn else 2e-2 + torch.testing.assert_close(out.float(), ref, rtol=tol, atol=tol) + + +@pytest.mark.skipif(not _is_sm100f(), reason="trtllm-gen decode kernels are SM100/SM103 only") +def test_staged_subpage_table_is_used_only_when_the_factor_matches(): + """prepare() stages one expansion for every dense layer, so the kernel must + take it when the factor agrees and expand its own when it does not; either + way the answer is the same.""" + subpages_per_slot = 9 + seq_lens = [PAGE_SIZE * 3, PAGE_SIZE + 5] + q, pool, block_table, seq_lens_t = _make_pool_inputs( + seq_lens, 8, 1, 1, subpages_per_slot, torch.bfloat16, seed=11 + ) + staged = torch.empty( + (block_table.shape[0], 2, block_table.shape[1]), device="cuda", dtype=torch.int32 + ) + write_subpage_block_table(block_table, subpages_per_slot, staged) + + args = (q, pool, subpages_per_slot, block_table, seq_lens_t, 1) + expanded_here = _run_dense(*args) + from_staged = _run_dense(*args, staged_subpage_table=staged, staged_subpages_per_slot=9) + # A stale factor must be ignored rather than trusted: this table addresses + # the wrong sub-pages entirely. + wrong = torch.zeros_like(staged) + ignored = _run_dense(*args, staged_subpage_table=wrong, staged_subpages_per_slot=2) + + torch.testing.assert_close(from_staged.float(), expanded_here.float(), rtol=1e-3, atol=1e-3) + torch.testing.assert_close(ignored.float(), expanded_here.float(), rtol=1e-3, atol=1e-3) + + +@pytest.mark.skipif(not _is_sm100f(), reason="trtllm-gen decode kernels are SM100/SM103 only") +def test_cuda_graph_replay_tracks_inputs(): + seq_lens = [PAGE_SIZE * 2, PAGE_SIZE * 3] + q, pool, block_table, seq_lens_t = _make_pool_inputs( + seq_lens, 8, 1, 1, 9, torch.bfloat16, seed=7 + ) + stub = _StubManager(pool, 9) + out = torch.zeros_like(q, dtype=torch.bfloat16) + + def run(): + minimax_m3_trtllm_gen_dense_decode( + q, + stub, + 0, + block_table, + seq_lens_t, + sm_scale=SM_SCALE, + output=out, + decode_query_len=1, + max_seq_len=PAGE_SIZE * 3, + max_num_requests=2, + ) + + run() # warm the arena and the counter buffer before capture + torch.cuda.synchronize() + graph = torch.cuda.CUDAGraph() + stream = torch.cuda.Stream() + stream.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(stream): + with torch.cuda.graph(graph): + run() + torch.cuda.current_stream().wait_stream(stream) + + q.copy_(torch.randn_like(q, dtype=torch.float32).to(torch.bfloat16)) + graph.replay() + torch.cuda.synchronize() + replayed = out.clone() + + out.zero_() + run() + torch.cuda.synchronize() + torch.testing.assert_close(replayed.float(), out.float(), rtol=2e-2, atol=2e-2) + + +# -------------------------------------------------------------------------- +# Gating +# -------------------------------------------------------------------------- + + +def test_declines_unsupported_geometry(): + """The reason is a string because it is logged; what matters here is that + each unsupported geometry is caught before the kernel is reached.""" + pool = torch.zeros((1, 1, PAGE_SIZE, 64), device="cuda", dtype=torch.bfloat16) + assert "head_dim 64" in dense_decode_unsupported_reason(_StubManager(pool, 2), 64) + + class _NoPool: + pass + + assert "flat sub-page pool" in dense_decode_unsupported_reason(_NoPool(), HEAD_DIM) + + +def test_accepts_the_m3_geometry(): + pytest.importorskip("flashinfer") + pool = torch.zeros((1, 1, PAGE_SIZE, HEAD_DIM), device="cuda", dtype=torch.bfloat16) + assert dense_decode_unsupported_reason(_StubManager(pool, 2), HEAD_DIM) is None diff --git a/tests/unittest/_torch/attention/sparse/test_minimax_m3_index_decode_score.py b/tests/unittest/_torch/attention/sparse/test_minimax_m3_index_decode_score.py new file mode 100644 index 000000000000..343d562ffb2f --- /dev/null +++ b/tests/unittest/_torch/attention/sparse/test_minimax_m3_index_decode_score.py @@ -0,0 +1,472 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-License-Identifier: Apache-2.0 +# PyTorch oracle vendored from vLLM (Apache-2.0), _reference_decode_index_score: +# https://github.com/vllm-project/vllm/blob/6f91edf96d3f3272945809c04702380053bff4de/tests/kernels/attention/test_minimax_m3.py#L188 +"""Correctness tests for the CuTe DSL MiniMax-M3 indexer decode scorer. + +The PyTorch oracle is ported from the vLLM reference linked in the file header +(v0.26.1rc0-77-g6f91edf96). +""" + +import pytest +import torch + +from tensorrt_llm._torch.attention_backend.sparse.minimax_m3.msa_indexer import _cutedsl_score +from tensorrt_llm._torch.attention_backend.sparse.minimax_m3.msa_utils import ( + MSA_REQUIRED_TOPK, + build_kv_page_indices, + msa_package_available, + select_blocks_from_maxscore, +) +from tensorrt_llm._utils import get_sm_version + +PAGE_SIZE = 128 +HEAD_DIM = 128 + +skip_not_sm100 = pytest.mark.skipif( + get_sm_version() not in (100, 103), + reason="CuTe DSL MiniMax-M3 index decode scoring requires SM100/SM103.", +) + +pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") + + +def _runner(): + """Return the CuTe DSL runner class, skipping the test if unavailable.""" + pytest.importorskip("cutlass") + from tensorrt_llm._torch.custom_ops import cute_dsl_custom_ops + + runner = getattr(cute_dsl_custom_ops, "CuteDSLMiniMaxM3IndexDecodeScoreRunner", None) + if runner is None: + pytest.skip("CuTe DSL custom ops are not registered in this build.") + return runner + + +def _reference_decode_index_score( + idx_q: torch.Tensor, + index_kv_cache: torch.Tensor, + block_table: torch.Tensor, + seq_lens: torch.Tensor, + decode_query_len: int, + score_block_stride: int, +) -> torch.Tensor: + """Per-block causal max of Q.K, in the kernel's [head, token, block] layout.""" + total_q, num_idx_heads, _ = idx_q.shape + out = torch.full( + (num_idx_heads, total_q, score_block_stride), + -float("inf"), + device=idx_q.device, + dtype=torch.float32, + ) + for req_id, seq_len in enumerate(seq_lens.tolist()): + num_blocks = (seq_len + PAGE_SIZE - 1) // PAGE_SIZE + if num_blocks == 0: + continue + token_start = req_id * decode_query_len + q = idx_q[token_start : token_start + decode_query_len].float() + pages = block_table[req_id, :num_blocks] + k = index_kv_cache[pages].reshape(num_blocks * PAGE_SIZE, -1).float() + score = torch.einsum("qhd,kd->hqk", q, k) + q_pos = seq_len - decode_query_len + torch.arange(decode_query_len, device=idx_q.device) + k_pos = torch.arange(k.shape[0], device=idx_q.device) + score.masked_fill_(k_pos[None, :] > q_pos[:, None], -float("inf")) + out[:, token_start : token_start + decode_query_len, :num_blocks] = ( + score.reshape(num_idx_heads, decode_query_len, num_blocks, PAGE_SIZE).max(dim=3).values + ) + return out + + +def _make_inputs(seq_lens, *, dtype, num_heads, decode_query_len, seed=0): + """Build (idx_q, index_k_cache, block_table, seq_lens_dev, score, expected).""" + generator = torch.Generator(device="cuda").manual_seed(seed) + seq_lens_dev = torch.tensor(seq_lens, device="cuda", dtype=torch.int32) + batch = len(seq_lens) + total_q = batch * decode_query_len + max_blocks = (max(seq_lens) + PAGE_SIZE - 1) // PAGE_SIZE + # Mirror the plan's max_k_tiles alignment so the tests exercise a score + # buffer wider than any single request's block count. + score_block_stride = ((max_blocks + 15) // 16) * 16 + num_pages = batch * max_blocks + + # Shuffled pages: a kernel that ignored the block table and walked pages + # linearly would still pass with an identity mapping. + block_table = torch.randperm(num_pages, device="cuda", generator=generator).to(torch.int32) + block_table = block_table.reshape(batch, max_blocks) + + idx_q = torch.randn(total_q, num_heads, HEAD_DIM, device="cuda", generator=generator).to(dtype) + index_k_cache = torch.randn( + num_pages, PAGE_SIZE, HEAD_DIM, device="cuda", generator=generator + ).to(dtype) + + # Production hands the kernel a transposed view of the selector's + # [heads, blocks, tokens] buffer, so build it the same way here. + backing = torch.full( + (num_heads, score_block_stride, total_q), + -float("inf"), + device="cuda", + dtype=torch.float32, + ) + score = backing.transpose(1, 2) + + expected = _reference_decode_index_score( + idx_q, index_k_cache, block_table, seq_lens_dev, decode_query_len, score_block_stride + ) + return idx_q, index_k_cache, block_table, seq_lens_dev, backing, score, expected + + +def _assert_scores_close(actual, expected): + """Compare fp32 scores. + + Both sides accumulate 128 exactly-representable products in fp32, but the + kernel's k-split order differs from the reference matmul's, so a few ULP of + drift is expected. + """ + torch.testing.assert_close(actual, expected, rtol=1e-5, atol=1e-4) + + +@skip_not_sm100 +@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float8_e4m3fn]) +@pytest.mark.parametrize( + ("num_heads", "decode_query_len"), + [(1, 1), (1, 3), (4, 1), (4, 8)], +) +@pytest.mark.parametrize( + "seq_lens", + [ + # Exactly one block, and one block plus a partial. + [128], + [1, 128, 129], + # Non-multiples of the page size, several requests. + [1025, 4097], + [300, 1500, 2049, 4096, 5000, 7777, 8192, 9001], + ], + ids=["one-block", "short-mixed", "two-req", "batch8"], +) +def test_index_decode_score_matches_reference(dtype, num_heads, decode_query_len, seq_lens): + if num_heads * decode_query_len > 32: + pytest.skip("BLOCK_Q must not exceed 32.") + # A request cannot have fewer KV positions than it has query tokens. + seq_lens = [max(s, decode_query_len) for s in seq_lens] + idx_q, k_cache, block_table, seq_lens_dev, _, score, expected = _make_inputs( + seq_lens, dtype=dtype, num_heads=num_heads, decode_query_len=decode_query_len + ) + + torch.ops.trtllm.cute_dsl_minimax_m3_index_decode_score( + idx_q, k_cache, block_table, seq_lens_dev, score, decode_query_len + ) + _assert_scores_close(score, expected) + + +@skip_not_sm100 +@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float8_e4m3fn]) +def test_index_decode_score_multi_round_split_k(dtype): + """A request longer than split_k pages makes the per-CTA block loop iterate. + + split_k is 256, so 33000 tokens is 258 blocks and CTAs 0 and 1 each handle + two blocks while the rest handle one. + """ + seq_lens = [33000, 32768 + 1] + idx_q, k_cache, block_table, seq_lens_dev, _, score, expected = _make_inputs( + seq_lens, dtype=dtype, num_heads=1, decode_query_len=1, seed=3 + ) + assert (max(seq_lens) + PAGE_SIZE - 1) // PAGE_SIZE > 256 + + torch.ops.trtllm.cute_dsl_minimax_m3_index_decode_score( + idx_q, k_cache, block_table, seq_lens_dev, score, 1 + ) + _assert_scores_close(score, expected) + + +@skip_not_sm100 +def test_index_decode_score_transposed_view_matches_contiguous(): + """The transposed selector view must produce the same values as a direct write. + + This is the zero-copy trick the indexer relies on: the kernel writes + [head, token, block] into a buffer that is contiguous as + [head, block, token], so the selector reads it without a transpose. + """ + seq_lens = [1025, 4097, 300] + idx_q, k_cache, block_table, seq_lens_dev, backing, transposed, expected = _make_inputs( + seq_lens, dtype=torch.bfloat16, num_heads=4, decode_query_len=1, seed=11 + ) + contiguous = torch.full_like(expected, -float("inf")) + + torch.ops.trtllm.cute_dsl_minimax_m3_index_decode_score( + idx_q, k_cache, block_table, seq_lens_dev, transposed, 1 + ) + torch.ops.trtllm.cute_dsl_minimax_m3_index_decode_score( + idx_q, k_cache, block_table, seq_lens_dev, contiguous, 1 + ) + + assert backing.is_contiguous() + assert not transposed.is_contiguous() + assert torch.equal(transposed, contiguous) + _assert_scores_close(contiguous, expected) + + +@skip_not_sm100 +@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float8_e4m3fn]) +def test_index_decode_score_feeds_selector(dtype): + """End to end: the produced scores must select the same blocks as the oracle.""" + seq_lens = [129, 1025, 4097, 8192] + idx_q, k_cache, block_table, seq_lens_dev, backing, score, expected = _make_inputs( + seq_lens, dtype=dtype, num_heads=1, decode_query_len=1, seed=5 + ) + n_valid = torch.tensor( + [(s + PAGE_SIZE - 1) // PAGE_SIZE for s in seq_lens], device="cuda", dtype=torch.int32 + ) + + torch.ops.trtllm.cute_dsl_minimax_m3_index_decode_score( + idx_q, k_cache, block_table, seq_lens_dev, score, 1 + ) + + actual_topk = select_blocks_from_maxscore( + backing, topk=MSA_REQUIRED_TOPK, n_valid_blocks=n_valid, init_blocks=0, local_blocks=1 + ) + expected_topk = select_blocks_from_maxscore( + expected.transpose(1, 2).contiguous(), + topk=MSA_REQUIRED_TOPK, + n_valid_blocks=n_valid, + init_blocks=0, + local_blocks=1, + ) + assert torch.equal(actual_topk, expected_topk) + + +@skip_not_sm100 +@pytest.mark.skipif(not msa_package_available(), reason="fmha_sm100 (MSA submodule) required") +@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float8_e4m3fn]) +def test_index_decode_score_matches_msa_proxy(dtype): + """A/B against the fmha_sm100 proxy pass the CuTe DSL scorer replaces. + + Both sides report the per-block max of raw Q.K: the proxy's max_score is + read off the MMA accumulator before the softmax scale, which in + output_maxscore mode is never applied, so the values compare directly. + """ + from tensorrt_llm._torch.attention_backend.sparse.minimax_m3.msa_indexer import _proxy_max_score + + seq_lens = [1025, 4097, 300] + idx_q, k_cache, block_table, seq_lens_dev, backing, score, _ = _make_inputs( + seq_lens, dtype=dtype, num_heads=1, decode_query_len=1, seed=17 + ) + sm_scale = HEAD_DIM**-0.5 + batch = len(seq_lens) + n_valid_list = [(s + PAGE_SIZE - 1) // PAGE_SIZE for s in seq_lens] + n_valid = torch.tensor(n_valid_list, device="cuda", dtype=torch.int32) + + torch.ops.trtllm.cute_dsl_minimax_m3_index_decode_score( + idx_q, k_cache, block_table, seq_lens_dev, score, 1 + ) + + qo_lens_cpu = torch.ones(batch, dtype=torch.int32) + kv_lens_cpu = torch.tensor(seq_lens, dtype=torch.int32) + kv_indices = build_kv_page_indices(block_table.cpu(), kv_lens_cpu, PAGE_SIZE).cuda() + msa_score = _proxy_max_score( + idx_q, + k_cache.unsqueeze(1), + qo_lens_cpu=qo_lens_cpu, + kv_lens_cpu=kv_lens_cpu, + qo_offset_cpu=kv_lens_cpu - qo_lens_cpu, + kv_indices=kv_indices, + sm_scale=sm_scale, + causal=True, + ) + + # The proxy plan sizes its block dim independently of the score buffer, and + # entries past a token's valid count are undefined on both sides. + width = min(backing.shape[1], msa_score.shape[1]) + for token, num_valid in enumerate(n_valid_list): + assert num_valid <= width + torch.testing.assert_close( + backing[:, :num_valid, token], + msa_score[:, :num_valid, token], + rtol=2e-2, + atol=2e-2, + ) + + kwargs = dict(topk=MSA_REQUIRED_TOPK, n_valid_blocks=n_valid, init_blocks=0, local_blocks=1) + assert torch.equal( + select_blocks_from_maxscore(backing[:, :width].contiguous(), **kwargs), + select_blocks_from_maxscore(msa_score[:, :width].contiguous(), **kwargs), + ) + + +@skip_not_sm100 +@pytest.mark.skipif(not msa_package_available(), reason="fmha_sm100 (MSA submodule) required") +@pytest.mark.parametrize("head_major_output", [False, True]) +def test_mixed_batch_split_selects_the_same_blocks_as_the_whole_batch_proxy(head_major_output): + """A mixed batch must select the same blocks however it was scored. + + The generation rows move onto the CuTe DSL scorer while the context row + stays on the fmha_sm100 proxy, now planned over its row alone, so the two + halves are scored into separate buffers and their tables joined. The gate is + that the joined table is the one the whole-batch proxy would have produced. + """ + from tensorrt_llm._torch.attention_backend.sparse.minimax_m3.common import MiniMaxM3SparseConfig + from tensorrt_llm._torch.attention_backend.sparse.minimax_m3.msa_indexer import MsaIndexer + + _runner() + + # One context request prefilling a fresh 300-token prompt, then three decode + # rows. Their KV lengths span 2 to 33 blocks, so top-k of 16 is a real choice + # for the longest of them. + qo_lens_cpu = torch.tensor([300, 1, 1, 1], dtype=torch.int32) + kv_lens_cpu = torch.tensor([300, 1025, 4097, 130], dtype=torch.int32) + qo_offset_cpu = kv_lens_cpu - qo_lens_cpu + ctx_rows, ctx_tokens = 1, 300 + batch, total_q = int(qo_lens_cpu.shape[0]), int(qo_lens_cpu.sum()) + + generator = torch.Generator(device="cuda").manual_seed(29) + max_blocks = int((kv_lens_cpu.max().item() + PAGE_SIZE - 1) // PAGE_SIZE) + num_pages = batch * max_blocks + block_table = ( + torch.randperm(num_pages, device="cuda", generator=generator) + .to(torch.int32) + .reshape(batch, max_blocks) + ) + # Four index heads over two KV heads, so the amax reduce to KV-head + # granularity runs on both halves and the two output layouts differ. + num_index_heads, num_kv_heads = 4, 2 + idx_q = torch.randn( + total_q, num_index_heads, HEAD_DIM, device="cuda", generator=generator + ).bfloat16() + idx_k_paged = torch.randn( + num_pages, PAGE_SIZE, HEAD_DIM, device="cuda", generator=generator + ).bfloat16()[:, None] + kv_indices = build_kv_page_indices(block_table.cpu(), kv_lens_cpu, PAGE_SIZE).cuda() + + indexer = MsaIndexer( + MiniMaxM3SparseConfig( + num_q_heads=8, + num_kv_heads=num_kv_heads, + head_dim=HEAD_DIM, + num_index_heads=num_index_heads, + sparse_index_dim=HEAD_DIM, + block_size=PAGE_SIZE, + topk=MSA_REQUIRED_TOPK, + ) + ) + common = dict( + idx_sm_scale=HEAD_DIM**-0.5, + kv_indices=kv_indices, + qo_lens_cpu=qo_lens_cpu, + kv_lens_cpu=kv_lens_cpu, + qo_offset_cpu=qo_offset_cpu, + head_major_output=head_major_output, + ) + + # No score buffer, so the scorer is not even attempted and the proxy plans + # the whole batch inline: the pre-split behaviour of every mixed step. + reference = indexer.select_blocks(idx_q, idx_k_paged, **common) + + # Mirrors the plan's max_k_tiles alignment, so the span's buffer is wider + # than any one request's block count, as it is in production. + max_score = torch.full( + (num_index_heads, ((max_blocks + 15) // 16) * 16, total_q - ctx_tokens), + -float("inf"), + device="cuda", + dtype=torch.float32, + ) + split = indexer.select_blocks( + idx_q, + idx_k_paged, + max_score=max_score, + block_table=block_table[ctx_rows:], + seq_lens_cuda=kv_lens_cpu[ctx_rows:].cuda(), + decode_query_len=1, + require_cutedsl=True, + gen_token_first=ctx_tokens, + ctx_rows=ctx_rows, + **common, + ) + + assert torch.equal(split, reference) + # The Triton sparse decode kernel reads the table head-major, so a joined + # table must permute to a contiguous view exactly as an unjoined one does. + assert split.permute(1, 0, 2).is_contiguous() is head_major_output + + +@skip_not_sm100 +def test_index_decode_score_cuda_graph_replay_tracks_inputs(): + seq_lens = [1025, 4097] + idx_q, k_cache, block_table, seq_lens_dev, backing, score, _ = _make_inputs( + seq_lens, dtype=torch.bfloat16, num_heads=1, decode_query_len=1, seed=23 + ) + + def run(): + torch.ops.trtllm.cute_dsl_minimax_m3_index_decode_score( + idx_q, k_cache, block_table, seq_lens_dev, score, 1 + ) + + # Warm up outside capture so the JIT compile never lands inside it. + for _ in range(3): + run() + torch.cuda.synchronize() + + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + run() + + # Shrink one request and shuffle its pages; the replay must track both. + new_seq_lens = [513, 4097] + seq_lens_dev.copy_(torch.tensor(new_seq_lens, device="cuda", dtype=torch.int32)) + block_table.copy_(block_table.flip(1)) + backing.fill_(-float("inf")) + graph.replay() + torch.cuda.synchronize() + + expected = _reference_decode_index_score( + idx_q, + k_cache, + block_table, + torch.tensor(new_seq_lens, device="cuda", dtype=torch.int32), + 1, + backing.shape[1], + ) + _assert_scores_close(score, expected) + + +@skip_not_sm100 +@pytest.mark.parametrize( + ("kwargs", "reason"), + [ + ({"num_heads": 8, "max_decode_query_len": 8}, "BLOCK_Q above 32"), + ({"page_size": 64}, "unsupported page size"), + ({"head_dim": 64}, "unsupported head dim"), + ({"q_dtype": torch.float16}, "unsupported dtype"), + ], +) +def test_index_decode_score_declines_unsupported_geometry(kwargs, reason): + runner = _runner() + supported = dict( + q_dtype=torch.bfloat16, + num_heads=1, + head_dim=HEAD_DIM, + page_size=PAGE_SIZE, + max_decode_query_len=1, + ) + assert runner.is_supported(**supported), "baseline geometry must be supported" + assert not runner.is_supported(**{**supported, **kwargs}), reason + + +@skip_not_sm100 +def test_cutedsl_score_helper_falls_back_on_unsupported_geometry(): + """_cutedsl_score must report False rather than raise, so the caller can + fall back to the fmha_sm100 proxy.""" + _runner() + batch, num_heads, dql = 2, 1, 1 + idx_q = torch.randn(batch * dql, num_heads, 64, device="cuda", dtype=torch.bfloat16) + k_paged = torch.randn(4, 1, PAGE_SIZE, 64, device="cuda", dtype=torch.bfloat16) + block_table = torch.zeros(batch, 2, device="cuda", dtype=torch.int32) + seq_lens = torch.full((batch,), 128, device="cuda", dtype=torch.int32) + max_score = torch.zeros(num_heads, 2, batch * dql, device="cuda").transpose(1, 2) + + assert not _cutedsl_score( + idx_q, + k_paged, + max_score, + block_table=block_table, + seq_lens_cuda=seq_lens, + decode_query_len=dql, + ) diff --git a/tests/unittest/_torch/attention/sparse/test_minimax_m3_msa_backend.py b/tests/unittest/_torch/attention/sparse/test_minimax_m3_msa_backend.py index aa1ebf4f718a..34ad72226453 100644 --- a/tests/unittest/_torch/attention/sparse/test_minimax_m3_msa_backend.py +++ b/tests/unittest/_torch/attention/sparse/test_minimax_m3_msa_backend.py @@ -19,6 +19,9 @@ from tensorrt_llm._torch.attention_backend.sparse.minimax_m3.msa_scatter import ( fused_write_layer_caches, ) +from tensorrt_llm._torch.attention_backend.sparse.minimax_m3.msa_utils import ( + msa_ported_decode_active, +) from tensorrt_llm._torch.attention_backend.sparse.utils import _resolve_minimax_m3_backend_cls from tensorrt_llm.llmapi.llm_args import MiniMaxM3SparseAttentionConfig @@ -149,6 +152,116 @@ def test_msa_metadata_rejects_undersized_max_score_buffer(): ) +MAX_NUM_SEQUENCES = 8 +MAX_BLOCKS_PER_SEQ = 64 + + +class _RecordingBuffers: + """The graph buffer pool, recording what each buffer was reserved as.""" + + def __init__(self): + self.requested = {} + + def get_buffer(self, tensor_shape, dtype, cache_name, capture_graph): + self.requested[cache_name] = (tuple(tensor_shape), dtype, capture_graph) + return torch.zeros(tensor_shape, device="cuda", dtype=dtype) + + +def _buffer_metadata(**manager_fields): + """Metadata ready for _create_msa_buffers, under capture.""" + metadata_cls = MiniMaxM3MsaSparseAttention.Metadata + metadata = metadata_cls.__new__(metadata_cls) + metadata.kv_cache_manager = SimpleNamespace( + max_blocks_per_seq=MAX_BLOCKS_PER_SEQ, + tokens_per_block=128, + get_index_k_buffer=lambda layer_idx, kv_layout=None: None, + **manager_fields, + ) + metadata.is_cuda_graph = True + metadata.cuda_graph_buffers = _RecordingBuffers() + metadata.max_num_sequences = MAX_NUM_SEQUENCES + metadata.max_num_tokens = 512 + # No sparse params, so the fmha_sm100 proxy scratch is skipped and this + # exercises only the layer-invariant buffers. + metadata._msa_params = None + return metadata + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +def test_msa_buffers_include_graph_stable_block_table(): + """The 2-D page table and per-request length the ported decode kernels take + must come from the graph buffer pool at the manager's worst-case geometry, + so their addresses survive capture.""" + metadata = _buffer_metadata() + + metadata._create_msa_buffers() + + assert metadata._msa_buffers_ready + assert metadata.msa_block_table.shape == (MAX_NUM_SEQUENCES, MAX_BLOCKS_PER_SEQ) + assert metadata.msa_seq_lens_cuda.shape == (MAX_NUM_SEQUENCES,) + # Reserved from the graph pool at the worst-case geometry, alongside the + # flat page table the fmha_sm100 path uses. + requested = metadata.cuda_graph_buffers.requested + assert requested["msa_block_table"] == ( + (MAX_NUM_SEQUENCES, MAX_BLOCKS_PER_SEQ), + torch.int32, + True, + ) + assert requested["msa_seq_lens_cuda"] == ((MAX_NUM_SEQUENCES,), torch.int32, True) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +@pytest.mark.parametrize( + ("factors", "expected"), + [((9, 9, 9), 9), ((9, 4, 9), 0)], + ids=["uniform-pool", "groups-disagree"], +) +def test_msa_buffers_stage_the_subpage_table_only_for_a_uniform_pool(factors, expected): + """The sub-page expansion is hoisted out of the dense layers into prepare(), + which runs before any layer is named. That is sound only where every layer + of the pool packs the same number of sub-pages per slot; where they + disagree, no table is staged and each dense layer expands its own. + """ + metadata = _buffer_metadata( + layer_offsets=dict.fromkeys(range(len(factors)), 0), + get_kv_subpage_pool=lambda layer_idx, kv_layout="HND": (None, factors[layer_idx]), + ) + + metadata._create_msa_buffers() + + assert metadata._msa_subpages_per_slot == expected + if expected == 0: + assert metadata.msa_subpage_block_table is None + assert "msa_subpage_block_table" not in metadata.cuda_graph_buffers.requested + else: + # One K row and one V row per slot, at the same worst-case geometry as + # the slot table it expands. + assert metadata.cuda_graph_buffers.requested["msa_subpage_block_table"] == ( + (MAX_NUM_SEQUENCES, 2, MAX_BLOCKS_PER_SEQ), + torch.int32, + True, + ) + + +def test_msa_subpage_rows_slice_the_generation_span(): + """A mixed step hands the dense kernel only the span's rows, and its block + table has to be sliced the same way the slot table is. The factor travels + with it so the kernel can tell a stale staging from its own geometry.""" + metadata_cls = MiniMaxM3MsaSparseAttention.Metadata + metadata = metadata_cls.__new__(metadata_cls) + metadata.msa_subpage_block_table = torch.arange(4 * 2 * 3, dtype=torch.int32).reshape(4, 2, 3) + metadata._msa_subpages_per_slot = 9 + + table, factor = metadata.msa_subpage_rows(2, 4) + assert factor == 9 + assert torch.equal(table, metadata.msa_subpage_block_table[2:4]) + + # Nothing staged: the caller expands its own layer's table, and the 0 + # factor is what tells it to. + metadata.msa_subpage_block_table = None + assert metadata.msa_subpage_rows(2, 4) == (None, 0) + + def test_msa_proxy_max_score_view_is_contiguous_over_stable_store(): """The proxy view fed to fmha_sm100 must be contiguous in the exact [num_index_heads, plan_max_k_tiles, num_tokens] shape the kernel writes, @@ -172,6 +285,11 @@ def test_msa_proxy_max_score_view_is_contiguous_over_stable_store(): with pytest.raises(ValueError, match=r"msa_max_score backing store"): metadata.msa_proxy_max_score_view(num_index_heads, worst_k, max_batch + 1) + # So is an empty one. Both writers address the view by block id, so a zero + # extent is not a small view but writes past the end of one. + with pytest.raises(ValueError, match=r"no block extent"): + metadata.msa_proxy_max_score_view(num_index_heads, 0, max_batch) + def test_msa_index_k_uses_hnd_cache_view_and_writer(): metadata_cls = MiniMaxM3MsaSparseAttention.Metadata @@ -482,6 +600,364 @@ def test_msa_scratch_sizing_covers_spec_verify_tokens(): ) +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +def test_lazily_allocated_scratch_publishes_the_bound_it_used(monkeypatch): + """The scratch is normally sized in _create_msa_buffers, but a metadata + built without sparse params allocates it here on first use. Either way the + worst-case bound has to be published: msa_proxy_max_score_view shapes the + view from it, including on a step that skips the proxy plan and so never + computes a bound of its own. + """ + from tensorrt_llm._torch.attention_backend.sparse.minimax_m3 import msa_backend + + # Both stand in for the fmha_sm100 submodule, which need not be built to + # test what is done with the number it returns. + monkeypatch.setattr(msa_backend, "require_msa_module", lambda: None) + monkeypatch.setattr(msa_backend, "_worst_case_proxy_max_k_tiles", lambda *a, **kw: 32) + + metadata_cls = MiniMaxM3MsaSparseAttention.Metadata + metadata = metadata_cls.__new__(metadata_cls) + metadata.kv_cache_manager = SimpleNamespace() + metadata.cuda_graph_buffers = None + metadata.max_num_sequences = 2 + metadata.max_num_tokens = 8 + metadata.msa_max_score = None + metadata._msa_worst_case_max_k_tiles = 0 + + metadata._ensure_msa_decode_scratch_buffers( + num_index_heads=4, + max_batch=2, + capture_graph=False, + required_max_k_tiles=16, + ) + + assert metadata.msa_worst_case_max_k_tiles == 32 + # The store was sized against that bound, so a view shaped by it fits. + assert metadata.msa_proxy_max_score_view(4, 32, 8).shape == (4, 32, 8) + + +def _resolution_metadata( + *, num_contexts=0, qo_lens=(1, 1), kv_lens=(9, 11), is_cuda_graph=False, page_size=128 +): + """Metadata with just enough state for _resolve_decode_kernels. + + The resolver reads only host-side facts, so seq_lens/kv_lens are enough to + drive the real msa_*_cpu length properties; no cache pool is needed. + """ + metadata_cls = MiniMaxM3MsaSparseAttention.Metadata + metadata = metadata_cls.__new__(metadata_cls) + metadata._msa_params = MiniMaxM3SparseAttentionConfig( + implementation="msa" + ).to_sparse_metadata_params() + metadata.mapping = None + # Assigned behind the seq_lens property, whose setter would stage a device + # copy the resolver never reads; num_seqs derives from it. The num_contexts + # setter then runs on_update() over both, as it does in a real step. + metadata._seq_lens = torch.tensor(qo_lens, dtype=torch.int32) + metadata.num_contexts = num_contexts + metadata.kv_lens = torch.tensor(kv_lens, dtype=torch.int32) + metadata.kv_cache_params = None + metadata.is_cuda_graph = is_cuda_graph + metadata._msa_captured_resolution = None + metadata.kv_cache_manager = SimpleNamespace( + tokens_per_block=page_size, + indexer_kv_dtype="bf16", + # Present, so the trtllm-gen dense support check passes. + get_kv_subpage_pool=lambda: None, + ) + metadata.msa_block_table = torch.zeros(len(qo_lens), 4, dtype=torch.int32) + metadata.msa_seq_lens_cuda = torch.zeros(len(qo_lens), dtype=torch.int32) + return metadata + + +def _force_cutedsl_supported(monkeypatch): + """Force the CuTe DSL geometry verdict, which otherwise needs an SM100 host. + + The resolver raises on an unsupported geometry, so without this every + resolution test would be a test of the runner's availability. + """ + monkeypatch.setattr( + MiniMaxM3MsaSparseAttention.Metadata, + "_cutedsl_indexer_supported", + lambda self, **kw: True, + ) + + +def test_resolve_decode_kernels_commits_on_uniform_decode(monkeypatch): + """A uniform pure-decode step resolves a span over the whole batch, which is + what lets prepare() skip the fmha_sm100 plans entirely.""" + _force_cutedsl_supported(monkeypatch) + metadata = _resolution_metadata() + + metadata._resolve_decode_kernels() + + assert msa_ported_decode_active(metadata) is True + assert metadata._msa_runs_no_fmha() is True + assert metadata.msa_decode_query_len == 1 + assert metadata.msa_max_kv_len == 11 + # The whole batch is the span, so nothing is left for fmha_sm100. + span = metadata.msa_decode_span + assert (span.row_first, span.row_last) == (0, 2) + assert (span.token_first, span.token_last) == (0, 2) + assert span.is_mixed is False + + +def test_resolve_decode_kernels_commits_the_generation_span_of_a_mixed_step(monkeypatch): + """A context request does not disqualify the step. + + The generation requests are the batch's row and token suffix, so the ported + kernels take that span and fmha_sm100 keeps the context prefix. + """ + _force_cutedsl_supported(monkeypatch) + # Two context requests (7 and 5 query tokens, the first a chunk of a long + # prompt) ahead of two decode rows. + metadata = _resolution_metadata(num_contexts=2, qo_lens=(7, 5, 1, 1), kv_lens=(4096, 5, 40, 33)) + + metadata._resolve_decode_kernels() + + span = metadata.msa_decode_span + assert (span.row_first, span.row_last) == (2, 4) + assert (span.token_first, span.token_last) == (12, 14) + assert span.query_len == 1 + assert span.is_mixed is True + assert msa_ported_decode_active(metadata) is True + # fmha_sm100 still runs the context prefix, so its page table stays live. + assert metadata._msa_runs_no_fmha() is False + # The trtllm-gen scheduling bound must come from the span's own rows: the + # 4096-token context row here would inflate a whole-batch maximum by 100x. + assert metadata.msa_max_kv_len == 40 + + +def test_resolve_decode_kernels_resolves_no_span_for_a_pure_prefill(monkeypatch): + """A step with no generation row has nothing for the ported kernels, and + fmha_sm100 keeps every plan and the page table they read.""" + _force_cutedsl_supported(monkeypatch) + metadata = _resolution_metadata(num_contexts=2, qo_lens=(5, 7), kv_lens=(5, 7)) + + metadata._resolve_decode_kernels() + + assert metadata.msa_decode_span is None + assert metadata.msa_decode_query_len is None + assert msa_ported_decode_active(metadata) is False + assert metadata._msa_runs_no_fmha() is False + + +def test_a_span_without_its_buffers_is_not_active(): + """The ported kernels address the page table and per-request lengths + directly, so a metadata carrying a query length but neither (the standalone + kernel tests, which never run prepare()) has not resolved a span.""" + metadata = SimpleNamespace( + msa_decode_query_len=1, msa_block_table=None, msa_seq_lens_cuda=torch.zeros(1) + ) + assert msa_ported_decode_active(metadata) is False + + metadata.msa_block_table = torch.zeros(1) + assert msa_ported_decode_active(metadata) is True + + metadata.msa_seq_lens_cuda = None + assert msa_ported_decode_active(metadata) is False + + +@pytest.mark.parametrize( + ("num_contexts", "qo_lens", "kv_lens"), + [ + # Ragged decode: the ported kernels' token -> request mapping breaks. + (0, (1, 3), (9, 11)), + # Ragged generation rows behind a context request. The context request + # is fine, but these rows still have no single query length. + (1, (5, 1, 2), (5, 9, 11)), + ], + ids=["ragged-decode", "ragged-mixed"], +) +def test_resolve_decode_kernels_raises_on_ragged_generation_rows( + monkeypatch, num_contexts, qo_lens, kv_lens +): + """There is no fmha_sm100 decode path left to fall back to, so a span the + ported kernels cannot serve has to surface rather than cost the step its + decode throughput silently.""" + _force_cutedsl_supported(monkeypatch) + metadata = _resolution_metadata(num_contexts=num_contexts, qo_lens=qo_lens, kv_lens=kv_lens) + + with pytest.raises(NotImplementedError, match=r"one query length"): + metadata._resolve_decode_kernels() + + assert metadata.msa_decode_span is None + + +def test_resolve_decode_kernels_raises_without_the_dense_subpage_pool(monkeypatch): + """trtllm-gen needs the flat sub-page pool, and its dense plan is gone, so a + manager without one cannot serve the dense layers at all.""" + _force_cutedsl_supported(monkeypatch) + metadata = _resolution_metadata() + del metadata.kv_cache_manager.get_kv_subpage_pool + + with pytest.raises(NotImplementedError, match=r"sub-page pool"): + metadata._resolve_decode_kernels() + + +def test_resolve_decode_kernels_raises_when_the_scorer_declines_the_geometry(monkeypatch): + """Same for the indexer: the proxy pass over the span is gone with it.""" + monkeypatch.setattr( + MiniMaxM3MsaSparseAttention.Metadata, + "_cutedsl_indexer_supported", + lambda self, **kw: False, + ) + metadata = _resolution_metadata() + + with pytest.raises(NotImplementedError, match=r"CuTe DSL indexer scorer"): + metadata._resolve_decode_kernels() + + +def test_fmha_plan_rows_narrow_to_the_context_prefix(monkeypatch): + """The plans must cover exactly the rows fmha_sm100 still runs. + + on_update_kv_lens patches a plan's length mirrors against the requests it + was built from, so a plan that claimed the whole batch while only the + context prefix ran would write the wrong lengths into the kernel. + """ + _force_cutedsl_supported(monkeypatch) + + mixed = _resolution_metadata(num_contexts=2, qo_lens=(7, 5, 1, 1), kv_lens=(4096, 5, 40, 33)) + mixed._msa_live_batch = 4 + mixed._resolve_decode_kernels() + # The span took the generation suffix, so the plans cover the prefix. + assert mixed._msa_fmha_plan_rows() == (0, 2) + + decode = _resolution_metadata() + decode._msa_live_batch = 2 + decode._resolve_decode_kernels() + # Nothing is left to plan on a pure-decode step the kernels fully own. + assert decode._msa_fmha_plan_rows() is None + + prefill = _resolution_metadata(num_contexts=2, qo_lens=(5, 7), kv_lens=(5, 7)) + prefill._msa_live_batch = 2 + prefill._resolve_decode_kernels() + # No span, so fmha_sm100 runs every row and is planned over all of them. + assert prefill._msa_fmha_plan_rows() == (0, 2) + + +def test_resolution_must_not_change_under_a_captured_graph(monkeypatch): + """The kernels inside a captured graph are fixed, so a later step that + resolves differently would stage inputs for a kernel that never runs.""" + _force_cutedsl_supported(monkeypatch) + metadata = _resolution_metadata(is_cuda_graph=True) + + def _step(): + metadata._resolve_decode_kernels() + metadata._check_capture_stable_resolution() + + _step() + # Same inputs: the replay agrees with the capture. + _step() + + # A replay whose batch turned mixed. The graph was captured with no + # fmha_sm100 plans at all, so the context prefix this step resolves has + # nothing to run under. + metadata._seq_lens = torch.tensor([5, 1, 1], dtype=torch.int32) + metadata.kv_lens = torch.tensor([5, 9, 11], dtype=torch.int32) + metadata.num_contexts = 1 + metadata.msa_block_table = torch.zeros(3, 4, dtype=torch.int32) + metadata.msa_seq_lens_cuda = torch.zeros(3, dtype=torch.int32) + with pytest.raises(RuntimeError, match=r"changed under a captured CUDA graph"): + _step() + + +def test_indexer_raises_when_a_committed_cutedsl_scorer_declines(): + """prepare() left the proxy plan covering nothing but this step's context + prefix, so a decline has no fallback for the generation span and must + surface instead of silently reading a stale page table.""" + from tensorrt_llm._torch.attention_backend.sparse.minimax_m3.common import MiniMaxM3SparseConfig + from tensorrt_llm._torch.attention_backend.sparse.minimax_m3.msa_indexer import MsaIndexer + + indexer = MsaIndexer( + MiniMaxM3SparseConfig( + num_q_heads=8, + num_kv_heads=1, + head_dim=128, + num_index_heads=4, + sparse_index_dim=128, + block_size=128, + topk=16, + ) + ) + # block_table/seq_lens left None, so the scorer cannot even be attempted. + with pytest.raises(RuntimeError, match=r"CuTe DSL indexer scorer declined the span"): + indexer.select_blocks( + torch.zeros(2, 4, 128), + torch.zeros(4, 1, 128, 128), + idx_sm_scale=1.0, + kv_indices=torch.zeros(4, dtype=torch.int32), + max_score=torch.zeros(4, 8, 2), + require_cutedsl=True, + ) + + +@pytest.mark.parametrize("head_major", [False, True]) +def test_combined_topk_table_preserves_the_requested_backing(head_major): + """Joining the two halves of a mixed step's table must not change its layout. + + The Triton sparse decode kernel reads the top-k table head-major, so a + joined table has to permute to a contiguous [num_kv_heads, total_q, topk] + exactly as the selector's own output does; a token-major join would silently + hand the kernel a strided view where production hands it a dense one. + """ + from tensorrt_llm._torch.attention_backend.sparse.minimax_m3.msa_indexer import ( + _combined_topk_table, + ) + + num_kv_heads, topk = 2, 16 + ctx = torch.arange(5 * num_kv_heads * topk, dtype=torch.int32).reshape(5, num_kv_heads, topk) + gen = -ctx[:3] - 1 + + combined = _combined_topk_table(ctx, gen, head_major=head_major) + + assert combined.shape == (8, num_kv_heads, topk) + assert torch.equal(combined[:5], ctx) + assert torch.equal(combined[5:], gen) + assert combined.permute(1, 0, 2).is_contiguous() is head_major + assert combined.is_contiguous() is not head_major + + +def test_paged_gqa_raises_when_a_committed_dense_step_declines(): + """The mirror of the indexer guard on the attention side. The step resolved + a span and dropped the dense plan, so a call site that finds the geometry + unsupported has nothing left to fall back to.""" + from tensorrt_llm._torch.attention_backend.fmha.msa_sparse_gqa import run_msa_paged_gqa + + num_heads, head_dim, num_pages, page_size = 8, 128, 4, 16 + attention = MiniMaxM3MsaSparseAttention.__new__(MiniMaxM3MsaSparseAttention) + attention.layer_idx = 3 + attention.head_dim = head_dim + attention.num_heads = num_heads + attention.q_scaling = 1.0 + + metadata = SimpleNamespace( + # No get_kv_subpage_pool, so trtllm-gen declines at the call site. + kv_cache_manager=SimpleNamespace( + get_buffers=lambda layer_idx, kv_layout=None: torch.zeros( + num_pages, 2, 1, page_size, head_dim + ) + ), + # A resolved pure-decode span, so the decline is the only thing that + # can send this call to fmha_sm100. + msa_decode_query_len=1, + msa_block_table=torch.zeros(2, 1, dtype=torch.int32), + msa_seq_lens_cuda=torch.zeros(2, dtype=torch.int32), + ) + + with pytest.raises(RuntimeError, match=r"skipped the fmha_sm100 dense plan"): + run_msa_paged_gqa( + attention, + torch.zeros(2, num_heads * head_dim), + None, + None, + metadata, + torch.zeros(2, num_heads * head_dim), + kv_block_indexes=None, + plan=None, + ) + + def test_per_token_valid_blocks_multi_token_decode(): """Spec-verify decode rows expose one entry per query TOKEN, walking the causal ladder within the verify window.""" @@ -615,8 +1091,8 @@ def _reference_scatter_write(k_cache, v_cache, idx_cache, slots, k, v, idx_k): @pytest.mark.parametrize("num_kv_heads", [1, 4]) @pytest.mark.parametrize("with_idx", [True, False]) def test_fused_scatter_matches_reference(cache_dtype, num_kv_heads, with_idx): - """The fused per-layer cache scatter must match the legacy write_kv_slots - path exactly on production-shaped inputs: non-contiguous HND cache views + """The fused per-layer cache scatter must match the per-cache write_kv_slots + writes exactly on production-shaped inputs: non-contiguous HND cache views carved from a pooled allocation and strided source rows sliced from a fused projection, including the bf16 -> fp8 cache cast. Asserting on the whole pool also catches stray writes outside the targeted slots.""" @@ -656,3 +1132,160 @@ def test_fused_scatter_matches_reference(cache_dtype, num_kv_heads, with_idx): torch.testing.assert_close(pool.to(torch.float32), ref_pool.to(torch.float32)) torch.testing.assert_close(idx_pool, ref_idx_pool) + + +def _mixed_batch_sparse_gqa_case(*, page_size, head_dim, num_kv_heads, group, topk, seed): + """A one-context-plus-three-decode batch for run_msa_paged_gqa. + + Returns the attention stub, the metadata fields both runs share, q, the + per-query top-k table, and the batch's context token count. Pages are + shuffled so a kernel that ignored the block table and indexed the cache by + logical block would not pass. + """ + from tensorrt_llm._torch.attention_backend.sparse.minimax_m3.msa_utils import ( + build_kv_page_indices, + per_token_valid_blocks, + ) + + generator = torch.Generator(device="cuda").manual_seed(seed) + num_heads = num_kv_heads * group + # Row 0 prefills a fresh 260-token prompt; rows 1-3 decode one token each. + qo_lens_cpu = torch.tensor([260, 1, 1, 1], dtype=torch.int32) + kv_lens_cpu = torch.tensor([260, 300, 1500, 129], dtype=torch.int32) + qo_offset_cpu = kv_lens_cpu - qo_lens_cpu + batch = int(qo_lens_cpu.shape[0]) + total_q = int(qo_lens_cpu.sum()) + max_blocks = int((kv_lens_cpu.max().item() + page_size - 1) // page_size) + num_pages = batch * max_blocks + + block_table = ( + torch.randperm(num_pages, device="cuda", generator=generator) + .to(torch.int32) + .reshape(batch, max_blocks) + ) + pool = torch.randn( + num_pages, + 2, + num_kv_heads, + page_size, + head_dim, + device="cuda", + generator=generator, + dtype=torch.float32, + ).to(torch.bfloat16) + + q = torch.randn( + total_q, num_heads * head_dim, device="cuda", generator=generator, dtype=torch.float32 + ).to(torch.bfloat16) + + # Select each token's earliest valid blocks, ascending with a -1 tail, as + # the indexer emits them. Deterministic, and valid for the context rows, + # whose causal extent grows token by token. + n_valid = per_token_valid_blocks( + qo_lens_cpu, kv_lens_cpu, qo_offset_cpu, causal=True, block_size=page_size + ) + table = torch.full((total_q, num_kv_heads, topk), -1, dtype=torch.int32) + for token, valid in enumerate(n_valid.tolist()): + real = min(topk, max(int(valid), 0)) + table[token, :, :real] = torch.arange(real, dtype=torch.int32) + # Head-major backing, so the .permute(1, 0, 2) in run_msa_paged_gqa is the + # zero-copy view it is in production. + head_major = table.permute(1, 0, 2).contiguous().cuda() + + attention = MiniMaxM3MsaSparseAttention.__new__(MiniMaxM3MsaSparseAttention) + attention.layer_idx = 0 + attention.head_dim = head_dim + attention.num_heads = num_heads + attention.q_scaling = 1.0 + + fields = dict( + kv_cache_manager=SimpleNamespace( + tokens_per_block=page_size, + get_buffers=lambda layer_idx, kv_layout=None: pool, + ), + msa_block_table=block_table, + msa_seq_lens_cuda=kv_lens_cpu.cuda(), + msa_kv_indices=build_kv_page_indices(block_table.cpu(), kv_lens_cpu, page_size).cuda(), + msa_qo_lens_cpu=qo_lens_cpu, + msa_kv_lens_cpu=kv_lens_cpu, + msa_qo_offset_cpu=qo_offset_cpu, + msa_max_kv_len=int(kv_lens_cpu[1:].max()), + max_num_requests=batch, + ) + return attention, fields, q, head_major.permute(1, 0, 2), int(qo_lens_cpu[0]) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +def test_mixed_batch_generation_span_matches_the_whole_batch_msa_path(): + """Splitting a mixed batch by phase must not change any row's answer. + + The generation rows move off fmha_sm100 and onto the Triton sparse decode + kernel while the context rows stay behind under a context-only plan, so the + correctness gate is that both halves still agree with a whole-batch + fmha_sm100 run, which a metadata carrying no span still takes. This is the + only test that covers which kernel produced which output rows. + """ + from tensorrt_llm._torch.attention_backend.fmha.msa_sparse_gqa import run_msa_paged_gqa + from tensorrt_llm._torch.attention_backend.sparse.minimax_m3.msa_backend import _MsaDecodeSpan + from tensorrt_llm._torch.attention_backend.sparse.minimax_m3.msa_utils import ( + MSA_REQUIRED_TOPK, + msa_package_available, + ) + from tensorrt_llm._utils import get_sm_version + + if not msa_package_available(): + pytest.skip("fmha_sm100 (MSA submodule) required") + if get_sm_version() not in (100, 103): + pytest.skip("fmha_sm100 requires SM100/SM103") + + page_size, head_dim = 128, 128 + attention, fields, q, kv_block_indexes, num_ctx_tokens = _mixed_batch_sparse_gqa_case( + page_size=page_size, + head_dim=head_dim, + num_kv_heads=1, + group=8, + topk=MSA_REQUIRED_TOPK, + seed=61, + ) + total_q = int(q.shape[0]) + + def run(**resolution): + output = torch.zeros_like(q) + run_msa_paged_gqa( + attention, + q, + None, + None, + SimpleNamespace(**fields, **resolution), + output, + kv_block_indexes=kv_block_indexes, + plan=None, + ) + torch.cuda.synchronize() + return output.view(total_q, attention.num_heads, head_dim).float() + + reference = run(msa_decode_span=None) + split = run( + # A property of the span on real metadata, and what + # msa_ported_decode_active reads; this fake carries both. + msa_decode_query_len=1, + msa_decode_span=_MsaDecodeSpan( + row_first=1, + row_last=4, + token_first=num_ctx_tokens, + token_last=total_q, + query_len=1, + ), + ) + + assert torch.isfinite(split).all() + # The context prefix runs on fmha_sm100 either way, but under a 1-row plan + # rather than a 4-row one, so its work partitioning differs. + torch.testing.assert_close( + split[:num_ctx_tokens], reference[:num_ctx_tokens], rtol=1e-2, atol=1e-2 + ) + # The generation rows change kernel outright, so they carry the wider + # tolerance the Triton-vs-fmha_sm100 A/B uses elsewhere. + torch.testing.assert_close( + split[num_ctx_tokens:], reference[num_ctx_tokens:], rtol=6e-2, atol=6e-2 + ) diff --git a/tests/unittest/_torch/attention/sparse/test_minimax_m3_sparse_attn_decode.py b/tests/unittest/_torch/attention/sparse/test_minimax_m3_sparse_attn_decode.py new file mode 100644 index 000000000000..583318ca4fc2 --- /dev/null +++ b/tests/unittest/_torch/attention/sparse/test_minimax_m3_sparse_attn_decode.py @@ -0,0 +1,360 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-License-Identifier: Apache-2.0 +# PyTorch oracle derived from vLLM (Apache-2.0), _reference_sparse_attn: +# https://github.com/vllm-project/vllm/blob/6f91edf96d3f3272945809c04702380053bff4de/tests/kernels/attention/test_minimax_m3.py#L755 +"""Correctness tests for the Triton MiniMax-M3 sparse block decode attention. + +The PyTorch oracle follows the vLLM reference linked in the file header +(v0.26.1rc0-77-g6f91edf96): softmax over exactly the selected blocks, each +truncated at the query token's own causal extent. +""" + +import pytest +import torch + +from tensorrt_llm._torch.attention_backend.sparse.minimax_m3.msa_utils import ( + MSA_REQUIRED_TOPK, + build_kv_page_indices, + msa_package_available, +) +from tensorrt_llm._torch.attention_backend.sparse.minimax_m3.triton_sparse_decode import ( + SPARSE_BLOCK_SIZE, + minimax_m3_sparse_attn_decode, + resolve_num_topk_chunks, +) +from tensorrt_llm._utils import get_sm_version + +PAGE_SIZE = SPARSE_BLOCK_SIZE +HEAD_DIM = 128 + +pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") + +skip_not_sm100 = pytest.mark.skipif( + get_sm_version() not in (100, 103), + reason="fmha_sm100 A/B comparison requires SM100/SM103.", +) + + +def _reference_sparse_decode( + q: torch.Tensor, + k_paged: torch.Tensor, + v_paged: torch.Tensor, + topk_idx: torch.Tensor, # [num_kv_heads, total_q, topk] + block_table: torch.Tensor, + seq_lens: torch.Tensor, + sm_scale: float, + decode_query_len: int, +) -> torch.Tensor: + """fp32 softmax over the selected blocks, truncated at each token's extent.""" + total_q, num_heads, head_dim = q.shape + num_kv_heads = k_paged.shape[1] + group = num_heads // num_kv_heads + max_topk = topk_idx.shape[-1] + out = torch.zeros(total_q, num_heads, head_dim, device=q.device, dtype=torch.float32) + k_f32, v_f32 = k_paged.float(), v_paged.float() + positions = torch.arange(PAGE_SIZE, device=q.device) + + for req, seq_len in enumerate(seq_lens.tolist()): + for intra in range(decode_query_len): + token = req * decode_query_len + intra + kv_len = max(seq_len - decode_query_len + intra + 1, 0) + num_blocks = (kv_len + PAGE_SIZE - 1) // PAGE_SIZE + real_topk = min(max_topk, num_blocks) + if real_topk == 0: + continue + for kv_head in range(num_kv_heads): + blocks = topk_idx[kv_head, token, :real_topk].tolist() + pages = [int(block_table[req, int(b)]) for b in blocks] + keys = torch.cat([k_f32[p, kv_head] for p in pages]) + values = torch.cat([v_f32[p, kv_head] for p in pages]) + valid = torch.cat([int(b) * PAGE_SIZE + positions < kv_len for b in blocks]) + q_rows = q[token, kv_head * group : (kv_head + 1) * group].float() + logits = (q_rows @ keys.T) * sm_scale + logits = logits.masked_fill(~valid[None, :], -float("inf")) + probs = torch.softmax(logits, dim=-1) + out[token, kv_head * group : (kv_head + 1) * group] = probs @ values + return out + + +def _random_topk( + seq_lens, + *, + num_kv_heads: int, + decode_query_len: int, + topk: int, + generator: torch.Generator, +) -> torch.Tensor: + """Ascending block ids with -1 tail padding, as the selector emits them.""" + total_q = len(seq_lens) * decode_query_len + table = torch.full((num_kv_heads, total_q, topk), -1, device="cuda", dtype=torch.int32) + for req, seq_len in enumerate(seq_lens): + for intra in range(decode_query_len): + token = req * decode_query_len + intra + kv_len = max(seq_len - decode_query_len + intra + 1, 0) + num_blocks = (kv_len + PAGE_SIZE - 1) // PAGE_SIZE + real_topk = min(topk, num_blocks) + if real_topk == 0: + continue + for kv_head in range(num_kv_heads): + perm = torch.randperm(num_blocks, device="cuda", generator=generator) + chosen = perm[:real_topk].sort().values + table[kv_head, token, :real_topk] = chosen.to(torch.int32) + return table + + +def _make_inputs( + seq_lens, + *, + kv_dtype, + q_dtype=torch.bfloat16, + num_kv_heads=1, + group=8, + decode_query_len=1, + topk=MSA_REQUIRED_TOPK, + seed=0, +): + generator = torch.Generator(device="cuda").manual_seed(seed) + batch = len(seq_lens) + total_q = batch * decode_query_len + num_heads = num_kv_heads * group + max_blocks = max(1, (max(seq_lens) + PAGE_SIZE - 1) // PAGE_SIZE) + num_pages = batch * max_blocks + + # Shuffled pages: a kernel that ignored the block table and indexed the + # cache by logical block would still pass with an identity mapping. + block_table = torch.randperm(num_pages, device="cuda", generator=generator) + block_table = block_table.to(torch.int32).reshape(batch, max_blocks) + seq_lens_dev = torch.tensor(seq_lens, device="cuda", dtype=torch.int32) + + q = torch.randn( + total_q, num_heads, HEAD_DIM, device="cuda", generator=generator, dtype=torch.float32 + ).to(q_dtype) + kv = torch.randn( + num_pages, + 2, + num_kv_heads, + PAGE_SIZE, + HEAD_DIM, + device="cuda", + generator=generator, + dtype=torch.float32, + ).to(kv_dtype) + # Non-contiguous K/V views of one coalesced pool, exactly as msa_paged_kv + # hands them over. + k_paged, v_paged = kv[:, 0], kv[:, 1] + + topk_idx = _random_topk( + seq_lens, + num_kv_heads=num_kv_heads, + decode_query_len=decode_query_len, + topk=topk, + generator=generator, + ) + return q, k_paged, v_paged, topk_idx, block_table, seq_lens_dev + + +def _run(q, k_paged, v_paged, topk_idx, block_table, seq_lens, decode_query_len, **kwargs): + out = torch.empty_like(q) + minimax_m3_sparse_attn_decode( + q, + k_paged, + v_paged, + topk_idx, + block_table, + seq_lens, + sm_scale=HEAD_DIM**-0.5, + output=out, + decode_query_len=decode_query_len, + **kwargs, + ) + return out + + +@pytest.mark.parametrize("kv_dtype", [torch.bfloat16, torch.float8_e4m3fn]) +@pytest.mark.parametrize( + ("num_kv_heads", "group", "decode_query_len"), + [(1, 8, 1), (1, 8, 2), (2, 16, 1), (4, 4, 3)], +) +@pytest.mark.parametrize( + "seq_lens", + [ + [128], + [1, 128, 129], + [1025, 4097], + [300, 1500, 2049, 4096, 5000, 7777], + ], + ids=["one-block", "short-mixed", "two-req", "batch6"], +) +def test_sparse_decode_matches_reference(kv_dtype, num_kv_heads, group, decode_query_len, seq_lens): + seq_lens = [max(s, decode_query_len) for s in seq_lens] + inputs = _make_inputs( + seq_lens, + kv_dtype=kv_dtype, + num_kv_heads=num_kv_heads, + group=group, + decode_query_len=decode_query_len, + ) + out = _run(*inputs, decode_query_len) + expected = _reference_sparse_decode( + *inputs, sm_scale=HEAD_DIM**-0.5, decode_query_len=decode_query_len + ) + torch.testing.assert_close(out.float(), expected, rtol=3e-2, atol=3e-2) + + +@pytest.mark.parametrize("num_topk_chunks", [1, 2, 4, 8, 16]) +def test_sparse_decode_split_k_invariant(num_topk_chunks): + """Flash-decoding must merge to the same answer for any split-K factor. + + A wrong LSE merge shows up here and nowhere else, because the default + chunk count for a small batch is usually large enough to hide it. + """ + seq_lens = [4097, 300, 8192] + inputs = _make_inputs(seq_lens, kv_dtype=torch.float8_e4m3fn, num_kv_heads=2, group=8, seed=7) + reference = _run(*inputs, 1, num_topk_chunks=1) + out = _run(*inputs, 1, num_topk_chunks=num_topk_chunks) + torch.testing.assert_close(out.float(), reference.float(), rtol=1e-2, atol=1e-2) + + +def test_sparse_decode_ignores_padded_topk_entries(): + """Rows whose valid block count is below topk must not read the -1 tail. + + Every request here is one or two blocks long against topk=16, so all but a + couple of entries are -1; dereferencing them would fault or corrupt. + """ + seq_lens = [1, 64, 128, 129, 200] + inputs = _make_inputs(seq_lens, kv_dtype=torch.bfloat16, seed=13) + topk_idx = inputs[3] + assert (topk_idx == -1).any() + + out = _run(*inputs, 1) + expected = _reference_sparse_decode(*inputs, sm_scale=HEAD_DIM**-0.5, decode_query_len=1) + torch.testing.assert_close(out.float(), expected, rtol=3e-2, atol=3e-2) + + +def test_sparse_decode_zero_length_rows_are_zero_not_nan(): + """CUDA-graph padding rows attend nothing; they must emit zeros. + + A NaN here would be discarded from the padded output but would still reach + the residual stream and the tensor-parallel all-reduce. + """ + seq_lens = [1024, 0, 512, 0] + inputs = _make_inputs(seq_lens, kv_dtype=torch.float8_e4m3fn, seed=21) + out = _run(*inputs, 1) + + assert torch.isfinite(out).all() + assert torch.equal(out[1], torch.zeros_like(out[1])) + assert torch.equal(out[3], torch.zeros_like(out[3])) + assert out[0].abs().sum() > 0 + + +def test_sparse_decode_accepts_token_major_topk_table(): + """The kernel reads the top-k table by stride, so either backing works.""" + seq_lens = [1025, 4097] + q, k_paged, v_paged, head_major, block_table, seq_lens_dev = _make_inputs( + seq_lens, kv_dtype=torch.bfloat16, num_kv_heads=2, group=8, seed=31 + ) + token_major = head_major.permute(1, 0, 2).contiguous().permute(1, 0, 2) + assert head_major.is_contiguous() and not token_major.is_contiguous() + + args = (q, k_paged, v_paged) + out_hm = _run(*args, head_major, block_table, seq_lens_dev, 1) + out_tm = _run(*args, token_major, block_table, seq_lens_dev, 1) + assert torch.equal(out_hm, out_tm) + + +def test_sparse_decode_cuda_graph_replay_tracks_inputs(): + """Replay must recompute from the live buffers, not reuse captured values.""" + seq_lens = [2048, 4097, 900] + q, k_paged, v_paged, topk_idx, block_table, seq_lens_dev = _make_inputs( + seq_lens, kv_dtype=torch.float8_e4m3fn, num_kv_heads=2, group=8, seed=41 + ) + out = torch.empty_like(q) + + def run(): + minimax_m3_sparse_attn_decode( + q, + k_paged, + v_paged, + topk_idx, + block_table, + seq_lens_dev, + sm_scale=HEAD_DIM**-0.5, + output=out, + decode_query_len=1, + num_topk_chunks=4, + ) + + # Warm up on a side stream so the Triton JIT and the scratch arena are + # settled before capture, which forbids both. + side = torch.cuda.Stream() + side.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(side): + run() + torch.cuda.current_stream().wait_stream(side) + + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + run() + + generator = torch.Generator(device="cuda").manual_seed(99) + q.copy_(torch.randn(q.shape, device="cuda", generator=generator, dtype=torch.float32)) + graph.replay() + torch.cuda.synchronize() + + expected = _reference_sparse_decode( + q, k_paged, v_paged, topk_idx, block_table, seq_lens_dev, HEAD_DIM**-0.5, 1 + ) + torch.testing.assert_close(out.float(), expected, rtol=3e-2, atol=3e-2) + + +@skip_not_sm100 +@pytest.mark.skipif(not msa_package_available(), reason="fmha_sm100 (MSA submodule) required") +def test_sparse_decode_matches_msa_kernel(): + """A/B against the fmha_sm100 sparse GQA path this kernel replaces.""" + from tensorrt_llm._torch.attention_backend.fmha.msa_sparse_gqa import run_msa_sparse_gqa + + seq_lens = [1025, 4097, 300, 8192] + q, k_paged, v_paged, topk_idx, block_table, seq_lens_dev = _make_inputs( + seq_lens, kv_dtype=torch.float8_e4m3fn, num_kv_heads=1, group=8, seed=53 + ) + sm_scale = HEAD_DIM**-0.5 + batch = len(seq_lens) + + triton_out = _run(q, k_paged, v_paged, topk_idx, block_table, seq_lens_dev, 1) + + qo_lens_cpu = torch.ones(batch, dtype=torch.int32) + kv_lens_cpu = torch.tensor(seq_lens, dtype=torch.int32) + kv_indices = build_kv_page_indices(block_table.cpu(), kv_lens_cpu, PAGE_SIZE).cuda() + msa_out = torch.empty_like(q) + run_msa_sparse_gqa( + q.to(torch.float8_e4m3fn), + k_paged, + v_paged, + topk_idx.permute(1, 0, 2).contiguous(), + kv_indices=kv_indices, + sm_scale=sm_scale, + qo_lens_cpu=qo_lens_cpu, + kv_lens_cpu=kv_lens_cpu, + qo_offset_cpu=kv_lens_cpu - qo_lens_cpu, + causal=True, + head_dim=HEAD_DIM, + out=msa_out, + use_fp8=True, + ) + + # fmha_sm100 quantizes q to E4M3 while the Triton kernel keeps it in + # bf16, so the two differ by q's quantization error alone. + torch.testing.assert_close(triton_out.float(), msa_out.float(), rtol=6e-2, atol=6e-2) + + +@pytest.mark.parametrize( + ("total_q", "num_kv_heads", "max_topk"), + [(1, 1, 16), (8, 1, 16), (64, 2, 16), (512, 4, 16), (4096, 8, 16)], +) +def test_resolve_num_topk_chunks_is_shape_only_power_of_two(total_q, num_kv_heads, max_topk): + """The split-K factor is frozen by shape, so a captured graph keeps it.""" + chunks = resolve_num_topk_chunks(total_q, num_kv_heads, max_topk) + assert 1 <= chunks <= max_topk + assert chunks & (chunks - 1) == 0 + assert chunks == resolve_num_topk_chunks(total_q, num_kv_heads, max_topk) diff --git a/tests/unittest/_torch/thop/parallel_hw_agnostic/test_minimax_m3_fp8_horizontal_producer.py b/tests/unittest/_torch/thop/parallel_hw_agnostic/test_minimax_m3_fp8_horizontal_producer.py index a05fef3915ec..1f6c41998b83 100644 --- a/tests/unittest/_torch/thop/parallel_hw_agnostic/test_minimax_m3_fp8_horizontal_producer.py +++ b/tests/unittest/_torch/thop/parallel_hw_agnostic/test_minimax_m3_fp8_horizontal_producer.py @@ -63,9 +63,9 @@ def test_minimax_m3_horizontal_producer_matches_separate_producers(num_tokens): slots = (torch.arange(num_tokens, dtype=torch.int32, device="cuda") * 37) % ( (num_pages - 1) * 128 ) - # Keep the parity reference slots valid: the legacy separate main-K/V - # producer does not support negative slots. Negative-slot handling is - # exercised below using horizontal eager execution versus graph replay. + # Keep the parity reference slots valid: the separate main-K/V producer + # this is compared against does not support negative slots. Negative-slot + # handling is exercised below, eager execution versus graph replay. rope_cache = _rope_cache(max(256, num_tokens)) main_width = (num_heads_q + 2 * num_kv_heads) * 128