diff --git a/tensorrt_llm/_torch/attention/ATTENTION_DEVELOPER_GUIDE.md b/tensorrt_llm/_torch/attention/ATTENTION_DEVELOPER_GUIDE.md index e56bbb5fa24a..6b2dd2003d4f 100644 --- a/tensorrt_llm/_torch/attention/ATTENTION_DEVELOPER_GUIDE.md +++ b/tensorrt_llm/_torch/attention/ATTENTION_DEVELOPER_GUIDE.md @@ -373,8 +373,9 @@ The FMHA package is split by role: [vendored-source lifecycle](../../../3rdparty/vendor-sources.md). Land upstream-worthy changes in FlashInfer and update the vendor lock; keep only TRT-LLM-specific adaptations in the persistent patch. -- `fmha/msa_sparse_gqa.py` integrates the packaged SM100/SM103 block-sparse - GQA implementation. +- `fmha/msa_prefill.py` integrates the packaged SM100/SM103 block-sparse GQA + implementation for the context phase, and `fmha/msa_decode.py` runs the + MiniMax-M3 decode kernels for the generation phase. - `fmha/flashinfer_sparse_mla.py` implements the FlashInfer SM120/SM121 sparse MLA FMHA library. - `fmha/flashinfer_trtllm_gen.py` implements the FlashInfer trtllm-gen FMHA diff --git a/tensorrt_llm/_torch/attention/backends/fmha/__init__.py b/tensorrt_llm/_torch/attention/backends/fmha/__init__.py index ac8cd20f5f6e..04d92f69630a 100644 --- a/tensorrt_llm/_torch/attention/backends/fmha/__init__.py +++ b/tensorrt_llm/_torch/attention/backends/fmha/__init__.py @@ -19,7 +19,8 @@ from .flashinfer_sparse_mla import FlashInferSparseMlaFmha from .flashinfer_trtllm_gen import FlashInferTrtllmGenFmha from .interface import Fmha, FmhaPhase -from .msa_sparse_gqa import MsaSparseGqaFmha +from .msa_decode import MsaDecodeFmha +from .msa_prefill import MsaPrefillFmha from .phased import FmhaParams, PhasedFmha from .prims_ts import PrimsTSFmha from .registry import DEFAULT_FMHA_LIBS, FMHA_LIBS, FmhaCls, get_enabled_fmha_lib_classes @@ -37,7 +38,8 @@ "FmhaCls", "FmhaParams", "FmhaPhase", - "MsaSparseGqaFmha", + "MsaDecodeFmha", + "MsaPrefillFmha", "PhasedFmha", "PrimsTSFmha", "TritonCustomMaskFmha", diff --git a/tensorrt_llm/_torch/attention/backends/fmha/msa_decode.py b/tensorrt_llm/_torch/attention/backends/fmha/msa_decode.py new file mode 100644 index 000000000000..c812a3ec5830 --- /dev/null +++ b/tensorrt_llm/_torch/attention/backends/fmha/msa_decode.py @@ -0,0 +1,243 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Generation-phase FMHA for MiniMax-M3, on kernels built for a decode shape. + +MsaPrefillFmha's fmha_sm100 kernel schedules a generation row like a context +row, a single query token occupying a 128-row Q tile. This library takes the +generation phase instead and dispatches by layer: a sparse layer to the Triton +block-sparse decode kernel over the blocks the indexer selected, a dense layer +to trtllm-gen over the full page table. + +Both need a uniform query length across the generation rows and a geometry they +support, and neither has a fallback: prepare() settles the query length per +step as metadata.msa_decode_span, and ensure_msa_available and +_validate_decode_kernel_support settle the geometry once per run. +""" + +from __future__ import annotations + +import math +from typing import TYPE_CHECKING, Optional + +import torch + +from ..sparse.minimax_m3.kernels.msa_utils import is_msa_layer, msa_paged_kv, write_msa_phase_kv +from ..sparse.minimax_m3.kernels.trtllm_gen_dense_decode import ( + DenseDecodeWorkspaceLayout, + dense_decode_workspace_layout, + minimax_m3_trtllm_gen_dense_decode, + split_dense_decode_workspace, +) +from .interface import FmhaPhase +from .phased import FmhaParams, PhasedFmha + +if TYPE_CHECKING: + from tensorrt_llm._torch.attention.backends.interface import AttentionForwardArgs + from tensorrt_llm._torch.attention.backends.trtllm import ( + TrtllmAttention, + TrtllmAttentionMetadata, + ) + + +class MsaDecodeFmha(PhasedFmha): + """MiniMax-M3 generation attention on the Triton and trtllm-gen kernels. + + Generation only: the context phase is MsaPrefillFmha's, and a mixed batch + is served by the two together through CombinedFmha. run_context is left to + the base class, which refuses it. + """ + + def __init__(self, attn: "TrtllmAttention"): + super().__init__(attn) + # Where this layer's trtllm-gen scratch sits inside the shared + # attention workspace, settled by prepare_workspace each step. + self._dense_layout: Optional[DenseDecodeWorkspaceLayout] = None + + @classmethod + def is_available(cls, attn: "TrtllmAttention") -> bool: + return is_msa_layer(attn) + + def is_supported( + self, + q: torch.Tensor, + k: Optional[torch.Tensor], + v: Optional[torch.Tensor], + metadata: "TrtllmAttentionMetadata", + forward_args: "AttentionForwardArgs", + *, + phase: Optional[FmhaPhase] = None, + ) -> bool: + # Every generation row is this library's and every context row + # MsaPrefillFmha's, whatever the step looks like, so each library + # claims its one phase and their partition of the phases is total. The + # phase-less query asks for the whole step, which neither can serve + # alone: a mixed step belongs to the two together through CombinedFmha. + return phase is FmhaPhase.GENERATION + + def prepare_workspace( + self, + q: torch.Tensor, + k: Optional[torch.Tensor], + v: Optional[torch.Tensor], + metadata: "TrtllmAttentionMetadata", + forward_args: "AttentionForwardArgs", + workspace: torch.Tensor, + ) -> None: + if forward_args.sparse_runtime_params.sparse_attn_indices is not None: + # A sparse layer; the Triton kernel takes its split-K scratch from + # the arena itself, sized by the grid it just chose. + return + self._reserve_dense_workspace(q, metadata, workspace) + + def _reserve_dense_workspace( + self, + q: torch.Tensor, + metadata: "TrtllmAttentionMetadata", + workspace: torch.Tensor, + ) -> None: + """Size the shared attention workspace for this layer's trtllm-gen scratch. + + Taking the scratch from the workspace every FMHA library shares is what + puts it in a captured graph's memory pool. Growing that workspace + mid-capture would allocate behind the recorded kernels, so it is + refused. + """ + # The manager has the pool: _validate_decode_kernel_support refused the + # run without it, so no phase would have reached here. + kv_pool, _ = metadata.kv_cache_manager.get_kv_subpage_pool(self.attn.layer_idx, "HND") + q_dtype = torch.float8_e4m3fn if kv_pool.dtype == torch.float8_e4m3fn else q.dtype + layout = dense_decode_workspace_layout( + q_dtype=q_dtype, + num_heads=self.attn.num_heads, + head_dim=self.attn.head_dim, + num_kv_heads=int(kv_pool.shape[1]), + max_num_requests=int(metadata.max_num_requests), + device=q.device, + ) + current_bytes = workspace.numel() * workspace.element_size() + if current_bytes < layout.total_bytes: + if metadata.is_cuda_graph and torch.cuda.is_current_stream_capturing(): + raise RuntimeError( + "The attention CUDA graph workspace holds " + f"{current_bytes} bytes, fewer than the {layout.total_bytes} " + "MiniMax-M3 dense decode needs. The scratch is sized by the " + "head geometry alone, so this should not move." + ) + workspace.resize_((math.ceil(layout.total_bytes / workspace.element_size()),)) + self._dense_layout = layout + + def run_generation(self, params: FmhaParams) -> None: + metadata = params.meta + span = metadata.msa_decode_span + # Both kernels below map a query token to its request through the span, + # while the phase params carry PhasedFmha's own derivation of the same + # boundary. Checking them against each other rejects a step whose + # generation rows were never described and one where the two disagree. + phase = (params.seq_offset, params.input_seq_length) + if span != phase: + raise RuntimeError( + "MsaDecodeFmha ran on a generation phase its decode span does " + f"not describe: the span is {span}, while the phase starts at " + f"row {params.seq_offset} with {params.input_seq_length} query " + "tokens per request." + ) + write_msa_phase_kv( + params.attn, + params.key_input, + params.value_input, + metadata, + params.fwd.attention_input_type, + token_offset=params.token_offset, + ) + row_first = params.seq_offset + row_last = row_first + metadata.num_generations + block_table = metadata.msa_block_table[row_first:row_last] + seq_lens = metadata.msa_seq_lens_cuda[row_first:row_last] + + kv_block_indexes = params.fwd.sparse_runtime_params.sparse_attn_indices + if kv_block_indexes is not None: + self._run_sparse(params, kv_block_indexes, block_table, seq_lens) + else: + self._run_dense(params, block_table, seq_lens) + + def _run_sparse( + self, + params: FmhaParams, + kv_block_indexes: torch.Tensor, + block_table: torch.Tensor, + seq_lens: torch.Tensor, + ) -> None: + # Function-local: this module is on the import path of every + # attention.backends.trtllm import, and the kernel pulls in Triton. + from ..sparse.minimax_m3.kernels.triton_sparse_decode import minimax_m3_sparse_attn_decode + + attn = params.attn + head_dim = attn.head_dim + num_tokens = params.num_tokens + k_paged, v_paged = msa_paged_kv(params.meta.kv_cache_manager, attn.layer_idx) + # q may still be FP8 from a fused producer; the kernel widens it + # in-register, so it is passed through as it arrives. + minimax_m3_sparse_attn_decode( + params.attention_input.view(num_tokens, attn.num_heads, head_dim), + k_paged, + v_paged, + # The kernel reads the top-k table head-major and the indexer + # builds it that way for every step, so this is a view. It reads + # every stride, so a mixed step's strided suffix works too. + kv_block_indexes[params.token_offset : params.token_offset + num_tokens].permute( + 1, 0, 2 + ), + block_table, + seq_lens, + sm_scale=(head_dim**-0.5) / float(attn.q_scaling), + output=params.context_buf.view(num_tokens, attn.num_heads, head_dim), + decode_query_len=params.input_seq_length, + ) + + def _run_dense( + self, + params: FmhaParams, + block_table: torch.Tensor, + seq_lens: torch.Tensor, + ) -> None: + """Attend the full page table on trtllm-gen, for a dense M3 layer. + + This is the kernel FlashInferTrtllmGenFmha runs, called directly rather + than through that library, which cannot address M3's pool. See the + trtllm_gen_dense_decode module docstring for why. + """ + attn = params.attn + metadata = params.meta + head_dim = attn.head_dim + num_tokens = params.num_tokens + row_first = params.seq_offset + # 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_table, staged_factor = metadata.msa_subpage_rows( + row_first, row_first + metadata.num_generations + ) + if self._dense_layout is None: + raise RuntimeError( + "MiniMax-M3 dense decode ran before prepare_workspace sized its " + "scratch out of the shared attention workspace." + ) + workspace, counters = split_dense_decode_workspace(params.workspace, self._dense_layout) + minimax_m3_trtllm_gen_dense_decode( + params.attention_input.view(num_tokens, attn.num_heads, head_dim), + metadata.kv_cache_manager, + attn.layer_idx, + block_table, + seq_lens, + sm_scale=(head_dim**-0.5) / float(attn.q_scaling), + output=params.context_buf.view(num_tokens, attn.num_heads, head_dim), + decode_query_len=params.input_seq_length, + 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, + workspace=workspace, + counters=counters, + ) + + +__all__ = ["MsaDecodeFmha"] diff --git a/tensorrt_llm/_torch/attention/backends/fmha/msa_prefill.py b/tensorrt_llm/_torch/attention/backends/fmha/msa_prefill.py new file mode 100644 index 000000000000..a0a42b87a031 --- /dev/null +++ b/tensorrt_llm/_torch/attention/backends/fmha/msa_prefill.py @@ -0,0 +1,252 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""MiniMax-M3 context-phase FMHA backed by MSA's fmha_sm100 kernel. + +MsaPrefillFmha wraps the fmha_sm100 paged sparse GQA kernel and participates +in the standard TrtllmAttention.forward dispatch loop. The owning MiniMax-M3 +MSA attention layer runs an MsaIndexer to select the per-query KV blocks and +publishes them on forward_args.sparse_runtime_params; this class attends over +them. + +It serves the context phase alone. fmha_sm100 schedules a generation row like +a context row, and a mixed batch cannot split the two apart (see +_mixed_batch_split in fmha_sm100/api.py), so the generation phase is +MsaDecodeFmha's outright. +""" + +from __future__ import annotations + +from typing import TYPE_CHECKING, Optional + +import torch + +from ..sparse.minimax_m3.kernels.msa_utils import ( + MSA_REQUIRED_HEAD_DIM, + is_msa_layer, + msa_paged_kv, + require_msa_module, + write_msa_phase_kv, +) +from .interface import FmhaPhase +from .phased import FmhaParams, PhasedFmha + +if TYPE_CHECKING: + from tensorrt_llm._torch.attention.backends.interface import AttentionForwardArgs + from tensorrt_llm._torch.attention.backends.trtllm import ( + TrtllmAttention, + TrtllmAttentionMetadata, + ) + + +def run_msa_sparse_gqa( + q: torch.Tensor, + k_paged: torch.Tensor, + v_paged: torch.Tensor, + kv_block_indexes: Optional[torch.Tensor] = None, + *, + kv_indices: torch.Tensor, + sm_scale: float, + qo_lens_cpu: Optional[torch.Tensor] = None, + kv_lens_cpu: Optional[torch.Tensor] = None, + qo_offset_cpu: Optional[torch.Tensor] = None, + causal: bool = True, + head_dim: int = MSA_REQUIRED_HEAD_DIM, + plan: Optional[tuple] = None, + out: Optional[torch.Tensor] = None, + use_fp8: bool = False, +) -> None: + """Run fmha_sm100 paged GQA (plan/run split). + + `kv_block_indexes`: if set, sparse top-k mode (fixed `kv_block_num=topk`); + if None, dense mode attending all pages in `kv_indices`. + `plan`: prebuilt execution plan; if None, built inline from the CPU length + tensors (the prebuilt plan for a step prepare() staged, inline for a test). + `out`: destination buffer the kernel writes in place. + `use_fp8`: FP8 KV cache. The caller must pass FP8 `q` to match the FP8 paged + K/V, since the kernel variant shares one dtype across q/k/v. Also selects the + FP8 AOT kernels for an inline sparse-prefill plan. + """ + fmha_sm100 = require_msa_module() + + if q.dim() != 3: + raise ValueError( + f"MSA paged GQA expects q [total_q, num_qo_heads, head_dim]; got {tuple(q.shape)}." + ) + if q.shape[-1] != head_dim: + raise NotImplementedError(f"MSA paged GQA supports head_dim={head_dim}; got {q.shape[-1]}.") + if k_paged.dim() != 4 or v_paged.dim() != 4: + raise ValueError( + "MSA paged GQA expects paged KV [num_pages, num_kv_heads, page_size, head_dim]; " + f"got k={tuple(k_paged.shape)}, v={tuple(v_paged.shape)}." + ) + if k_paged.shape != v_paged.shape: + raise ValueError( + f"MSA paged GQA requires k and v to share shape; " + f"got k={tuple(k_paged.shape)}, v={tuple(v_paged.shape)}." + ) + + if plan is None: + # kv_block_num is planned only for the sparse (block-indexed) path; + # dense paged GQA leaves it unset and attends the full page table. + kv_block_num = int(kv_block_indexes.shape[-1]) if kv_block_indexes is not None else -1 + plan = fmha_sm100.fmha_sm100_plan( + qo_lens_cpu, + kv_lens_cpu, + int(q.shape[1]), # num query heads. + num_kv_heads=int(k_paged.shape[1]), + qo_offset=qo_offset_cpu, + page_size=int(k_paged.shape[2]), + kv_block_num=kv_block_num, + causal=causal, + num_kv_splits=1, + use_fp8_kvcache=use_fp8, + ) + fmha_sm100.fmha_sm100( + q, + k_paged, + v_paged, + plan, + kv_indices=kv_indices, + kv_block_indexes=kv_block_indexes, + out=out, + sm_scale=sm_scale, + output_maxscore=False, + ) + + +def run_msa_prefill_gqa( + attn: "TrtllmAttention", + q: torch.Tensor, + metadata: "TrtllmAttentionMetadata", + output: torch.Tensor, + *, + kv_block_indexes: Optional[torch.Tensor], + plan: Optional[tuple], + row_first: int, + num_rows: int, +) -> None: + """Run paged GQA over one row range into output in place. + + Shared by the sparse layers (kv_block_indexes is the per-query top-k table + for these rows, with the sparse plan) and the dense layers + (kv_block_indexes None, with the dense plan, attending the full page + table). + + `q` and `output` are already the phase's token slice. `row_first` and + `num_rows` are its batch rows, which narrow the host length tensors an + inline plan would read. + """ + head_dim = attn.head_dim + num_tokens = int(q.shape[0]) + if num_tokens == 0: + return + q_view = q.view(num_tokens, attn.num_heads, head_dim) + out_view = output.view(num_tokens, attn.num_heads, head_dim) + k_paged, v_paged = msa_paged_kv(metadata.kv_cache_manager, attn.layer_idx) + sm_scale = (head_dim**-0.5) / float(attn.q_scaling) + + # 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. + use_fp8 = k_paged.dtype == torch.float8_e4m3fn + if use_fp8: + q_view = q_view.to(torch.float8_e4m3fn) + + def rows_of(lens: Optional[torch.Tensor]) -> Optional[torch.Tensor]: + """Narrow a per-request host length tensor to this phase's rows. + + Slicing keeps the pinned backing, so an inline plan still stages these + with non-blocking copies. + """ + if lens is None: + return None + return lens[row_first : row_first + num_rows] + + run_msa_sparse_gqa( + q_view, + k_paged, + v_paged, + kv_block_indexes, + kv_indices=metadata.msa_kv_indices, + sm_scale=sm_scale, + 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, + use_fp8=use_fp8, + ) + + +class MsaPrefillFmha(PhasedFmha): + """SM100 paged GQA FMHA powered by MSA's fmha_sm100 kernel. + + Handles the context phase of every MiniMax-M3 MSA layer. Sparse layers pass + the indexer's selected KV block indices on + forward_args.sparse_runtime_params.sparse_attn_indices and attend those + blocks; dense layers leave the indices None and attend the full page table. + Requires head_dim 128 and 4-D HND paged K/V. + + The generation phase is MsaDecodeFmha's, so run_generation is left to the + base class, which refuses it. + """ + + @classmethod + def is_available(cls, attn: "TrtllmAttention") -> bool: + return is_msa_layer(attn) + + def is_supported( + self, + q: torch.Tensor, + k: Optional[torch.Tensor], + v: Optional[torch.Tensor], + metadata: "TrtllmAttentionMetadata", + forward_args: "AttentionForwardArgs", + *, + phase: Optional[FmhaPhase] = None, + ) -> bool: + # A step's context rows are this library's and its generation rows + # MsaDecodeFmha's, whatever the step looks like, so each library claims + # its one phase and their partition of the phases is total. The + # phase-less query asks for the whole step, which neither can serve + # alone: a mixed step belongs to the two together through CombinedFmha. + return phase is FmhaPhase.CONTEXT + + def run_context(self, params: FmhaParams) -> None: + metadata = params.meta + write_msa_phase_kv( + params.attn, + params.key_input, + params.value_input, + metadata, + params.fwd.attention_input_type, + token_offset=params.token_offset, + ) + # Sparse layers attend the per-query top-k blocks with the sparse plan; + # dense layers leave the indices None and attend the full page table + # with the dense plan. + kv_block_indexes = params.fwd.sparse_runtime_params.sparse_attn_indices + is_sparse_layer = kv_block_indexes is not None + if is_sparse_layer: + kv_block_indexes = kv_block_indexes[ + params.token_offset : params.token_offset + params.num_tokens + ] + run_msa_prefill_gqa( + params.attn, + params.attention_input, + metadata, + params.context_buf, + kv_block_indexes=kv_block_indexes, + plan=( + metadata.msa_prefill_gqa_plan + if is_sparse_layer + else metadata.msa_prefill_dense_plan + ), + row_first=params.seq_offset, + num_rows=metadata.num_contexts, + ) + + +__all__ = ["MsaPrefillFmha", "run_msa_prefill_gqa", "run_msa_sparse_gqa"] diff --git a/tensorrt_llm/_torch/attention/backends/fmha/msa_sparse_gqa.py b/tensorrt_llm/_torch/attention/backends/fmha/msa_sparse_gqa.py deleted file mode 100644 index ee85aaf02e0d..000000000000 --- a/tensorrt_llm/_torch/attention/backends/fmha/msa_sparse_gqa.py +++ /dev/null @@ -1,249 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 -"""Block-sparse GQA FMHA backed by MSA's fmha_sm100 kernel. - -MsaSparseGqaFmha wraps the fmha_sm100 paged sparse GQA kernel and -participates in the standard TrtllmAttention.forward dispatch loop. The -owning MiniMax-M3 MSA attention layer runs an MsaIndexer to select the -per-query KV blocks and publishes them on forward_args.sparse_runtime_params; -this class attends over them. -""" - -from __future__ import annotations - -from typing import TYPE_CHECKING, Optional - -import torch - -from tensorrt_llm._utils import is_sm_100f - -from .interface import Fmha - -if TYPE_CHECKING: - from tensorrt_llm._torch.attention.backends.interface import AttentionForwardArgs - from tensorrt_llm._torch.attention.backends.trtllm import ( - TrtllmAttention, - TrtllmAttentionMetadata, - ) - - -def run_msa_sparse_gqa( - q: torch.Tensor, - k_paged: torch.Tensor, - v_paged: torch.Tensor, - kv_block_indexes: Optional[torch.Tensor] = None, - *, - kv_indices: torch.Tensor, - sm_scale: float, - qo_lens_cpu: Optional[torch.Tensor] = None, - kv_lens_cpu: Optional[torch.Tensor] = None, - qo_offset_cpu: Optional[torch.Tensor] = None, - causal: bool = True, - head_dim: int = 128, - plan: Optional[tuple] = None, - out: Optional[torch.Tensor] = None, - use_fp8: bool = False, -) -> None: - """Run fmha_sm100 paged GQA (plan/run split). - - `kv_block_indexes`: if set, sparse top-k mode (fixed `kv_block_num=topk`); - if None, dense mode attending all pages in `kv_indices`. - `plan`: prebuilt execution plan; if None, built inline from the CPU length - tensors (eager prefill/tests vs. CUDA-graph decode). - `out`: destination buffer the kernel writes in place. - `use_fp8`: FP8 KV cache. The caller must pass FP8 `q` to match the FP8 paged - K/V, since the kernel variant shares one dtype across q/k/v. Also selects the - FP8 AOT kernels for an inline sparse-prefill plan; no-op for the decode planner. - """ - from tensorrt_llm._torch.attention.backends.sparse.minimax_m3.msa_utils import ( - require_msa_module, - ) - - fmha_sm100 = require_msa_module() - - if q.dim() != 3: - raise ValueError( - f"MsaSparseGqaFmha expects q [total_q, num_qo_heads, head_dim]; got {tuple(q.shape)}." - ) - if q.shape[-1] != head_dim: - raise NotImplementedError( - f"MsaSparseGqaFmha supports head_dim={head_dim}; got {q.shape[-1]}." - ) - if k_paged.dim() != 4 or v_paged.dim() != 4: - raise ValueError( - "MsaSparseGqaFmha expects paged KV [num_pages, num_kv_heads, page_size, head_dim]; " - f"got k={tuple(k_paged.shape)}, v={tuple(v_paged.shape)}." - ) - if k_paged.shape != v_paged.shape: - raise ValueError( - f"MsaSparseGqaFmha requires k and v to share shape; " - f"got k={tuple(k_paged.shape)}, v={tuple(v_paged.shape)}." - ) - - if plan is None: - # kv_block_num is planned only for the sparse (block-indexed) path; - # dense paged GQA leaves it unset and attends the full page table. - kv_block_num = int(kv_block_indexes.shape[-1]) if kv_block_indexes is not None else -1 - plan = fmha_sm100.fmha_sm100_plan( - qo_lens_cpu, - kv_lens_cpu, - int(q.shape[1]), # num query heads. - num_kv_heads=int(k_paged.shape[1]), - qo_offset=qo_offset_cpu, - page_size=int(k_paged.shape[2]), - kv_block_num=kv_block_num, - causal=causal, - num_kv_splits=1, - use_fp8_kvcache=use_fp8, - ) - fmha_sm100.fmha_sm100( - q, - k_paged, - v_paged, - plan, - kv_indices=kv_indices, - kv_block_indexes=kv_block_indexes, - out=out, - sm_scale=sm_scale, - output_maxscore=False, - ) - - -def run_msa_paged_gqa( - attn: "TrtllmAttention", - q: torch.Tensor, - k: Optional[torch.Tensor], - v: Optional[torch.Tensor], - metadata: "TrtllmAttentionMetadata", - output: torch.Tensor, - *, - kv_block_indexes: Optional[torch.Tensor], - plan: Optional[tuple], -) -> None: - """Write the new-token main K/V, then run paged GQA into output in place. - - Shared by the sparse layers (kv_block_indexes is the per-query top-k table, - 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. - """ - from tensorrt_llm._torch.attention.backends.sparse.minimax_m3.msa_utils import ( - msa_paged_kv, - write_msa_main_kv, - ) - - layer_idx = attn.layer_idx - head_dim = attn.head_dim - kv_cache_manager = metadata.kv_cache_manager - num_tokens = int(q.shape[0]) - if k is not None and v is not None: - write_msa_main_kv( - kv_cache_manager, layer_idx, metadata.msa_out_cache_loc[:num_tokens], k, v - ) - - q_view = q.view(num_tokens, attn.num_heads, head_dim) - out_view = output.view(num_tokens, attn.num_heads, head_dim) - k_paged, v_paged = msa_paged_kv(kv_cache_manager, layer_idx) - sm_scale = (head_dim**-0.5) / float(attn.q_scaling) - - # 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. - use_fp8 = k_paged.dtype == torch.float8_e4m3fn - if use_fp8: - q_view = q_view.to(torch.float8_e4m3fn) - - run_msa_sparse_gqa( - q_view, - k_paged, - v_paged, - kv_block_indexes, - 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, - causal=True, - head_dim=head_dim, - plan=plan, - out=out_view, - use_fp8=use_fp8, - ) - - -class MsaSparseGqaFmha(Fmha): - """SM100 paged GQA FMHA powered by MSA's fmha_sm100 kernel. - - Handles every MiniMax-M3 MSA layer. Sparse layers pass the indexer's - selected KV block indices on forward_args.sparse_runtime_params.sparse_attn_indices - 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. - """ - - @classmethod - def is_available(cls, attn: Optional["TrtllmAttention"] = None) -> bool: - if ( - attn is not None - and getattr(attn, "skip_correction_threshold", 0.0) > 0.0 - and not cls.supports_skip_correction - ): - return False - - # fmha_sm100 runs only on the SM100 family and is packaged in the - # wheel, so it is unavailable off SM100 or without the wheel. - # Imported lazily because the minimax_m3 package init imports the trtllm - # attention classes, which a module-scope import here would cycle with. - from tensorrt_llm._torch.attention.backends.sparse.minimax_m3.msa_utils import ( - msa_package_available, - ) - - if not is_sm_100f() or not msa_package_available(): - return False - # Only the MiniMax-M3 MSA layer uses this library. Matching the lowered - # sparse algorithm lets FmhaManager construction add it to that layer - # alone, so no custom library discovery is needed. - return attn.sparse_params is not None and attn.sparse_params.algorithm == "minimax_m3" - - def forward( - self, - q: torch.Tensor, - k: Optional[torch.Tensor], - v: Optional[torch.Tensor], - metadata: "TrtllmAttentionMetadata", - forward_args: "AttentionForwardArgs", - ) -> None: - output = forward_args.output - if output is None: - raise RuntimeError(f"{type(self).__name__} requires an output buffer.") - - # Sparse layers attend the per-query top-k blocks with the sparse plan; - # dense layers leave the indices None and attend the full page table - # with the dense plan. - kv_block_indexes = forward_args.sparse_runtime_params.sparse_attn_indices - if kv_block_indexes is not None: - plan = metadata.msa_decode_gqa_plan - if plan is None: - plan = getattr(metadata, "msa_eager_gqa_plan", None) - else: - plan = metadata.msa_decode_dense_plan - if plan is None: - plan = getattr(metadata, "msa_eager_dense_plan", None) - run_msa_paged_gqa( - self.attn, - q, - k, - v, - metadata, - output, - kv_block_indexes=kv_block_indexes, - plan=plan, - ) - - -__all__ = ["MsaSparseGqaFmha"] diff --git a/tensorrt_llm/_torch/attention/backends/fmha/phased.py b/tensorrt_llm/_torch/attention/backends/fmha/phased.py index 08b09177e330..560d3c73f344 100644 --- a/tensorrt_llm/_torch/attention/backends/fmha/phased.py +++ b/tensorrt_llm/_torch/attention/backends/fmha/phased.py @@ -55,6 +55,12 @@ class FmhaParams: cyclic_attention_window_size: int = 0 num_tokens: int = 0 seq_offset: int = 0 + # First query token of this phase on the axis of the q handed to the + # library, which covers the whole batch only where both phases share one + # tensor. The phase tensors above are already sliced by it; a library that + # indexes a separate per-token input, such as a sparse block table, needs + # it to take the matching slice. + token_offset: int = 0 tokens_per_block: int = 64 kv_factor: int = 0 total_num_blocks: int = 0 @@ -247,6 +253,7 @@ def forward( params.max_past_kv_length = max_past_kv_len params.num_tokens = num_ctx_tokens params.seq_offset = seq_offset + params.token_offset = token_offset params.input_seq_length = max_context_q_len params.batch_size = num_seqs params.num_requests = num_seqs @@ -290,6 +297,7 @@ def forward( params.max_past_kv_length = max_past_kv_len params.num_tokens = num_gen_tokens params.seq_offset = seq_offset + params.token_offset = token_offset params.input_seq_length = input_seq_length params.batch_size = num_seqs params.num_requests = num_seqs // metadata.beam_width diff --git a/tensorrt_llm/_torch/attention/backends/fmha/registry.py b/tensorrt_llm/_torch/attention/backends/fmha/registry.py index 4e545924e63e..fa4cf822d218 100644 --- a/tensorrt_llm/_torch/attention/backends/fmha/registry.py +++ b/tensorrt_llm/_torch/attention/backends/fmha/registry.py @@ -18,8 +18,11 @@ from .cute_dsl_mla import CuteDslMlaFmha from .fallback import FallbackFmha +from .flashinfer_sparse_mla import FlashInferSparseMlaFmha from .flashinfer_trtllm_gen import FlashInferTrtllmGenFmha from .interface import Fmha +from .msa_decode import MsaDecodeFmha +from .msa_prefill import MsaPrefillFmha from .prims_ts import PrimsTSFmha from .triton_custom_mask import TritonCustomMaskFmha @@ -27,19 +30,15 @@ def init_fmha_libs() -> dict[str, "FmhaCls"]: - """Build the ordered FMHA library registry. - - Backend classes are imported inside this factory rather than at module - scope, so backends can import trtllm attention classes at module scope - without an import cycle. - """ - from .flashinfer_sparse_mla import FlashInferSparseMlaFmha - from .msa_sparse_gqa import MsaSparseGqaFmha - + """Build the ordered FMHA library registry.""" return { "triton_custom_mask": TritonCustomMaskFmha, "cute_dsl_mla": CuteDslMlaFmha, - "msa_sparse_gqa": MsaSparseGqaFmha, + # A pair, not a priority: msa_decode serves a MiniMax-M3 MSA layer's + # generation phase and msa_prefill its context phase, and neither will + # take the other's. + "msa_decode": MsaDecodeFmha, + "msa_prefill": MsaPrefillFmha, "flashinfer_sparse_mla": FlashInferSparseMlaFmha, "prims_ts": PrimsTSFmha, "flashinfer_trtllm_gen": FlashInferTrtllmGenFmha, diff --git a/tensorrt_llm/_torch/attention/backends/sparse/minimax_m3/__init__.py b/tensorrt_llm/_torch/attention/backends/sparse/minimax_m3/__init__.py index e58370c410cb..d7c144884b42 100644 --- a/tensorrt_llm/_torch/attention/backends/sparse/minimax_m3/__init__.py +++ b/tensorrt_llm/_torch/attention/backends/sparse/minimax_m3/__init__.py @@ -16,10 +16,6 @@ KV-slot writer, block-priority sentinels, and the paged-cache slot mapping builder shared by both backends. - * :mod:`.msa_utils` -- MSA-only (fmha_sm100) helpers: import guard, - kernel precondition constants, HND paged-cache - adapters, main-KV writer, page-table builder, - valid-block counting, and top-k selection. * :mod:`.triton_kernels` -- OpenAI Triton kernels (per-block max score, masked softmax for sparse GQA). * :mod:`.triton_backend` -- the Triton reference algorithm (vectorized @@ -35,24 +31,41 @@ selection submodule. * :mod:`.msa_availability`-- SM100 and fmha_sm100 gating for the MSA path. + * :mod:`.kernels` -- the kernels themselves and the paged-cache + write they share with the backends, imported + directly by the FMHA libraries that drive + them. This package's public surface re-exports the names callers historically imported from ``...sparse.minimax_m3`` so external importers (the model code, ``sparse.utils``, focused tests) keep -working unchanged. +working unchanged. The re-export is lazy (PEP 562 __getattr__) so that +importing :mod:`.kernels` from an FMHA library does not pull in +:mod:`.msa_backend`, which subclasses TrtllmAttention and would close a cycle +back through the FMHA registry. """ -# The dense Triton oracle in the model imports these paged-cache helpers, so -# they stay importable from the package. They are package-private and are not +import importlib + +# The dense Triton oracle in the model imports the two paged-cache helpers, so +# they stay reachable from the package. They are package-private and are not # part of __all__. Every other backend/metadata/config symbol is imported # directly from its defining submodule by the code that needs it. -from .cache_manager import MiniMaxM3KVCacheManagerV2 -from .msa_backend import MiniMaxM3MsaSparseAttention -from .triton_backend import ( - MiniMaxM3SparseRuntimeBackend, - _gather_paged_batched, - _write_main_kv_slots_to_pool, -) +_LAZY_EXPORTS = { + "MiniMaxM3KVCacheManagerV2": ".cache_manager", + "MiniMaxM3MsaSparseAttention": ".msa_backend", + "MiniMaxM3SparseRuntimeBackend": ".triton_backend", + "_gather_paged_batched": ".triton_backend", + "_write_main_kv_slots_to_pool": ".triton_backend", +} + + +def __getattr__(name: str): + module = _LAZY_EXPORTS.get(name) + if module is None: + raise AttributeError(f"module {__name__!r} has no attribute {name!r}") + return getattr(importlib.import_module(module, __name__), name) + __all__ = [ "MiniMaxM3KVCacheManagerV2", diff --git a/tensorrt_llm/_torch/attention/backends/sparse/minimax_m3/cache_manager.py b/tensorrt_llm/_torch/attention/backends/sparse/minimax_m3/cache_manager.py index 043acaaaa73b..e65b168c5d14 100644 --- a/tensorrt_llm/_torch/attention/backends/sparse/minimax_m3/cache_manager.py +++ b/tensorrt_llm/_torch/attention/backends/sparse/minimax_m3/cache_manager.py @@ -25,7 +25,7 @@ from __future__ import annotations -from typing import List, Optional, Sequence +from typing import List, Optional, Sequence, Tuple import torch @@ -308,24 +308,15 @@ 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( + def _kv_slot_geometry( 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.) + ) -> Tuple[int, torch.dtype, int, int, List[int]]: + """Resolve one layer's position in the coalesced K/V pool. - 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``. + Returns (addr_key, torch_dtype, num_slots, scale, page_shape), where + scale is the number of equal-sized sub-pages a slot packs and + page_shape is one sub-page's shape in kv_layout. This layer's K is + sub-page 0 and its V sub-page 1, counting from addr_key. When omitted, ``kv_layout`` follows the selected sparse backend. """ if kv_layout is None: @@ -334,7 +325,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] @@ -345,13 +336,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}." ) @@ -378,27 +369,66 @@ 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``. + """ + 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/backends/sparse/minimax_m3/common.py b/tensorrt_llm/_torch/attention/backends/sparse/minimax_m3/common.py index 94c2c7be7974..5151fa2d4b6d 100644 --- a/tensorrt_llm/_torch/attention/backends/sparse/minimax_m3/common.py +++ b/tensorrt_llm/_torch/attention/backends/sparse/minimax_m3/common.py @@ -4,8 +4,8 @@ Both the Triton reference and the MSA (fmha_sm100) path share these backend-neutral pieces: the lowered parameter and per-rank kernel config -bundles, block-priority sentinels, KV-slot writers, and the paged-cache -slot mapping builder. MSA-only helpers live in :mod:`.msa_utils`. +bundles, block-priority sentinels, and the paged-cache slot mapping builder. +MSA-only helpers live in :mod:`.kernels.msa_utils`. """ from __future__ import annotations @@ -19,6 +19,10 @@ from ..params import SparseMetadataParams, SparseParams +# Re-exported so this module stays the one place the backends reach for +# paged-cache plumbing, wherever it is defined. +from .kernels.paged_cache import write_kv_slots + if TYPE_CHECKING: from tensorrt_llm.mapping import Mapping @@ -159,43 +163,6 @@ def from_sparse_params( ) -def write_kv_slots( - cache: torch.Tensor, - out_cache_loc: torch.Tensor, - values: torch.Tensor, - *, - layout: Literal["NHD", "HND"] = "NHD", -) -> None: - """Write per-token values into a K, V, or index-K cache at given slots. - - Handles a 3-D flat-slot cache and a 4-D paged view. `layout` sets the paged - axis order: "NHD" is [num_pages, tokens_per_block, num_heads, channel], - "HND" is [num_pages, num_heads, tokens_per_block, channel]. The paged view - is non-contiguous, so the slot id is split into (page, within) and written - by multi-dim assignment. `values` is always [num_tokens, num_heads, channel]. - - Callers must provide valid slots for every live token. The production M3 - mapping satisfies this contract: ``get_block_ids_per_seq`` canonicalizes - padded ``BAD_PAGE_INDEX`` entries before ``build_paged_kv_slot_mapping`` - selects only the allocated live-token positions. - """ - with torch.no_grad(): - if cache.ndim >= 4: - token_axis = 2 if layout == "HND" else 1 - tokens_per_block = int(cache.shape[token_axis]) - out_long = out_cache_loc.to(torch.long) - page = out_long // tokens_per_block - within = out_long % tokens_per_block - if layout == "HND": - # Advanced indices on dims 0 and 2 broadcast to [num_tokens] and - # move front, giving a [num_tokens, num_heads, channel] target. - cache[page, :, within, :] = values.to(cache.dtype) - else: - cache[page, within] = values.to(cache.dtype) - else: - cache.index_copy_(0, out_cache_loc.to(torch.long), values.to(cache.dtype)) - - class PagedKvSlotMapping(NamedTuple): """One step's paged-cache slot mapping (see build_paged_kv_slot_mapping).""" diff --git a/tensorrt_llm/_torch/attention/backends/sparse/minimax_m3/kernels/__init__.py b/tensorrt_llm/_torch/attention/backends/sparse/minimax_m3/kernels/__init__.py new file mode 100644 index 000000000000..f9351009f2a4 --- /dev/null +++ b/tensorrt_llm/_torch/attention/backends/sparse/minimax_m3/kernels/__init__.py @@ -0,0 +1,11 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""MiniMax-M3 attention kernels and the cache plumbing they need. + +Nothing here may import a sibling module of the parent package: those modules +import attention.backends.trtllm, so reaching them from an FMHA library would +close an import cycle. + +Submodules are not re-exported: triton_sparse_decode pulls in Triton, and this +package sits on the import path of every attention.backends.trtllm import. +""" diff --git a/tensorrt_llm/_torch/attention/backends/sparse/minimax_m3/msa_utils.py b/tensorrt_llm/_torch/attention/backends/sparse/minimax_m3/kernels/msa_utils.py similarity index 80% rename from tensorrt_llm/_torch/attention/backends/sparse/minimax_m3/msa_utils.py rename to tensorrt_llm/_torch/attention/backends/sparse/minimax_m3/kernels/msa_utils.py index 6bde824d04d8..28716a6dece5 100644 --- a/tensorrt_llm/_torch/attention/backends/sparse/minimax_m3/msa_utils.py +++ b/tensorrt_llm/_torch/attention/backends/sparse/minimax_m3/kernels/msa_utils.py @@ -11,9 +11,10 @@ import torch -from tensorrt_llm._utils import maybe_pin_memory +from tensorrt_llm._torch.attention.backends.interface import AttentionInputType +from tensorrt_llm._utils import is_sm_100f, maybe_pin_memory -from .common import write_kv_slots +from .paged_cache import write_kv_slots # fmha_sm100 ships only head_dim 128 variants and the MiniMax-M3 checkpoint # selects topk 16. Callers enforce these early so a misconfiguration fails @@ -22,6 +23,17 @@ MSA_REQUIRED_HEAD_DIM = 128 +def is_msa_layer(attn) -> bool: + """Whether this layer's attention is served by the MiniMax-M3 MSA kernels.""" + sparse_params = attn.sparse_params + return ( + is_sm_100f() + and sparse_params is not None + and sparse_params.algorithm == "minimax_m3" + and getattr(sparse_params, "implementation", None) == "msa" + ) + + def _install_msa_cutlass_compatibility() -> None: """Provide legacy CuTe aliases still referenced by the packaged MSA sources.""" try: @@ -111,6 +123,45 @@ def write_msa_main_kv( ) +def write_msa_phase_kv( + attn, + k: Optional[torch.Tensor], + v: Optional[torch.Tensor], + metadata, + attention_input_type, + *, + token_offset: int, +) -> None: + """Write one phase's new-token main K/V for one layer. + + The decode kernels and fmha_sm100 read the paged cache directly, so a + phase's new-token K/V has to be resident before its kernel runs. Each FMHA + library writes the rows it is about to attend, so a mixed batch's two + phases cover the step between them and neither repeats the other's write. + + k and v are the phase's token slice, and token_offset its first token on + the step's token axis, which is what msa_out_cache_loc is indexed by. + """ + if attention_input_type != AttentionInputType.mixed: + raise NotImplementedError( + "MiniMax-M3 MSA attention requires the mixed attention input type, " + f"but this call passed {attention_input_type!r}. Its cache slots and " + "top-k block table are indexed by whole-step token position." + ) + if k is None or v is None: + return + num_tokens = int(k.shape[0]) + if num_tokens == 0: + return + write_msa_main_kv( + metadata.kv_cache_manager, + attn.layer_idx, + metadata.msa_out_cache_loc[token_offset : token_offset + num_tokens], + k, + v, + ) + + def build_kv_page_indices( block_ids_cpu: torch.Tensor, kv_lens_cpu: torch.Tensor, @@ -219,4 +270,5 @@ def select_blocks_from_maxscore( "require_msa_module", "select_blocks_from_maxscore", "write_msa_main_kv", + "write_msa_phase_kv", ] diff --git a/tensorrt_llm/_torch/attention/backends/sparse/minimax_m3/kernels/paged_cache.py b/tensorrt_llm/_torch/attention/backends/sparse/minimax_m3/kernels/paged_cache.py new file mode 100644 index 000000000000..3a20070d56c3 --- /dev/null +++ b/tensorrt_llm/_torch/attention/backends/sparse/minimax_m3/kernels/paged_cache.py @@ -0,0 +1,49 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Paged-cache writes shared by the MiniMax-M3 kernels and backends.""" + +from __future__ import annotations + +from typing import Literal + +import torch + + +def write_kv_slots( + cache: torch.Tensor, + out_cache_loc: torch.Tensor, + values: torch.Tensor, + *, + layout: Literal["NHD", "HND"] = "NHD", +) -> None: + """Write per-token values into a K, V, or index-K cache at given slots. + + Handles a 3-D flat-slot cache and a 4-D paged view. `layout` sets the paged + axis order: "NHD" is [num_pages, tokens_per_block, num_heads, channel], + "HND" is [num_pages, num_heads, tokens_per_block, channel]. The paged view + is non-contiguous, so the slot id is split into (page, within) and written + by multi-dim assignment. `values` is always [num_tokens, num_heads, channel]. + + Callers must provide valid slots for every live token. The production M3 + mapping satisfies this contract: ``get_block_ids_per_seq`` canonicalizes + padded ``BAD_PAGE_INDEX`` entries before ``build_paged_kv_slot_mapping`` + selects only the allocated live-token positions. + """ + with torch.no_grad(): + if cache.ndim >= 4: + token_axis = 2 if layout == "HND" else 1 + tokens_per_block = int(cache.shape[token_axis]) + out_long = out_cache_loc.to(torch.long) + page = out_long // tokens_per_block + within = out_long % tokens_per_block + if layout == "HND": + # Advanced indices on dims 0 and 2 broadcast to [num_tokens] and + # move front, giving a [num_tokens, num_heads, channel] target. + cache[page, :, within, :] = values.to(cache.dtype) + else: + cache[page, within] = values.to(cache.dtype) + else: + cache.index_copy_(0, out_cache_loc.to(torch.long), values.to(cache.dtype)) + + +__all__ = ["write_kv_slots"] diff --git a/tensorrt_llm/_torch/attention/backends/sparse/minimax_m3/triton_sparse_decode.py b/tensorrt_llm/_torch/attention/backends/sparse/minimax_m3/kernels/triton_sparse_decode.py similarity index 97% rename from tensorrt_llm/_torch/attention/backends/sparse/minimax_m3/triton_sparse_decode.py rename to tensorrt_llm/_torch/attention/backends/sparse/minimax_m3/kernels/triton_sparse_decode.py index 107e5102717e..4e00b384b857 100644 --- a/tensorrt_llm/_torch/attention/backends/sparse/minimax_m3/triton_sparse_decode.py +++ b/tensorrt_llm/_torch/attention/backends/sparse/minimax_m3/kernels/triton_sparse_decode.py @@ -325,6 +325,14 @@ def minimax_m3_sparse_attn_decode( f"MiniMax-M3 sparse decode requires page_size={SPARSE_BLOCK_SIZE}; " f"got {int(k_paged.shape[2])}." ) + if num_heads % num_kv_heads: + # The kernel covers heads pid_kh * gqa_group_size + [0, group), so a + # head count that is not a whole number of groups leaves the tail heads + # unwritten and the merge kernel reads uninitialized partials. + raise ValueError( + f"MiniMax-M3 sparse decode requires num_heads ({num_heads}) to be a " + f"multiple of num_kv_heads ({num_kv_heads})." + ) 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 diff --git a/tensorrt_llm/_torch/attention/backends/sparse/minimax_m3/trtllm_gen_dense_decode.py b/tensorrt_llm/_torch/attention/backends/sparse/minimax_m3/kernels/trtllm_gen_dense_decode.py similarity index 61% rename from tensorrt_llm/_torch/attention/backends/sparse/minimax_m3/trtllm_gen_dense_decode.py rename to tensorrt_llm/_torch/attention/backends/sparse/minimax_m3/kernels/trtllm_gen_dense_decode.py index 7e395088fb0b..b477dc06a8d7 100644 --- a/tensorrt_llm/_torch/attention/backends/sparse/minimax_m3/trtllm_gen_dense_decode.py +++ b/tensorrt_llm/_torch/attention/backends/sparse/minimax_m3/kernels/trtllm_gen_dense_decode.py @@ -3,9 +3,9 @@ """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. +they only run there because MsaPrefillFmha 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 @@ -21,7 +21,7 @@ from __future__ import annotations import functools -from typing import Optional +from typing import NamedTuple, Optional, Tuple import torch @@ -45,6 +45,11 @@ def _counter_size(num_heads: int, max_num_requests: int, device_index: int) -> i return int(_get_multi_ctas_kv_counter_size(num_heads, max_num_requests, multi_processor_count)) +def _device_index(device: torch.device) -> int: + """Device ordinal of device, resolving the current one where it has none.""" + return device.index if device.index is not None else torch.cuda.current_device() + + def _counter_buffer( device: torch.device, num_heads: int, max_num_requests: int, reserve: bool ) -> torch.Tensor: @@ -55,9 +60,8 @@ def _counter_buffer( 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)], + [_counter_size(num_heads, max_num_requests, _device_index(device))], torch.uint8, buffer_name="m3_trtllm_gen_kv_counters", reserve_buffer=reserve, @@ -83,6 +87,88 @@ def _workspace(q_dtype: torch.dtype, num_heads: int, head_dim: int, num_kv_heads return int(layout["trtllm_gen_workspace_size"]) +def _workspace_buffer(size: int, reserve: bool) -> torch.Tensor: + return get_memory_buffers().get_buffer( + [size], + torch.uint8, + buffer_name="m3_trtllm_gen_workspace", + reserve_buffer=reserve, + ) + + +class DenseDecodeWorkspaceLayout(NamedTuple): + """Byte layout of the scratch one dense decode call needs. + + The trtllm-gen slab sits at the front of the buffer and the multi-CTA KV + counters behind it, so a caller with one byte buffer can serve both. + """ + + workspace_bytes: int + # Aligned start of the counter block. + counter_offset: int + total_bytes: int + + +# The alignment torch's own allocations start at, so a counter view taken at +# this offset is as aligned as a standalone buffer would be. +_COUNTER_OFFSET_ALIGNMENT = 256 + + +def dense_decode_workspace_layout( + *, + q_dtype: torch.dtype, + num_heads: int, + head_dim: int, + num_kv_heads: int, + max_num_requests: int, + device: torch.device, +) -> DenseDecodeWorkspaceLayout: + """Lay out the scratch a dense decode call needs inside one buffer. + + Both parts are sized by the head geometry and the request bound alone, + none of which move during a run, so a caller that knows those before the + phase runs can size its buffer there. See split_dense_decode_workspace for + the views this describes. + """ + workspace_bytes = _workspace(q_dtype, num_heads, head_dim, num_kv_heads) + counter_bytes = _counter_size(num_heads, max_num_requests, _device_index(device)) + counter_offset = ( + (workspace_bytes + _COUNTER_OFFSET_ALIGNMENT - 1) + // _COUNTER_OFFSET_ALIGNMENT + * _COUNTER_OFFSET_ALIGNMENT + ) + return DenseDecodeWorkspaceLayout( + workspace_bytes=workspace_bytes, + counter_offset=counter_offset, + total_bytes=counter_offset + counter_bytes, + ) + + +def split_dense_decode_workspace( + buffer: torch.Tensor, layout: DenseDecodeWorkspaceLayout +) -> Tuple[torch.Tensor, torch.Tensor]: + """Take the slab and the zeroed KV counters out of one byte buffer. + + The counters are zeroed on every call because the buffer is shared with + the other FMHA libraries, any of which may have written it. + """ + if buffer.dim() != 1 or buffer.element_size() != 1: + raise ValueError( + "MiniMax-M3 dense decode scratch must come from a flat byte buffer; " + f"got {tuple(buffer.shape)} of {buffer.dtype}." + ) + if int(buffer.numel()) < layout.total_bytes: + raise ValueError( + f"MiniMax-M3 dense decode was handed a {int(buffer.numel())}-byte " + f"buffer, but this call needs {layout.total_bytes}. It was sized " + "from a different head geometry than the kernel sees." + ) + workspace_bytes = buffer.view(torch.uint8) + counters = workspace_bytes[layout.counter_offset : layout.total_bytes] + counters.zero_() + return workspace_bytes[: layout.workspace_bytes], counters + + def subpage_block_table( block_table: torch.Tensor, subpages_per_slot: int, reserve: bool = False ) -> torch.Tensor: @@ -148,6 +234,8 @@ def minimax_m3_trtllm_gen_dense_decode( max_num_requests: int, staged_subpage_table: Optional[torch.Tensor] = None, staged_subpages_per_slot: int = 0, + workspace: Optional[torch.Tensor] = None, + counters: Optional[torch.Tensor] = None, enable_pdl: bool = True, ) -> None: """Full-context decode attention through trtllm-gen, in place into output. @@ -155,6 +243,12 @@ def minimax_m3_trtllm_gen_dense_decode( 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. + + workspace and counters are the scratch a caller already carved out of its + own buffer (see split_dense_decode_workspace); both are taken from the + arena here when a caller has not, as the standalone kernel tests do. A + supplied buffer is checked against the size this call needs, since the + caller sizes it from a geometry it derives for itself. """ import flashinfer @@ -168,12 +262,24 @@ def minimax_m3_trtllm_gen_dense_decode( 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, - ) + workspace_size = _workspace(q.dtype, num_heads, int(q.shape[2]), int(kv_pool.shape[1])) + if workspace is None: + workspace = _workspace_buffer(workspace_size, reserve) + elif int(workspace.numel()) < workspace_size: + raise ValueError( + f"MiniMax-M3 dense decode was handed a {int(workspace.numel())}-byte " + f"trtllm-gen workspace, but this call needs {workspace_size}. It was " + "reserved from a different head geometry than the kernel sees." + ) + counter_size = _counter_size(num_heads, max_num_requests, _device_index(q.device)) + if counters is None: + counters = _counter_buffer(q.device, num_heads, max_num_requests, reserve) + elif int(counters.numel()) < counter_size: + raise ValueError( + f"MiniMax-M3 dense decode was handed {int(counters.numel())} bytes of " + f"multi-CTA KV counters, but this call needs {counter_size}. They were " + "reserved from a different head or request bound than the kernel sees." + ) 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) @@ -199,18 +305,16 @@ def minimax_m3_trtllm_gen_dense_decode( cum_seq_lens_q=None, kv_cache_sf=None, # M3 stores unscaled E4M3 uses_shared_paged_kv_idx=False, - multi_ctas_kv_counter_buffer=_counter_buffer( - q.device, num_heads, max_num_requests, reserve - ), + multi_ctas_kv_counter_buffer=counters, ) @functools.lru_cache(maxsize=1) -def _flashinfer_available() -> bool: +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. + The verdict below is consulted by ensure_msa_available and again per dense + layer, so the import statement would otherwise be on the hot path. """ try: import flashinfer # noqa: F401 @@ -222,21 +326,26 @@ def _flashinfer_available() -> bool: 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. + Takes head_dim rather than a query tensor so the metadata's up-front + validation 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(): + if not flashinfer_available(): return "flashinfer is not installed." return None __all__ = [ + "DenseDecodeWorkspaceLayout", "dense_decode_unsupported_reason", + "dense_decode_workspace_layout", + "flashinfer_available", "minimax_m3_trtllm_gen_dense_decode", + "split_dense_decode_workspace", "subpage_block_table", "uniform_subpages_per_slot", "write_subpage_block_table", diff --git a/tensorrt_llm/_torch/attention/backends/sparse/minimax_m3/msa_availability.py b/tensorrt_llm/_torch/attention/backends/sparse/minimax_m3/msa_availability.py index 6e423da224d2..f76443e2e1ee 100644 --- a/tensorrt_llm/_torch/attention/backends/sparse/minimax_m3/msa_availability.py +++ b/tensorrt_llm/_torch/attention/backends/sparse/minimax_m3/msa_availability.py @@ -2,17 +2,19 @@ # SPDX-License-Identifier: Apache-2.0 """Availability checks for the MiniMax-M3 MSA sparse attention kernels. -The MSA kernels are provided by the fmha_sm100 package bundled with -TensorRT-LLM and run only on the SM100 architecture family (SM100 and SM103). -These helpers gate backend selection so a request for the MSA path fails early -with a clear message on unsupported systems. +The MSA path runs prefill on the fmha_sm100 kernels bundled with TensorRT-LLM +and decode on the CuTe DSL indexer scorer plus the Triton and trtllm-gen decode +kernels. Both halves are required: there is no fallback from one to the other, +so a missing package makes the path unavailable rather than slower. These +helpers gate backend selection so that fails early with a clear message. """ from __future__ import annotations from tensorrt_llm._utils import get_sm_version, is_sm_100f -from .msa_utils import msa_package_available +from .kernels.msa_utils import msa_package_available +from .kernels.trtllm_gen_dense_decode import flashinfer_available # fmha_sm100 runs on the SM100 architecture family (SM100 and SM103). Other # architectures, including SM120, are not supported. @@ -21,6 +23,10 @@ def ensure_msa_available() -> None: """Raise RuntimeError if the MSA sparse attention path cannot run here.""" + # Function-local: msa_indexer reaches the trtllm attention classes through + # this package's init, which a module-scope import here would cycle with. + from .msa_indexer import cutedsl_score_runner + if not msa_package_available(): raise RuntimeError( f"MiniMax-M3 MSA sparse attention requires the {MSA_PACKAGE} kernels " @@ -32,6 +38,18 @@ def ensure_msa_available() -> None: "MiniMax-M3 MSA sparse attention requires an SM100 or SM103 device, " f"but the current device reports SM version {sm_version}." ) + if cutedsl_score_runner() is None: + raise RuntimeError( + "MiniMax-M3 MSA sparse attention scores decode steps on the CuTe DSL " + "indexer kernel, which requires the nvidia-cutlass-dsl package. " + "Install it or select the 'triton' implementation." + ) + if not flashinfer_available(): + raise RuntimeError( + "MiniMax-M3 MSA sparse attention runs its dense layers through the " + "trtllm-gen decode kernel, which requires flashinfer. Install it or " + "select the 'triton' implementation." + ) __all__ = [ diff --git a/tensorrt_llm/_torch/attention/backends/sparse/minimax_m3/msa_backend.py b/tensorrt_llm/_torch/attention/backends/sparse/minimax_m3/msa_backend.py index 93b0d9af8f5d..bc55aa6cf51f 100644 --- a/tensorrt_llm/_torch/attention/backends/sparse/minimax_m3/msa_backend.py +++ b/tensorrt_llm/_torch/attention/backends/sparse/minimax_m3/msa_backend.py @@ -5,7 +5,8 @@ * MiniMaxM3MsaSparseAttention subclasses TrtllmAttention and reuses its inherited forward, overriding only the sparse hooks and owning an MsaIndexer. - * The main sparse GQA runs through the registered MsaSparseGqaFmha. + * The main attention runs through the registered MsaPrefillFmha for context + rows and MsaDecodeFmha for generation rows. * The indexer calls fmha_sm100 directly to produce the per-query selected block indices, which the model layer threads through forward_args.sparse_backend_args. @@ -18,15 +19,15 @@ needed here. The classes subclass TrtllmAttention and TrtllmAttentionMetadata, imported at -module scope. This is cycle-free because the fmha registry defers its -MsaSparseGqaFmha import (see fmha/registry.py), so trtllm's import chain does -not reach this module. +module scope. That is cycle-free only because the dependency runs one way: the +two FMHA libraries reach the kernels through ...minimax_m3.kernels and never +load this module, so trtllm's import chain does not come back here. """ from __future__ import annotations from dataclasses import dataclass -from typing import Optional, Tuple +from typing import NamedTuple, Optional, Tuple import torch @@ -34,6 +35,7 @@ from tensorrt_llm._torch.attention.backends.trtllm import TrtllmAttention, TrtllmAttentionMetadata from tensorrt_llm._utils import maybe_pin_memory from tensorrt_llm.bindings import DataType +from tensorrt_llm.models.modeling_utils import QuantConfig from .common import ( MiniMaxM3SparseConfig, @@ -41,14 +43,19 @@ build_paged_kv_slot_mapping, write_kv_slots, ) -from .msa_indexer import MsaIndexer -from .msa_utils import ( +from .kernels.msa_utils import ( MSA_REQUIRED_HEAD_DIM, MSA_REQUIRED_TOPK, build_kv_page_indices, per_token_valid_blocks, require_msa_module, ) +from .kernels.trtllm_gen_dense_decode import ( + dense_decode_unsupported_reason, + uniform_subpages_per_slot, + write_subpage_block_table, +) +from .msa_indexer import MsaIndexer, cutedsl_score_runner def _cache_device(meta) -> torch.device: @@ -57,7 +64,8 @@ def _cache_device(meta) -> torch.device: if kv_cache_manager is not None: try: return kv_cache_manager.get_buffers(0).device - except Exception: + except (AttributeError, IndexError, KeyError): + # A manager that exposes no layer-0 buffer, as in a focused test. pass return torch.device(f"cuda:{torch.cuda.current_device()}") @@ -89,140 +97,43 @@ def _worst_case_proxy_max_k_tiles( return int(proxy_plan[3]["max_k_tiles"]) -# Per-step fmha_sm100 plan tensors that must live in CUDA-graph-stable buffers. -# At num_kv_splits=1 the plan carries no split-KV workspaces, and -# cute_workspace_buffer is the vendor's cached scratch (kept by reference, not -# copied). -_MSA_PLAN_STABLE_KEYS = ( - "packed_work_range", - "packed_work_info", - "qo_segment_offsets", - "kv_segment_offsets", - "kv_page_indptr", - "qo_segment_lens", - "kv_segment_lens", - "qo_offset", -) -_MSA_PLAN_INT64_KEYS = ("packed_work_range", "packed_work_info") -# fmha_sm100 sizes packed_work_info at 131072 * max(num_kv_splits, 1); forcing -# num_kv_splits=1 pins this worklist width. -_MSA_PACKED_WORK_INFO_LEN = 131072 -_MSA_SPLIT_KV_KEYS = ( - "kv_tile_begin_indices", - "kv_tile_end_indices", - "kv_split_indices", - "num_kv_splits_per_row", - "workspace_o", - "workspace_lse", -) - - -class _MsaGraphSafePlan: - """CUDA-graph-stable mirror of one fmha_sm100 decode plan. +class MsaDecodeSpan(NamedTuple): + """Generation rows served by MiniMax-M3's dedicated decode kernels. - Owns fixed device buffers for the per-step plan worklists. refresh() copies - a freshly built plan into them and returns a plan tuple pointing at the - stable buffers, so the captured fmha_sm100 run reads addresses that do not - change across replays. Mirrors FlashInfer's fixed indptr/indices buffers. - - Only valid at num_kv_splits=1: the plan then has no split-KV workspaces - (refresh() asserts this), and cute_workspace_buffer and the scalar fields - pass through unchanged. + A named carrier for the pair each of them indexes by, in the style of + PagedKvSlotMapping in common.py. See msa_decode_span. """ - def __init__(self, metadata, name: str, *, max_batch: int, num_ctas: int, capture_graph: bool): - buffers = metadata.cuda_graph_buffers - self._buf = {} - # Set by refresh(), read through the plan property. - self._plan: Optional[tuple] = None - # cute_workspace_buffer must keep a fixed address across steps for the - # captured graph to replay correctly. Pin it on first use and fail if - # it moves. - self._ws_ptr: Optional[int] = None - for key in _MSA_PLAN_STABLE_KEYS: - if key == "packed_work_range": - shape = (num_ctas,) - elif key == "packed_work_info": - shape = (_MSA_PACKED_WORK_INFO_LEN,) - elif key in ("qo_segment_offsets", "kv_segment_offsets", "kv_page_indptr"): - shape = (max_batch + 1,) - else: - shape = (max_batch,) - dtype = torch.int64 if key in _MSA_PLAN_INT64_KEYS else torch.int32 - self._buf[key] = metadata.get_empty( - buffers, - shape, - cache_name=f"{name}_{key}", - dtype=dtype, - capture_graph=capture_graph, - ) - - @property - def plan(self) -> Optional[tuple]: - """The current graph-safe plan tuple, or None if no decode plan is live.""" - return self._plan - - def reset(self) -> None: - """Drop the live plan tuple (e.g. for a prefill/mixed or captured step).""" - self._plan = None - - def refresh(self, plan_tuple) -> tuple: - has_mixed, split, batch, decode, prefill = plan_tuple - if has_mixed: - raise RuntimeError( - "MSA decode expects a single (non-mixed) fmha_sm100 plan; a decode " - "batch must be pure decode." - ) - for key in _MSA_SPLIT_KV_KEYS: - if decode.get(key) is not None: - raise RuntimeError( - f"MSA decode plan used split-KV workspace {key!r}; num_kv_splits=1 " - "is required for graph-safe decode." - ) - ws = decode.get("cute_workspace_buffer") - if ws is not None: - if self._ws_ptr is None: - self._ws_ptr = ws.data_ptr() - elif ws.data_ptr() != self._ws_ptr: - raise RuntimeError( - "cute_workspace_buffer moved across steps; the fmha_sm100 plan " - "is not CUDA-graph safe." - ) - rebuilt = dict(decode) - for key in _MSA_PLAN_STABLE_KEYS: - src = decode.get(key) - if src is None: - continue - n = int(src.shape[0]) - dst = self._buf[key] - if n > dst.shape[0]: - raise ValueError( - f"MSA plan buffer {key} ({dst.shape[0]}) is smaller than the plan tensor ({n})." - ) - dst[:n].copy_(src, non_blocking=True) - rebuilt[key] = dst[:n] - self._plan = (has_mixed, split, batch, rebuilt, prefill) - return self._plan + # First generation row, which is also the step's context row count. + row_first: int + # Uniform query token count per generation request. + query_len: int @dataclass(init=False) class MiniMaxM3MsaSparseAttentionMetadata(TrtllmAttentionMetadata): """TrtllmAttentionMetadata for MiniMax-M3 MSA sparse layers. + A step is prepared for a fixed division of labour: its context rows run on + fmha_sm100 through MsaPrefillFmha, and its generation rows on MiniMax-M3's + dedicated decode kernels through MsaDecodeFmha. msa_decode_span says where + the second range begins. Neither library chooses, so the staging below is + unambiguous. + 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). 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, - read only while building plans in prepare() (outside capture), so they - need no graph-stable storage. Plans are built in _build_step_plans: - pure-decode batches use the graph-safe owners (msa_decode_*_plan) while - prefill/mixed batches keep plain eager tuples (msa_eager_*_plan). + msa_qo_offset_cpu) are host properties of the base seq_lens/kv_lens, read + only while building plans in prepare() (outside capture), so they need no + graph-stable storage. The plans themselves (msa_prefill_*_plan) cover the + context rows alone, and a step carrying those is never captured, so they + need none either. """ # Graph-stable buffers; consumers slice to the live count at the call @@ -231,124 +142,214 @@ 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 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 staged rather than derived at the call site. + 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 # _msa_buffers_ready gates the once-only device buffers; # _msa_fields_ready marks that the current step's buffers are populated. _msa_buffers_ready: bool = False _msa_fields_ready: bool = False - # Sparse geometry the decode plans need. + # Sparse geometry the plans need. _msa_params: Optional[MiniMaxM3SparseMetadataParams] = None - # Plan owners, created lazily when the decode plans are first built and - # reused across steps. Each owns its graph-safe plan buffers and the - # current refreshed plan tuple. - _msa_proxy_plan: Optional["_MsaGraphSafePlan"] = None - _msa_gqa_plan: Optional["_MsaGraphSafePlan"] = None - _msa_dense_plan: Optional["_MsaGraphSafePlan"] = None - # Eager (prefill/mixed) plans, plain tuples with no graph-stable buffers - # since prefill runs eagerly and is never CUDA-graph captured. Built once - # per step in prepare() and reused by every layer. - _msa_eager_proxy_plan: Optional[tuple] = None - _msa_eager_gqa_plan: Optional[tuple] = None - _msa_eager_dense_plan: Optional[tuple] = None - # Eager (prefill/mixed) per-token valid-block count. It is layer-invariant - # (a function of qo/kv lengths and page size), so it is computed on the host - # and staged to the device once per step via a non-blocking copy_, then - # reused by every sparse layer's indexer. _msa_eager_n_valid_buf is the - # persistent backing store for the view. - _msa_eager_n_valid_buf: Optional[torch.Tensor] = None - _msa_eager_n_valid_blocks: Optional[torch.Tensor] = None + # This step's fmha_sm100 plans, plain tuples with no graph-stable buffers + # because they cover the context rows alone and a step carrying those is + # never captured. Built once per step in prepare() and reused by every + # layer. + _msa_prefill_proxy_plan: Optional[tuple] = None + _msa_prefill_gqa_plan: Optional[tuple] = None + _msa_prefill_dense_plan: Optional[tuple] = None + # Per-token valid-block count for the prefill-side indexer proxy. It is + # layer-invariant (a function of qo/kv lengths and page size), so it is + # computed on the host and staged to the device once per step via a + # non-blocking copy_, then reused by every sparse layer's indexer. + # _msa_prefill_n_valid_buf is the persistent backing store for the view. + _msa_prefill_n_valid_buf: Optional[torch.Tensor] = None + _msa_prefill_n_valid_blocks: Optional[torch.Tensor] = None + # This step's per-request host lengths, staged by _stage_host_lengths. + _msa_qo_lens_cpu: Optional[torch.Tensor] = None + _msa_kv_lens_cpu: Optional[torch.Tensor] = None + _msa_qo_offset_cpu: Optional[torch.Tensor] = None + # Set once per step by _set_decode_span(), ahead of every other preparation + # step; see msa_decode_span. + _msa_decode_span: Optional[MsaDecodeSpan] = None + # See msa_max_kv_len. + _msa_max_kv_len: int = 0 + # See msa_worst_case_max_k_tiles. + _msa_worst_case_max_k_tiles: int = 0 def __post_init__(self) -> None: super().__post_init__() params = self.sparse_metadata_params self._msa_params = params if isinstance(params, MiniMaxM3SparseMetadataParams) else None self._create_msa_buffers() + self._validate_decode_kernel_support() @property def msa_qo_lens_cpu(self) -> Optional[torch.Tensor]: - """Per-request query length (host int32), from the base seq_lens. - - Pinned where pinning helps, as with the other two length properties: - the planners stage them to the device with non-blocking copies, which - degrade to a synchronous staging copy from pageable memory. - """ - seq_lens = self.seq_lens - if seq_lens is None: - return None - out = seq_lens[: self.num_seqs] - if out.dtype != torch.int32: - out = out.to(torch.int32) - return maybe_pin_memory(out) + """Per-request query length (host int32), from the base seq_lens.""" + return self._msa_qo_lens_cpu @property def msa_kv_lens_cpu(self) -> Optional[torch.Tensor]: """Per-request KV length, cached plus new tokens (host int32).""" - kv_lens = getattr(self, "kv_lens", None) - if self.seq_lens is None or kv_lens is None: - return None - out = kv_lens[: self.num_seqs] - if out.dtype != torch.int32: - out = out.to(torch.int32) - return maybe_pin_memory(out) + return self._msa_kv_lens_cpu @property def msa_qo_offset_cpu(self) -> Optional[torch.Tensor]: """Per-request causal offset (kv_len - qo_len), the cached prefix length.""" - qo = self.msa_qo_lens_cpu - kv = self.msa_kv_lens_cpu - if qo is None or kv is None: - return None - return maybe_pin_memory(kv - qo) + return self._msa_qo_offset_cpu + + def _stage_host_lengths(self) -> None: + """Build this step's per-request host length tensors, once. + + Every planner and both FMHA libraries read these, several of them once + per layer, and each build slices, casts to int32, and pins. Pinning is + what lets the planners stage them with non-blocking copies instead of + copying out of pageable memory, so it is worth paying for once a step. + """ + + def as_pinned_int32(lens: torch.Tensor) -> torch.Tensor: + rows = lens[: self.num_seqs] + return maybe_pin_memory(rows.to(torch.int32) if rows.dtype != torch.int32 else rows) + + seq_lens = self.seq_lens + kv_lens = getattr(self, "kv_lens", None) + self._msa_qo_lens_cpu = None if seq_lens is None else as_pinned_int32(seq_lens) + self._msa_kv_lens_cpu = ( + None if seq_lens is None or kv_lens is None else as_pinned_int32(kv_lens) + ) + self._msa_qo_offset_cpu = ( + None + if self._msa_kv_lens_cpu is None + else maybe_pin_memory(self._msa_kv_lens_cpu - self._msa_qo_lens_cpu) + ) + + # The attention plans cover the context rows; the proxy plan covers + # whatever rows the CuTe DSL scorer did not take. See _msa_attn_plan_rows + # and _msa_proxy_plan_rows. + @property + def msa_prefill_proxy_plan(self) -> Optional[tuple]: + """Prebuilt indexer proxy plan for this step's fmha_sm100 rows.""" + return self._msa_prefill_proxy_plan @property - def msa_decode_proxy_plan(self) -> Optional[tuple]: - """Proxy (max-score) plan tuple, or None outside decode.""" - plan = self._msa_proxy_plan - return plan.plan if plan is not None else None + def msa_prefill_gqa_plan(self) -> Optional[tuple]: + """Prebuilt context-phase sparse GQA plan.""" + return self._msa_prefill_gqa_plan @property - def msa_decode_gqa_plan(self) -> Optional[tuple]: - """Sparse GQA plan tuple, or None outside decode.""" - plan = self._msa_gqa_plan - return plan.plan if plan is not None else None + def msa_prefill_dense_plan(self) -> Optional[tuple]: + """Prebuilt context-phase dense GQA plan.""" + return self._msa_prefill_dense_plan @property - def msa_decode_dense_plan(self) -> Optional[tuple]: - """Dense GQA plan tuple, shared by dense layers 0 to 2.""" - plan = self._msa_dense_plan - return plan.plan if plan is not None else None + def msa_decode_span(self) -> Optional[MsaDecodeSpan]: + """This step's generation rows, or None where it has none. + + A batch is ordered context-first, so the generation requests are the + row suffix [row_first, num_seqs) and their query tokens the matching + token suffix. Those rows run on MiniMax-M3's dedicated decode kernels, + which address query tokens by the uniform per-request query_len; see + _set_decode_span. + """ + return self._msa_decode_span @property - def msa_eager_proxy_plan(self) -> Optional[tuple]: - """Prebuilt indexer proxy plan for the eager (prefill/mixed) path.""" - return self._msa_eager_proxy_plan + def msa_decode_query_len(self) -> Optional[int]: + """Uniform per-request query length over this step's generation rows.""" + span = self._msa_decode_span + return span.query_len if span is not None else None @property - def msa_eager_gqa_plan(self) -> Optional[tuple]: - """Prebuilt sparse GQA plan for the eager (prefill/mixed) path.""" - return self._msa_eager_gqa_plan + def msa_max_kv_len(self) -> int: + """Staged max KV length over this step's generation rows. + + A scheduling upper bound for the decode kernels, taken over the + generation rows alone so a long context request cannot inflate it. + """ + return self._msa_max_kv_len @property - def msa_eager_dense_plan(self) -> Optional[tuple]: - """Prebuilt dense GQA plan for the eager (prefill/mixed) path.""" - return self._msa_eager_dense_plan + 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_prefill_n_valid_blocks(self) -> Optional[torch.Tensor]: + """Device int32 valid-block count for the fmha_sm100 proxy rows, or None + where the step has none (a pure-decode step or a structural test).""" + return self._msa_prefill_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. - The eager GQA and dense plans must pass use_fp8_kvcache so the inline - sparse-prefill path selects the FP8 AOT kernels; it is a no-op for the - decode planner. Mirrors the k_paged.dtype check in run_msa_paged_gqa. + The GQA and dense plans must pass use_fp8_kvcache so the inline + sparse-prefill path selects the FP8 AOT kernels. Mirrors the k_paged + dtype check in run_msa_prefill_gqa. """ kv_cache_manager = self.kv_cache_manager return kv_cache_manager is not None and kv_cache_manager.dtype == DataType.FP8 + def _validate_decode_kernel_support(self) -> None: + """Require the decode kernels to accept this run's cache geometry. + + The generation phase runs on them alone, so the geometry is settled + once against the manager this metadata was built for. Doing it per step + would only offer the choice of running the wrong kernel. It belongs + here rather than in the attention's own validation because the manager, + not the layer, fixes the page size, the index dtype and the sub-page + pool. + """ + params = self._msa_params + kv_cache_manager = self.kv_cache_manager + if params is None or kv_cache_manager is None: + # No MSA geometry to check, as for a structural test's metadata. + return + page_size = int(kv_cache_manager.tokens_per_block) + if not self._cutedsl_indexer_supported( + num_index_heads=params.num_index_heads, + page_size=page_size, + # One query token per generation request; see msa_decode_span. + decode_query_len=1, + ): + raise RuntimeError( + "The MiniMax-M3 CuTe DSL indexer scorer does not support this " + f"configuration: {params.num_index_heads} index heads, page size " + f"{page_size}, index dtype {self._msa_index_kv_dtype()}." + ) + dense_unsupported = dense_decode_unsupported_reason(kv_cache_manager, MSA_REQUIRED_HEAD_DIM) + if dense_unsupported is not None: + raise RuntimeError( + "The MiniMax-M3 dense layers run on the trtllm-gen decode kernel, " + f"but {dense_unsupported}" + ) + def _create_msa_buffers(self) -> None: """Allocate the CUDA-graph-stable MSA device buffers. @@ -382,6 +383,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, + ) # The proxy scratch needs the fmha_sm100 plan geometry. This metadata # exists only for the MSA backend, whose selection already required the # kernels, so a failed import here is a hard error rather than a reason @@ -395,6 +421,7 @@ def _create_msa_buffers(self) -> None: kv_cache_manager=kv_cache_manager, max_batch=max_num_sequences, ) + self._msa_worst_case_max_k_tiles = int(max_k_tiles) self._alloc_msa_proxy_scratch( num_index_heads=params.num_index_heads, max_batch=max_num_sequences, @@ -469,6 +496,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_batch=max_batch, @@ -476,50 +504,192 @@ def _ensure_msa_decode_scratch_buffers( capture_graph=capture_graph, ) - def _ensure_eager_n_valid_buffer(self, total_q: int, device: torch.device) -> torch.Tensor: - """Return a persistent device int32 buffer for the eager valid-block count. + def _ensure_prefill_n_valid_buffer(self, total_q: int, device: torch.device) -> torch.Tensor: + """Return a persistent device int32 buffer for the valid-block count. - The eager path is never CUDA-graph captured, so a plain device tensor, - grown on demand and reused across steps, is sufficient. It is sized to - the worst-case per-step query-token count. + A step carrying context rows is never CUDA-graph captured, so a plain + device tensor, grown on demand and reused across steps, is sufficient. + It is sized to the worst-case per-step query-token count. """ - buf = self._msa_eager_n_valid_buf + buf = self._msa_prefill_n_valid_buf if buf is None or buf.numel() < total_q or buf.device != device: cap = max(int(total_q), int(getattr(self, "max_num_tokens", 0) or 0), 1) buf = torch.empty(cap, dtype=torch.int32, device=device) - self._msa_eager_n_valid_buf = buf + self._msa_prefill_n_valid_buf = buf return buf def prepare(self) -> None: super().prepare() + self._check_beam_width() + # Everything below reads these. + self._stage_host_lengths() + # Set first: both _build_msa_fields and _build_step_plans skip the + # fmha_sm100 preparation the decode kernels replace. + self._set_decode_span() self._build_msa_fields() + self._check_capture_is_pure_decode() self._build_step_plans() + def _check_beam_width(self) -> None: + """Fail on beam search, which every MSA site assumes away. + + The decode kernels take one row per request, while a beam batch holds + beam_width rows, so the block table and lengths handed to them would + cover only the first 1 / beam_width of the batch. + """ + if self.beam_width != 1: + raise NotImplementedError( + "MiniMax-M3 MSA attention does not support beam search, but this " + f"step has beam_width={self.beam_width}. Use beam_width=1 or the " + "non-MSA MiniMax-M3 backend." + ) + + def _set_decode_span(self) -> None: + """Describe this step's generation rows, ahead of any preparation work. + + The span is a description of the batch, not a choice between kernels: + the generation rows always run on MiniMax-M3's dedicated decode kernels + and the context rows always run on fmha_sm100. Whether those kernels + can serve the run at all is settled once, up front, by + ensure_msa_available and _validate_decode_kernel_support, so there is + no per-step verdict here for the FMHA libraries to disagree about. + + The one property of the rows themselves that has to hold is a single + positive query length across them, which the kernels derive the request + id from, so a batch without it is rejected rather than served. + """ + self._msa_decode_span = None + self._msa_max_kv_len = 0 + qo_lens_cpu = self.msa_qo_lens_cpu + kv_lens_cpu = self.msa_kv_lens_cpu + if qo_lens_cpu is None or kv_lens_cpu is None: + 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 describe. + 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_max > 1: + raise NotImplementedError( + "MiniMax-M3 MSA attention does not support speculative decoding " + "(multiple query tokens per decode step): generation rows " + f"[{row_first}, {row_last}) carry up to {qo_max} query tokens. " + "Disable speculative decoding or use the non-MSA MiniMax-M3 backend." + ) + if qo_min != qo_max or qo_max <= 0: + raise RuntimeError( + "MiniMax-M3 MSA attention needs one positive query length across a " + f"step's generation rows, which the decode kernels derive the " + f"request id from, but rows [{row_first}, {row_last}) carry " + f"{gen_qo_lens.tolist()}." + ) + # Staged, i.e. before the overlap scheduler's correction, which only + # shrinks lengths. That keeps it a valid upper bound even when it is + # baked into a CUDA graph. + self._msa_max_kv_len = int(kv_lens_cpu[row_first:].max()) + self._msa_decode_span = MsaDecodeSpan(row_first=row_first, query_len=qo_max) + + def _check_capture_is_pure_decode(self) -> None: + """Fail if a CUDA graph step carries context rows. + + A context row is planned eagerly, as a plain tuple of per-step tensors + (see _build_step_plans), so a graph that captured one would replay + fmha_sm100 against the addresses of a step that has passed. Decode + needs no such plan, which is why capture is confined to it. + """ + if self.is_cuda_graph and int(self.num_contexts or 0) > 0: + raise RuntimeError( + "MiniMax-M3 MSA attention captured a CUDA graph for a step with " + f"{int(self.num_contexts)} context requests. Only pure-decode steps " + "are graph-safe here; see _build_step_plans." + ) + + def _msa_runs_no_fmha(self) -> bool: + """Whether nothing this step reaches fmha_sm100. + + When True its whole per-step preparation is dead: the plans and the + flattened msa_kv_indices page table. That is a pure-decode step, since + the decode kernels then own every row. A mixed step never qualifies, as + fmha_sm100 still runs the context prefix. + """ + span = self._msa_decode_span + return span is not None and span.row_first == 0 + + def _msa_proxy_plan_rows(self) -> Optional[Tuple[int, int]]: + """Batch rows this step's fmha_sm100 indexer proxy plan must cover. + + The indexer is not split by phase: it runs once per sparse layer over + the whole batch, so its plan covers whatever the CuTe DSL scorer did + not take. + + * pure prefill, so no span: the whole batch, the proxy scoring it all; + * mixed: the context prefix only; + * pure decode: None, no rows left for the proxy. + """ + span = self._msa_decode_span + if span is None: + return (0, int(self.num_seqs)) + return (0, span.row_first) if span.row_first > 0 else None + + def _msa_attn_plan_rows(self) -> Optional[Tuple[int, int]]: + """Batch rows the fmha_sm100 attention plans must cover. + + The context prefix, which is the whole of what fmha_sm100 attends: the + generation rows are the decode kernels' and are never planned for it. + """ + num_contexts = int(self.num_contexts or 0) + return (0, num_contexts) if num_contexts > 0 else None + + def _msa_index_kv_dtype(self) -> torch.dtype: + """dtype of the paged index-K cache, which index Q is cast to. + + The CuTe DSL scorer requires index Q and K to match, and run_indexer + casts Q to the cache dtype, so the cache decides what the scorer sees. + """ + 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), + ) + ) + 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, head counts, topk, page size), never on the layer, so they are built - once here and reused by every layer: - - * Pure-decode batches mirror the plans into the CUDA-graph-stable - _MsaGraphSafePlan buffers (surfaced by msa_decode_*_plan), because - decode is captured and the plan worklists must keep a fixed address - across replays. - * 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. + once here and reused by every layer. + + Each plan covers only the rows fmha_sm100 still runs, per + _msa_proxy_plan_rows and _msa_attn_plan_rows, which leaves a pure + decode step nothing to plan: its attention is the decode kernels' and + its block selection the CuTe DSL scorer's. That is what keeps the plans + off the CUDA-graph path, since a captured step is a pure-decode one + (_check_capture_is_pure_decode), and so lets them be plain tuples of + per-step tensors. """ - # Drop any plan tuples from the previous step; the msa_decode_*_plan and - # msa_eager_*_plan properties then report None until rebuilt below. - for plan in (self._msa_proxy_plan, self._msa_gqa_plan, self._msa_dense_plan): - if plan is not None: - plan.reset() - self._msa_eager_proxy_plan = None - self._msa_eager_gqa_plan = None - self._msa_eager_dense_plan = None - self._msa_eager_n_valid_blocks = None + self._msa_prefill_proxy_plan = None + self._msa_prefill_gqa_plan = None + self._msa_prefill_dense_plan = None + self._msa_prefill_n_valid_blocks = None if not self._msa_fields_ready: return # Geometry is captured in __post_init__; skip when it is unavailable. @@ -527,130 +697,91 @@ def _build_step_plans(self) -> None: if params is None: return num_index_heads = params.num_index_heads - num_q_heads, num_kv_heads = params.sharded_head_counts(self.mapping) - topk = params.topk - - fmha_sm100 = require_msa_module() qo_lens_cpu = self.msa_qo_lens_cpu kv_lens_cpu = self.msa_kv_lens_cpu qo_offset_cpu = self.msa_qo_offset_cpu if qo_lens_cpu is None or kv_lens_cpu is None or qo_offset_cpu is None: return batch = int(qo_lens_cpu.shape[0]) - device = _cache_device(self) page_size = int(self.kv_cache_manager.tokens_per_block) - capture_graph = self.is_cuda_graph - max_batch = int(self.max_num_sequences) - # A decode batch is pure generation (no context requests). Only that - # path is CUDA-graph captured and uses the graph-stable plan buffers. - is_decode = int(self.num_contexts or 0) == 0 + if self._msa_runs_no_fmha(): + # Pure decode: no plan to build, but the scorer still writes its + # scores into the proxy scratch and reads the valid-block count. + self._ensure_msa_decode_scratch_buffers( + num_index_heads=num_index_heads, + max_batch=int(self.max_num_sequences), + capture_graph=self.is_cuda_graph, + # No proxy plan, so the worst case is the only bound available. + required_max_k_tiles=self._msa_worst_case_max_k_tiles, + ) + n_valid = per_token_valid_blocks( + qo_lens_cpu, kv_lens_cpu, qo_offset_cpu, causal=True, block_size=page_size + ) + self.msa_n_valid_blocks[:batch].copy_(n_valid.to(torch.int32), non_blocking=True) + return + + fmha_sm100 = require_msa_module() + num_q_heads, num_kv_heads = params.sharded_head_counts(self.mapping) # The main-attention GQA and dense plans need use_fp8_kvcache so the - # eager (inline sparse-prefill) kernel selection matches an FP8 paged - # cache; it is a no-op for the decode planner. The proxy runs over the - # bf16 index-K cache, so it never needs the flag. + # inline sparse-prefill kernel selection matches an FP8 paged cache. + # The proxy runs over the bf16 index-K cache, so it never needs the + # flag. use_fp8 = self._msa_main_kv_is_fp8() + def plan_for(rows: Optional[Tuple[int, int]], **plan_kwargs) -> Optional[tuple]: + """Plan one site over the given rows, or None where it has none. + + Slicing keeps the length tensors' pinned backing, so a single-phase + plan stages as cheaply as a whole-batch one. + """ + if rows is None: + return None + first, last = rows + whole = (first, last) == (0, batch) + 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, + self._msa_prefill_proxy_plan = plan_for( + self._msa_proxy_plan_rows(), + 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, + attn_rows = self._msa_attn_plan_rows() + # Sparse layers: kv_block_num=topk limits attention to top-k blocks. + self._msa_prefill_gqa_plan = plan_for( + attn_rows, + 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, + kv_block_num=params.topk, 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 layers: no kv_block_num, so the full page table is attended. + self._msa_prefill_dense_plan = plan_for( + attn_rows, + 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, ) - - if not is_decode: - # Prefill and mixed batches run eagerly, so keep the plain plan - # tuples and leave the graph-safe owners reset. - self._msa_eager_proxy_plan = proxy_plan - self._msa_eager_gqa_plan = gqa_plan - self._msa_eager_dense_plan = dense_plan - # Stage the valid-block count to the device once for the whole step - # (see _msa_eager_n_valid_blocks). - n_valid_host = per_token_valid_blocks( - qo_lens_cpu, kv_lens_cpu, qo_offset_cpu, causal=True, block_size=page_size - ) - total_q = int(n_valid_host.shape[0]) - if total_q > 0: - dev_buf = self._ensure_eager_n_valid_buffer(total_q, device) - dev_buf[:total_q].copy_(n_valid_host.to(torch.int32), non_blocking=True) - self._msa_eager_n_valid_blocks = dev_buf[:total_q] - return - - required_max_k_tiles = int(proxy_plan[3]["max_k_tiles"]) - 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, - ) - - # Allocate the graph-safe plan owners once per metadata; later steps - # only refresh their contents below. - if self._msa_proxy_plan is None: - num_ctas = torch.cuda.get_device_properties(device).multi_processor_count - self._msa_proxy_plan = _MsaGraphSafePlan( - self, - "msa_proxy_plan", - max_batch=max_batch, - num_ctas=num_ctas, - capture_graph=capture_graph, - ) - self._msa_gqa_plan = _MsaGraphSafePlan( - self, - "msa_gqa_plan", - max_batch=max_batch, - num_ctas=num_ctas, - capture_graph=capture_graph, - ) - self._msa_dense_plan = _MsaGraphSafePlan( - self, - "msa_dense_plan", - max_batch=max_batch, - num_ctas=num_ctas, - capture_graph=capture_graph, - ) - - # 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) - - n_valid = per_token_valid_blocks( + # Stage the valid-block count to the device once for the whole step + # (see msa_prefill_n_valid_blocks). + n_valid_host = per_token_valid_blocks( qo_lens_cpu, kv_lens_cpu, qo_offset_cpu, causal=True, block_size=page_size ) - self.msa_n_valid_blocks[:batch].copy_(n_valid.to(torch.int32), non_blocking=True) + total_q = int(n_valid_host.shape[0]) + if total_q > 0: + dev_buf = self._ensure_prefill_n_valid_buffer(total_q, _cache_device(self)) + dev_buf[:total_q].copy_(n_valid_host.to(torch.int32), non_blocking=True) + self._msa_prefill_n_valid_blocks = dev_buf[:total_q] def _build_msa_fields(self) -> None: """Populate the MSA cache-write buffers for this step. @@ -676,14 +807,6 @@ def _build_msa_fields(self) -> None: cache_device = _cache_device(self) page_size = int(kv_cache_manager.tokens_per_block) - is_prefill = int(self.num_contexts or 0) > 0 - if not is_prefill and int(qo_lens_cpu.max().item()) > 1: - raise NotImplementedError( - "MiniMax-M3 MSA attention does not support speculative decoding " - "(multiple query tokens per decode step). Disable speculative " - "decoding or use the non-MSA MiniMax-M3 backend." - ) - # Built in prepare() (outside capture), so these transients are # fine: forwards read only the persistent buffers filled below. # qo_offset is the prefix length, so one build covers prefill @@ -696,25 +819,57 @@ def _build_msa_fields(self) -> None: device=cache_device, ) 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 decode 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 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, outside capture into a graph-stable + # buffer as with 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], + ) self._msa_fields_ready = True def msa_idx_k_cache(self, layer_idx: int) -> torch.Tensor: @@ -746,6 +901,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( @@ -814,6 +976,35 @@ def _validate_msa_preconditions(self) -> None: f"MSA backend requires topk={MSA_REQUIRED_TOPK}, got {config.topk}." ) + def update_quant_config(self, new_quant_config: Optional[QuantConfig]) -> None: + """Build the FMHA manager, then require the MSA pair among its libraries. + + The base class defers this past __init__ when weight creation is + deferred, and reruns it whenever the quant config lands, so the pair + check belongs here rather than with the other preconditions. + """ + super().update_quant_config(new_quant_config) + self._validate_fmha_pair() + + def _validate_fmha_pair(self) -> None: + """Require both halves of the MSA pair on this layer. + + MsaPrefillFmha serves the context phase and MsaDecodeFmha the + generation phase, and neither will take the other's. Losing one, as a + TLLM_FMHA_LIBS subset would, leaves that phase to a library that + refuses it. Checking as they are built reports it once per layer rather + than on the first step that happens to carry those rows. + """ + present = {type(fmha).__name__ for fmha in self._fmha_manager.fmha_libs} + missing = sorted({"MsaPrefillFmha", "MsaDecodeFmha"} - present) + if missing: + raise RuntimeError( + f"MiniMax-M3 MSA attention layer {self.layer_idx} is missing the FMHA " + f"{'library' if len(missing) == 1 else 'libraries'} {', '.join(missing)}. " + "The two are a pair covering one phase each; enable both (msa_prefill " + "and msa_decode) or none." + ) + @classmethod def support_fused_rope(cls) -> bool: # The MiniMax-M3 model layer applies partial RoPE to the main and @@ -832,17 +1023,17 @@ def run_indexer( The model layer runs this before forward and threads the result through forward_args.sparse_backend_args. Returns [total_q, num_kv_heads, topk]. - Decode uses the prebuilt graph-safe proxy plan; prefill and mixed - batches use the prebuilt eager proxy plan. + The generation rows are scored by the CuTe DSL kernel and any context + rows by the fmha_sm100 proxy pass, over the plan prepare() built. """ 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]) + # The span says how the scoring is split between the CuTe DSL kernel + # and the fmha_sm100 proxy pass. + span = metadata.msa_decode_span # Preserve split column views without allowing an implicit copy. The # scorer and cache writer below both honor their source strides. - head_major_output = ( - int(metadata.num_contexts or 0) > 0 and int(metadata.num_generations or 0) == 0 - ) idx_q_view = idx_q.view(num_tokens, config.num_index_heads, config.sparse_index_dim) idx_k_cache = metadata.msa_idx_k_cache(self.layer_idx) configured_for_fp8 = self.indexer_kv_dtype == "fp8" @@ -874,26 +1065,59 @@ def run_indexer( # The fused production path arrives here with E4M3 Q and an already # populated cache; the BF16 path writes its live K above. - # 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. - proxy_plan = metadata.msa_decode_proxy_plan - if proxy_plan is not None: - # 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"]) + # Inputs for the CuTe DSL scorer, which takes this step's generation + # span. Left None on a pure-prefill step, which has no span, so the + # proxy plan scores the whole batch instead. gen_first is the span's + # first query token: the scorer takes [gen_first, num_tokens) over rows + # [ctx_rows, row_last), the proxy the context prefix ahead of both. + block_table = None + seq_lens_cuda = None + decode_query_len = None + gen_first = 0 + ctx_rows = 0 + if span is not None: + ctx_rows = span.row_first + decode_query_len = span.query_len + row_last = int(metadata.num_seqs) + # Derived from the row count and the uniform query length, as + # PhasedFmha derives the attention phase's token offset, so the + # scorer and the decode kernels agree on the boundary. + gen_first = num_tokens - decode_query_len * (row_last - ctx_rows) + 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, over the scratch prepare() staged for whichever + # scorer owns the rows: a span starting at row 0 leaves the CuTe DSL + # kernel every row and the graph-stable valid-block buffer, anything + # else keeps a proxy plan and the per-step count. When neither is + # present (a standalone test that skips prepare) select_blocks plans + # inline and computes the valid-block count itself. + proxy_plan = metadata.msa_prefill_proxy_plan + if proxy_plan is None and span is not None and span.row_first == 0: + # No proxy plan to read max_k_tiles from for the contiguous score + # view, so it is shaped to the worst case, which the scorer + # accepts: it takes every score stride at runtime. max_score = metadata.msa_proxy_max_score_view( - config.num_index_heads, plan_max_k_tiles, num_tokens + config.num_index_heads, metadata.msa_worst_case_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 - n_valid_blocks = metadata.msa_eager_n_valid_blocks + n_valid_blocks = metadata.msa_prefill_n_valid_blocks if n_valid_blocks is not None: n_valid_blocks = n_valid_blocks[:num_tokens] + # The scorer fills the buffer it is handed, shaped to the span's + # tokens alone: 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 span is not None + else None + ) return self.indexer.select_blocks( idx_q_view, idx_k_cache, @@ -905,7 +1129,12 @@ def run_indexer( proxy_plan=proxy_plan, max_score=max_score, n_valid_blocks=n_valid_blocks, - head_major_output=head_major_output, + require_cutedsl=span is not None, + 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( diff --git a/tensorrt_llm/_torch/attention/backends/sparse/minimax_m3/msa_indexer.py b/tensorrt_llm/_torch/attention/backends/sparse/minimax_m3/msa_indexer.py index 7ce9cdd854da..ea08907e34af 100644 --- a/tensorrt_llm/_torch/attention/backends/sparse/minimax_m3/msa_indexer.py +++ b/tensorrt_llm/_torch/attention/backends/sparse/minimax_m3/msa_indexer.py @@ -18,7 +18,7 @@ import torch -from .msa_utils import ( +from .kernels.msa_utils import ( MSA_REQUIRED_TOPK, per_token_valid_blocks, require_msa_module, @@ -155,6 +155,33 @@ 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` backs the result + the way select_blocks_from_maxscore backs its own output, so the combined + table permutes to a contiguous [num_kv_heads, total_q, topk] as an unsplit + one does. + """ + 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) + # Only a mixed step splits the table and a mixed step is never captured, so + # this allocation cannot land in a graph's memory pool. + 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", @@ -202,61 +229,173 @@ def select_blocks( proxy_plan: Optional[tuple] = None, 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. + Plan/run split, mirroring the sparse GQA: production passes a + precomputed device `n_valid_blocks`, a prebuilt `proxy_plan` wherever + any row is left to the proxy, and, on a pure-decode step, the + preallocated `max_score` the captured region runs into. + + `block_table`, `seq_lens_cuda` and `decode_query_len` put the CuTe DSL + scorer on this step's generation span in place of the fmha_sm100 proxy + pass; leaving them unset (the standalone kernel tests) runs the proxy + over the whole batch. `require_cutedsl` says prepare() narrowed the + proxy plan to the context prefix, so a decline has no fallback. + + `gen_token_first` and `ctx_rows` mark where that span starts: the + scorer takes query tokens [gen_token_first, total_q) and rows + [ctx_rows, batch), the proxy the context prefix ahead of both, and both + are 0 on a pure-decode step. The halves score into separate buffers, + since fmha_sm100 writes a contiguous [heads, k_tiles, tokens] block and + cannot fill a slice of the scorer's, then their tables are joined. """ - config = self.config + page_size = int(idx_k_paged.shape[2]) + gen_first = int(gen_token_first) + + 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 + ): + # The scorer emits raw Q.K rather than idx_sm_scale * Q.K, as the + # fmha_sm100 proxy does; ranking and the +inf init/local forcing in + # select_blocks_from_maxscore are invariant under a positive scale. + 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 proxy_plan is None: - max_score = _proxy_max_score( + if require_cutedsl and not scored: + raise RuntimeError( + "MiniMax-M3 prepare() narrowed the fmha_sm100 proxy plan to this " + "step's context prefix, but the CuTe DSL indexer scorer declined " + "its generation span. There is no proxy pass left to score it; " + "the scorer's geometry is required up front by " + "_validate_decode_kernel_support, so this should not happen." + ) + + # The scorer took nothing, so the proxy runs every token under the + # whole-batch plan it was handed. + 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, - ) - 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, + idx_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( qo_lens_cpu, kv_lens_cpu, qo_offset_cpu, causal=True, - block_size=int(idx_k_paged.shape[2]), + block_size=page_size, + ) + + gen_table = self._select(max_score, n_valid_blocks[gen_first:]) + if gen_first == 0: + return gen_table + # The proxy scores the context prefix into its own buffer, under the + # plan prepare() 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], + ) + return _combined_topk_table(ctx_table, gen_table, head_major=True) + + 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, + ) -> torch.Tensor: + """Reduce scores to KV-head granularity and take the top-k blocks. + + Always into a head-major backing, whatever the step: the Triton sparse + decode kernel reads the table head-major and takes every generation + row, while fmha_sm100 reads whatever strides it is handed. + """ 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, - head_major_output=head_major_output, + init_blocks=self.config.init_blocks, + local_blocks=self.config.local_blocks, + head_major_output=True, ) -__all__ = ["MsaIndexer"] +__all__ = ["MsaIndexer", "cutedsl_score_runner"] diff --git a/tests/microbenchmarks/minimax_m3_index_decode_score.py b/tests/microbenchmarks/minimax_m3_index_decode_score.py index eb1d344a552a..1a25db0269c1 100644 --- a/tests/microbenchmarks/minimax_m3_index_decode_score.py +++ b/tests/microbenchmarks/minimax_m3_index_decode_score.py @@ -16,11 +16,11 @@ import torch import tensorrt_llm._torch.custom_ops # noqa: F401 -from tensorrt_llm._torch.attention.backends.sparse.minimax_m3.msa_indexer import _proxy_max_score -from tensorrt_llm._torch.attention.backends.sparse.minimax_m3.msa_utils import ( +from tensorrt_llm._torch.attention.backends.sparse.minimax_m3.kernels.msa_utils import ( build_kv_page_indices, msa_package_available, ) +from tensorrt_llm._torch.attention.backends.sparse.minimax_m3.msa_indexer import _proxy_max_score PAGE_SIZE = 128 HEAD_DIM = 128 @@ -28,15 +28,7 @@ def _flat_page_table(block_table: torch.Tensor, kv_lens_cpu: torch.Tensor) -> torch.Tensor: """Flatten a block table into the per-request page ids fmha_sm100 consumes.""" - batch, max_pages = block_table.shape - intra = torch.arange(PAGE_SIZE, dtype=torch.int32) - req_to_token = (block_table.cpu().to(torch.int32) * PAGE_SIZE).unsqueeze(2) + intra - return build_kv_page_indices( - req_to_token.reshape(batch, max_pages * PAGE_SIZE), - torch.arange(batch, dtype=torch.int32), - kv_lens_cpu, - PAGE_SIZE, - ) + return build_kv_page_indices(block_table.cpu().to(torch.int32), kv_lens_cpu, PAGE_SIZE) def _time_us(fn, warmup: int = 20, iters: int = 100) -> float: diff --git a/tests/unittest/_torch/attention/sparse/msa/test_minimax_m3_dense_decode.py b/tests/unittest/_torch/attention/sparse/msa/test_minimax_m3_dense_decode.py index 1eddb226516a..980b5bd0734d 100644 --- a/tests/unittest/_torch/attention/sparse/msa/test_minimax_m3_dense_decode.py +++ b/tests/unittest/_torch/attention/sparse/msa/test_minimax_m3_dense_decode.py @@ -28,9 +28,11 @@ import pytest import torch -from tensorrt_llm._torch.attention.backends.sparse.minimax_m3.trtllm_gen_dense_decode import ( +from tensorrt_llm._torch.attention.backends.sparse.minimax_m3.kernels.trtllm_gen_dense_decode import ( dense_decode_unsupported_reason, + dense_decode_workspace_layout, minimax_m3_trtllm_gen_dense_decode, + split_dense_decode_workspace, subpage_block_table, uniform_subpages_per_slot, write_subpage_block_table, @@ -40,7 +42,7 @@ PAGE_SIZE = 32 HEAD_DIM = 128 -# The bmm1 scale, spelled as run_msa_paged_gqa spells it at q_scaling 1. +# The bmm1 scale, spelled as run_msa_prefill_gqa spells it at q_scaling 1. SM_SCALE = HEAD_DIM**-0.5 @@ -49,6 +51,45 @@ def _is_sm100f() -> bool: return major == 10 and minor in (0, 3) +# -------------------------------------------------------------------------- +# Scratch layout +# -------------------------------------------------------------------------- + + +def test_dense_decode_scratch_is_carved_out_of_one_buffer(): + """Both halves of the scratch must fit the total the layout reported. + + MsaDecodeFmha grows the shared attention workspace to that total and then + splits it, so an offset or a size that disagreed with it would hand the + kernel a view running off the end of the workspace. + """ + layout = dense_decode_workspace_layout( + q_dtype=torch.bfloat16, + num_heads=8, + head_dim=HEAD_DIM, + num_kv_heads=1, + max_num_requests=4, + device=torch.device("cuda"), + ) + + assert layout.counter_offset >= layout.workspace_bytes + assert layout.total_bytes > layout.counter_offset + + buffer = torch.full((layout.total_bytes,), 7, dtype=torch.int8, device="cuda") + workspace, counters = split_dense_decode_workspace(buffer, layout) + + assert workspace.numel() == layout.workspace_bytes + assert counters.numel() == layout.total_bytes - layout.counter_offset + assert workspace.data_ptr() == buffer.data_ptr() + assert counters.data_ptr() == buffer.data_ptr() + layout.counter_offset + # Only the counters are cleared; the slab is the kernel's own to initialize. + assert int(counters.sum()) == 0 + assert int(workspace[0]) == 7 + + with pytest.raises(ValueError, match="but this call needs"): + split_dense_decode_workspace(buffer[:-1], layout) + + # -------------------------------------------------------------------------- # Block-table expansion # -------------------------------------------------------------------------- diff --git a/tests/unittest/_torch/attention/sparse/msa/test_minimax_m3_index_decode_score.py b/tests/unittest/_torch/attention/sparse/msa/test_minimax_m3_index_decode_score.py index ad5c213bae55..c9945d39d48c 100644 --- a/tests/unittest/_torch/attention/sparse/msa/test_minimax_m3_index_decode_score.py +++ b/tests/unittest/_torch/attention/sparse/msa/test_minimax_m3_index_decode_score.py @@ -12,13 +12,13 @@ import pytest import torch -from tensorrt_llm._torch.attention.backends.sparse.minimax_m3.msa_indexer import _cutedsl_score -from tensorrt_llm._torch.attention.backends.sparse.minimax_m3.msa_utils import ( +from tensorrt_llm._torch.attention.backends.sparse.minimax_m3.kernels.msa_utils import ( MSA_REQUIRED_TOPK, build_kv_page_indices, msa_package_available, select_blocks_from_maxscore, ) +from tensorrt_llm._torch.attention.backends.sparse.minimax_m3.msa_indexer import _cutedsl_score from tensorrt_llm._utils import get_sm_version PAGE_SIZE = 128 @@ -38,7 +38,7 @@ def _flat_page_table(block_table: torch.Tensor, kv_lens_cpu: torch.Tensor) -> to The production helper concatenates the valid prefix of each request's block-id row according to its KV length. """ - return build_kv_page_indices(block_table.cpu(), kv_lens_cpu, PAGE_SIZE) + return build_kv_page_indices(block_table.cpu().to(torch.int32), kv_lens_cpu, PAGE_SIZE) def _runner(): diff --git a/tests/unittest/_torch/attention/sparse/msa/test_minimax_m3_msa_selector.py b/tests/unittest/_torch/attention/sparse/msa/test_minimax_m3_msa_selector.py index 8fb321c3bc48..38e7569cdfe0 100644 --- a/tests/unittest/_torch/attention/sparse/msa/test_minimax_m3_msa_selector.py +++ b/tests/unittest/_torch/attention/sparse/msa/test_minimax_m3_msa_selector.py @@ -9,7 +9,7 @@ _INIT_SCORE, _LOCAL_SCORE, ) -from tensorrt_llm._torch.attention.backends.sparse.minimax_m3.msa_utils import ( +from tensorrt_llm._torch.attention.backends.sparse.minimax_m3.kernels.msa_utils import ( select_blocks_from_maxscore, ) @@ -194,7 +194,7 @@ def test_fused_selector_head_major_output_through_msa_q2k_consumer(): if torch.cuda.get_device_capability()[0] != 10: pytest.skip("SM100 (Blackwell) required") - from tensorrt_llm._torch.attention.backends.sparse.minimax_m3.msa_utils import ( + from tensorrt_llm._torch.attention.backends.sparse.minimax_m3.kernels.msa_utils import ( msa_package_available, ) diff --git a/tests/unittest/_torch/attention/sparse/msa/test_minimax_m3_sparse_attn_decode.py b/tests/unittest/_torch/attention/sparse/msa/test_minimax_m3_sparse_attn_decode.py index 4e8496d610a4..5a59ca65d700 100644 --- a/tests/unittest/_torch/attention/sparse/msa/test_minimax_m3_sparse_attn_decode.py +++ b/tests/unittest/_torch/attention/sparse/msa/test_minimax_m3_sparse_attn_decode.py @@ -13,12 +13,12 @@ import pytest import torch -from tensorrt_llm._torch.attention.backends.sparse.minimax_m3.msa_utils import ( +from tensorrt_llm._torch.attention.backends.sparse.minimax_m3.kernels.msa_utils import ( MSA_REQUIRED_TOPK, build_kv_page_indices, msa_package_available, ) -from tensorrt_llm._torch.attention.backends.sparse.minimax_m3.triton_sparse_decode import ( +from tensorrt_llm._torch.attention.backends.sparse.minimax_m3.kernels.triton_sparse_decode import ( SPARSE_BLOCK_SIZE, minimax_m3_sparse_attn_decode, resolve_num_topk_chunks, @@ -42,7 +42,7 @@ def _flat_page_table(block_table: torch.Tensor, kv_lens_cpu: torch.Tensor) -> to The production helper concatenates the valid prefix of each request's block-id row according to its KV length. """ - return build_kv_page_indices(block_table.cpu(), kv_lens_cpu, PAGE_SIZE) + return build_kv_page_indices(block_table.cpu().to(torch.int32), kv_lens_cpu, PAGE_SIZE) def _reference_sparse_decode( @@ -375,7 +375,7 @@ def run(): @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.backends.fmha.msa_sparse_gqa import run_msa_sparse_gqa + from tensorrt_llm._torch.attention.backends.fmha.msa_prefill 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( diff --git a/tests/unittest/_torch/attention/sparse/msa/test_msa_backend.py b/tests/unittest/_torch/attention/sparse/msa/test_msa_backend.py index 62884e61a626..bba4c7918d58 100644 --- a/tests/unittest/_torch/attention/sparse/msa/test_msa_backend.py +++ b/tests/unittest/_torch/attention/sparse/msa/test_msa_backend.py @@ -9,6 +9,7 @@ """ import sys +import weakref from types import ModuleType, SimpleNamespace from unittest.mock import Mock @@ -19,7 +20,11 @@ MiniMaxM3KVCacheManagerV2, MiniMaxM3MsaSparseAttention, ) -from tensorrt_llm._torch.attention.backends.sparse.minimax_m3.msa_utils import msa_paged_kv +from tensorrt_llm._torch.attention.backends.sparse.minimax_m3.kernels.msa_utils import ( + MSA_REQUIRED_TOPK, + msa_paged_kv, +) +from tensorrt_llm._torch.attention.backends.sparse.minimax_m3.msa_backend import MsaDecodeSpan from tensorrt_llm._torch.attention.backends.sparse.registry import _resolve_minimax_m3_backend_cls from tensorrt_llm._torch.pyexecutor.kv_cache.kv_cache_manager_v2 import KVCacheManagerV2 from tensorrt_llm.bindings import DataType @@ -27,7 +32,7 @@ def test_msa_package_availability_installs_cutlass_compatibility_aliases(monkeypatch): - from tensorrt_llm._torch.attention.backends.sparse.minimax_m3.msa_utils import ( + from tensorrt_llm._torch.attention.backends.sparse.minimax_m3.kernels.msa_utils import ( msa_package_available, ) @@ -53,7 +58,7 @@ def test_msa_package_availability_installs_cutlass_compatibility_aliases(monkeyp def test_msa_import_preserves_cute_compile_option_selection() -> None: - from tensorrt_llm._torch.attention.backends.sparse.minimax_m3.msa_utils import ( + from tensorrt_llm._torch.attention.backends.sparse.minimax_m3.kernels.msa_utils import ( msa_package_available, ) @@ -200,6 +205,114 @@ 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 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,) + 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): + """prepare() stages the sub-page expansion before any layer is named, which + 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, so its sub-page + 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, @@ -223,13 +336,18 @@ 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) + @pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA") def test_msa_paged_kv_preserves_tma_compatible_outer_stride() -> None: if torch.cuda.get_device_capability()[0] != 10: pytest.skip("SM100 (Blackwell) required") - from tensorrt_llm._torch.attention.backends.sparse.minimax_m3.msa_utils import ( + from tensorrt_llm._torch.attention.backends.sparse.minimax_m3.kernels.msa_utils import ( msa_package_available, ) @@ -271,7 +389,7 @@ def test_msa_paged_hnd_input_materializes_unaligned_outer_stride() -> None: if torch.cuda.get_device_capability()[0] != 10: pytest.skip("SM100 (Blackwell) required") - from tensorrt_llm._torch.attention.backends.sparse.minimax_m3.msa_utils import ( + from tensorrt_llm._torch.attention.backends.sparse.minimax_m3.kernels.msa_utils import ( msa_package_available, ) @@ -422,16 +540,17 @@ def select_blocks(self, idx_q, idx_k_cache, **kwargs): attention.indexer = FakeIndexer() class FakeMetadata: - msa_decode_proxy_plan = None - msa_eager_proxy_plan = (False, 0, 2, {}, None) - msa_eager_all_blocks_empty = False - msa_eager_n_valid_blocks = torch.ones(2, dtype=torch.int32, device="cuda") + # A chunked-prefill step, whose rows are all fmha_sm100's, so the + # indexer takes the proxy path and needs no decode span. + msa_decode_span = None + msa_prefill_proxy_plan = (False, 0, 2, {}, None) + msa_prefill_n_valid_blocks = torch.ones(2, dtype=torch.int32, device="cuda") msa_kv_indices = torch.arange(2, dtype=torch.int32, device="cuda") msa_qo_lens_cpu = torch.ones(2, dtype=torch.int32) msa_kv_lens_cpu = torch.full((2,), 128, dtype=torch.int32) msa_qo_offset_cpu = torch.full((2,), 127, dtype=torch.int32) - num_contexts = 0 - num_generations = 2 + num_contexts = 2 + num_generations = 0 def __init__(self, dtype: torch.dtype) -> None: backing = torch.empty(2 * 7, 1, 128, 128, dtype=dtype, device="cuda") @@ -503,27 +622,47 @@ def msa_idx_k_cache(self, layer_idx: int) -> torch.Tensor: @pytest.mark.parametrize( - ("num_contexts", "num_generations", "expected_head_major"), - [(2, 0, True), (1, 1, False), (0, 2, False)], + ( + "num_contexts", + "num_generations", + "expected_gen_first", + "expected_ctx_rows", + "expected_query_len", + ), + [(2, 0, 0, 0, None), (1, 1, 1, 1, 1), (0, 2, 0, 0, 1)], + ids=["prefill", "mixed", "decode"], ) -def test_run_indexer_routes_head_major_output_by_batch_mode( +def test_run_indexer_hands_the_indexer_this_steps_generation_span( num_contexts: int, num_generations: int, - expected_head_major: bool, + expected_gen_first: int, + expected_ctx_rows: int, + expected_query_len: int | None, ) -> None: - num_tokens, num_index_heads, sparse_index_dim = 3, 4, 128 + """The scorer split follows the batch, and both halves must agree on it. + + The CuTe DSL scorer takes the generation span and the fmha_sm100 proxy the + context prefix ahead of it, so the first query token of the span is derived + from the row count and the uniform query length exactly as PhasedFmha + derives the attention phase's token offset. + + A pure-prefill step reports no split at all rather than one covering every + row: it has no span, so the proxy plan scores the whole batch and there is + no boundary for the two halves to disagree about. + """ + num_tokens, num_index_heads, sparse_index_dim = num_contexts + num_generations, 4, 128 captured = {} class FakeIndexer: def select_blocks(self, *args: object, **kwargs: object) -> torch.Tensor: del args - captured["head_major_output"] = kwargs["head_major_output"] + captured.update(kwargs) return torch.zeros(num_tokens, 1, 16, dtype=torch.int32) class FakeMetadata: - msa_decode_proxy_plan = None - msa_eager_proxy_plan = ("eager",) - msa_eager_n_valid_blocks = torch.ones(num_tokens, dtype=torch.int32) + msa_prefill_n_valid_blocks = torch.ones(num_tokens, dtype=torch.int32) + msa_n_valid_blocks = torch.ones(num_tokens, dtype=torch.int32) + msa_worst_case_max_k_tiles = 8 msa_kv_indices = torch.arange(num_tokens, dtype=torch.int32) msa_qo_lens_cpu = torch.tensor([num_tokens], dtype=torch.int32) msa_kv_lens_cpu = torch.tensor([num_tokens], dtype=torch.int32) @@ -532,6 +671,17 @@ class FakeMetadata: def __init__(self) -> None: self.num_contexts = num_contexts self.num_generations = num_generations + self.num_seqs = num_contexts + num_generations + # The suffix of single-token generation rows prepare() would + # describe, which is empty for a pure-prefill step. + self.msa_decode_span = ( + MsaDecodeSpan(row_first=num_contexts, query_len=1) if num_generations > 0 else None + ) + # Planned for exactly the context rows, so a pure-decode step has + # no plan at all. + self.msa_prefill_proxy_plan = ("prefill",) if num_contexts > 0 else None + self.msa_block_table = torch.zeros(self.num_seqs, 4, dtype=torch.int32) + self.msa_seq_lens_cuda = torch.zeros(self.num_seqs, dtype=torch.int32) self.idx_k_cache = torch.empty( num_tokens, 1, @@ -547,6 +697,11 @@ def msa_idx_k_cache(self, layer_idx: int) -> torch.Tensor: del layer_idx return self.idx_k_cache + def msa_proxy_max_score_view( + self, num_index_heads: int, max_k_tiles: int, tokens: int + ) -> torch.Tensor: + return torch.zeros(num_index_heads, max_k_tiles, tokens) + attention = SimpleNamespace( layer_idx=0, m3_config=SimpleNamespace( @@ -567,7 +722,11 @@ def msa_idx_k_cache(self, layer_idx: int) -> torch.Tensor: ) assert result.shape == (num_tokens, 1, 16) - assert captured["head_major_output"] is expected_head_major + assert captured["gen_token_first"] == expected_gen_first + assert captured["ctx_rows"] == expected_ctx_rows + # Set together with the rest of the scorer's inputs, so it says whether the + # step has a generation span for the scorer to take at all. + assert captured["decode_query_len"] == expected_query_len @pytest.mark.parametrize( @@ -583,12 +742,12 @@ def test_msa_proxy_max_score_strided_index_k_matches_packed( if torch.cuda.get_device_capability()[0] != 10: pytest.skip("SM100 (Blackwell) required") + from tensorrt_llm._torch.attention.backends.sparse.minimax_m3.kernels.msa_utils import ( + msa_package_available, + ) from tensorrt_llm._torch.attention.backends.sparse.minimax_m3.msa_indexer import ( _proxy_max_score, ) - from tensorrt_llm._torch.attention.backends.sparse.minimax_m3.msa_utils import ( - msa_package_available, - ) if not msa_package_available(): pytest.skip("fmha_sm100 (MSA) not importable") @@ -651,7 +810,7 @@ def test_build_kv_page_indices_matches_first_slot_of_each_page(): row holds at its page boundaries, since both use the manager's tokens_per_block as the page size. Rows are ragged (0-padded block ids, global and non-contiguous) and one request has no KV at all.""" - from tensorrt_llm._torch.attention.backends.sparse.minimax_m3.msa_utils import ( + from tensorrt_llm._torch.attention.backends.sparse.minimax_m3.kernels.msa_utils import ( build_kv_page_indices, ) @@ -732,3 +891,492 @@ def get_block_ids_per_seq(self, request_ids): ) for b, slot in enumerate(padded.out_cache_loc.tolist()): assert slot in req_to_token[b].tolist() + + +@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.backends.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 bound 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, 2).shape == (4, 32, 2) + + +def _span_metadata(*, num_contexts=0, qo_lens=(1, 1), kv_lens=(9, 11), is_cuda_graph=False): + """Metadata with just enough state for _set_decode_span. + + The span is a description of the batch's rows, 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.mapping = None + # Assigned behind the seq_lens property, whose setter would stage a device + # copy nothing here 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.is_cuda_graph = is_cuda_graph + # prepare() stages these before the span, and the span reads them. + metadata._stage_host_lengths() + return metadata + + +def test_the_decode_span_of_a_pure_decode_step_is_the_whole_batch(): + """A pure-decode step leaves fmha_sm100 nothing, which is what lets + prepare() skip its plans and its page table entirely.""" + metadata = _span_metadata() + + metadata._set_decode_span() + + assert metadata.msa_decode_span == (0, 1) + assert metadata.msa_decode_query_len == 1 + assert metadata.msa_max_kv_len == 11 + assert metadata._msa_runs_no_fmha() is True + + +def test_the_decode_span_of_a_mixed_step_is_its_generation_suffix(): + """A context request does not move the generation rows off their kernels. + + The generation requests are the batch's row suffix, so the decode kernels + take that span and fmha_sm100 keeps the context prefix. + """ + # Two context requests (7 and 5 query tokens, the first a chunk of a long + # prompt) ahead of two decode rows. + metadata = _span_metadata(num_contexts=2, qo_lens=(7, 5, 1, 1), kv_lens=(4096, 5, 40, 33)) + + metadata._set_decode_span() + + assert metadata.msa_decode_span == (2, 1) + # 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_a_pure_prefill_step_has_no_decode_span(): + """A step with no generation row has nothing for the decode kernels, and + fmha_sm100 keeps every plan and the page table they read.""" + metadata = _span_metadata(num_contexts=2, qo_lens=(5, 7), kv_lens=(5, 7)) + + metadata._set_decode_span() + + assert metadata.msa_decode_span is None + assert metadata.msa_decode_query_len is None + assert metadata._msa_runs_no_fmha() is False + + +def test_plan_rows_narrow_to_the_rows_fmha_sm100_still_runs(): + """The plans must cover exactly those rows. + + A plan is built from the host lengths of the requests it covers, so one that + claimed the whole batch while only the context prefix ran would schedule the + kernel over generation rows nothing dispatches to it. The attention plan + covers the context rows because that is the only phase MsaPrefillFmha runs; + the indexer's proxy plan covers whatever the CuTe DSL scorer did not take, + which is the same range plus a pure-prefill step's generation-free batch. + """ + mixed = _span_metadata(num_contexts=2, qo_lens=(7, 5, 1, 1), kv_lens=(4096, 5, 40, 33)) + mixed._set_decode_span() + # The span took the generation suffix, so what is left is the prefix. + assert mixed._msa_proxy_plan_rows() == (0, 2) + assert mixed._msa_attn_plan_rows() == (0, 2) + + decode = _span_metadata() + decode._set_decode_span() + # Nothing is left to plan on a pure-decode step the kernels fully own. + assert decode._msa_proxy_plan_rows() is None + assert decode._msa_attn_plan_rows() is None + + prefill = _span_metadata(num_contexts=2, qo_lens=(5, 7), kv_lens=(5, 7)) + prefill._set_decode_span() + # No generation row, so the proxy scores every row. + assert prefill._msa_proxy_plan_rows() == (0, 2) + assert prefill._msa_attn_plan_rows() == (0, 2) + + +@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.backends.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 _phase_libraries(): + """One instance of each library, sharing a stub layer. + + Built without __init__ because neither library's state is under test here; + _attn_ref is what Fmha.attn reads. + """ + from tensorrt_llm._torch.attention.backends.fmha.msa_decode import MsaDecodeFmha + from tensorrt_llm._torch.attention.backends.fmha.msa_prefill import MsaPrefillFmha + + attention = MiniMaxM3MsaSparseAttention.__new__(MiniMaxM3MsaSparseAttention) + attention.layer_idx = 3 + attn_ref = weakref.ref(attention) + libraries = [] + for cls in (MsaDecodeFmha, MsaPrefillFmha): + library = cls.__new__(cls) + library._attn_ref = attn_ref + libraries.append(library) + return attention, libraries + + +def _generation_params( + attention, metadata, *, seq_offset, input_seq_length, token_offset=0, key_input=None +): + """Phase params as PhasedFmha.forward builds them for a generation phase.""" + from tensorrt_llm._torch.attention.backends.fmha.phased import FmhaParams + from tensorrt_llm._torch.attention.backends.interface import AttentionInputType + + params = FmhaParams( + attn=attention, + meta=metadata, + fwd=SimpleNamespace( + sparse_runtime_params=SimpleNamespace(sparse_attn_indices=None), + attention_input_type=AttentionInputType.mixed, + ), + workspace=torch.zeros(1), + ) + params.seq_offset = seq_offset + params.input_seq_length = input_seq_length + params.token_offset = token_offset + params.key_input = key_input + params.value_input = key_input + return params + + +def test_decode_fmha_runs_the_phase_its_span_describes(): + """The agreeing case is the one production takes, so the check must let it + through to the dispatch by layer type.""" + attention, (decode, _) = _phase_libraries() + metadata = SimpleNamespace( + msa_decode_span=(2, 1), + num_generations=2, + msa_block_table=torch.zeros(4, 4, dtype=torch.int32), + msa_seq_lens_cuda=torch.zeros(4, dtype=torch.int32), + ) + params = _generation_params(attention, metadata, seq_offset=2, input_seq_length=1) + dispatched = [] + decode._run_dense = lambda params, block_table, seq_lens: dispatched.append( + (block_table.shape[0], seq_lens.shape[0]) + ) + + decode.run_generation(params) + + # The span's rows, not the whole batch: a dense layer attends the page + # table of the generation requests alone. + assert dispatched == [(2, 2)] + + +def test_each_phase_writes_the_cache_slots_of_the_rows_it_attends(): + """A mixed step's cache write is split with its phases, not repeated. + + A phase that wrote from the head of msa_out_cache_loc instead of its own + token offset would land the generation rows' K/V in the context rows' + slots. + """ + from tensorrt_llm._torch.attention.backends.interface import AttentionInputType + from tensorrt_llm._torch.attention.backends.sparse.minimax_m3.kernels.msa_utils import ( + write_msa_phase_kv, + ) + + page_size, num_kv_heads, head_dim = 8, 1, 4 + buffers = torch.zeros(2, 2, num_kv_heads, page_size, head_dim) + metadata = SimpleNamespace( + kv_cache_manager=SimpleNamespace(get_buffers=lambda layer_idx, kv_layout=None: buffers), + msa_out_cache_loc=torch.tensor([1, 2, page_size + 3], dtype=torch.int32), + ) + attention = SimpleNamespace(layer_idx=0) + k = torch.arange(1, 3 * head_dim + 1, dtype=torch.float32).reshape(3, head_dim) + v = -k + + # Two context tokens, then the one-token generation row behind them. + write_msa_phase_kv(attention, k[:2], v[:2], metadata, AttentionInputType.mixed, token_offset=0) + write_msa_phase_kv(attention, k[2:], v[2:], metadata, AttentionInputType.mixed, token_offset=2) + + k_cache, v_cache = buffers[:, 0], buffers[:, 1] + torch.testing.assert_close(k_cache[0, 0, 1], k[0]) + torch.testing.assert_close(k_cache[0, 0, 2], k[1]) + torch.testing.assert_close(k_cache[1, 0, 3], k[2]) + torch.testing.assert_close(v_cache[1, 0, 3], v[2]) + # Only the three named slots were touched. + assert int((k_cache != 0).sum()) == 3 * head_dim + + +@pytest.mark.parametrize( + "num_contexts, num_generations, expected", + [ + (1, 0, "MsaPrefillFmha"), + (0, 1, "MsaDecodeFmha"), + (1, 1, "CombinedFmha"), + ], +) +def test_the_pair_takes_one_phase_each_and_a_mixed_step_together( + num_contexts, num_generations, expected +): + """Neither library may claim a whole step, only its own phase. + + FmhaManager asks each library about the whole step before asking about a + phase, and a library that answered yes to the first would then be handed + the phase it refuses. Claiming one phase each is what leaves a mixed step + to CombinedFmha, which is the only way its generation rows reach the + decode kernels. + """ + from tensorrt_llm._torch.attention.backends.fmha.manager import FmhaManager + from tensorrt_llm._torch.attention.backends.interface import ( + AttentionForwardArgs, + AttentionInputType, + ) + + attention, libraries = _phase_libraries() + # CombinedFmha's PhasedFmha.__init__ reads the layer's head geometry. + attention.is_mla_enable = False + attention.kv_lora_rank = None + attention.v_head_dim = None + attention.head_dim = 128 + manager = FmhaManager.__new__(FmhaManager) + manager.fmha_libs = libraries + + q = torch.empty((num_contexts + num_generations, 4)) + metadata = SimpleNamespace( + num_contexts=num_contexts, + num_generations=num_generations, + num_ctx_tokens=num_contexts, + use_spec_decoding=False, + ) + forward_args = AttentionForwardArgs(attention_input_type=AttentionInputType.mixed) + + # The whole-step query, which the manager asks first. + assert not any( + library.is_supported(q, None, None, metadata, forward_args) for library in libraries + ) + + selected = manager._select_uncached(attention, q, None, None, metadata, forward_args) + + assert type(selected).__name__ == expected + if expected == "CombinedFmha": + decode, prefill = libraries + assert selected._get_context_impl() is prefill + assert selected._get_generation_impl() is decode + + +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_prefill_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.backends.sparse.minimax_m3.kernels.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_prefill_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 go to the Triton sparse decode kernel while the context + rows stay on fmha_sm100 under a context-only plan, so the correctness gate + is that both halves still agree with a whole-batch fmha_sm100 run. + """ + from tensorrt_llm._torch.attention.backends.fmha.msa_prefill import run_msa_prefill_gqa + from tensorrt_llm._torch.attention.backends.sparse.minimax_m3.kernels.msa_utils import ( + msa_package_available, + msa_paged_kv, + ) + from tensorrt_llm._torch.attention.backends.sparse.minimax_m3.kernels.triton_sparse_decode import ( + minimax_m3_sparse_attn_decode, + ) + 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]) + metadata = SimpleNamespace(**fields) + num_heads = attention.num_heads + sm_scale = head_dim**-0.5 + + def run_gqa(output, *, token_first, token_last, row_first, num_rows): + """One fmha_sm100 call, as MsaPrefillFmha.run_context makes it.""" + run_msa_prefill_gqa( + attention, + q[token_first:token_last], + metadata, + output[token_first:token_last], + kv_block_indexes=kv_block_indexes[token_first:token_last], + plan=None, + row_first=row_first, + num_rows=num_rows, + ) + + # Reference: one whole-batch fmha_sm100 call over every row, which is what + # the split has to reproduce. + reference = torch.zeros_like(q) + run_gqa(reference, token_first=0, token_last=total_q, row_first=0, num_rows=4) + + # Split: the context prefix keeps fmha_sm100 under a one-row plan and the + # generation suffix goes to the Triton kernel, as MsaDecodeFmha dispatches + # it. + split = torch.zeros_like(q) + run_gqa(split, token_first=0, token_last=num_ctx_tokens, row_first=0, num_rows=1) + k_paged, v_paged = msa_paged_kv(metadata.kv_cache_manager, attention.layer_idx) + num_gen_tokens = total_q - num_ctx_tokens + minimax_m3_sparse_attn_decode( + q[num_ctx_tokens:].view(num_gen_tokens, num_heads, head_dim), + k_paged, + v_paged, + kv_block_indexes[num_ctx_tokens:].permute(1, 0, 2), + metadata.msa_block_table[1:4], + metadata.msa_seq_lens_cuda[1:4], + sm_scale=sm_scale, + output=split[num_ctx_tokens:].view(num_gen_tokens, num_heads, head_dim), + decode_query_len=1, + ) + torch.cuda.synchronize() + + reference = reference.view(total_q, num_heads, head_dim).float() + split = split.view(total_q, num_heads, head_dim).float() + + 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_sparse_mqa_gqa.py b/tests/unittest/_torch/attention/sparse/test_sparse_mqa_gqa.py index c6a8860a8a5f..df7cf235fe26 100644 --- a/tests/unittest/_torch/attention/sparse/test_sparse_mqa_gqa.py +++ b/tests/unittest/_torch/attention/sparse/test_sparse_mqa_gqa.py @@ -39,7 +39,7 @@ from utils.util import getSMVersion import tensorrt_llm -from tensorrt_llm._torch.attention.backends.fmha.msa_sparse_gqa import run_msa_sparse_gqa +from tensorrt_llm._torch.attention.backends.fmha.msa_prefill import run_msa_sparse_gqa from tensorrt_llm._torch.attention.backends.interface import ( AttentionForwardArgs, AttentionRuntimeFeatures, @@ -47,7 +47,9 @@ from tensorrt_llm._torch.attention.backends.sparse.dsa.kernels import ( triton_convert_req_index_to_global_index, ) -from tensorrt_llm._torch.attention.backends.sparse.minimax_m3.msa_utils import msa_package_available +from tensorrt_llm._torch.attention.backends.sparse.minimax_m3.kernels.msa_utils import ( + msa_package_available, +) from tensorrt_llm._torch.attention.backends.sparse.params import SparseParams from tensorrt_llm._torch.attention.backends.trtllm import ( TrtllmAttention,