From cea5d2d18c1402cfe9562158cf073f2e811d9b7c Mon Sep 17 00:00:00 2001 From: WeiHaocheng <20514172+WeiHaocheng@users.noreply.github.com> Date: Mon, 29 Jun 2026 05:50:59 -0700 Subject: [PATCH 1/4] [None][feat] Integrate M3 sparse attention kernels form MSA Signed-off-by: WeiHaocheng <20514172+WeiHaocheng@users.noreply.github.com> --- .../_torch/attention_backend/fmha/__init__.py | 8 + .../attention_backend/fmha/block_sparse.py | 166 ++++ .../attention_backend/fmha/indexer_proxy.py | 160 ++++ .../attention_backend/fmha/interface.py | 28 +- .../attention_backend/fmha/msa_proxy_mqa.py | 150 ++++ .../attention_backend/fmha/msa_sparse_gqa.py | 162 ++++ .../_torch/attention_backend/fmha/registry.py | 12 + .../sparse/minimax_m3/__init__.py | 10 + .../sparse/minimax_m3/metadata.py | 7 + .../sparse/minimax_m3/msa_backend.py | 794 ++++++++++++++++++ .../_torch/attention_backend/sparse/utils.py | 25 +- tensorrt_llm/llmapi/llm_args.py | 9 + .../sparse/test_minimax_m3_msa_backend.py | 553 ++++++++++++ 13 files changed, 2074 insertions(+), 10 deletions(-) create mode 100644 tensorrt_llm/_torch/attention_backend/fmha/block_sparse.py create mode 100644 tensorrt_llm/_torch/attention_backend/fmha/indexer_proxy.py create mode 100644 tensorrt_llm/_torch/attention_backend/fmha/msa_proxy_mqa.py create mode 100644 tensorrt_llm/_torch/attention_backend/fmha/msa_sparse_gqa.py create mode 100644 tensorrt_llm/_torch/attention_backend/sparse/minimax_m3/msa_backend.py create mode 100644 tests/unittest/_torch/attention/sparse/test_minimax_m3_msa_backend.py diff --git a/tensorrt_llm/_torch/attention_backend/fmha/__init__.py b/tensorrt_llm/_torch/attention_backend/fmha/__init__.py index 1c3981abcf91..c4a54755de97 100644 --- a/tensorrt_llm/_torch/attention_backend/fmha/__init__.py +++ b/tensorrt_llm/_torch/attention_backend/fmha/__init__.py @@ -13,13 +13,18 @@ # See the License for the specific language governing permissions and # limitations under the License. +from .block_sparse import BlockSparseFmha from .fallback import FallbackFmha from .flashinfer_trtllm_gen import FlashInferTrtllmGenFmha +from .indexer_proxy import IndexerProxyFmha from .interface import Fmha +from .msa_proxy_mqa import MsaProxyMqaFmha +from .msa_sparse_gqa import MsaSparseGqaFmha from .phased import FmhaParams, PhasedFmha from .registry import DEFAULT_FMHA_LIBS, FMHA_LIBS, FmhaCls, get_enabled_fmha_lib_classes __all__ = [ + "BlockSparseFmha", "DEFAULT_FMHA_LIBS", "FMHA_LIBS", "FallbackFmha", @@ -27,6 +32,9 @@ "Fmha", "FmhaCls", "FmhaParams", + "IndexerProxyFmha", + "MsaProxyMqaFmha", + "MsaSparseGqaFmha", "PhasedFmha", "get_enabled_fmha_lib_classes", ] diff --git a/tensorrt_llm/_torch/attention_backend/fmha/block_sparse.py b/tensorrt_llm/_torch/attention_backend/fmha/block_sparse.py new file mode 100644 index 000000000000..ebb72fcb049e --- /dev/null +++ b/tensorrt_llm/_torch/attention_backend/fmha/block_sparse.py @@ -0,0 +1,166 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Abstract base for block-sparse paged-KV FMHA backends. + +Sparse-attention algorithms typically split work into two phases: + + 1. A predictor pass that produces a per-(query, KV head) list of + "selected KV blocks". See + :class:`tensorrt_llm._torch.attention_backend.fmha.indexer_proxy.IndexerProxyFmha` + for the FMHA library family that implements that phase. + 2. A *block-sparse* main attention pass that consumes the selected + block indices and runs the actual attention on a paged KV cache, + skipping unselected blocks. + +:class:`BlockSparseFmha` is the abstract base for the FMHA libraries +that implement phase (2). Like :class:`IndexerProxyFmha`, they live in +the same :data:`FMHA_LIBS` registry as standard main-attention FMHA +backends (FlashInfer trtllm-gen, fallback) so the same +``TLLM_FMHA_LIBS`` env var selects them; they opt out of the standard +:meth:`TrtllmAttention.forward` dispatch loop by returning ``False`` +from :meth:`is_supported` because their input contract +(``kv_block_indexes``, sparse-attention metadata) does not fit the +standard :class:`AttentionForwardArgs` signature, and they are invoked +directly by sparse-attention attention backends that have access to +those extra inputs. + +See +:class:`tensorrt_llm._torch.attention_backend.fmha.msa_sparse_gqa.MsaSparseGqaFmha` +for the canonical concrete implementation (MSA's ``fmha_sm100`` +sparse GQA kernel). +""" + +from __future__ import annotations + +from abc import abstractmethod +from typing import TYPE_CHECKING, Optional + +import torch + +from .interface import Fmha + +if TYPE_CHECKING: + from tensorrt_llm._torch.attention_backend.interface import AttentionForwardArgs + from tensorrt_llm._torch.attention_backend.trtllm import TrtllmAttentionMetadata + + +class BlockSparseFmha(Fmha): + """Abstract FMHA backend that consumes ``kv_block_indexes``. + + Block-sparse backends accept a per-query list of selected KV block + indices (produced by a sparse predictor; see + :class:`IndexerProxyFmha`) and run paged GQA attention restricted + to those blocks. They are invoked directly by sparse-attention + backends that own the predictor output, not by the standard + :meth:`TrtllmAttention.forward` dispatch loop. The standard + dispatch loop is opted out of via :meth:`is_supported` returning + ``False``. + + Concrete subclasses must implement :meth:`forward_block_sparse`, + which has a stable, dedicated signature carrying the + sparse-attention metadata that does not fit the standard + :class:`AttentionForwardArgs`. Sparse-attention backends locate + concrete subclasses via :func:`get_enabled_fmha_lib_classes` + filtered to subclasses of :class:`BlockSparseFmha` and call + :meth:`forward_block_sparse` directly. + """ + + @abstractmethod + def forward_block_sparse( + self, + q: torch.Tensor, + k_paged: torch.Tensor, + v_paged: torch.Tensor, + kv_block_indexes: torch.Tensor, + *, + qo_lens_cpu: torch.Tensor, + kv_lens_cpu: torch.Tensor, + qo_offset_cpu: Optional[torch.Tensor], + kv_indices: torch.Tensor, + sm_scale: float, + causal: bool, + ) -> torch.Tensor: + """Run block-sparse paged GQA attention. + + Parameters + ---------- + q : torch.Tensor + Shape ``[total_q, num_qo_heads, head_dim]`` (bf16/fp16). + k_paged : torch.Tensor + Paged K cache in HND layout + ``[num_pages, num_kv_heads, page_size, head_dim]``. + v_paged : torch.Tensor + Paged V cache, same shape as ``k_paged``. + kv_block_indexes : torch.Tensor + Shape ``[total_q, num_kv_heads, topk]``, dtype int32, + ascending per row with ``-1`` padding at the tail. Encodes + the per-query subset of KV blocks selected by the + preceding sparse predictor. + qo_lens_cpu, kv_lens_cpu : torch.Tensor + Shape ``[batch]``, dtype int32, on CPU. Per-request Q/O + and KV lengths. + qo_offset_cpu : torch.Tensor, optional + Shape ``[batch]``, dtype int32, on CPU. Per-request causal + offset (i.e. prefix length). Ignored when ``causal=False``. + kv_indices : torch.Tensor + Shape ``[sum_pages_across_batch]``, dtype int32, on the + cache device. Flattened paged-KV page table. + sm_scale : float + Softmax scale. + causal : bool + Whether to apply a causal mask. + + Returns + ------- + torch.Tensor + Shape ``[total_q, num_qo_heads, head_dim]``, dtype + bfloat16. The attention output over the selected KV blocks. + """ + ... + + def is_supported( + self, + q: torch.Tensor, + k: Optional[torch.Tensor], + v: Optional[torch.Tensor], + metadata: "TrtllmAttentionMetadata", + forward_args: "AttentionForwardArgs", + ) -> bool: + # Block-sparse backends consume sparse-attention metadata + # (kv_block_indexes, paged HND KV, per-batch lens) that does + # not fit AttentionForwardArgs. They are invoked directly by + # sparse-attention attention backends; returning False keeps + # us out of the standard TrtllmAttention.forward dispatch loop. + return False + + def forward( + self, + q: torch.Tensor, + k: Optional[torch.Tensor], + v: Optional[torch.Tensor], + metadata: "TrtllmAttentionMetadata", + forward_args: "AttentionForwardArgs", + ) -> None: + raise NotImplementedError( + f"{type(self).__name__} is a block-sparse FMHA backend; it is " + "invoked via forward_block_sparse() by sparse-attention " + "backends, not by the standard FMHA dispatch path. Locate it " + "via get_enabled_fmha_lib_classes() filtered to subclasses of " + "BlockSparseFmha." + ) + + +__all__ = ["BlockSparseFmha"] diff --git a/tensorrt_llm/_torch/attention_backend/fmha/indexer_proxy.py b/tensorrt_llm/_torch/attention_backend/fmha/indexer_proxy.py new file mode 100644 index 000000000000..06423c93ab6c --- /dev/null +++ b/tensorrt_llm/_torch/attention_backend/fmha/indexer_proxy.py @@ -0,0 +1,160 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Abstract base for indexer-style FMHA backends. + +Sparse-attention predictors (the MiniMax-M3 indexer, future Top-k +selectors) need a fast 'score every KV block against an MQA query' +pass over a paged KV cache. The score tensor is then fed into a top-k +selector to produce the sparse block indices consumed by the main +attention. + +That proxy attention is structurally a regular FMHA call -- it takes +``Q/K/V`` over a paged KV cache, runs causal attention, etc. -- but +its output is a ``max_score`` tensor rather than an attention output. +The :class:`IndexerProxyFmha` base class lets multiple proxy +implementations (MSA's ``fmha_sm100``, a future Triton path, etc.) +live in the same :data:`FMHA_LIBS` registry as main-attention FMHA +backends. They opt out of the main-attention dispatch loop by +returning ``False`` from :meth:`is_supported` and instead expose a +custom :meth:`forward_proxy` entry point that callers (sparse +indexers) invoke directly after looking the class up in the registry. + +See :class:`tensorrt_llm._torch.attention_backend.fmha.msa_proxy_mqa.MsaProxyMqaFmha` +for the canonical concrete implementation. +""" + +from __future__ import annotations + +from abc import abstractmethod +from typing import TYPE_CHECKING, Optional + +import torch + +from .interface import Fmha + +if TYPE_CHECKING: + from tensorrt_llm._torch.attention_backend.interface import AttentionForwardArgs + from tensorrt_llm._torch.attention_backend.trtllm import TrtllmAttentionMetadata + + +class IndexerProxyFmha(Fmha): + """Abstract FMHA backend that produces per-block max scores. + + Indexer-style backends are owned by sparse-attention indexers, not + by :class:`TrtllmAttention`. They are constructed without an + ``attn`` argument and never dispatched by the standard + ``TrtllmAttention.forward`` loop -- :meth:`is_supported` always + returns ``False``, so the loop skips past them when iterating + :data:`FMHA_LIBS`. + + Concrete subclasses must implement :meth:`forward_proxy`, which + has a custom (and stable) signature tailored to the proxy-MQA use + case. Indexers locate concrete subclasses via the standard + :func:`get_enabled_fmha_lib_classes` helper, filter by + ``issubclass(IndexerProxyFmha)`` and ``is_available()``, and call + ``forward_proxy`` directly. The Fmha registry remains the single + source of truth for which proxy implementations are reachable on + this build. + + Future expansion: this base may grow companion methods for other + indexer compute primitives (e.g. block-score reductions). The + no-op :meth:`is_supported` keeps the implementation outside the + main-attention dispatch path regardless of how many additional + methods are added. + """ + + @abstractmethod + def forward_proxy( + self, + idx_q: torch.Tensor, + idx_k_paged: torch.Tensor, + *, + qo_lens_cpu: torch.Tensor, + kv_lens_cpu: torch.Tensor, + qo_offset_cpu: Optional[torch.Tensor], + kv_indices: torch.Tensor, + sm_scale: float, + causal: bool, + ) -> torch.Tensor: + """Compute a per-(qo_head, kv_tile) max-score tensor. + + Parameters + ---------- + idx_q : torch.Tensor + Shape ``[total_q, num_qo_heads, head_dim]`` (bf16/fp16). + idx_k_paged : torch.Tensor + Paged KV in HND layout + ``[num_pages, num_kv_heads, page_size, head_dim]``. For the + canonical MQA proxy, ``num_kv_heads == 1`` and ``idx_k`` is + broadcast across every QO head during scoring. + qo_lens_cpu, kv_lens_cpu : torch.Tensor + Shape ``[batch]``, dtype int32, on CPU. Per-request Q/O + and KV lengths. + qo_offset_cpu : torch.Tensor, optional + Shape ``[batch]``, dtype int32, on CPU. Per-request causal + offset (i.e. prefix length). Ignored when ``causal=False``. + kv_indices : torch.Tensor + Shape ``[sum_pages_across_batch]``, dtype int32, on the + cache device. Flattened paged-KV page table. + sm_scale : float + Softmax scale applied to the QK scores prior to the + per-block max reduction. + causal : bool + Whether to apply a causal mask. Prefill batches typically + use ``True``; pure-decode batches use ``False``. + + Returns + ------- + torch.Tensor + Shape ``[num_qo_heads, max_k_tiles, total_q]``, dtype + float32. Out-of-range tile slots are padded with ``-inf`` + so a subsequent top-k selector can ignore them. + """ + ... + + def is_supported( + self, + q: torch.Tensor, + k: Optional[torch.Tensor], + v: Optional[torch.Tensor], + metadata: "TrtllmAttentionMetadata", + forward_args: "AttentionForwardArgs", + ) -> bool: + # Indexer-style backends never participate in the standard + # TrtllmAttention.forward dispatch loop. Returning False keeps + # us out of `for fmha in self.fmha_libs: if fmha.is_supported` + # so we don't have to spuriously claim/refuse main attention + # work. + return False + + def forward( + self, + q: torch.Tensor, + k: Optional[torch.Tensor], + v: Optional[torch.Tensor], + metadata: "TrtllmAttentionMetadata", + forward_args: "AttentionForwardArgs", + ) -> None: + raise NotImplementedError( + f"{type(self).__name__} is an indexer-style proxy FMHA backend; " + "it produces per-block max scores via forward_proxy() and is " + "not driven by the standard FMHA dispatch path. Sparse-attention " + "indexers should locate it via get_enabled_fmha_lib_classes() " + "filtered to subclasses of IndexerProxyFmha." + ) + + +__all__ = ["IndexerProxyFmha"] diff --git a/tensorrt_llm/_torch/attention_backend/fmha/interface.py b/tensorrt_llm/_torch/attention_backend/fmha/interface.py index f14f26964b25..5a24382d4adf 100644 --- a/tensorrt_llm/_torch/attention_backend/fmha/interface.py +++ b/tensorrt_llm/_torch/attention_backend/fmha/interface.py @@ -29,20 +29,40 @@ class Fmha(ABC): - """Common runtime contract for TRT-LLM attention FMHA libraries.""" + """Common runtime contract for TRT-LLM attention FMHA libraries. - def __init__(self, attn: "TrtllmAttention"): - self._attn_ref: weakref.ReferenceType["TrtllmAttention"] = weakref.ref(attn) + Most FMHA backends are owned by a :class:`TrtllmAttention` layer and + are driven by :meth:`TrtllmAttention.forward`. A small subset + (indexer-style proxy FMHA libraries used by sparse-attention + predictors) live in the same registry but have no owning attention + layer; they pass ``None`` for ``attn``. The :meth:`attn` property + raises only when callers actually dereference the owner, so the + no-owner subclasses never trip it as long as they do not call into + ``self.attn``. + """ + + def __init__(self, attn: Optional["TrtllmAttention"] = None): + self._attn_ref: Optional[weakref.ReferenceType["TrtllmAttention"]] = ( + weakref.ref(attn) if attn is not None else None + ) @property def attn(self) -> "TrtllmAttention": + if self._attn_ref is None: + raise RuntimeError( + f"{type(self).__name__} was constructed without an owning " + "TrtllmAttention instance. This typically means an " + "indexer-style FMHA backend was asked for a property " + "that only makes sense for main-attention FMHA " + "backends." + ) attn = self._attn_ref() if attn is None: raise RuntimeError("The owning TrtllmAttention instance has been garbage collected.") return attn @classmethod - def is_available(cls, attn: "TrtllmAttention") -> bool: + def is_available(cls, attn: Optional["TrtllmAttention"] = None) -> bool: return True def is_supported( diff --git a/tensorrt_llm/_torch/attention_backend/fmha/msa_proxy_mqa.py b/tensorrt_llm/_torch/attention_backend/fmha/msa_proxy_mqa.py new file mode 100644 index 000000000000..837a258241e3 --- /dev/null +++ b/tensorrt_llm/_torch/attention_backend/fmha/msa_proxy_mqa.py @@ -0,0 +1,150 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Proxy MQA FMHA backed by MSA's ``fmha_sm100`` kernel. + +This module provides :class:`MsaProxyMqaFmha`, an +:class:`IndexerProxyFmha` implementation that wraps MSA's +``fmha_sm100`` dense FMHA in ``output_maxscore`` mode for use by +sparse-attention indexers (currently the MiniMax-M3 indexer). + +The kernel is SM100-only and the ``fmha_sm100`` Python package is an +optional external dependency (https://github.com/MiniMax-AI/MSA). On +hosts where either precondition is missing, :meth:`is_available` +returns ``False`` so the registry skips the class. +""" + +from __future__ import annotations + +from typing import Optional + +import torch + +from tensorrt_llm.logger import logger + +from .indexer_proxy import IndexerProxyFmha + + +class MsaProxyMqaFmha(IndexerProxyFmha): + """SM100 MQA proxy FMHA powered by MSA's ``fmha_sm100`` kernel. + + The MiniMax-M3 sparse-attention indexer (and, in the future, + other sparse predictors that need to score every KV block against + an MQA query) selects this backend when ``fmha_sm100`` is + importable and the current device exposes compute capability + 10 (SM100 family). + + Hard requirements (checked at runtime in :meth:`forward_proxy`): + * ``idx_q`` head dim is 128 -- the only ``fmha_sm100`` variant + shipped today. + * ``idx_k_paged`` is 4-D HND with ``num_kv_heads == 1`` (MQA). + """ + + HEAD_DIM = 128 + + @classmethod + def is_available(cls, attn=None) -> bool: + try: + import fmha_sm100 # noqa: F401 + except ImportError: + logger.debug("MsaProxyMqaFmha is unavailable: fmha_sm100 package not installed.") + return False + if not torch.cuda.is_available(): + logger.debug("MsaProxyMqaFmha is unavailable: no CUDA device.") + return False + try: + major, _ = torch.cuda.get_device_capability() + except Exception: + return False + if major != 10: + logger.debug( + "MsaProxyMqaFmha is unavailable: requires SM100 (compute capability 10.x), " + f"got compute capability {major}.x." + ) + return False + return True + + def forward_proxy( + self, + idx_q: torch.Tensor, + idx_k_paged: torch.Tensor, + *, + qo_lens_cpu: torch.Tensor, + kv_lens_cpu: torch.Tensor, + qo_offset_cpu: Optional[torch.Tensor], + kv_indices: torch.Tensor, + sm_scale: float, + causal: bool, + ) -> torch.Tensor: + """Run the MQA proxy FMHA and return ``[num_qo_heads, max_k_tiles, total_q]``. + + The implementation follows MSA's documented two-call pattern: + ``fmha_sm100_plan`` builds the per-shape plan (CUDA-graph-stable + worklists, KV-split workspaces) and ``fmha_sm100`` runs the + kernel with ``output_o=False, output_maxscore=True`` so only + the score tensor is materialized. + """ + # Imported here (not at module top) so the registry can still + # advertise the class on hosts where fmha_sm100 is absent -- + # is_available() handles the off-host case. + import fmha_sm100 + + if idx_q.dim() != 3: + raise ValueError( + "MsaProxyMqaFmha expects idx_q with shape [total_q, num_qo_heads, head_dim]; " + f"got {tuple(idx_q.shape)}." + ) + if idx_q.shape[-1] != self.HEAD_DIM: + raise NotImplementedError( + f"MsaProxyMqaFmha currently supports head_dim={self.HEAD_DIM}; " + f"got {idx_q.shape[-1]}." + ) + if idx_k_paged.dim() != 4 or idx_k_paged.shape[1] != 1: + raise ValueError( + "MsaProxyMqaFmha expects MQA paged KV " + "[num_pages, 1, page_size, head_dim]; " + f"got {tuple(idx_k_paged.shape)}." + ) + if idx_k_paged.shape[-1] != self.HEAD_DIM: + raise NotImplementedError( + f"MsaProxyMqaFmha currently supports head_dim={self.HEAD_DIM}; " + f"got idx_k_paged head_dim={idx_k_paged.shape[-1]}." + ) + + page_size = int(idx_k_paged.shape[2]) + proxy_plan = fmha_sm100.fmha_sm100_plan( + qo_lens_cpu, + kv_lens_cpu, + idx_q.shape[1], # num_qo_heads (= num_index_heads for M3) + num_kv_heads=1, + qo_offset=qo_offset_cpu, + page_size=page_size, + output_maxscore=True, + causal=causal, + ) + _, max_score = fmha_sm100.fmha_sm100( + idx_q, + idx_k_paged, + idx_k_paged, # v passthrough -- proxy ignores V via output_o=False + proxy_plan, + kv_indices=kv_indices, + output_o=False, + output_maxscore=True, + sm_scale=sm_scale, + ) + return max_score + + +__all__ = ["MsaProxyMqaFmha"] diff --git a/tensorrt_llm/_torch/attention_backend/fmha/msa_sparse_gqa.py b/tensorrt_llm/_torch/attention_backend/fmha/msa_sparse_gqa.py new file mode 100644 index 000000000000..c014e8f2354f --- /dev/null +++ b/tensorrt_llm/_torch/attention_backend/fmha/msa_sparse_gqa.py @@ -0,0 +1,162 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Block-sparse GQA FMHA backed by MSA's ``fmha_sm100`` kernel. + +This module provides :class:`MsaSparseGqaFmha`, a +:class:`BlockSparseFmha` implementation that wraps MSA's +``fmha_sm100`` paged sparse GQA kernel. It consumes the +``kv_block_indexes`` produced by an upstream proxy + top-k pass (see +:class:`MsaProxyMqaFmha`) and runs the main attention on only the +selected KV blocks. + +The kernel is SM100-only and the ``fmha_sm100`` Python package is an +optional external dependency (https://github.com/MiniMax-AI/MSA). On +hosts where either precondition is missing, :meth:`is_available` +returns ``False`` so the registry skips the class. +""" + +from __future__ import annotations + +from typing import Optional + +import torch + +from tensorrt_llm.logger import logger + +from .block_sparse import BlockSparseFmha + + +class MsaSparseGqaFmha(BlockSparseFmha): + """SM100 block-sparse GQA FMHA powered by MSA's ``fmha_sm100`` kernel. + + Consumes ``kv_block_indexes`` (typically produced by + :class:`MsaProxyMqaFmha` + ``sparse_topk_select`` upstream) and + runs paged GQA attention over the selected blocks. Used by + :mod:`tensorrt_llm._torch.attention_backend.sparse.minimax_m3.msa_backend` + for the MiniMax-M3 sparse layers' main attention pass. + + Hard requirements (checked at runtime in + :meth:`forward_block_sparse`): + * ``q``/``k``/``v`` head dim is 128 -- the only ``fmha_sm100`` + variant shipped today. + * ``k_paged`` and ``v_paged`` are 4-D HND paged caches with + matching ``num_kv_heads`` and ``page_size``. + """ + + HEAD_DIM = 128 + + @classmethod + def is_available(cls, attn=None) -> bool: + try: + import fmha_sm100 # noqa: F401 + except ImportError: + logger.debug("MsaSparseGqaFmha is unavailable: fmha_sm100 package not installed.") + return False + if not torch.cuda.is_available(): + logger.debug("MsaSparseGqaFmha is unavailable: no CUDA device.") + return False + try: + major, _ = torch.cuda.get_device_capability() + except Exception: + return False + if major != 10: + logger.debug( + "MsaSparseGqaFmha is unavailable: requires SM100 (compute capability 10.x), " + f"got compute capability {major}.x." + ) + return False + return True + + def forward_block_sparse( + self, + q: torch.Tensor, + k_paged: torch.Tensor, + v_paged: torch.Tensor, + kv_block_indexes: torch.Tensor, + *, + qo_lens_cpu: torch.Tensor, + kv_lens_cpu: torch.Tensor, + qo_offset_cpu: Optional[torch.Tensor], + kv_indices: torch.Tensor, + sm_scale: float, + causal: bool, + ) -> torch.Tensor: + """Run block-sparse paged GQA and return ``[total_q, num_qo_heads, head_dim]``. + + Follows MSA's two-call pattern: ``fmha_sm100_plan`` builds the + per-shape sparse plan (with ``kv_block_num`` derived from + ``kv_block_indexes.shape[-1]``) and ``fmha_sm100`` runs the + kernel with the block indices threaded through. + """ + # Imported here (not at module top) so the registry can still + # advertise the class on hosts where fmha_sm100 is absent -- + # is_available() handles the off-host case. + import fmha_sm100 + + if q.dim() != 3: + raise ValueError( + "MsaSparseGqaFmha expects q with shape [total_q, num_qo_heads, head_dim]; " + f"got {tuple(q.shape)}." + ) + if q.shape[-1] != self.HEAD_DIM: + raise NotImplementedError( + f"MsaSparseGqaFmha currently supports head_dim={self.HEAD_DIM}; got {q.shape[-1]}." + ) + if k_paged.dim() != 4 or v_paged.dim() != 4: + raise ValueError( + "MsaSparseGqaFmha expects paged KV with shape " + "[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 k_paged.shape[-1] != self.HEAD_DIM: + raise NotImplementedError( + f"MsaSparseGqaFmha currently supports head_dim={self.HEAD_DIM}; " + f"got k_paged head_dim={k_paged.shape[-1]}." + ) + + num_qo_heads = int(q.shape[1]) + num_kv_heads = int(k_paged.shape[1]) + page_size = int(k_paged.shape[2]) + + sparse_plan = fmha_sm100.fmha_sm100_plan( + qo_lens_cpu, + kv_lens_cpu, + num_qo_heads, + num_kv_heads=num_kv_heads, + qo_offset=qo_offset_cpu, + page_size=page_size, + kv_block_num=int(kv_block_indexes.shape[-1]), + causal=causal, + ) + out, _ = fmha_sm100.fmha_sm100( + q, + k_paged, + v_paged, + sparse_plan, + kv_indices=kv_indices, + kv_block_indexes=kv_block_indexes, + sm_scale=sm_scale, + output_maxscore=False, + ) + return out + + +__all__ = ["MsaSparseGqaFmha"] diff --git a/tensorrt_llm/_torch/attention_backend/fmha/registry.py b/tensorrt_llm/_torch/attention_backend/fmha/registry.py index 97467657e9b5..0ec9c0a79f10 100644 --- a/tensorrt_llm/_torch/attention_backend/fmha/registry.py +++ b/tensorrt_llm/_torch/attention_backend/fmha/registry.py @@ -19,12 +19,24 @@ from .fallback import FallbackFmha from .flashinfer_trtllm_gen import FlashInferTrtllmGenFmha from .interface import Fmha +from .msa_proxy_mqa import MsaProxyMqaFmha +from .msa_sparse_gqa import MsaSparseGqaFmha FmhaCls: TypeAlias = type[Fmha] FMHA_LIBS: dict[str, FmhaCls] = { "flashinfer_trtllm_gen": FlashInferTrtllmGenFmha, "fallback": FallbackFmha, + # Indexer-style proxy FMHA. Returns False from is_supported() so the + # main-attention dispatch loop ignores it; sparse-attention indexers + # locate it via get_enabled_fmha_lib_classes() filtered to + # subclasses of IndexerProxyFmha. + "msa_proxy_mqa": MsaProxyMqaFmha, + # Block-sparse main-attention FMHA. Consumes kv_block_indexes + # produced by an upstream proxy + top-k pass. Same opt-out pattern + # as the proxy: returns False from is_supported() and is invoked + # directly by sparse-attention backends via forward_block_sparse(). + "msa_sparse_gqa": MsaSparseGqaFmha, } DEFAULT_FMHA_LIBS: tuple[str, ...] = tuple(FMHA_LIBS) diff --git a/tensorrt_llm/_torch/attention_backend/sparse/minimax_m3/__init__.py b/tensorrt_llm/_torch/attention_backend/sparse/minimax_m3/__init__.py index c248ffd91119..e3f1a79745e3 100644 --- a/tensorrt_llm/_torch/attention_backend/sparse/minimax_m3/__init__.py +++ b/tensorrt_llm/_torch/attention_backend/sparse/minimax_m3/__init__.py @@ -58,6 +58,12 @@ get_minimax_m3_attention_metadata_cls, replace_metadata, ) +from .msa_backend import ( + get_minimax_m3_attention_backend_cls_with_msa, + get_minimax_m3_msa_attention_backend_cls, + minimax_m3_msa_sparse_decode, + minimax_m3_msa_sparse_prefill, +) __all__ = [ "MiniMaxM3KVCacheManagerV2", @@ -68,8 +74,12 @@ "allocate_minimax_m3_static_buffers", "build_runtime_metadata_from_kv_manager", "get_minimax_m3_attention_backend_cls", + "get_minimax_m3_attention_backend_cls_with_msa", "get_minimax_m3_attention_metadata_cls", "get_minimax_m3_kv_cache_manager_cls", + "get_minimax_m3_msa_attention_backend_cls", + "minimax_m3_msa_sparse_decode", + "minimax_m3_msa_sparse_prefill", "minimax_m3_sparse_decode", "minimax_m3_sparse_prefill", "replace_metadata", diff --git a/tensorrt_llm/_torch/attention_backend/sparse/minimax_m3/metadata.py b/tensorrt_llm/_torch/attention_backend/sparse/minimax_m3/metadata.py index 925ca6c0189b..790e7694cf9b 100644 --- a/tensorrt_llm/_torch/attention_backend/sparse/minimax_m3/metadata.py +++ b/tensorrt_llm/_torch/attention_backend/sparse/minimax_m3/metadata.py @@ -43,6 +43,13 @@ class MiniMaxM3SparseParams(SparseParams): local_blocks: int = 1 score_type: str = "max" disable_index_value: bool = True + # When True, the layer dispatches the sparse forward through the + # MSA-backed FMHA runtime (``fmha_sm100`` + ``sparse_topk_select``) + # instead of the in-tree Triton + SDPA reference path. The MSA stack + # is only available on SM100 and requires the external + # ``fmha_sm100`` package; the layer raises a descriptive error if + # it is requested without those preconditions met. + use_msa: bool = False @dataclass(frozen=True) diff --git a/tensorrt_llm/_torch/attention_backend/sparse/minimax_m3/msa_backend.py b/tensorrt_llm/_torch/attention_backend/sparse/minimax_m3/msa_backend.py new file mode 100644 index 000000000000..6cfb24938ad9 --- /dev/null +++ b/tensorrt_llm/_torch/attention_backend/sparse/minimax_m3/msa_backend.py @@ -0,0 +1,794 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""MSA-backed FMHA runtime for MiniMax-M3 sparse attention. + +Implements the same two-step sparse attention as +:mod:`tensorrt_llm._torch.attention_backend.sparse.minimax_m3.backend`, +but lowers the per-block max score, top-k selection, and sparse GQA to +the external ``fmha_sm100`` kernels (a.k.a. MSA, MiniMax Sparse +Attention; see https://github.com/MiniMax-AI/MSA) instead of the +in-tree Triton + SDPA reference path. + +The flow per forward call is: + + 1. **Index proxy pass.** Run a dense MQA FMHA over the index branch + with ``output_maxscore=True`` and ``num_kv_heads=1``. ``fmha_sm100`` + returns ``(None, max_score)`` where ``max_score`` has shape + ``[num_index_heads, max_k_tiles, total_qo_len]``. + 2. **Block selection.** Feed ``max_score`` into ``sparse_topk_select`` + to produce ascending per-row top-k block indices with ``-1`` + padding. The kernel currently fixes ``topk = 16`` and the + MiniMax-M3 checkpoint matches that default exactly. + 3. **Sparse GQA.** Run a second ``fmha_sm100`` over the main K/V + branch, passing the block indices via ``kv_block_indexes``. + +This module deliberately keeps the cache-layout adapter explicit so +the MSA backend is selectable per-layer without disturbing the +existing Triton path. The MSA package is imported lazily; if it is +absent the backend raises a descriptive error at construction time. +""" + +from __future__ import annotations + +import functools +from typing import TYPE_CHECKING, Optional, Tuple + +import torch + +from .backend import _write_main_kv_slots, get_minimax_m3_attention_backend_cls +from .metadata import ( + MiniMaxM3SparseAttentionMetadata, + MiniMaxM3SparseConfig, + ensure_metadata_on_device, +) + +if TYPE_CHECKING: + from .metadata import MiniMaxM3SparseParams + + +# MSA's ``sparse_topk_select`` kernel only ships a topk=16 path today, +# and ``fmha_sm100`` only ships head_dim=128 variants. Enforce these +# preconditions early so layer construction fails with a clear message +# rather than a cryptic shape error from inside the MSA JIT. +_MSA_REQUIRED_TOPK = 16 +_MSA_REQUIRED_HEAD_DIM = 128 + + +def _require_msa_module(): + """Lazy-import ``fmha_sm100`` and raise a clear error on failure. + + Returns the imported module. The import is guarded so that the + MSA backend can be advertised in the config schema even on hosts + where MSA is not installed; the error only fires when a sparse + layer actually attempts to dispatch the MSA path. + """ + try: + import fmha_sm100 # noqa: F401 + except ImportError as exc: # pragma: no cover - install-time error + raise RuntimeError( + "MiniMax-M3 MSA backend requires the external `fmha_sm100` " + "package (MSA: https://github.com/MiniMax-AI/MSA). Install " + "it with `pip install fmha_sm100`, or unset " + "`sparse_use_msa` in the sparse attention config to fall " + "back to the Triton reference path." + ) from exc + return fmha_sm100 + + +# --------------------------------------------------------------------------- +# Cache layout adapters +# --------------------------------------------------------------------------- +# +# TRT-LLM's :class:`KVCacheManagerV2` stores the main K/V cache as a 5-D +# pool ``[num_pages, kv_factor, tokens_per_block, num_kv_heads, head_dim]`` +# with NHD ordering. ``fmha_sm100`` expects paged K/V tensors with +# layout ``[num_pages, num_kv_heads, page_size, head_dim]`` (HND). +# Sparse layers also carry a side index-K cache whose flat-slot layout +# is ``[num_slots, 1, sparse_index_dim]`` or the 4-D paged variant +# ``[num_pages, tokens_per_block, 1, sparse_index_dim]``. +# +# The helpers below produce MSA-compatible views without mutating the +# pool's underlying storage. They permute and call ``.contiguous()`` +# once per call; that is the per-call cost we accept until the cache +# manager grows a native HND view. + + +def _cache_view_to_msa_paged(cache_view: torch.Tensor) -> torch.Tensor: + """Convert ``[num_pages, page_size, num_kv_heads, head_dim]`` -> HND. + + The 4-D ``cache_view`` is the non-contiguous slice ``kv_pool[:, k]`` + (``k=0`` for K, ``k=1`` for V) returned by ``KVCacheManagerV2``. + MSA's ``fmha_sm100`` expects ``[num_pages, num_kv_heads, page_size, + head_dim]``; we permute dims 1 and 2 and force a contiguous copy + so the kernel reads sequential memory. + + For flat-slot 3-D caches (focused unit tests) the function + interprets the cache as a single virtual page of size ``num_slots`` + and returns ``[1, num_kv_heads, num_slots, head_dim]``. That keeps + the test surface working without forcing every test to provide a + real paged cache. + """ + if cache_view.dim() == 4: + return cache_view.permute(0, 2, 1, 3).contiguous() + if cache_view.dim() == 3: + # Flat-slot fallback. Treat the cache as a single page so the + # MSA kernel's page-table indexing still resolves correctly. + return cache_view.permute(1, 0, 2).unsqueeze(0).contiguous() + raise ValueError( + f"Unsupported cache view rank {cache_view.dim()} for MSA paged conversion; " + f"expected 3 (flat-slot) or 4 (paged multi-dim)." + ) + + +def _idx_cache_to_msa_paged(idx_cache: torch.Tensor) -> torch.Tensor: + """Convert the side index-K cache to MSA's HND paged layout. + + Accepts: + * 3-D ``[num_slots, 1, sparse_index_dim]`` -> ``[1, 1, num_slots, sparse_index_dim]`` + * 4-D ``[num_pages, tokens_per_block, 1, sparse_index_dim]`` + -> ``[num_pages, 1, tokens_per_block, sparse_index_dim]`` + """ + if idx_cache.dim() == 4: + return idx_cache.permute(0, 2, 1, 3).contiguous() + if idx_cache.dim() == 3: + return idx_cache.permute(1, 0, 2).unsqueeze(0).contiguous() + raise ValueError(f"Unsupported index cache rank {idx_cache.dim()} for MSA paged conversion.") + + +def _page_size_from_view(cache_view: torch.Tensor) -> int: + if cache_view.dim() == 4: + return int(cache_view.shape[1]) + if cache_view.dim() == 3: + return int(cache_view.shape[0]) + raise ValueError(f"Unsupported cache view rank {cache_view.dim()} for page-size lookup.") + + +def _build_kv_indices_and_lens( + metadata: MiniMaxM3SparseAttentionMetadata, + page_size: int, +) -> Tuple[torch.Tensor, torch.Tensor]: + """Build ``kv_indices`` and ``kv_segment_lens`` for ``fmha_sm100``. + + ``kv_indices`` is the flattened per-request page table: + ``concat([pages_of(seq_0), pages_of(seq_1), ...])`` with dtype + int32. The pages of a sequence come from + ``metadata.req_to_token[slot_id, ::page_size] // page_size`` — i.e. + the page index of every block start. ``kv_segment_lens`` carries + the per-request effective KV length (already on the cache device + inside ``metadata``). + + The helper assumes ``page_size == metadata.req_to_token.stride[-1]`` + block geometry, which is the contract :class:`MiniMaxM3KVCacheManagerV2` + enforces (``tokens_per_block == sparse_block_size``). + """ + device = metadata.req_to_token.device + slot_ids_long = metadata.slot_ids.to(torch.long) + req_rows = metadata.req_to_token.index_select(0, slot_ids_long).to(torch.long) + batch = int(req_rows.shape[0]) + max_kv_len = int(req_rows.shape[1]) + seq_lens_cpu = metadata.seq_lens_cpu.to(torch.long).tolist() + + page_lists = [] + for b in range(batch): + kv_len = int(seq_lens_cpu[b]) + if kv_len <= 0: + continue + num_pages = (kv_len + page_size - 1) // page_size + # First slot of each page gives the page id (each block is a + # contiguous run of page_size slots, see KVCacheManagerV2). + # Clamp so trailing pages (rounded up beyond the request's + # block count) reuse the last valid block - they are masked out + # by seq_lens in the kernel. + max_page = max_kv_len // page_size + page_starts = torch.arange(num_pages, device=device, dtype=torch.long) * page_size + page_starts = page_starts.clamp_max(max(0, max_kv_len - 1)) + page_ids = req_rows[b].gather(0, page_starts) // page_size + page_ids = page_ids.clamp_min(0).clamp_max(max(0, max_page - 1)) + page_lists.append(page_ids.to(torch.int32)) + + if page_lists: + kv_indices = torch.cat(page_lists, dim=0) + else: + kv_indices = torch.empty(0, dtype=torch.int32, device=device) + + return kv_indices, metadata.seq_lens.to(torch.int32) + + +# --------------------------------------------------------------------------- +# Forward primitives +# --------------------------------------------------------------------------- + + +@functools.lru_cache(maxsize=1) +def _select_proxy_fmha_class(): + """Pick the first available :class:`IndexerProxyFmha` from the registry. + + Walks :func:`get_enabled_fmha_lib_classes` and returns the first + class that (a) is a subclass of :class:`IndexerProxyFmha` and (b) + reports :meth:`is_available` ``True``. Returns ``None`` if no proxy + backend is available; callers raise a descriptive error pointing + at the ``TLLM_FMHA_LIBS`` env var in that case. + + The result is cached because the registry contents only change at + process start (driven by the ``TLLM_FMHA_LIBS`` env var, also + cached). + """ + from ...fmha import IndexerProxyFmha, get_enabled_fmha_lib_classes + + for cls in get_enabled_fmha_lib_classes(): + if not issubclass(cls, IndexerProxyFmha): + continue + if cls.is_available(): + return cls + return None + + +def _msa_index_proxy_and_topk( + idx_q: torch.Tensor, + idx_k_paged: torch.Tensor, + *, + qo_lens_cpu: torch.Tensor, + kv_lens_cpu: torch.Tensor, + qo_offset_cpu: Optional[torch.Tensor], + kv_indices: torch.Tensor, + config: MiniMaxM3SparseConfig, + idx_sm_scale: float, + causal: bool, + init_blocks: int, + local_blocks: int, +) -> torch.Tensor: + """Run a proxy FMHA + top-k block selection. + + The proxy attention is delegated to an :class:`IndexerProxyFmha` + implementation looked up via the standard FMHA registry (defaults + to :class:`MsaProxyMqaFmha` when MSA is installed). The + group-reduction and ``sparse_topk_select`` step remain here because + they are MiniMax-M3-specific algorithmic choices, not generic + proxy concerns. + + Inputs + ------ + idx_q : ``[total_q, num_index_heads, sparse_index_dim]``, bf16/fp16. + idx_k_paged : ``[num_pages, 1, page_size, sparse_index_dim]`` + Index K cache in MSA's HND paged layout. The dense proxy pass + runs MQA (``num_kv_heads_dense=1``) so the single replicated + index-K head is consumed directly. + qo_lens_cpu, kv_lens_cpu : ``[batch]``, int32, on CPU. + qo_offset_cpu : optional ``[batch]`` int32 on CPU. Causal offset + per request; ``None`` defaults to ``kv_lens - qo_lens``. + kv_indices : ``[sum_pages]`` int32 on the cache device. + config : layer-invariant sparse config. + idx_sm_scale : softmax scale for the index attention. + causal : whether to apply causal masking (True for prefill, False + for pure-decode batches). + init_blocks, local_blocks : forced-include block counts. + + Returns + ------- + torch.Tensor + ``[total_q, num_kv_heads, topk]`` int32 ascending top-k block + indices with ``-1`` padding (MSA kernel contract). When + ``num_index_heads > num_kv_heads`` the per-block max score is + reduced across each KV head's index-head group before + ``sparse_topk_select``, mirroring the + ``score_type='max'`` reduction the reference path performs. + """ + # ``sparse_topk_select`` still lives in fmha_sm100 -- the + # top-k step is not factored into the FMHA registry today. + # (Future work: register it as a separate selector library.) + fmha_sm100 = _require_msa_module() + + proxy_cls = _select_proxy_fmha_class() + if proxy_cls is None: + raise RuntimeError( + "No IndexerProxyFmha backend is available. The MiniMax-M3 MSA " + "backend needs an indexer-style proxy FMHA enabled via " + "TLLM_FMHA_LIBS (e.g. 'msa_proxy_mqa') and the corresponding " + "external dependency installed (e.g. the `fmha_sm100` package " + "for `msa_proxy_mqa`). Set TLLM_FMHA_LIBS=msa_proxy_mqa or " + "leave TLLM_FMHA_LIBS at its default to enable all registered " + "backends." + ) + + proxy = proxy_cls() + max_score = proxy.forward_proxy( + 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=causal, + ) + + # ``max_score`` has shape ``[num_index_heads, max_k_tiles, total_q]``. + # Reduce across each KV head's index-head group via amax so the + # downstream ``sparse_topk_select`` picks blocks per KV head, which + # is exactly the granularity ``fmha_sm100``'s sparse FMHA expects. + if config.num_index_heads % config.num_kv_heads != 0: + raise ValueError( + f"num_index_heads ({config.num_index_heads}) must be divisible by " + f"num_kv_heads ({config.num_kv_heads}) for MSA group-max reduction." + ) + group = config.num_index_heads // config.num_kv_heads + if group > 1: + max_score_kv = max_score.view( + config.num_kv_heads, group, max_score.shape[1], max_score.shape[2] + ).amax(dim=1) + else: + max_score_kv = max_score + + # ``num_valid_pages`` is per-call so the kernel masks out the + # rounded-up tail tiles; we pass the maximum across the batch and + # rely on the kernel's ``idx >= num_valid_pages`` check. + page_size = int(idx_k_paged.shape[2]) + max_valid_pages = ( + int(((kv_lens_cpu + page_size - 1) // page_size).max().item()) if kv_lens_cpu.numel() else 0 + ) + if max_valid_pages <= 0: + # Degenerate batch (no KV) — return all-padded indices. + return torch.full( + (idx_q.shape[0], config.num_kv_heads, _MSA_REQUIRED_TOPK), + -1, + dtype=torch.int32, + device=idx_q.device, + ) + + return fmha_sm100.sparse_topk_select( + max_score_kv.contiguous(), + _MSA_REQUIRED_TOPK, + num_valid_pages=max_valid_pages, + force_begin_blocks=init_blocks, + force_end_blocks=local_blocks, + ) + + +@functools.lru_cache(maxsize=1) +def _select_block_sparse_fmha_class(): + """Pick the first available :class:`BlockSparseFmha` from the registry. + + Mirrors :func:`_select_proxy_fmha_class` for the main attention + pass: walks :func:`get_enabled_fmha_lib_classes` looking for any + class that (a) is a subclass of :class:`BlockSparseFmha` and (b) + reports :meth:`is_available` ``True``. Returns ``None`` if no + block-sparse backend is available; callers raise a descriptive + error pointing at the ``TLLM_FMHA_LIBS`` env var. + """ + from ...fmha import BlockSparseFmha, get_enabled_fmha_lib_classes + + for cls in get_enabled_fmha_lib_classes(): + if not issubclass(cls, BlockSparseFmha): + continue + if cls.is_available(): + return cls + return None + + +def _msa_sparse_attention( + q: torch.Tensor, + k_paged: torch.Tensor, + v_paged: torch.Tensor, + kv_block_indexes: torch.Tensor, + *, + qo_lens_cpu: torch.Tensor, + kv_lens_cpu: torch.Tensor, + qo_offset_cpu: Optional[torch.Tensor], + kv_indices: torch.Tensor, + sm_scale: float, + causal: bool, +) -> torch.Tensor: + """Run block-sparse paged GQA over the selected blocks. + + The kernel call is delegated to a :class:`BlockSparseFmha` + implementation looked up via the standard FMHA registry (defaults + to :class:`MsaSparseGqaFmha` when MSA is installed). + ``num_q_heads`` / ``num_kv_heads`` / ``page_size`` are derived from + the tensor shapes inside the backend so callers do not need to + pass them explicitly. + + Returns ``[total_q, num_q_heads, head_dim]`` bfloat16. + """ + sparse_cls = _select_block_sparse_fmha_class() + if sparse_cls is None: + raise RuntimeError( + "No BlockSparseFmha backend is available. The MiniMax-M3 MSA " + "backend needs a block-sparse FMHA enabled via TLLM_FMHA_LIBS " + "(e.g. 'msa_sparse_gqa') and the corresponding external " + "dependency installed (e.g. the `fmha_sm100` package for " + "`msa_sparse_gqa`). Set TLLM_FMHA_LIBS=msa_sparse_gqa or " + "leave TLLM_FMHA_LIBS at its default to enable all " + "registered backends." + ) + + sparse = sparse_cls() + return sparse.forward_block_sparse( + q, + k_paged, + v_paged, + kv_block_indexes, + qo_lens_cpu=qo_lens_cpu, + kv_lens_cpu=kv_lens_cpu, + qo_offset_cpu=qo_offset_cpu, + kv_indices=kv_indices, + sm_scale=sm_scale, + causal=causal, + ) + + +# --------------------------------------------------------------------------- +# Public forward entry points +# --------------------------------------------------------------------------- + + +def _qo_lens_offsets_from_metadata( + metadata: MiniMaxM3SparseAttentionMetadata, +) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """Extract CPU qo_lens / kv_lens / qo_offset tensors. + + Returns ``(qo_lens_cpu, kv_lens_cpu, qo_offset_cpu)`` all on CPU + with dtype int32. ``qo_offset_cpu`` is ``None`` for pure-decode + batches (causal=False) and is the per-request prefix length for + prefill batches. + """ + seq_lens_cpu = metadata.seq_lens_cpu.to(torch.int32) + if metadata.is_prefill: + if metadata.extend_seq_lens_cpu is None: + raise RuntimeError("prefill metadata requires extend_seq_lens_cpu") + qo_lens_cpu = torch.tensor(metadata.extend_seq_lens_cpu, dtype=torch.int32) + # ``prefix_lens`` is the causal offset. + if metadata.prefix_lens is None: + raise RuntimeError("prefill metadata requires prefix_lens") + qo_offset_cpu = metadata.prefix_lens.detach().to(device="cpu", dtype=torch.int32) + else: + batch = int(metadata.slot_ids.shape[0]) + qo_lens_cpu = torch.ones(batch, dtype=torch.int32) + qo_offset_cpu = (seq_lens_cpu - 1).to(torch.int32) + return qo_lens_cpu, seq_lens_cpu, qo_offset_cpu + + +def minimax_m3_msa_sparse_prefill( + q: torch.Tensor, + k_cache: torch.Tensor, + v_cache: torch.Tensor, + idx_q: torch.Tensor, + idx_k_cache: torch.Tensor, + metadata: MiniMaxM3SparseAttentionMetadata, + config: MiniMaxM3SparseConfig, + *, + sm_scale: Optional[float] = None, + idx_sm_scale: Optional[float] = None, +) -> torch.Tensor: + """MiniMax-M3 sparse prefill backed by MSA's ``fmha_sm100``. + + Inputs follow the same conventions as + :func:`minimax_m3_sparse_prefill`. The main differences are: + + * The K/V caches are accepted in TRT-LLM's pool layout and + permuted to MSA's HND paged layout internally. + * The index branch's ``idx_v_cache`` is unused: MSA's proxy pass + only consumes ``max_score``, so the M3 default + ``disable_index_value=True`` is the only supported mode. + * ``num_index_heads`` must be a multiple of ``num_kv_heads``; + the per-block max score is reduced with ``amax`` over each KV + head's index-head group before ``sparse_topk_select`` runs. + """ + if not metadata.is_prefill: + raise ValueError("MSA prefill entry called with decode metadata") + if metadata.q_batch_row is None or metadata.q_positions is None: + raise ValueError("prefill metadata requires q_batch_row / q_positions") + if config.head_dim != _MSA_REQUIRED_HEAD_DIM: + raise NotImplementedError( + f"MSA backend currently supports head_dim={_MSA_REQUIRED_HEAD_DIM}; " + f"got {config.head_dim}." + ) + if config.sparse_index_dim != _MSA_REQUIRED_HEAD_DIM: + raise NotImplementedError( + f"MSA backend requires sparse_index_dim={_MSA_REQUIRED_HEAD_DIM}; " + f"got {config.sparse_index_dim}." + ) + if config.topk != _MSA_REQUIRED_TOPK: + raise NotImplementedError( + f"MSA backend currently supports topk={_MSA_REQUIRED_TOPK}; got {config.topk}." + ) + + sm_scale = sm_scale if sm_scale is not None else config.head_dim**-0.5 + idx_sm_scale = idx_sm_scale if idx_sm_scale is not None else config.sparse_index_dim**-0.5 + + k_paged = _cache_view_to_msa_paged(k_cache) + v_paged = _cache_view_to_msa_paged(v_cache) + idx_k_paged = _idx_cache_to_msa_paged(idx_k_cache) + page_size = _page_size_from_view(k_cache) + if page_size != config.block_size: + raise ValueError( + f"MSA backend requires page_size == sparse_block_size; " + f"got page_size={page_size}, sparse_block_size={config.block_size}." + ) + + qo_lens_cpu, kv_lens_cpu, qo_offset_cpu = _qo_lens_offsets_from_metadata(metadata) + kv_indices, _ = _build_kv_indices_and_lens(metadata, page_size) + + kv_block_indexes = _msa_index_proxy_and_topk( + 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, + config=config, + idx_sm_scale=idx_sm_scale, + causal=True, + init_blocks=config.init_blocks, + local_blocks=config.local_blocks, + ) + + out = _msa_sparse_attention( + q, + k_paged, + v_paged, + kv_block_indexes, + qo_lens_cpu=qo_lens_cpu, + kv_lens_cpu=kv_lens_cpu, + qo_offset_cpu=qo_offset_cpu, + kv_indices=kv_indices, + sm_scale=sm_scale, + causal=True, + ) + + total_q = int(q.shape[0]) + return out.reshape(total_q, config.num_q_heads * config.head_dim).contiguous() + + +def minimax_m3_msa_sparse_decode( + q: torch.Tensor, + idx_q: torch.Tensor, + k_cache: torch.Tensor, + v_cache: torch.Tensor, + idx_k_cache: torch.Tensor, + metadata: MiniMaxM3SparseAttentionMetadata, + config: MiniMaxM3SparseConfig, + *, + sm_scale: Optional[float] = None, + idx_sm_scale: Optional[float] = None, +) -> torch.Tensor: + """Pure-decode MSA path. + + Decode is the ``qo_len == 1`` specialization of + :func:`minimax_m3_msa_sparse_prefill` with ``causal=False`` (the + new token attends to all cached positions through ``seq_lens - 1``, + no in-batch causality). ``MSA``'s ``fmha_sm100`` handles the + sub-32 query path internally via its decode kernel selection + (``_prefill_qlen_threshold(sparse=True) == 32``). + """ + if metadata.is_prefill: + raise ValueError("MSA decode entry called with prefill metadata") + if config.head_dim != _MSA_REQUIRED_HEAD_DIM: + raise NotImplementedError( + f"MSA backend currently supports head_dim={_MSA_REQUIRED_HEAD_DIM}; " + f"got {config.head_dim}." + ) + if config.sparse_index_dim != _MSA_REQUIRED_HEAD_DIM: + raise NotImplementedError( + f"MSA backend requires sparse_index_dim={_MSA_REQUIRED_HEAD_DIM}; " + f"got {config.sparse_index_dim}." + ) + if config.topk != _MSA_REQUIRED_TOPK: + raise NotImplementedError( + f"MSA backend currently supports topk={_MSA_REQUIRED_TOPK}; got {config.topk}." + ) + + sm_scale = sm_scale if sm_scale is not None else config.head_dim**-0.5 + idx_sm_scale = idx_sm_scale if idx_sm_scale is not None else config.sparse_index_dim**-0.5 + + k_paged = _cache_view_to_msa_paged(k_cache) + v_paged = _cache_view_to_msa_paged(v_cache) + idx_k_paged = _idx_cache_to_msa_paged(idx_k_cache) + page_size = _page_size_from_view(k_cache) + if page_size != config.block_size: + raise ValueError( + f"MSA backend requires page_size == sparse_block_size; " + f"got page_size={page_size}, sparse_block_size={config.block_size}." + ) + + qo_lens_cpu, kv_lens_cpu, qo_offset_cpu = _qo_lens_offsets_from_metadata(metadata) + kv_indices, _ = _build_kv_indices_and_lens(metadata, page_size) + + kv_block_indexes = _msa_index_proxy_and_topk( + 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, + config=config, + idx_sm_scale=idx_sm_scale, + causal=False, + init_blocks=config.init_blocks, + local_blocks=config.local_blocks, + ) + + out = _msa_sparse_attention( + q, + k_paged, + v_paged, + kv_block_indexes, + qo_lens_cpu=qo_lens_cpu, + kv_lens_cpu=kv_lens_cpu, + qo_offset_cpu=qo_offset_cpu, + kv_indices=kv_indices, + sm_scale=sm_scale, + causal=False, + ) + + batch = int(q.shape[0]) + return out.reshape(batch, config.num_q_heads * config.head_dim).contiguous() + + +# --------------------------------------------------------------------------- +# AttentionBackend wrapper +# --------------------------------------------------------------------------- + + +@functools.lru_cache(maxsize=1) +def get_minimax_m3_msa_attention_backend_cls(): + """Return :class:`MiniMaxM3MSARuntimeBackend` (lazy import). + + Subclasses the canonical Triton-backed + :class:`MiniMaxM3SparseRuntimeBackend` and overrides + :meth:`forward_sparse` so the per-layer dispatch goes through the + MSA prefill / decode primitives above. Construction validates that + the layer's config is compatible with the MSA kernel surface + (head_dim, sparse_index_dim, topk, disable_index_value) before any + forward call. + """ + base_cls = get_minimax_m3_attention_backend_cls() + + class MiniMaxM3MSARuntimeBackend(base_cls): # type: ignore[misc] + """MSA-backed ``MiniMaxM3SparseRuntimeBackend``. + + See module docstring for the two-step dispatch flow. The + backend reuses every parent-class lifecycle method + (``support_fused_rope``, ``__init__`` validation, the standard + ``forward`` keyword-routing) and only swaps the prefill / + decode primitives. + """ + + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + if not self.disable_index_value: + raise NotImplementedError( + "MSA backend currently requires disable_index_value=True " + "(MSA's proxy pass only consumes max_score; an index-V " + "path is not implemented yet)." + ) + # Validate kernel preconditions up-front so layer + # construction fails with a clear message rather than a + # cryptic shape mismatch on the first forward call. + if self.m3_config.head_dim != _MSA_REQUIRED_HEAD_DIM: + raise NotImplementedError( + f"MSA backend requires head_dim={_MSA_REQUIRED_HEAD_DIM}, " + f"got {self.m3_config.head_dim}." + ) + if self.m3_config.sparse_index_dim != _MSA_REQUIRED_HEAD_DIM: + raise NotImplementedError( + f"MSA backend requires sparse_index_dim={_MSA_REQUIRED_HEAD_DIM}, " + f"got {self.m3_config.sparse_index_dim}." + ) + if self.m3_config.topk != _MSA_REQUIRED_TOPK: + raise NotImplementedError( + f"MSA backend requires topk={_MSA_REQUIRED_TOPK}, got {self.m3_config.topk}." + ) + + def forward_sparse( + self, + *, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + idx_q: torch.Tensor, + idx_k: torch.Tensor, + idx_v: Optional[torch.Tensor], + k_cache: torch.Tensor, + v_cache: torch.Tensor, + idx_k_cache: torch.Tensor, + idx_v_cache: Optional[torch.Tensor], + out_cache_loc: torch.Tensor, + m3_metadata: "MiniMaxM3SparseAttentionMetadata", + sm_scale: Optional[float] = None, + idx_sm_scale: Optional[float] = None, + ) -> torch.Tensor: + """MSA-backed sparse forward. + + Mirrors the parent class's prologue (cache writes, metadata + device migration) and then dispatches to the MSA prefill / + decode primitives instead of the Triton path. The MSA + kernels read directly from the paged caches we just + populated; we do **not** gather K/V into per-batch padded + slabs. + """ + if idx_v is not None: + raise NotImplementedError( + "MSA backend does not consume idx_v (disable_index_value=True is required)." + ) + num_kv_heads = self.m3_config.num_kv_heads + head_dim = self.m3_config.head_dim + sparse_index_dim = self.m3_config.sparse_index_dim + num_idx_heads = self.m3_config.num_index_heads + + num_tokens = int(q.shape[0]) + q_view = q.view(num_tokens, self.num_heads, head_dim) + k_view = k.view(num_tokens, num_kv_heads, head_dim) + v_view = v.view(num_tokens, num_kv_heads, head_dim) + idx_q_view = idx_q.view(num_tokens, num_idx_heads, sparse_index_dim) + idx_k_view = idx_k.view(num_tokens, 1, sparse_index_dim) + + cache_device = k_cache.device + if any( + t is not None and t.device != cache_device + for t in ( + m3_metadata.req_to_token, + m3_metadata.slot_ids, + m3_metadata.seq_lens, + m3_metadata.prefix_lens, + m3_metadata.cu_seqlens_q, + m3_metadata.q_batch_row, + m3_metadata.q_positions, + ) + ): + m3_metadata = ensure_metadata_on_device(m3_metadata, cache_device) + + # KV writes use the same layout-aware helper as the + # Triton backend; the MSA kernels read from the same paged + # pool afterwards. + _write_main_kv_slots(k_cache, out_cache_loc, k_view) + _write_main_kv_slots(v_cache, out_cache_loc, v_view) + _write_main_kv_slots(idx_k_cache, out_cache_loc, idx_k_view) + + if m3_metadata.is_prefill: + return minimax_m3_msa_sparse_prefill( + q_view, + k_cache, + v_cache, + idx_q_view, + idx_k_cache, + m3_metadata, + self.m3_config, + sm_scale=sm_scale, + idx_sm_scale=idx_sm_scale, + ) + return minimax_m3_msa_sparse_decode( + q_view, + idx_q_view, + k_cache, + v_cache, + idx_k_cache, + m3_metadata, + self.m3_config, + sm_scale=sm_scale, + idx_sm_scale=idx_sm_scale, + ) + + return MiniMaxM3MSARuntimeBackend + + +def get_minimax_m3_attention_backend_cls_with_msa( + sparse_params: "MiniMaxM3SparseParams", +): + """Return the MSA backend class when ``sparse_params.use_msa`` is set. + + Falls back to the Triton-backed + :class:`MiniMaxM3SparseRuntimeBackend` otherwise. Centralises the + Triton-vs-MSA dispatch decision so :mod:`...sparse.utils` and the + model layer do not duplicate the predicate. + """ + if getattr(sparse_params, "use_msa", False): + return get_minimax_m3_msa_attention_backend_cls() + return get_minimax_m3_attention_backend_cls() + + +__all__ = [ + "get_minimax_m3_attention_backend_cls_with_msa", + "get_minimax_m3_msa_attention_backend_cls", + "minimax_m3_msa_sparse_decode", + "minimax_m3_msa_sparse_prefill", +] diff --git a/tensorrt_llm/_torch/attention_backend/sparse/utils.py b/tensorrt_llm/_torch/attention_backend/sparse/utils.py index 12e9b146c4ef..088239346db7 100644 --- a/tensorrt_llm/_torch/attention_backend/sparse/utils.py +++ b/tensorrt_llm/_torch/attention_backend/sparse/utils.py @@ -39,14 +39,29 @@ def get_sparse_attn_kv_cache_manager( ) +def _resolve_minimax_m3_backend_cls( + sparse_attention_config: "SparseAttentionConfig"): + """Pick the Triton or MSA-backed M3 backend class. + + Honours ``sparse_use_msa`` on the user-facing + :class:`MiniMaxM3SparseAttentionConfig`. Falls back to the Triton + reference path when the flag is unset, preserving the legacy + behaviour for callers that have not opted in. + """ + from .minimax_m3 import (get_minimax_m3_attention_backend_cls, + get_minimax_m3_msa_attention_backend_cls) + if getattr(sparse_attention_config, "sparse_use_msa", False): + return get_minimax_m3_msa_attention_backend_cls() + return get_minimax_m3_attention_backend_cls() + + def get_vanilla_sparse_attn_attention_backend( sparse_attention_config: "SparseAttentionConfig"): - from .minimax_m3 import get_minimax_m3_attention_backend_cls from .rocket import RocketVanillaAttention if sparse_attention_config.algorithm == "rocket": return RocketVanillaAttention elif sparse_attention_config.algorithm == "minimax_m3": - return get_minimax_m3_attention_backend_cls() + return _resolve_minimax_m3_backend_cls(sparse_attention_config) else: raise ValueError( f"Unsupported sparse attention algorithm in vanilla attention backend: {sparse_attention_config.algorithm}" @@ -59,7 +74,6 @@ def get_trtllm_sparse_attn_attention_backend( from .deepseek_v4 import DeepseekV4TrtllmAttention from .dsa import DSATrtllmAttention - from .minimax_m3 import get_minimax_m3_attention_backend_cls from .rocket import RocketTrtllmAttention if sparse_attention_config.algorithm == "rocket": return RocketTrtllmAttention @@ -75,7 +89,7 @@ def get_trtllm_sparse_attn_attention_backend( # `create_attention(...)` dispatch in `Attention.__init__` # returns an instantiable AttentionBackend under the trtllm # attention backend slot. - return get_minimax_m3_attention_backend_cls() + return _resolve_minimax_m3_backend_cls(sparse_attention_config) else: raise ValueError( f"Unsupported sparse attention algorithm in trtllm attention backend: {sparse_attention_config.algorithm}" @@ -84,9 +98,8 @@ def get_trtllm_sparse_attn_attention_backend( def get_flashinfer_sparse_attn_attention_backend( sparse_attention_config: "SparseAttentionConfig"): - from .minimax_m3 import get_minimax_m3_attention_backend_cls if sparse_attention_config.algorithm == "minimax_m3": - return get_minimax_m3_attention_backend_cls() + return _resolve_minimax_m3_backend_cls(sparse_attention_config) raise ValueError( f"Unsupported sparse attention algorithm in flashinfer attention backend: {sparse_attention_config.algorithm}" ) diff --git a/tensorrt_llm/llmapi/llm_args.py b/tensorrt_llm/llmapi/llm_args.py index 3093f2f1e34f..10d6155f4fcc 100644 --- a/tensorrt_llm/llmapi/llm_args.py +++ b/tensorrt_llm/llmapi/llm_args.py @@ -644,6 +644,14 @@ class MiniMaxM3SparseAttentionConfig(BaseSparseAttentionConfig): default=True, description="If True, skip the index V branch (M3 checkpoint default).", ) + sparse_use_msa: bool = Field( + default=False, + description=("If True, route the sparse forward through the MSA-backed " + "FMHA runtime (``fmha_sm100`` + ``sparse_topk_select``) " + "instead of the in-tree Triton + SDPA reference path. " + "Requires SM100 and the external ``fmha_sm100`` package " + "(https://github.com/MiniMax-AI/MSA) to be importable."), + ) def supports_backend(self, backend: str) -> bool: return backend == "pytorch" @@ -664,6 +672,7 @@ def to_sparse_params(self, **kwargs): local_blocks=self.sparse_local_blocks, score_type=self.sparse_score_type, disable_index_value=self.sparse_disable_index_value, + use_msa=self.sparse_use_msa, ) diff --git a/tests/unittest/_torch/attention/sparse/test_minimax_m3_msa_backend.py b/tests/unittest/_torch/attention/sparse/test_minimax_m3_msa_backend.py new file mode 100644 index 000000000000..d7dab1d5432f --- /dev/null +++ b/tests/unittest/_torch/attention/sparse/test_minimax_m3_msa_backend.py @@ -0,0 +1,553 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Unit tests for the MSA-backed MiniMax-M3 sparse attention runtime. + +The MSA stack (``fmha_sm100`` package) is SM100-only and not installed +in the standard CI runners. These tests therefore split into two +tiers: + + * **Static plumbing tests** (always run): config -> params -> backend + class wiring; layout adapter shapes; metadata derivation. These do + not touch the MSA kernels and run anywhere PyTorch + the TRT-LLM + Python imports succeed. + * **Live kernel tests** (skipped when ``fmha_sm100`` is unavailable + or no SM100 GPU is present): end-to-end forward parity check + against the in-tree Triton reference path. +""" + +from __future__ import annotations + +import importlib +from typing import List + +import pytest +import torch + +# Module-under-test imports are kept lazy so the file imports cleanly +# on hosts where heavy TRT-LLM C++ extensions are absent (the rest of +# the testsuite uses the same pattern). +sparse_minimax_m3 = pytest.importorskip("tensorrt_llm._torch.attention_backend.sparse.minimax_m3") +msa_backend = pytest.importorskip( + "tensorrt_llm._torch.attention_backend.sparse.minimax_m3.msa_backend" +) +m3_backend = pytest.importorskip("tensorrt_llm._torch.attention_backend.sparse.minimax_m3.backend") + + +# --------------------------------------------------------------------------- +# Param + config plumbing +# --------------------------------------------------------------------------- + + +def test_sparse_attention_config_threads_use_msa_into_params(): + """``sparse_use_msa=True`` must land on the lowered SparseParams.""" + from tensorrt_llm.llmapi.llm_args import MiniMaxM3SparseAttentionConfig + + cfg = MiniMaxM3SparseAttentionConfig(sparse_use_msa=True) + params = cfg.to_sparse_params() + assert params.use_msa is True + + cfg_default = MiniMaxM3SparseAttentionConfig() + params_default = cfg_default.to_sparse_params() + assert params_default.use_msa is False + + +def test_backend_dispatch_picks_msa_when_flag_set(): + """``utils._resolve_minimax_m3_backend_cls`` honours sparse_use_msa.""" + from tensorrt_llm._torch.attention_backend.sparse.utils import _resolve_minimax_m3_backend_cls + from tensorrt_llm.llmapi.llm_args import MiniMaxM3SparseAttentionConfig + + triton_cls = sparse_minimax_m3.get_minimax_m3_attention_backend_cls() + msa_cls = sparse_minimax_m3.get_minimax_m3_msa_attention_backend_cls() + + triton_cfg = MiniMaxM3SparseAttentionConfig(sparse_use_msa=False) + assert _resolve_minimax_m3_backend_cls(triton_cfg) is triton_cls + + msa_cfg = MiniMaxM3SparseAttentionConfig(sparse_use_msa=True) + resolved = _resolve_minimax_m3_backend_cls(msa_cfg) + assert resolved is msa_cls + # MSA class subclasses the Triton class so the model layer's + # isinstance check still accepts it. + assert issubclass(resolved, triton_cls) + + +def test_get_with_msa_helper_routes_on_params_flag(): + """``get_minimax_m3_attention_backend_cls_with_msa`` matches the dispatch.""" + triton_cls = sparse_minimax_m3.get_minimax_m3_attention_backend_cls() + msa_cls = sparse_minimax_m3.get_minimax_m3_msa_attention_backend_cls() + + no_msa = sparse_minimax_m3.MiniMaxM3SparseConfig # noqa: F841 (just exercise import) + from tensorrt_llm._torch.attention_backend.sparse.minimax_m3.metadata import ( + MiniMaxM3SparseParams, + ) + + triton_params = MiniMaxM3SparseParams(use_msa=False) + msa_params = MiniMaxM3SparseParams(use_msa=True) + assert msa_backend.get_minimax_m3_attention_backend_cls_with_msa(triton_params) is triton_cls + assert msa_backend.get_minimax_m3_attention_backend_cls_with_msa(msa_params) is msa_cls + + +# --------------------------------------------------------------------------- +# Cache layout adapter +# --------------------------------------------------------------------------- + + +def test_cache_view_to_msa_paged_4d_layout(): + """4-D pool views permute (page_size, num_kv_heads) -> HND.""" + num_pages, tokens_per_block, num_kv_heads, head_dim = 3, 128, 2, 128 + cache_view = torch.randn( + num_pages, tokens_per_block, num_kv_heads, head_dim, dtype=torch.bfloat16 + ) + paged = msa_backend._cache_view_to_msa_paged(cache_view) + assert paged.shape == (num_pages, num_kv_heads, tokens_per_block, head_dim) + assert paged.is_contiguous() + # Values must agree element-wise: paged[p, h, t, d] == view[p, t, h, d]. + expected = cache_view.permute(0, 2, 1, 3).contiguous() + assert torch.equal(paged, expected) + + +def test_cache_view_to_msa_paged_3d_flat_slot_layout(): + """3-D flat-slot caches collapse into a single virtual page.""" + num_slots, num_kv_heads, head_dim = 64, 2, 128 + cache_view = torch.randn(num_slots, num_kv_heads, head_dim, dtype=torch.bfloat16) + paged = msa_backend._cache_view_to_msa_paged(cache_view) + assert paged.shape == (1, num_kv_heads, num_slots, head_dim) + assert paged.is_contiguous() + expected = cache_view.permute(1, 0, 2).unsqueeze(0).contiguous() + assert torch.equal(paged, expected) + + +def test_idx_cache_to_msa_paged_handles_both_ranks(): + """Index cache adapter mirrors the K-cache one but for single-head data.""" + num_pages, tokens_per_block, sparse_index_dim = 2, 128, 128 + paged_4d = torch.randn(num_pages, tokens_per_block, 1, sparse_index_dim, dtype=torch.bfloat16) + out_4d = msa_backend._idx_cache_to_msa_paged(paged_4d) + assert out_4d.shape == (num_pages, 1, tokens_per_block, sparse_index_dim) + assert torch.equal(out_4d, paged_4d.permute(0, 2, 1, 3).contiguous()) + + flat_3d = torch.randn(num_pages * tokens_per_block, 1, sparse_index_dim, dtype=torch.bfloat16) + out_3d = msa_backend._idx_cache_to_msa_paged(flat_3d) + assert out_3d.shape == (1, 1, flat_3d.shape[0], sparse_index_dim) + + +def test_cache_view_to_msa_paged_rejects_other_ranks(): + with pytest.raises(ValueError, match="rank"): + msa_backend._cache_view_to_msa_paged(torch.empty(4, 8)) + with pytest.raises(ValueError, match="rank"): + msa_backend._idx_cache_to_msa_paged(torch.empty(2)) + + +# --------------------------------------------------------------------------- +# Metadata derivation +# --------------------------------------------------------------------------- + + +def _build_metadata_for_test( + *, + is_prefill: bool, + seq_lens: List[int], + extend_seq_lens: List[int] | None = None, + page_size: int = 128, +): + """Construct a minimal :class:`MiniMaxM3SparseAttentionMetadata`. + + Builds a stable, contiguous slot layout so the page-index gather in + :func:`_build_kv_indices_and_lens` produces deterministic block ids. + """ + from tensorrt_llm._torch.attention_backend.sparse.minimax_m3.metadata import ( + MiniMaxM3SparseAttentionMetadata, + ) + + batch = len(seq_lens) + max_kv_len = max(((s + page_size - 1) // page_size) * page_size for s in seq_lens) + req_to_token = torch.arange(batch * max_kv_len, dtype=torch.int32).view(batch, max_kv_len) + slot_ids = torch.arange(batch, dtype=torch.int32) + seq_lens_t = torch.tensor(seq_lens, dtype=torch.int32) + + if is_prefill: + assert extend_seq_lens is not None + prefix_lens = torch.tensor( + [seq_lens[b] - extend_seq_lens[b] for b in range(batch)], dtype=torch.int32 + ) + cu_q = [0] + for x in extend_seq_lens: + cu_q.append(cu_q[-1] + x) + cu_seqlens_q = torch.tensor(cu_q, dtype=torch.int32) + meta = MiniMaxM3SparseAttentionMetadata( + is_prefill=True, + req_to_token=req_to_token, + slot_ids=slot_ids, + seq_lens=seq_lens_t, + seq_lens_cpu=seq_lens_t, + prefix_lens=prefix_lens, + cu_seqlens_q=cu_seqlens_q, + extend_seq_lens_cpu=list(extend_seq_lens), + ) + else: + meta = MiniMaxM3SparseAttentionMetadata( + is_prefill=False, + req_to_token=req_to_token, + slot_ids=slot_ids, + seq_lens=seq_lens_t, + seq_lens_cpu=seq_lens_t, + ) + meta.prepare() + return meta + + +def test_qo_lens_offsets_prefill_uses_extend_lens_and_prefix(): + meta = _build_metadata_for_test( + is_prefill=True, + seq_lens=[256, 384], + extend_seq_lens=[128, 128], + ) + qo_lens_cpu, kv_lens_cpu, qo_offset_cpu = msa_backend._qo_lens_offsets_from_metadata(meta) + assert qo_lens_cpu.dtype == torch.int32 + assert qo_lens_cpu.tolist() == [128, 128] + assert kv_lens_cpu.tolist() == [256, 384] + # Causal offset is per-request prefix length (kv - extend). + assert qo_offset_cpu.tolist() == [128, 256] + + +def test_qo_lens_offsets_decode_uses_unit_q_and_kv_minus_one(): + meta = _build_metadata_for_test(is_prefill=False, seq_lens=[200, 300]) + qo_lens_cpu, kv_lens_cpu, qo_offset_cpu = msa_backend._qo_lens_offsets_from_metadata(meta) + assert qo_lens_cpu.tolist() == [1, 1] + assert kv_lens_cpu.tolist() == [200, 300] + # Decode causal offset = kv_len - 1 (the new token's position). + assert qo_offset_cpu.tolist() == [199, 299] + + +def test_build_kv_indices_packs_per_request_pages(): + """``kv_indices`` is the concatenated per-request page table.""" + page_size = 128 + meta = _build_metadata_for_test( + is_prefill=False, + seq_lens=[page_size * 2, page_size * 3 - 5], + page_size=page_size, + ) + kv_indices, kv_lens = msa_backend._build_kv_indices_and_lens(meta, page_size) + + # Request 0 owns slots [0, 384) so its pages are {0, 1, 2}, but only + # 2 pages are needed (seq_len=256). Request 1 starts at slot 384, + # so its first page is page id 3. + # The helper builds page indices by ``req_to_token[b, page_starts] + # // page_size`` where the slot ids are contiguous arange. + # For request 0 (seq=256): page_starts=[0, 128], slots=[0, 128], + # page ids=[0, 1]. + # For request 1 (seq=379, ceil=3 pages): page_starts=[0, 128, 256] + # mapped to req_rows[1] = [384..767], slots=[384, 512, 640], + # page ids=[3, 4, 5]. + assert kv_indices.dtype == torch.int32 + assert kv_indices.tolist() == [0, 1, 3, 4, 5] + assert kv_lens.tolist() == [256, 379] + + +# --------------------------------------------------------------------------- +# Backend class construction +# --------------------------------------------------------------------------- + + +def test_msa_backend_rejects_disable_index_value_false(): + """Layer construction surfaces the unsupported-mode error early.""" + from tensorrt_llm._torch.attention_backend.sparse.minimax_m3.metadata import ( + MiniMaxM3SparseParams, + ) + + msa_cls = sparse_minimax_m3.get_minimax_m3_msa_attention_backend_cls() + bad_params = MiniMaxM3SparseParams(disable_index_value=False, use_msa=True) + with pytest.raises(NotImplementedError, match="disable_index_value"): + msa_cls( + layer_idx=0, + num_heads=8, + head_dim=128, + num_kv_heads=1, + sparse_params=bad_params, + ) + + +def test_msa_backend_validates_required_dims(): + """Layer construction surfaces head_dim / sparse_index_dim mismatches.""" + from tensorrt_llm._torch.attention_backend.sparse.minimax_m3.metadata import ( + MiniMaxM3SparseParams, + ) + + msa_cls = sparse_minimax_m3.get_minimax_m3_msa_attention_backend_cls() + bad_params_head_dim = MiniMaxM3SparseParams(use_msa=True) + # head_dim != 128 must trip the precondition (sparse_index_dim is + # the default 128 in MiniMaxM3SparseParams). + with pytest.raises(NotImplementedError, match="head_dim"): + msa_cls( + layer_idx=0, + num_heads=8, + head_dim=64, + num_kv_heads=1, + sparse_params=bad_params_head_dim, + ) + + bad_params_topk = MiniMaxM3SparseParams(use_msa=True, topk=8) + with pytest.raises(NotImplementedError, match="topk"): + msa_cls( + layer_idx=0, + num_heads=8, + head_dim=128, + num_kv_heads=1, + sparse_params=bad_params_topk, + ) + + +def test_msa_proxy_mqa_registered_in_fmha_libs(): + """The proxy MQA FMHA must live in the standard ``FMHA_LIBS`` registry. + + Keeps the indexer's dispatch path symmetric with the main-attention + FMHA backends (FlashInfer trtllm-gen, fallback): the same env var + (``TLLM_FMHA_LIBS``) governs which proxy implementations are + reachable, and a regression that drops the class from the registry + breaks the M3 sparse forward at runtime. + """ + from tensorrt_llm._torch.attention_backend.fmha import ( + FMHA_LIBS, + IndexerProxyFmha, + MsaProxyMqaFmha, + ) + + assert "msa_proxy_mqa" in FMHA_LIBS + assert FMHA_LIBS["msa_proxy_mqa"] is MsaProxyMqaFmha + assert issubclass(MsaProxyMqaFmha, IndexerProxyFmha) + + +def test_msa_proxy_mqa_opts_out_of_main_dispatch(): + """is_supported must return False for the main-attention dispatch loop. + + Confirms the indexer-style FMHA never claims work intended for + main FMHA backends: a regression here would cause TrtllmAttention + to call ``MsaProxyMqaFmha.forward`` and crash with + ``NotImplementedError``. + """ + from tensorrt_llm._torch.attention_backend.fmha import MsaProxyMqaFmha + + fmha = MsaProxyMqaFmha() # constructed without an owning attn + # Pass placeholder args -- is_supported should reject regardless. + assert fmha.is_supported(None, None, None, None, None) is False + with pytest.raises(NotImplementedError, match="forward_proxy"): + fmha.forward(None, None, None, None, None) + + +def test_select_proxy_fmha_class_picks_msa_when_available(monkeypatch): + """``_select_proxy_fmha_class`` returns MsaProxyMqaFmha when it is_available.""" + from tensorrt_llm._torch.attention_backend.fmha import MsaProxyMqaFmha + + # Force is_available to True and confirm the lookup picks MsaProxyMqaFmha. + monkeypatch.setattr(MsaProxyMqaFmha, "is_available", classmethod(lambda cls, attn=None: True)) + # Clear the lru_cache on the selector so monkeypatch takes effect. + msa_backend._select_proxy_fmha_class.cache_clear() + try: + cls = msa_backend._select_proxy_fmha_class() + finally: + msa_backend._select_proxy_fmha_class.cache_clear() + assert cls is MsaProxyMqaFmha + + +def test_msa_sparse_gqa_registered_in_fmha_libs(): + """The block-sparse main FMHA must live in the standard ``FMHA_LIBS`` registry. + + Same rationale as the proxy registration test: ``TLLM_FMHA_LIBS`` + is the single env var governing which sparse-FMHA implementations + are reachable, and a regression that drops the class breaks the + M3 sparse forward at runtime. + """ + from tensorrt_llm._torch.attention_backend.fmha import ( + FMHA_LIBS, + BlockSparseFmha, + MsaSparseGqaFmha, + ) + + assert "msa_sparse_gqa" in FMHA_LIBS + assert FMHA_LIBS["msa_sparse_gqa"] is MsaSparseGqaFmha + assert issubclass(MsaSparseGqaFmha, BlockSparseFmha) + + +def test_msa_sparse_gqa_opts_out_of_main_dispatch(): + """is_supported must return False for the main-attention dispatch loop.""" + from tensorrt_llm._torch.attention_backend.fmha import MsaSparseGqaFmha + + fmha = MsaSparseGqaFmha() # no owning attn + assert fmha.is_supported(None, None, None, None, None) is False + with pytest.raises(NotImplementedError, match="forward_block_sparse"): + fmha.forward(None, None, None, None, None) + + +def test_select_block_sparse_fmha_class_picks_msa_when_available(monkeypatch): + """``_select_block_sparse_fmha_class`` returns MsaSparseGqaFmha when available.""" + from tensorrt_llm._torch.attention_backend.fmha import MsaSparseGqaFmha + + monkeypatch.setattr(MsaSparseGqaFmha, "is_available", classmethod(lambda cls, attn=None: True)) + msa_backend._select_block_sparse_fmha_class.cache_clear() + try: + cls = msa_backend._select_block_sparse_fmha_class() + finally: + msa_backend._select_block_sparse_fmha_class.cache_clear() + assert cls is MsaSparseGqaFmha + + +def test_select_block_sparse_fmha_class_returns_none_when_unavailable(monkeypatch): + """``_select_block_sparse_fmha_class`` returns None when no backend is available.""" + from tensorrt_llm._torch.attention_backend.fmha import ( + BlockSparseFmha, + get_enabled_fmha_lib_classes, + ) + + for cls in get_enabled_fmha_lib_classes(): + if issubclass(cls, BlockSparseFmha): + monkeypatch.setattr(cls, "is_available", classmethod(lambda cls, attn=None: False)) + msa_backend._select_block_sparse_fmha_class.cache_clear() + try: + assert msa_backend._select_block_sparse_fmha_class() is None + finally: + msa_backend._select_block_sparse_fmha_class.cache_clear() + + +def test_select_proxy_fmha_class_returns_none_when_unavailable(monkeypatch): + """``_select_proxy_fmha_class`` returns None when no backend is_available.""" + # Force every indexer-style class in the registry to report unavailable. + from tensorrt_llm._torch.attention_backend.fmha import ( + IndexerProxyFmha, + get_enabled_fmha_lib_classes, + ) + + for cls in get_enabled_fmha_lib_classes(): + if issubclass(cls, IndexerProxyFmha): + monkeypatch.setattr(cls, "is_available", classmethod(lambda cls, attn=None: False)) + msa_backend._select_proxy_fmha_class.cache_clear() + try: + assert msa_backend._select_proxy_fmha_class() is None + finally: + msa_backend._select_proxy_fmha_class.cache_clear() + + +def test_require_msa_module_raises_when_absent(monkeypatch): + """The lazy import wrapper produces a descriptive error.""" + + # Force the import to fail by sabotaging ``sys.modules``. + import sys + + original = sys.modules.pop("fmha_sm100", None) + try: + # Also prevent re-import from a possible install. + monkeypatch.setattr( + "builtins.__import__", + _make_import_blocklist({"fmha_sm100"}, fallback=__import__), + ) + with pytest.raises(RuntimeError, match="fmha_sm100"): + msa_backend._require_msa_module() + finally: + if original is not None: + sys.modules["fmha_sm100"] = original + + +def _make_import_blocklist(blocked, *, fallback): + """Return an ``__import__`` shim that raises for names in ``blocked``.""" + + def _shim(name, globals=None, locals=None, fromlist=(), level=0): + if name in blocked or name.split(".")[0] in blocked: + raise ImportError(f"blocked: {name}") + return fallback(name, globals, locals, fromlist, level) + + return _shim + + +# --------------------------------------------------------------------------- +# Live kernel parity (SM100 + fmha_sm100 only) +# --------------------------------------------------------------------------- + + +def _msa_available() -> bool: + if not torch.cuda.is_available(): + return False + major, _ = torch.cuda.get_device_capability() + if major != 10: # SM100 family + return False + try: + importlib.import_module("fmha_sm100") + return True + except ImportError: + return False + + +@pytest.mark.skipif(not _msa_available(), reason="MSA fmha_sm100 + SM100 GPU required") +def test_msa_prefill_runs_end_to_end(): + """Smoke test: the MSA prefill path executes without raising. + + A bit-exact parity check vs the Triton reference is left to the + Minimax-M3 integration tests because the two paths differ in their + top-k semantics (union-OR vs amax-then-topk). This smoke test + confirms the public entry points compose correctly: cache layout + adapters, ``fmha_sm100_plan``, ``sparse_topk_select``, and the + second-stage ``fmha_sm100`` call wire up to a kernel launch and + return a tensor of the expected shape. + """ + from tensorrt_llm._torch.attention_backend.sparse.minimax_m3.metadata import ( + MiniMaxM3SparseConfig, + ) + + device = torch.device("cuda") + page_size = 128 + seq_lens = [256, 384] + extend_seq_lens = [128, 128] + num_q_heads, num_kv_heads, head_dim = 8, 1, 128 + num_index_heads, sparse_index_dim = 4, 128 + + config = MiniMaxM3SparseConfig( + num_q_heads=num_q_heads, + num_kv_heads=num_kv_heads, + head_dim=head_dim, + num_index_heads=num_index_heads, + sparse_index_dim=sparse_index_dim, + block_size=page_size, + topk=16, + init_blocks=0, + local_blocks=1, + ) + + meta = _build_metadata_for_test( + is_prefill=True, + seq_lens=seq_lens, + extend_seq_lens=extend_seq_lens, + page_size=page_size, + ) + # Move tensors to device for the live kernel run. + meta.req_to_token = meta.req_to_token.to(device) + meta.slot_ids = meta.slot_ids.to(device) + meta.seq_lens = meta.seq_lens.to(device) + meta.prefix_lens = meta.prefix_lens.to(device) + meta.cu_seqlens_q = meta.cu_seqlens_q.to(device) + if meta.q_batch_row is not None: + meta.q_batch_row = meta.q_batch_row.to(device) + if meta.q_positions is not None: + meta.q_positions = meta.q_positions.to(device) + + total_q = sum(extend_seq_lens) + q = torch.randn(total_q, num_q_heads, head_dim, dtype=torch.bfloat16, device=device) + idx_q = torch.randn( + total_q, num_index_heads, sparse_index_dim, dtype=torch.bfloat16, device=device + ) + num_pages = max(((s + page_size - 1) // page_size) for s in seq_lens) * len(seq_lens) + k_cache = torch.randn( + num_pages, page_size, num_kv_heads, head_dim, dtype=torch.bfloat16, device=device + ) + v_cache = torch.randn_like(k_cache) + idx_k_cache = torch.randn( + num_pages, page_size, 1, sparse_index_dim, dtype=torch.bfloat16, device=device + ) + + out = msa_backend.minimax_m3_msa_sparse_prefill( + q, + k_cache, + v_cache, + idx_q, + idx_k_cache, + meta, + config, + ) + assert out.shape == (total_q, num_q_heads * head_dim) + assert out.dtype == torch.bfloat16 + assert torch.isfinite(out).all() From 6d73982e1ed97a8da37087aa0895d0af7278fa48 Mon Sep 17 00:00:00 2001 From: WeiHaocheng <20514172+WeiHaocheng@users.noreply.github.com> Date: Thu, 2 Jul 2026 02:37:19 -0700 Subject: [PATCH 2/4] [None][fix] Fix parallel config of Minimax M3 MXFP8 test Signed-off-by: WeiHaocheng <20514172+WeiHaocheng@users.noreply.github.com> --- .../defs/accuracy/test_llm_api_pytorch.py | 12 +++++++++--- .../integration/test_lists/qa/llm_function_core.txt | 2 +- 2 files changed, 10 insertions(+), 4 deletions(-) diff --git a/tests/integration/defs/accuracy/test_llm_api_pytorch.py b/tests/integration/defs/accuracy/test_llm_api_pytorch.py index 641ed8c11ca4..f97f9382d7bf 100644 --- a/tests/integration/defs/accuracy/test_llm_api_pytorch.py +++ b/tests/integration/defs/accuracy/test_llm_api_pytorch.py @@ -7261,16 +7261,20 @@ def test_auto_dtype(self, tp_size, ep_size): task = GSM8K(self.MODEL_NAME) task.evaluate(llm) - @pytest.mark.skip_less_device(8) + @pytest.mark.skip_less_device(4) @pytest.mark.skip_less_device_memory(140000) - @parametrize_with_ids("tp_size,ep_size", [(8, 8)]) + @parametrize_with_ids("tp_size,ep_size", [(4, 4)]) def test_mxfp8(self, tp_size, ep_size): # MXFP8 checkpoint: weights are MXFP8 (e4m3 + UE8M0 1x32 block # scales) with MXFP8 dynamic activations; the KV cache stays in # BF16 and the sparse attention path is unchanged from BF16. model_name = "MiniMaxAI/MiniMax-M3-MXFP8" model_path = f"{llm_models_root()}/MiniMax-M3-MXFP8" - kv_cache_config = KvCacheConfig(free_gpu_memory_fraction=0.6, + # Halving TP from the BF16 reference (TP=8) doubles per-rank + # model + KV footprint; cap KV cache at 0.4 of free memory and + # constrain batch / token budget so the runtime allocator stays + # under the PyTorch cap. + kv_cache_config = KvCacheConfig(free_gpu_memory_fraction=0.4, enable_block_reuse=False) sparse_attention_config = MiniMaxM3SparseAttentionConfig() with LLM(model_path, @@ -7279,6 +7283,8 @@ def test_mxfp8(self, tp_size, ep_size): kv_cache_config=kv_cache_config, sparse_attention_config=sparse_attention_config, max_seq_len=4096, + max_batch_size=32, + max_num_tokens=4096, trust_remote_code=True) as llm: assert llm.args.quant_config.quant_algo == QuantAlgo.MXFP8 task = MMLU(model_name) diff --git a/tests/integration/test_lists/qa/llm_function_core.txt b/tests/integration/test_lists/qa/llm_function_core.txt index d2a91b70a423..823474494227 100644 --- a/tests/integration/test_lists/qa/llm_function_core.txt +++ b/tests/integration/test_lists/qa/llm_function_core.txt @@ -664,7 +664,7 @@ accuracy/test_llm_api_pytorch.py::TestLlama3_3_70BInstruct::test_nvfp4_tp4[torch accuracy/test_llm_api_pytorch.py::TestMiniMaxM2::test_4gpus[attention_dp=False-cuda_graph=True-overlap_scheduler=True-tp_size=4-ep_size=4] accuracy/test_llm_api_pytorch.py::TestMiniMaxM2_5::test_4gpus[attention_dp=False-cuda_graph=True-overlap_scheduler=True-tp_size=4-ep_size=4] accuracy/test_llm_api_pytorch.py::TestMiniMaxM3::test_auto_dtype[tp_size=8-ep_size=8] TIMEOUT (180) -accuracy/test_llm_api_pytorch.py::TestMiniMaxM3::test_mxfp8[tp_size=8-ep_size=8] TIMEOUT (180) +accuracy/test_llm_api_pytorch.py::TestMiniMaxM3::test_mxfp8[tp_size=4-ep_size=4] TIMEOUT (180) accuracy/test_llm_api_pytorch.py::TestMinistral8BInstruct::test_auto_dtype accuracy/test_llm_api_pytorch.py::TestMinistral8BInstruct::test_fp8 accuracy/test_llm_api_pytorch.py::TestMistralLarge3_675B::test_fp8[latency_moe_deepgemm] From 4619ef6e425cfcfa5c6570360610cac8d4ba9fa0 Mon Sep 17 00:00:00 2001 From: WeiHaocheng <20514172+WeiHaocheng@users.noreply.github.com> Date: Mon, 6 Jul 2026 02:26:12 -0700 Subject: [PATCH 3/4] [TRTLLM-14019][feat] Support MSA kernel with enable cuda graph Signed-off-by: WeiHaocheng <20514172+WeiHaocheng@users.noreply.github.com> --- requirements.txt | 4 + .../attention_backend/fmha/msa_proxy_mqa.py | 8 + .../attention_backend/fmha/msa_sparse_gqa.py | 1 + .../sparse/minimax_m3/cache_manager.py | 20 +- .../minimax_m3/decode_wrapper/__init__.py | 21 + .../minimax_m3/decode_wrapper/dispatch.py | 497 ++++++++++++++++++ .../sparse/minimax_m3/decode_wrapper/topk.py | 97 ++++ .../minimax_m3/decode_wrapper/worklist.py | 86 +++ .../sparse/minimax_m3/metadata.py | 155 +++++- .../sparse/minimax_m3/msa_backend.py | 369 ++++++++++--- .../sparse/minimax_m3/msa_plan_cache.py | 284 ++++++++++ .../_torch/models/modeling_minimaxm3.py | 38 ++ .../test_minimax_m3_decode_driver_vs_msa.py | 467 ++++++++++++++++ 13 files changed, 1973 insertions(+), 74 deletions(-) create mode 100644 tensorrt_llm/_torch/attention_backend/sparse/minimax_m3/decode_wrapper/__init__.py create mode 100644 tensorrt_llm/_torch/attention_backend/sparse/minimax_m3/decode_wrapper/dispatch.py create mode 100644 tensorrt_llm/_torch/attention_backend/sparse/minimax_m3/decode_wrapper/topk.py create mode 100644 tensorrt_llm/_torch/attention_backend/sparse/minimax_m3/decode_wrapper/worklist.py create mode 100644 tensorrt_llm/_torch/attention_backend/sparse/minimax_m3/msa_plan_cache.py create mode 100644 tests/unittest/_torch/attention/sparse/test_minimax_m3_decode_driver_vs_msa.py diff --git a/requirements.txt b/requirements.txt index a243f9b26ab5..dcee1e744592 100644 --- a/requirements.txt +++ b/requirements.txt @@ -55,6 +55,10 @@ peft>=0.18.1,<0.19.0 patchelf einops flashinfer-python==0.6.12 +# MiniMax Sparse Attention (MSA) kernels for MiniMax-M3 (sparse_use_msa=true). +# SM100-only, JIT-compiled on first use; MIT-licensed. Not yet on PyPI, so pin +# a commit; pip fetches the vendored CUTLASS headers via git submodules. +fmha_sm100 @ git+https://github.com/MiniMax-AI/MSA.git@e2ebe7656649f619af0ad1d457b534283034655e opencv-python-headless xgrammar==0.1.32 llguidance==0.7.29 diff --git a/tensorrt_llm/_torch/attention_backend/fmha/msa_proxy_mqa.py b/tensorrt_llm/_torch/attention_backend/fmha/msa_proxy_mqa.py index 837a258241e3..70c3d6162d7b 100644 --- a/tensorrt_llm/_torch/attention_backend/fmha/msa_proxy_mqa.py +++ b/tensorrt_llm/_torch/attention_backend/fmha/msa_proxy_mqa.py @@ -95,6 +95,13 @@ def forward_proxy( worklists, KV-split workspaces) and ``fmha_sm100`` runs the kernel with ``output_o=False, output_maxscore=True`` so only the score tensor is materialized. + + This wrapper serves the eager prefill path only. Decode runs + through the in-tree graph-safe driver + (``sparse.minimax_m3.decode_wrapper``) because + ``fmha_sm100_plan`` is CUDA-graph-hostile (unpinned H2D + staging, per-call device allocations, device-side cost sweep + with ``.tolist()``). """ # Imported here (not at module top) so the registry can still # advertise the class on hosts where fmha_sm100 is absent -- @@ -133,6 +140,7 @@ def forward_proxy( page_size=page_size, output_maxscore=True, causal=causal, + num_kv_splits=1, ) _, max_score = fmha_sm100.fmha_sm100( idx_q, diff --git a/tensorrt_llm/_torch/attention_backend/fmha/msa_sparse_gqa.py b/tensorrt_llm/_torch/attention_backend/fmha/msa_sparse_gqa.py index c014e8f2354f..b4ab275536f7 100644 --- a/tensorrt_llm/_torch/attention_backend/fmha/msa_sparse_gqa.py +++ b/tensorrt_llm/_torch/attention_backend/fmha/msa_sparse_gqa.py @@ -145,6 +145,7 @@ def forward_block_sparse( page_size=page_size, kv_block_num=int(kv_block_indexes.shape[-1]), causal=causal, + num_kv_splits=1, ) out, _ = fmha_sm100.fmha_sm100( q, diff --git a/tensorrt_llm/_torch/attention_backend/sparse/minimax_m3/cache_manager.py b/tensorrt_llm/_torch/attention_backend/sparse/minimax_m3/cache_manager.py index b9badca6ea35..45341136baee 100644 --- a/tensorrt_llm/_torch/attention_backend/sparse/minimax_m3/cache_manager.py +++ b/tensorrt_llm/_torch/attention_backend/sparse/minimax_m3/cache_manager.py @@ -154,7 +154,17 @@ def __init__( # then from ``sparse_attn_config``, then from the M3 checkpoint # convention (layers 0..2 dense, 3..N-1 sparse, # disable_index_value=True, sparse_index_dim=128). - sparse_attn_config = kwargs.get("sparse_attn_config") + # The pyexecutor constructs KV cache managers with the kwarg + # named ``sparse_attention_config`` (see + # ``_util.py:_create_kv_cache_manager``); accept the short + # spelling too for direct test construction. Reading only the + # short name silently yielded ``use_msa=False`` in production, + # so ``prepare()`` never pre-built MSA plans and the captured + # decode fell back to in-forward planning (the frozen-plan CUDA + # graph bug). + sparse_attn_config = kwargs.get("sparse_attn_config") or kwargs.get( + "sparse_attention_config" + ) num_layers = kwargs.get("num_layers") if sparse_index_dim is None: @@ -173,6 +183,14 @@ def __init__( self.sparse_layer_ids = sorted(int(i) for i in sparse_layer_ids) self.disable_index_value_layer_ids = set(int(i) for i in disable_index_value_layer_ids) self.sparse_index_dim = int(sparse_index_dim) + # Surface whether the runtime dispatches through the MSA-backed + # FMHA (``fmha_sm100``) path so ``MiniMaxM3AttentionMetadata.prepare()`` + # can pre-build the MSA plan objects outside the CUDA graph + # capture window. Read from the sparse-attention config; the + # config's ``sparse_use_msa`` field is the single source of truth + # for backend dispatch (see + # :func:`tensorrt_llm._torch.attention_backend.sparse.utils._resolve_minimax_m3_backend_cls`). + self.use_msa = bool(getattr(sparse_attn_config, "sparse_use_msa", False)) super().__init__(*args, **kwargs) diff --git a/tensorrt_llm/_torch/attention_backend/sparse/minimax_m3/decode_wrapper/__init__.py b/tensorrt_llm/_torch/attention_backend/sparse/minimax_m3/decode_wrapper/__init__.py new file mode 100644 index 000000000000..94712b9d5031 --- /dev/null +++ b/tensorrt_llm/_torch/attention_backend/sparse/minimax_m3/decode_wrapper/__init__.py @@ -0,0 +1,21 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""CUDA-graph-safe decode wrapper for MiniMax-M3 sparse attention. + +Drives the external MSA (``fmha_sm100``) SM100 kernels with launch +arguments assembled from device tensors, so decode steps can be +captured into CUDA graphs and replayed correctly as sequence state +advances. Public surface: + +* :func:`proxy_mqa_decode` — proxy MQA (indexer) pass, one KV head, + per-KV-block max-score output. +* :func:`sparse_gqa_decode` — block-sparse GQA main pass over the + top-k selected KV blocks. +""" + +from .dispatch import proxy_mqa_decode, sparse_gqa_decode + +__all__ = [ + "proxy_mqa_decode", + "sparse_gqa_decode", +] diff --git a/tensorrt_llm/_torch/attention_backend/sparse/minimax_m3/decode_wrapper/dispatch.py b/tensorrt_llm/_torch/attention_backend/sparse/minimax_m3/decode_wrapper/dispatch.py new file mode 100644 index 000000000000..25621310a25c --- /dev/null +++ b/tensorrt_llm/_torch/attention_backend/sparse/minimax_m3/decode_wrapper/dispatch.py @@ -0,0 +1,497 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Graph-safe decode driver for MiniMax-M3 sparse attention. + +Replaces MSA's host-centric ``fmha_sm100_plan`` / ``fmha_sm100`` driver +(``api.py``) for the decode path while invoking the *same* JIT-compiled +SM100 kernel binaries via ``fmha_sm100.jit.get_fmha_variant``: the MSA +plan bakes host-side values into the launch, which CUDA graph replays +would freeze, so only the driver is replaced and the kernels are +reused. + +Design contract: + +* No plan/run split — every call assembles launch args directly. +* Everything per-step-varying is a device tensor: ``seq_lens``, + ``kv_page_indptr``, ``kv_indices``, ``kv_block_indexes``, the + ``max_score`` contents. +* Host-baked values are geometry / per-batch-size constants only: + head counts, pack factor, page size, ``max_k_tiles`` capacity, + worklists (a pure function of batch size for decode). +* Every method is callable inside a CUDA graph capture and yields + correct results at replay: no ``.item()`` / ``.cpu()`` / ``.tolist()``. +""" + +from __future__ import annotations + +import math +from dataclasses import dataclass +from typing import Dict, Tuple + +import torch + +from .topk import select_topk_blocks +from .worklist import build_decode_worklist + +# Mirrors fmha_sm100.jit._PACK_FACTORS. +_PACK_FACTORS = (1, 2, 4, 6, 8, 16) +_QO_TILE_SIZE = 128 +_WORKSPACE_BYTES = 32 * 1024 * 1024 + + +def _compute_pack_factor(max_qo_len: int, num_qo_heads: int, num_kv_heads: int) -> int: + """Verbatim port of ``fmha_sm100.api._compute_pack_factor``.""" + if num_kv_heads == -1: + return 1 + h_r = num_qo_heads // num_kv_heads + if h_r <= 1 or max_qo_len <= 0 or max_qo_len > 32: + return 1 + max_pf = 128 // max_qo_len + for pf in reversed(_PACK_FACTORS): + if pf <= max_pf and pf <= h_r and h_r % pf == 0: + return pf + return 1 + + +def _max_k_tiles_capacity(max_kv_len: int) -> int: + """MSA's ``max_k_tiles`` formula (``api.py:658``) at max capacity. + + Number of 128-token KV blocks rounded up to a multiple of 128 (the + kernel's max-score tile stride granularity). + """ + return math.ceil(math.ceil(max_kv_len / 128) / 128) * 128 + + +@dataclass(frozen=True) +class M3DecodeGeometry: + """Compile/alloc-time constants for one M3 layer family (per rank).""" + + num_q_heads: int + num_kv_heads: int + num_index_heads: int + head_dim: int + page_size: int + topk: int + init_blocks: int + local_blocks: int + max_batch: int + max_kv_len: int + + def __post_init__(self): + if self.head_dim != 128 or self.page_size != 128: + raise NotImplementedError( + "MSA SM100 decode kernels require head_dim=128 and page_size=128; " + f"got head_dim={self.head_dim}, page_size={self.page_size}." + ) + if self.num_q_heads % self.num_kv_heads != 0: + raise ValueError("num_q_heads must be divisible by num_kv_heads") + if self.num_index_heads % self.num_kv_heads != 0: + raise ValueError("num_index_heads must be divisible by num_kv_heads") + + +class M3DecodeKernelDriver: + """Persistent-state decode driver for one (device, geometry) pair. + + Allocates every buffer once at construction with max-capacity + geometry; per-call views are prefixes/strided views of those + buffers so their ``data_ptr()`` is stable across CUDA graph + replays. + """ + + def __init__(self, geometry: M3DecodeGeometry, device: torch.device): + import fmha_sm100 # noqa: F401 — hard dependency of this driver + from fmha_sm100.jit import _dlpack_dtype_code, get_fmha_variant + + self.geom = g = geometry + self.device = device + + # --- pack factors and packed head counts (decode: qo_len == 1) --- + self.pf_proxy = _compute_pack_factor(1, g.num_index_heads, 1) + self.heads_packed_proxy = g.num_index_heads // self.pf_proxy + self.pf_sparse = _compute_pack_factor(1, g.num_q_heads, g.num_kv_heads) + self.heads_packed_sparse = g.num_q_heads // self.pf_sparse + + self.max_k_tiles = _max_k_tiles_capacity(g.max_kv_len) + self.num_ctas = int(torch.cuda.get_device_properties(device).multi_processor_count) + + # --- kernel variant modules (JIT-compiled once, same binaries the + # MSA api path runs) ----------------------------------------- + bf16_code = _dlpack_dtype_code(torch.bfloat16) + # Proxy: OnlyScore (sparse_mode=2), single_wg, no split. + self._proxy_module = get_fmha_variant( + bf16_code, _QO_TILE_SIZE, True, 2, g.page_size, False, self.pf_proxy + ) + # Sparse GQA: Sparse (sparse_mode=0). + self._sparse_module = get_fmha_variant( + bf16_code, _QO_TILE_SIZE, True, 0, g.page_size, False, self.pf_sparse + ) + + # --- persistent buffers ----------------------------------------- + self.workspace_buffer = torch.empty(_WORKSPACE_BYTES, dtype=torch.uint8, device=device) + self._max_score_flat = torch.empty( + g.num_index_heads * self.max_k_tiles * g.max_batch, + dtype=torch.float32, + device=device, + ) + self._kv_block_indexes = torch.full( + (g.max_batch, g.num_kv_heads, g.topk), -1, dtype=torch.int32, device=device + ) + self._out = torch.empty( + g.max_batch, g.num_q_heads, g.head_dim, dtype=torch.bfloat16, device=device + ) + self._kv_segment_offsets = torch.zeros(g.max_batch + 1, dtype=torch.int32, device=device) + self._valid_pages = torch.zeros(g.max_batch, dtype=torch.int32, device=device) + self._qo_offset = torch.zeros(g.max_batch, dtype=torch.int32, device=device) + + # Per-batch-size constants, built lazily and cached (bounded by + # the distinct batch sizes seen: CUDA graph buckets + eager). + self._qo_const_cache: Dict[Tuple[int, int], Tuple[torch.Tensor, torch.Tensor]] = {} + self._worklist_cache: Dict[Tuple[int, int], Tuple[torch.Tensor, torch.Tensor]] = {} + + # ------------------------------------------------------------------ + # Cached per-batch-size constants (host work happens once per shape, + # outside any capture — callers warm shapes up before capturing). + # ------------------------------------------------------------------ + + def _qo_consts(self, batch: int, pack_factor: int) -> Tuple[torch.Tensor, torch.Tensor]: + """(qo_segment_lens, qo_segment_offsets) for packed decode lens.""" + key = (batch, pack_factor) + cached = self._qo_const_cache.get(key) + if cached is None: + lens = torch.full((batch,), pack_factor, dtype=torch.int32, device=self.device) + offsets = ( + (torch.arange(batch + 1, dtype=torch.int64) * pack_factor) + .to(torch.int32) + .to(self.device) + ) + cached = (lens, offsets) + self._qo_const_cache[key] = cached + return cached + + def _worklist(self, batch: int, num_packed_heads: int) -> Tuple[torch.Tensor, torch.Tensor]: + key = (batch, num_packed_heads) + cached = self._worklist_cache.get(key) + if cached is None: + cached = build_decode_worklist( + batch_size=batch, + num_packed_heads=num_packed_heads, + num_ctas=self.num_ctas, + device=self.device, + ) + self._worklist_cache[key] = cached + return cached + + def warmup_shapes(self, batch: int) -> None: + """Pre-build all per-batch-size constants for ``batch``. + + Call once per CUDA graph bucket before capture so no host-side + cache misses happen inside the captured region. + """ + self._qo_consts(batch, self.pf_proxy) + self._qo_consts(batch, self.pf_sparse) + self._worklist(batch, self.heads_packed_proxy) + self._worklist(batch, self.heads_packed_sparse) + + # ------------------------------------------------------------------ + # Shared per-call device-side metadata refresh + # ------------------------------------------------------------------ + + def _kv_offsets_view(self, seq_lens: torch.Tensor, batch: int) -> torch.Tensor: + """Cumulative KV lengths into the persistent buffer (device op).""" + view = self._kv_segment_offsets[: batch + 1] + torch.cumsum(seq_lens, 0, dtype=torch.int32, out=view[1:]) + return view + + def _qo_offset_view(self, seq_lens: torch.Tensor, batch: int) -> torch.Tensor: + """Per-request causal offset ``kv_len - 1`` (device op). + + The kernel's causal bound is inclusive (attend positions + ``<= offset + q_idx``); with one query token at position + ``kv_len - 1`` this unmasks exactly the ``kv_len`` cached + positions. ``kv_len`` itself would leak one stale slot from a + partially-filled last page in *sparse* mode, which has no + secondary seqlen clip (verified empirically in + ``test_minimax_m3_decode_driver_vs_msa.py`` hetero cases). + """ + view = self._qo_offset[:batch] + torch.sub(seq_lens, 1, out=view) + return view + + # ------------------------------------------------------------------ + # Proxy MQA pass (indexer): per-KV-block max scores + # ------------------------------------------------------------------ + + def proxy_max_score( + self, + idx_q: torch.Tensor, + idx_k_paged: torch.Tensor, + *, + seq_lens: torch.Tensor, + kv_page_indptr: torch.Tensor, + kv_indices: torch.Tensor, + sm_scale: float, + ) -> torch.Tensor: + """Dense MQA proxy pass, OnlyScore mode. + + Parameters + ---------- + idx_q : ``[batch, num_index_heads, 128]`` bf16 (decode: 1 token/req). + idx_k_paged : ``[num_pages, 1, page_size, 128]`` bf16 (HND). + seq_lens : ``[batch]`` int32 device — per-request KV length. + kv_page_indptr : ``[batch + 1]`` int32 device. + kv_indices : ``[total_pages]`` int32 device page table. + sm_scale : softmax scale (does not affect max-score ranking). + + Returns + ------- + ``[num_index_heads, max_k_tiles, batch]`` fp32 view into the + persistent max-score buffer; unwritten tiles are ``-inf``. + """ + g = self.geom + batch = idx_q.shape[0] + seq_lens = seq_lens[:batch] + + max_score = torch.as_strided( + self._max_score_flat, + (g.num_index_heads, self.max_k_tiles, batch), + (self.max_k_tiles * batch, batch, 1), + ) + max_score.fill_(float("-inf")) + + qo_lens, qo_offsets = self._qo_consts(batch, self.pf_proxy) + work_range, work_info = self._worklist(batch, self.heads_packed_proxy) + kv_offsets = self._kv_offsets_view(seq_lens, batch) + qo_offset = self._qo_offset_view(seq_lens, batch) + + self._proxy_module.run( + self.workspace_buffer, + idx_q, + idx_k_paged, + idx_k_paged, + qo_lens, + seq_lens, + qo_offsets, + kv_offsets, + work_range, + work_info, + None, # out (OnlyScore) + float(sm_scale), + 1.0, + 1.0, + 1.0, + 1.0, + self.pf_proxy, # max_qo_len after packing + qo_offset, # kv_len - 1: inclusive causal bound = last cached pos + 1, # num_kv_splits + None, + None, + None, # kv_tile_begin / end / split + None, + None, # workspace_o / workspace_lse + None, # num_kv_splits_per_row + _QO_TILE_SIZE, + kv_indices, + kv_page_indptr, + max_score, + self.max_k_tiles, + None, # kv_block_indexes + self.pf_proxy, + True, # qo_len_uniform + torch.cuda.current_stream().cuda_stream, + ) + return max_score + + # ------------------------------------------------------------------ + # Top-k block selection (device-driven) + # ------------------------------------------------------------------ + + def select_blocks( + self, + max_score: torch.Tensor, + *, + seq_lens: torch.Tensor, + ) -> torch.Tensor: + """Group-reduce index-head scores to KV heads and pick top-k blocks. + + Returns ``[batch, num_kv_heads, topk]`` int32 view into the + persistent ``kv_block_indexes`` buffer. + """ + g = self.geom + batch = max_score.shape[2] + seq_lens = seq_lens[:batch] + + group = g.num_index_heads // g.num_kv_heads + if group > 1: + max_score_kv = max_score.view(g.num_kv_heads, group, max_score.shape[1], batch).amax( + dim=1 + ) + else: + max_score_kv = max_score + + valid_pages = self._valid_pages[:batch] + torch.div(seq_lens + (g.page_size - 1), g.page_size, rounding_mode="floor", out=valid_pages) + + out = self._kv_block_indexes[:batch] + return select_topk_blocks( + max_score_kv, + valid_pages, + topk=g.topk, + init_blocks=g.init_blocks, + local_blocks=g.local_blocks, + out=out, + ) + + # ------------------------------------------------------------------ + # Sparse block-GQA main pass + # ------------------------------------------------------------------ + + def sparse_attention( + self, + q: torch.Tensor, + k_paged: torch.Tensor, + v_paged: torch.Tensor, + kv_block_indexes: torch.Tensor, + *, + seq_lens: torch.Tensor, + kv_page_indptr: torch.Tensor, + kv_indices: torch.Tensor, + sm_scale: float, + ) -> torch.Tensor: + """Block-sparse paged GQA over the selected KV blocks. + + Parameters + ---------- + q : ``[batch, num_q_heads, 128]`` bf16. + k_paged / v_paged : ``[num_pages, num_kv_heads, page_size, 128]`` bf16. + kv_block_indexes : ``[batch, num_kv_heads, topk]`` int32 ascending, + ``-1`` tail padded. + seq_lens / kv_page_indptr / kv_indices : as in ``proxy_max_score``. + + Returns + ------- + ``[batch, num_q_heads, 128]`` bf16 view into the persistent + output buffer. + """ + batch = q.shape[0] + seq_lens = seq_lens[:batch] + + out = self._out[:batch] + qo_lens, qo_offsets = self._qo_consts(batch, self.pf_sparse) + work_range, work_info = self._worklist(batch, self.heads_packed_sparse) + kv_offsets = self._kv_offsets_view(seq_lens, batch) + qo_offset = self._qo_offset_view(seq_lens, batch) + + self._sparse_module.run( + self.workspace_buffer, + q, + k_paged, + v_paged, + qo_lens, + seq_lens, + qo_offsets, + kv_offsets, + work_range, + work_info, + out, + float(sm_scale), + 1.0, + 1.0, + 1.0, + 1.0, + self.pf_sparse, # max_qo_len after packing + qo_offset, # kv_len - 1 (see _qo_offset_view) + 1, # num_kv_splits + None, + None, + None, + None, + None, + None, + _QO_TILE_SIZE, + kv_indices, + kv_page_indptr, + None, # max_score + -1, # max_k_tiles + kv_block_indexes, + self.pf_sparse, + True, # qo_len_uniform + torch.cuda.current_stream().cuda_stream, + ) + return out + + +# --------------------------------------------------------------------------- +# Module-level convenience API (driver cache + functional entry points) +# --------------------------------------------------------------------------- + +_driver_cache: Dict[Tuple, M3DecodeKernelDriver] = {} + + +def get_decode_driver(geometry: M3DecodeGeometry, device: torch.device) -> M3DecodeKernelDriver: + key = ( + geometry, + device.type, + device.index if device.index is not None else torch.cuda.current_device(), + ) + driver = _driver_cache.get(key) + if driver is None: + driver = M3DecodeKernelDriver(geometry, device) + _driver_cache[key] = driver + return driver + + +def proxy_mqa_decode( + idx_q: torch.Tensor, + idx_k_paged: torch.Tensor, + *, + geometry: M3DecodeGeometry, + seq_lens: torch.Tensor, + kv_page_indptr: torch.Tensor, + kv_indices: torch.Tensor, + sm_scale: float, +) -> torch.Tensor: + """Functional proxy pass: returns per-KV-block max scores.""" + driver = get_decode_driver(geometry, idx_q.device) + return driver.proxy_max_score( + idx_q, + idx_k_paged, + seq_lens=seq_lens, + kv_page_indptr=kv_page_indptr, + kv_indices=kv_indices, + sm_scale=sm_scale, + ) + + +def sparse_gqa_decode( + q: torch.Tensor, + k_paged: torch.Tensor, + v_paged: torch.Tensor, + kv_block_indexes: torch.Tensor, + *, + geometry: M3DecodeGeometry, + seq_lens: torch.Tensor, + kv_page_indptr: torch.Tensor, + kv_indices: torch.Tensor, + sm_scale: float, +) -> torch.Tensor: + """Functional sparse GQA pass over selected KV blocks.""" + driver = get_decode_driver(geometry, q.device) + return driver.sparse_attention( + q, + k_paged, + v_paged, + kv_block_indexes, + seq_lens=seq_lens, + kv_page_indptr=kv_page_indptr, + kv_indices=kv_indices, + sm_scale=sm_scale, + ) + + +__all__ = [ + "M3DecodeGeometry", + "M3DecodeKernelDriver", + "get_decode_driver", + "proxy_mqa_decode", + "sparse_gqa_decode", +] diff --git a/tensorrt_llm/_torch/attention_backend/sparse/minimax_m3/decode_wrapper/topk.py b/tensorrt_llm/_torch/attention_backend/sparse/minimax_m3/decode_wrapper/topk.py new file mode 100644 index 000000000000..023a79385ee9 --- /dev/null +++ b/tensorrt_llm/_torch/attention_backend/sparse/minimax_m3/decode_wrapper/topk.py @@ -0,0 +1,97 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Device-driven top-k KV-block selection for M3 sparse decode. + +Replaces ``fmha_sm100.sparse_topk_select`` on the decode path. The MSA +kernel takes ``num_valid_pages`` as a **single global host int** (the +max over the batch), which is both CUDA-graph-hostile (it varies per +step) and imprecise for heterogeneous batches (the forced "local" +window is anchored at the *global* last block instead of each row's +own last valid block). + +This implementation reads per-row valid page counts from a device +tensor, matching the in-tree Triton reference semantics +(``backend.py:_index_attention_and_select``: force first +``init_blocks`` and last ``local_blocks`` *valid* blocks per row, mask +invalid blocks, pick top-k, ascending indices, ``-1`` tail padding) +while staying pure torch ops — fully CUDA-graph-capturable. +""" + +from __future__ import annotations + +from typing import Optional + +import torch + +# Finite sentinel mirroring MSA's FLT_MAX forced-score marker (avoids +# inf arithmetic edge cases in torch.topk). +_FORCE_SCORE = torch.finfo(torch.float32).max +# Sort key sentinel that pushes -1 (invalid) entries to the tail while +# staying well inside int32. +_PAD_KEY = 0x40000000 + + +def select_topk_blocks( + max_score_kv: torch.Tensor, + valid_pages: torch.Tensor, + *, + topk: int, + init_blocks: int, + local_blocks: int, + out: Optional[torch.Tensor] = None, +) -> torch.Tensor: + """Select per-(token, kv-head) top-k KV block indices. + + Parameters + ---------- + max_score_kv : ``[num_kv_heads, max_k_tiles, total_q]`` float32. + Per-128-token-block max scores (already group-reduced to KV-head + granularity). Blocks beyond a row's valid count may hold any + value — they are masked here. + valid_pages : ``[total_q]`` int32/int64 device tensor. + Per-token count of valid KV blocks (``ceil(kv_len / page_size)``). + topk : number of blocks to select (M3: 16). + init_blocks : always include blocks ``[0, init_blocks)``. + local_blocks : always include blocks + ``[valid - local_blocks, valid)`` (per row). + out : optional ``[total_q, num_kv_heads, topk]`` int32 buffer. + + Returns + ------- + ``[total_q, num_kv_heads, topk]`` int32, ascending block indices, + ``-1`` padded at the tail (the sparse FMHA kernel's + ``kv_block_indexes`` contract). + """ + num_kv_heads, max_k_tiles, total_q = max_score_kv.shape + if max_k_tiles < topk: + raise ValueError(f"max_k_tiles ({max_k_tiles}) must be >= topk ({topk})") + + scores = max_score_kv.permute(2, 0, 1) # [total_q, H_kv, K] + k_idx = torch.arange(max_k_tiles, device=scores.device, dtype=torch.int64) + valid = valid_pages.to(torch.int64).view(total_q, 1, 1) + + forced = k_idx < init_blocks + if local_blocks > 0: + forced = forced | ((k_idx >= valid - local_blocks) & (k_idx < valid)) + else: + forced = forced.expand(total_q, 1, max_k_tiles) + + scores = torch.where(forced, scores.new_full((), _FORCE_SCORE), scores) + scores = torch.where(k_idx >= valid, scores.new_full((), float("-inf")), scores) + + top_scores, top_idx = torch.topk(scores, topk, dim=-1) + # Rows with fewer than topk valid blocks pick -inf slots: pad them. + top_idx = torch.where(top_scores == float("-inf"), -1, top_idx) + + # Ascending by block index with -1 at the tail. + sort_key = torch.where(top_idx < 0, _PAD_KEY, top_idx) + sort_key, _ = torch.sort(sort_key, dim=-1) + result = torch.where(sort_key == _PAD_KEY, -1, sort_key).to(torch.int32) + + if out is not None: + out.copy_(result) + return out + return result + + +__all__ = ["select_topk_blocks"] diff --git a/tensorrt_llm/_torch/attention_backend/sparse/minimax_m3/decode_wrapper/worklist.py b/tensorrt_llm/_torch/attention_backend/sparse/minimax_m3/decode_wrapper/worklist.py new file mode 100644 index 000000000000..0f00cb871342 --- /dev/null +++ b/tensorrt_llm/_torch/attention_backend/sparse/minimax_m3/decode_wrapper/worklist.py @@ -0,0 +1,86 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Static persistent-CTA worklists for the M3 sparse decode kernels. + +The MSA C++ SM100 FMHA kernel (``csrc/include/fmha_tile_scheduler.hpp``, +``HostPrecomputedTileScheduler``) walks a device worklist: + +* ``packed_work_range[cta] = end << 32 | start`` — the slice of + ``packed_work_info`` CTA ``cta`` owns; +* ``packed_work_info[i] = qo_tile << 32 | (head & 0xFFFF) << 16 | + (batch & 0xFFFF)`` — one (batch, packed-head, qo-tile) work item. + +MSA computes these with a device planner kernel +(``csrc/include/plan.cuh``) parameterised by per-step KV lengths for +load balancing. For **decode** (``qo_len == 1`` per request, +``num_kv_splits == 1``) the *set* of work items is a pure function of +``(batch_size, num_packed_heads)`` — KV lengths only affect which CTA +runs which item, never whether an item exists, and each item reads its +own KV bounds from device tensors at execution time. So the worklist +is a per-batch-size constant: build it once on the host, keep it in a +persistent device buffer, and CUDA graph replays stay correct for any +KV lengths. + +Item order is batch-major then head; CTA assignment is contiguous +chunks (equal cost per item is the decode regime). Per-item outputs +are independent, so assignment differences vs. MSA's greedy planner +cannot change results — only load balance. +""" + +from __future__ import annotations + +from typing import Tuple + +import torch + + +def build_decode_worklist( + *, + batch_size: int, + num_packed_heads: int, + num_ctas: int, + device: torch.device, +) -> Tuple[torch.Tensor, torch.Tensor]: + """Build ``(packed_work_range, packed_work_info)`` for one decode shape. + + Parameters + ---------- + batch_size : number of requests (== work-item batch dim; for the + per-token-expanded sparse pass this equals ``total_q``). + num_packed_heads : ``num_qo_heads // pack_factor`` — the head count + the kernel iterates after pack-GQA folding. + num_ctas : persistent grid width (SM count); also the length of + ``packed_work_range``. + device : CUDA device for the returned buffers. + + Returns + ------- + (packed_work_range ``[num_ctas]`` int64, packed_work_info ``[n]`` int64) + with ``n = batch_size * num_packed_heads`` (every item has + ``qo_tile == 0`` because decode packs at most 128 q rows per tile). + """ + if batch_size <= 0: + raise ValueError(f"batch_size must be positive, got {batch_size}") + if batch_size > 0xFFFF: + raise ValueError(f"batch_size {batch_size} exceeds the 16-bit work-item field") + if num_packed_heads <= 0 or num_packed_heads > 0xFFFF: + raise ValueError(f"num_packed_heads out of range: {num_packed_heads}") + + n_items = batch_size * num_packed_heads + + # Batch-major enumeration: item = b * H + h -> head h, batch b. + batch_idx = torch.arange(n_items, dtype=torch.int64) // num_packed_heads + head_idx = torch.arange(n_items, dtype=torch.int64) % num_packed_heads + work_info = (head_idx << 16) | batch_idx # qo_tile == 0 + + # Contiguous, maximally even split of [0, n_items) across CTAs. + # CTA c owns [c * n // C rounded, ...) via cumulative fair shares. + bounds = (torch.arange(num_ctas + 1, dtype=torch.int64) * n_items) // num_ctas + starts = bounds[:-1] + ends = bounds[1:] + work_range = (ends << 32) | starts + + return work_range.to(device), work_info.to(device) + + +__all__ = ["build_decode_worklist"] diff --git a/tensorrt_llm/_torch/attention_backend/sparse/minimax_m3/metadata.py b/tensorrt_llm/_torch/attention_backend/sparse/minimax_m3/metadata.py index 790e7694cf9b..6311fceff978 100644 --- a/tensorrt_llm/_torch/attention_backend/sparse/minimax_m3/metadata.py +++ b/tensorrt_llm/_torch/attention_backend/sparse/minimax_m3/metadata.py @@ -22,13 +22,17 @@ import dataclasses import functools +import os from dataclasses import dataclass, field -from typing import List, Literal, Optional, Tuple +from typing import TYPE_CHECKING, List, Literal, Optional, Tuple import torch from ..params import SparseParams +if TYPE_CHECKING: + from .msa_plan_cache import MsaPlanCache, MsaPlanCacheGeometry + @dataclass(frozen=True) class MiniMaxM3SparseParams(SparseParams): @@ -681,6 +685,85 @@ def build_runtime_metadata_from_kv_manager( return meta, out_cache_loc +def _build_msa_plans_for_metadata( + *, + m3_meta: "MiniMaxM3SparseAttentionMetadata", + geometry: "MsaPlanCacheGeometry", + cache_device: torch.device, + max_batch: int, + plan_cache: Optional["MsaPlanCache"], +) -> Tuple[Optional["MsaPlanCache"], Optional[dict]]: + """Refresh the persistent paged-KV staging for one scheduler step. + + Returns ``(plan_cache, msa_plans_dict)``. ``plan_cache`` owns the + persistent ``kv_indices`` / ``kv_page_indptr`` buffers (stable + ``data_ptr()`` across CUDA graph replays); ``msa_plans_dict`` is + the payload attached to ``self.minimax_m3["msa_plans"]`` so the + forward path can read the staged tables plus the per-request CPU + lens/offsets the prefill path consumes. + """ + # Local import so this file stays importable on hosts without MSA. + from .msa_plan_cache import MsaPlanCache + + # 1) Derive per-request CPU tensors (mirrors msa_backend. + # _qo_lens_offsets_from_metadata, kept in sync here so the + # forward sees the same values it would have built itself). + seq_lens_cpu = m3_meta.seq_lens_cpu.to(torch.int32) + batch = int(seq_lens_cpu.shape[0]) + if m3_meta.is_prefill: + if m3_meta.extend_seq_lens_cpu is None or m3_meta.prefix_lens is None: + # Prefill metadata is incomplete; skip staging. The sparse + # forward's eager fallback builds the page table in-forward + # (safe when outside capture). + return plan_cache, None + qo_lens_cpu = torch.tensor(m3_meta.extend_seq_lens_cpu, dtype=torch.int32) + qo_offset_cpu = m3_meta.prefix_lens.detach().to(device="cpu", dtype=torch.int32) + else: + qo_lens_cpu = torch.ones(batch, dtype=torch.int32) + qo_offset_cpu = (seq_lens_cpu - 1).to(torch.int32) + + # 2) Allocate the staging buffers lazily on the first call. Sizes + # are picked so any padded batch up to ``max_batch`` (from + # ``AttentionMetadata.max_num_sequences``) fits. + if plan_cache is None: + # kv_indices is max_batch * max_pages_per_seq; page_size is the + # sparse config's block_size (128 for M3) so max_pages_per_seq + # comes from req_to_token's max_kv_len column dimension. Use + # the current metadata's ``req_to_token`` as the size witness. + max_kv_len = int(m3_meta.req_to_token.shape[1]) + max_pages_per_seq = max(1, max_kv_len // int(geometry.block_size)) + max_kv_indices = max_batch * max_pages_per_seq + plan_cache = MsaPlanCache( + device=cache_device, + geometry=geometry, + max_batch=max_batch, + max_kv_indices=max_kv_indices, + ) + + # 3) Refresh the page table in-place into the persistent buffers. + plan_cache.build_from_metadata( + req_to_token=m3_meta.req_to_token, + slot_ids=m3_meta.slot_ids, + seq_lens=m3_meta.seq_lens, + seq_lens_cpu=m3_meta.seq_lens_cpu, + page_size=int(geometry.block_size), + ) + + msa_plans = { + "kv_indices": plan_cache.kv_indices, + "kv_page_indptr": plan_cache.kv_page_indptr, + "qo_lens_cpu": qo_lens_cpu, + "kv_lens_cpu": seq_lens_cpu, + "qo_offset_cpu": qo_offset_cpu, + "geometry": geometry, + # Capacity constants for the in-tree decode driver (stable + # across steps so the driver cache key stays constant). + "max_batch": int(plan_cache.max_batch), + "max_kv_len": int(m3_meta.req_to_token.shape[1]), + } + return plan_cache, msa_plans + + @functools.lru_cache(maxsize=1) def get_minimax_m3_attention_metadata_cls(): """Return :class:`MiniMaxM3AttentionMetadata` (lazy import). @@ -721,6 +804,20 @@ class MiniMaxM3AttentionMetadata(AttentionMetadata): # ``prepare()`` call decides to use them (``is_cuda_graph`` / # graph-stable mode). _m3_static_buffers: Optional[dict] = None + # MSA (fmha_sm100) plan cache with persistent stable buffers. + # Populated lazily on the first prepare() call that has both: + # * ``use_msa=True`` on the KV cache manager (see + # :class:`MiniMaxM3KVCacheManagerV2`), AND + # * geometry attached by the model layer's first sparse + # forward (``_msa_geometry`` -- see + # ``modeling_minimaxm3.py``). + # Both conditions are needed because the geometry is only known + # once a sparse layer runs, and this happens during eager warmup + # (before the CUDA graph capture pass). From the capture pass + # onwards, ``prepare()`` rebuilds the plans into the persistent + # buffers so the captured forward reads from stable addresses. + _msa_plan_cache: Optional["MsaPlanCache"] = None + _msa_geometry: Optional["MsaPlanCacheGeometry"] = None def _maybe_get_m3_static_buffers( self, cache_device: torch.device, kv_cache_manager @@ -798,6 +895,8 @@ def prepare(self) -> None: # memory and either produces wrong tokens or fires # ``Indexing.cu:1515`` ``srcIndex < srcSelectDimSize``. self.minimax_m3 = None + if os.environ.get("TLLM_M3_SYNC") == "pre": + torch.cuda.synchronize() # Production path: build the M3 metadata from the standard # AttentionMetadata fields. Requires kv_cache_manager + the @@ -908,6 +1007,60 @@ def prepare(self) -> None: "out_cache_loc": out_cache_loc, } + # -- MSA plan pre-build -- + # Runs OUTSIDE any CUDA graph capture window (prepare() is + # called from the model_engine's ``_prepare_inputs``, which + # sits between the scheduler and the captured forward). + # We only build plans when: + # * the KV cache manager was constructed with + # ``sparse_use_msa=True``; AND + # * the model layer has populated ``_msa_geometry`` on a + # prior eager forward call. + # If ``_msa_geometry`` is not yet set (first eager warmup + # pass), the MSA backend's forward falls back to its + # in-forward plan call. That path is safe outside capture + # and lets us bootstrap the geometry without a chicken-and- + # egg dependency. + use_msa = bool(getattr(kv_cache_manager, "use_msa", False)) + if os.environ.get("TLLM_M3_DEBUG_PREPARE") == "1": + import sys as _sys + + print( + f"[m3-prepare-debug] is_cuda_graph={getattr(self, 'is_cuda_graph', None)} " + f"use_msa={use_msa} inst_geom={self._msa_geometry is not None} " + f"batch={batch_size} is_extend={is_extend} " + f"kv_lens={kv_lens_cpu_list[:4]} num_cached={list(num_cached_per_seq)[:4]} " + f"m3_meta_id={id(m3_meta)} self_id={id(self)}", + file=_sys.stderr, + flush=True, + ) + geometry = self._msa_geometry + if geometry is None: + # Layer-constructor registration (always available once + # any MSA-backed sparse layer exists — in particular + # before CUDA graph capture, whose metadata instances + # never see the per-instance publication from the first + # eager forward). + from .msa_plan_cache import get_global_msa_geometry + + geometry = get_global_msa_geometry() + if use_msa and geometry is not None: + max_batch = int(getattr(self, "max_num_sequences", None) or self.max_num_requests) + self._msa_plan_cache, msa_plans = _build_msa_plans_for_metadata( + m3_meta=m3_meta, + geometry=geometry, + cache_device=cache_device, + max_batch=max_batch, + plan_cache=self._msa_plan_cache, + ) + if msa_plans is not None: + self.minimax_m3["msa_plans"] = msa_plans + # Route the same dict through the algorithm-side + # metadata so ``msa_backend.forward_sparse`` reads + # the staged tables without changing its call + # signature. + m3_meta.msa_plans = msa_plans + return MiniMaxM3AttentionMetadata diff --git a/tensorrt_llm/_torch/attention_backend/sparse/minimax_m3/msa_backend.py b/tensorrt_llm/_torch/attention_backend/sparse/minimax_m3/msa_backend.py index 6cfb24938ad9..7f81e200b10e 100644 --- a/tensorrt_llm/_torch/attention_backend/sparse/minimax_m3/msa_backend.py +++ b/tensorrt_llm/_torch/attention_backend/sparse/minimax_m3/msa_backend.py @@ -15,13 +15,17 @@ with ``output_maxscore=True`` and ``num_kv_heads=1``. ``fmha_sm100`` returns ``(None, max_score)`` where ``max_score`` has shape ``[num_index_heads, max_k_tiles, total_qo_len]``. - 2. **Block selection.** Feed ``max_score`` into ``sparse_topk_select`` - to produce ascending per-row top-k block indices with ``-1`` - padding. The kernel currently fixes ``topk = 16`` and the - MiniMax-M3 checkpoint matches that default exactly. + 2. **Block selection.** Reduce ``max_score`` to KV-head granularity + and select top-k blocks per query (ascending indices, ``-1`` + padded) with per-query valid-block masking and forced init/local + blocks. 3. **Sparse GQA.** Run a second ``fmha_sm100`` over the main K/V branch, passing the block indices via ``kv_block_indexes``. +Prefill runs this flow eagerly through the MSA plan/run API; pure +decode runs it through the CUDA-graph-safe in-tree driver +(:mod:`.decode_wrapper`). + This module deliberately keeps the cache-layout adapter explicit so the MSA backend is selectable per-layer without disturbing the existing Triton path. The MSA package is imported lazily; if it is @@ -35,7 +39,12 @@ import torch -from .backend import _write_main_kv_slots, get_minimax_m3_attention_backend_cls +from .backend import ( + _INIT_SCORE, + _LOCAL_SCORE, + _write_main_kv_slots, + get_minimax_m3_attention_backend_cls, +) from .metadata import ( MiniMaxM3SparseAttentionMetadata, MiniMaxM3SparseConfig, @@ -46,10 +55,11 @@ from .metadata import MiniMaxM3SparseParams -# MSA's ``sparse_topk_select`` kernel only ships a topk=16 path today, -# and ``fmha_sm100`` only ships head_dim=128 variants. Enforce these -# preconditions early so layer construction fails with a clear message -# rather than a cryptic shape error from inside the MSA JIT. +# ``fmha_sm100`` only ships head_dim=128 variants, and the in-tree +# block-selection/driver geometry is validated for topk=16 (the +# MiniMax-M3 checkpoint value). Enforce these preconditions early so +# layer construction fails with a clear message rather than a cryptic +# shape error from inside the MSA JIT. _MSA_REQUIRED_TOPK = 16 _MSA_REQUIRED_HEAD_DIM = 128 @@ -67,8 +77,9 @@ def _require_msa_module(): except ImportError as exc: # pragma: no cover - install-time error raise RuntimeError( "MiniMax-M3 MSA backend requires the external `fmha_sm100` " - "package (MSA: https://github.com/MiniMax-AI/MSA). Install " - "it with `pip install fmha_sm100`, or unset " + "package (MSA: https://github.com/MiniMax-AI/MSA; not on " + "PyPI). Install it with `pip install " + "'git+https://github.com/MiniMax-AI/MSA.git'`, or unset " "`sparse_use_msa` in the sparse attention config to fall " "back to the Triton reference path." ) from exc @@ -165,7 +176,6 @@ def _build_kv_indices_and_lens( slot_ids_long = metadata.slot_ids.to(torch.long) req_rows = metadata.req_to_token.index_select(0, slot_ids_long).to(torch.long) batch = int(req_rows.shape[0]) - max_kv_len = int(req_rows.shape[1]) seq_lens_cpu = metadata.seq_lens_cpu.to(torch.long).tolist() page_lists = [] @@ -174,16 +184,24 @@ def _build_kv_indices_and_lens( if kv_len <= 0: continue num_pages = (kv_len + page_size - 1) // page_size - # First slot of each page gives the page id (each block is a - # contiguous run of page_size slots, see KVCacheManagerV2). - # Clamp so trailing pages (rounded up beyond the request's - # block count) reuse the last valid block - they are masked out - # by seq_lens in the kernel. - max_page = max_kv_len // page_size + # First slot of each page gives the *global* page id into the paged + # cache (each block is a contiguous run of page_size slots, see + # KVCacheManagerV2). ``req_rows[b]`` already holds valid slot ids, + # so ``// page_size`` yields valid global page ids by construction. + # + # Do NOT clamp these to a per-request bound: ``max_kv_len // + # page_size`` is the per-request page count, not a global page-id + # bound. Clamping page ids to ``max_page - 1`` collapses the page + # table for every request whose pages exceed that count (i.e. every + # request after the first in a contiguous layout, and virtually all + # requests in production where block ids are global and + # non-contiguous), making the proxy FMHA read the wrong K/V and + # corrupting the block scores. (Ported from + # brb/feat/minimax_m3_mxfp8_msa commit 677bcb45e5.) page_starts = torch.arange(num_pages, device=device, dtype=torch.long) * page_size - page_starts = page_starts.clamp_max(max(0, max_kv_len - 1)) + # ``page_starts`` is bounded by ``(num_pages - 1) * page_size < kv_len`` + # so it never over-reads; no clamp needed on the read index either. page_ids = req_rows[b].gather(0, page_starts) // page_size - page_ids = page_ids.clamp_min(0).clamp_max(max(0, max_page - 1)) page_lists.append(page_ids.to(torch.int32)) if page_lists: @@ -223,6 +241,98 @@ class that (a) is a subclass of :class:`IndexerProxyFmha` and (b) return None +def _per_token_valid_blocks( + qo_lens_cpu: torch.Tensor, + kv_lens_cpu: torch.Tensor, + qo_offset_cpu: Optional[torch.Tensor], + *, + causal: bool, + block_size: int, +) -> torch.Tensor: + """Per-query number of valid KV blocks (causal-aware), on CPU. + + Ported from brb/feat/minimax_m3_mxfp8_msa commit 92d8405af5. + Expands per-request lens/offsets to a per-*token* ``[total_q]`` + tensor so block selection can honour each query token's own causal + extent — which ``sparse_topk_select``'s scalar ``num_valid_pages`` / + ``force_end_blocks`` cannot express. + """ + qo = qo_lens_cpu.to(torch.long) + kv = kv_lens_cpu.to(torch.long) + batch = int(qo.shape[0]) + total = int(qo.sum().item()) + if total == 0: + return torch.zeros(0, dtype=torch.long) + batch_row = torch.repeat_interleave(torch.arange(batch, dtype=torch.long), qo) + starts = torch.zeros(batch, dtype=torch.long) + if batch > 1: + starts[1:] = torch.cumsum(qo, 0)[:-1] + intra = torch.arange(total, dtype=torch.long) - starts[batch_row] + kv_per = kv[batch_row] + if causal: + if qo_offset_cpu is not None: + off = qo_offset_cpu.to(torch.long)[batch_row] + else: + off = (kv - qo)[batch_row] + eff = torch.minimum(off + intra + 1, kv_per) + else: + eff = kv_per + return (eff + block_size - 1) // block_size + + +def _select_blocks_from_maxscore( + max_score_kv: torch.Tensor, + *, + topk: int, + n_valid_blocks: torch.Tensor, + init_blocks: int, + local_blocks: int, +) -> torch.Tensor: + """Per-query block selection from per-KV-head block scores, in torch. + + Ported from brb/feat/minimax_m3_mxfp8_msa commit 92d8405af5. + Mirrors the reference ``backend._index_attention_and_select`` + selection (init/local forced blocks + per-query valid-block masking + + top-k) on the ``amax``-reduced per-KV-head scores + (``[num_kv_heads, n_blocks, total_q]``). Replaces + ``fmha_sm100.sparse_topk_select``, whose scalar ``num_valid_pages`` + / forced-block windows are batch-wide and therefore wrong for every + query shorter than the batch-longest. + + Returns ``[total_q, num_kv_heads, topk]`` int32, ascending block + indices with ``-1`` tail padding (unchanged downstream contract). + """ + num_kv_heads, n_blocks, total_q = max_score_kv.shape + device = max_score_kv.device + scores = max_score_kv.permute(2, 0, 1).to(torch.float32).clone() # [q, kv, blk] + block_ids = torch.arange(n_blocks, device=device, dtype=torch.long) + nvb = n_valid_blocks.to(device=device, dtype=torch.long) # [total_q] + + if init_blocks > 0: + init_mask = block_ids.view(1, 1, -1) < init_blocks + scores = torch.where(init_mask, torch.full_like(scores, _INIT_SCORE), scores) + if local_blocks > 0: + local_start = (nvb - local_blocks).clamp_min(0) # [total_q] + local_mask = (block_ids.view(1, -1) >= local_start.view(-1, 1)) & ( + block_ids.view(1, -1) < nvb.view(-1, 1) + ) # [total_q, n_blocks] + scores = torch.where(local_mask.unsqueeze(1), torch.full_like(scores, _LOCAL_SCORE), scores) + # Per-query replacement for the kernel's scalar num_valid_pages clamp. + block_valid = block_ids.view(1, -1) < nvb.view(-1, 1) # [total_q, n_blocks] + scores = scores.masked_fill(~block_valid.unsqueeze(1), float("-inf")) + + k = min(topk, n_blocks) + vals, idx = scores.topk(k=k, dim=-1) # [total_q, kv, k] + idx = torch.where(vals != float("-inf"), idx, torch.full_like(idx, -1)) + sort_key = torch.where(idx < 0, torch.full_like(idx, n_blocks), idx) + sort_key, _ = torch.sort(sort_key, dim=-1) + idx = torch.where(sort_key >= n_blocks, torch.full_like(sort_key, -1), sort_key) + if k < topk: + pad = torch.full((total_q, num_kv_heads, topk - k), -1, dtype=idx.dtype, device=device) + idx = torch.cat([idx, pad], dim=-1) + return idx.to(torch.int32) + + def _msa_index_proxy_and_topk( idx_q: torch.Tensor, idx_k_paged: torch.Tensor, @@ -273,11 +383,6 @@ def _msa_index_proxy_and_topk( ``sparse_topk_select``, mirroring the ``score_type='max'`` reduction the reference path performs. """ - # ``sparse_topk_select`` still lives in fmha_sm100 -- the - # top-k step is not factored into the FMHA registry today. - # (Future work: register it as a separate selector library.) - fmha_sm100 = _require_msa_module() - proxy_cls = _select_proxy_fmha_class() if proxy_cls is None: raise RuntimeError( @@ -319,14 +424,19 @@ def _msa_index_proxy_and_topk( else: max_score_kv = max_score - # ``num_valid_pages`` is per-call so the kernel masks out the - # rounded-up tail tiles; we pass the maximum across the batch and - # rely on the kernel's ``idx >= num_valid_pages`` check. + # Per-query valid-block counts + torch selection (replaces + # ``fmha_sm100.sparse_topk_select``, whose scalar num_valid_pages / + # forced windows are batch-wide — wrong for heterogeneous batches + # and for prefill where each token has its own causal extent). page_size = int(idx_k_paged.shape[2]) - max_valid_pages = ( - int(((kv_lens_cpu + page_size - 1) // page_size).max().item()) if kv_lens_cpu.numel() else 0 + n_valid_blocks = _per_token_valid_blocks( + qo_lens_cpu, + kv_lens_cpu, + qo_offset_cpu, + causal=causal, + block_size=page_size, ) - if max_valid_pages <= 0: + if n_valid_blocks.numel() == 0 or int(n_valid_blocks.max().item()) <= 0: # Degenerate batch (no KV) — return all-padded indices. return torch.full( (idx_q.shape[0], config.num_kv_heads, _MSA_REQUIRED_TOPK), @@ -335,12 +445,12 @@ def _msa_index_proxy_and_topk( device=idx_q.device, ) - return fmha_sm100.sparse_topk_select( - max_score_kv.contiguous(), - _MSA_REQUIRED_TOPK, - num_valid_pages=max_valid_pages, - force_begin_blocks=init_blocks, - force_end_blocks=local_blocks, + return _select_blocks_from_maxscore( + max_score_kv, + topk=_MSA_REQUIRED_TOPK, + n_valid_blocks=n_valid_blocks, + init_blocks=init_blocks, + local_blocks=local_blocks, ) @@ -416,6 +526,112 @@ def _msa_sparse_attention( ) +# --------------------------------------------------------------------------- +# In-tree graph-safe decode driver +# --------------------------------------------------------------------------- + + +def _intree_sparse_decode( + q: torch.Tensor, + idx_q: torch.Tensor, + k_paged: torch.Tensor, + v_paged: torch.Tensor, + idx_k_paged: torch.Tensor, + metadata: MiniMaxM3SparseAttentionMetadata, + config: MiniMaxM3SparseConfig, + *, + sm_scale: float, + idx_sm_scale: float, + page_size: int, +) -> torch.Tensor: + """Pure-decode path through the in-tree graph-safe driver. + + Replaces the MSA ``fmha_sm100_plan`` / ``fmha_sm100`` / + ``sparse_topk_select`` host driver with + :mod:`.decode_wrapper.dispatch` while running the same + JIT-compiled kernel binaries. Everything per-step-varying is read + from device tensors, so this function is CUDA-graph-capturable and + every replay tracks the current ``seq_lens`` / page tables (the + exact property the MSA host driver lacks). + """ + from .decode_wrapper.dispatch import M3DecodeGeometry, get_decode_driver + + batch = int(q.shape[0]) + seq_lens = metadata.seq_lens.to(torch.int32) + + msa_plans = getattr(metadata, "msa_plans", None) + if msa_plans is not None: + kv_indices = msa_plans["kv_indices"] + kv_page_indptr = msa_plans["kv_page_indptr"] + max_batch = int(msa_plans.get("max_batch") or 0) + max_kv_len = int(msa_plans.get("max_kv_len") or 0) + else: + # Eager fallback when prepare() did not pre-stage the page + # table (e.g. focused unit tests). Host-side work is fine here, + # but it must never run inside a capture — fail loudly instead + # of letting the unpinned H2D copy below produce a cryptic + # capture error (or worse, silently freeze stale values). + if torch.cuda.is_current_stream_capturing(): + raise RuntimeError( + "MiniMax-M3 in-tree decode reached the eager fallback during " + "CUDA graph capture: metadata.msa_plans was not pre-staged by " + "prepare(). This means the MSA geometry was not registered " + "before capture (see msa_plan_cache.set_global_msa_geometry)." + ) + kv_indices, _ = _build_kv_indices_and_lens(metadata, page_size) + num_pages_cpu = (metadata.seq_lens_cpu.to(torch.long) + page_size - 1) // page_size + kv_page_indptr = torch.zeros(batch + 1, dtype=torch.int32) + kv_page_indptr[1:] = num_pages_cpu.to(torch.int32).cumsum(0) + kv_page_indptr = kv_page_indptr.to(q.device, non_blocking=True) + max_batch = 0 + max_kv_len = 0 + + if max_batch <= 0: + # Stable power-of-two capacity so the driver cache key does not + # churn as eager batch sizes vary. + max_batch = max(64, 1 << (batch - 1).bit_length()) + if max_kv_len <= 0: + max_kv_len = int(metadata.req_to_token.shape[1]) + + geometry = M3DecodeGeometry( + num_q_heads=config.num_q_heads, + num_kv_heads=config.num_kv_heads, + num_index_heads=config.num_index_heads, + head_dim=config.head_dim, + page_size=page_size, + topk=config.topk, + init_blocks=config.init_blocks, + local_blocks=config.local_blocks, + max_batch=max_batch, + max_kv_len=max_kv_len, + ) + driver = get_decode_driver(geometry, q.device) + + max_score = driver.proxy_max_score( + idx_q, + idx_k_paged, + seq_lens=seq_lens, + kv_page_indptr=kv_page_indptr, + kv_indices=kv_indices, + sm_scale=idx_sm_scale, + ) + kv_block_indexes = driver.select_blocks(max_score, seq_lens=seq_lens) + out = driver.sparse_attention( + q, + k_paged, + v_paged, + kv_block_indexes, + seq_lens=seq_lens, + kv_page_indptr=kv_page_indptr, + kv_indices=kv_indices, + sm_scale=sm_scale, + ) + # ``out`` aliases the driver's persistent buffer; the caller + # consumes it immediately (o_proj input) before the next layer's + # dispatch overwrites it, which is stream-ordered and safe. + return out.reshape(batch, config.num_q_heads * config.head_dim) + + # --------------------------------------------------------------------------- # Public forward entry points # --------------------------------------------------------------------------- @@ -505,8 +721,15 @@ def minimax_m3_msa_sparse_prefill( f"got page_size={page_size}, sparse_block_size={config.block_size}." ) - qo_lens_cpu, kv_lens_cpu, qo_offset_cpu = _qo_lens_offsets_from_metadata(metadata) - kv_indices, _ = _build_kv_indices_and_lens(metadata, page_size) + msa_plans = getattr(metadata, "msa_plans", None) + if msa_plans is not None: + qo_lens_cpu = msa_plans["qo_lens_cpu"] + kv_lens_cpu = msa_plans["kv_lens_cpu"] + qo_offset_cpu = msa_plans["qo_offset_cpu"] + kv_indices = msa_plans["kv_indices"] + else: + qo_lens_cpu, kv_lens_cpu, qo_offset_cpu = _qo_lens_offsets_from_metadata(metadata) + kv_indices, _ = _build_kv_indices_and_lens(metadata, page_size) kv_block_indexes = _msa_index_proxy_and_topk( idx_q, @@ -551,14 +774,16 @@ def minimax_m3_msa_sparse_decode( sm_scale: Optional[float] = None, idx_sm_scale: Optional[float] = None, ) -> torch.Tensor: - """Pure-decode MSA path. + """Pure-decode path: always the in-tree graph-safe driver. Decode is the ``qo_len == 1`` specialization of - :func:`minimax_m3_msa_sparse_prefill` with ``causal=False`` (the - new token attends to all cached positions through ``seq_lens - 1``, - no in-batch causality). ``MSA``'s ``fmha_sm100`` handles the - sub-32 query path internally via its decode kernel selection - (``_prefill_qlen_threshold(sparse=True) == 32``). + :func:`minimax_m3_msa_sparse_prefill`. It runs the same + ``fmha_sm100`` kernel binaries as the prefill path but through + :mod:`.decode_wrapper.dispatch` — device-tensor launch + arguments, device-side top-k, CUDA-graph-capturable end to end. + The legacy MSA host-driver decode (``fmha_sm100_plan`` + + ``sparse_topk_select``) was removed after the in-tree driver + reached bit-parity in eager and passed GSM8K under CUDA graphs. """ if metadata.is_prefill: raise ValueError("MSA decode entry called with prefill metadata") @@ -590,39 +815,19 @@ def minimax_m3_msa_sparse_decode( f"got page_size={page_size}, sparse_block_size={config.block_size}." ) - qo_lens_cpu, kv_lens_cpu, qo_offset_cpu = _qo_lens_offsets_from_metadata(metadata) - kv_indices, _ = _build_kv_indices_and_lens(metadata, page_size) - - kv_block_indexes = _msa_index_proxy_and_topk( - 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, - config=config, - idx_sm_scale=idx_sm_scale, - causal=False, - init_blocks=config.init_blocks, - local_blocks=config.local_blocks, - ) - - out = _msa_sparse_attention( + return _intree_sparse_decode( q, + idx_q, k_paged, v_paged, - kv_block_indexes, - qo_lens_cpu=qo_lens_cpu, - kv_lens_cpu=kv_lens_cpu, - qo_offset_cpu=qo_offset_cpu, - kv_indices=kv_indices, + idx_k_paged, + metadata, + config, sm_scale=sm_scale, - causal=False, + idx_sm_scale=idx_sm_scale, + page_size=page_size, ) - batch = int(q.shape[0]) - return out.reshape(batch, config.num_q_heads * config.head_dim).contiguous() - # --------------------------------------------------------------------------- # AttentionBackend wrapper @@ -678,6 +883,26 @@ def __init__(self, *args, **kwargs): raise NotImplementedError( f"MSA backend requires topk={_MSA_REQUIRED_TOPK}, got {self.m3_config.topk}." ) + # Register the per-rank sparse geometry process-wide at + # construction time — before any forward, hence before any + # CUDA graph capture — so every metadata instance's + # ``prepare()`` (including the CUDA graph runner's separate + # instances) can pre-build the MSA plans / kv-indices + # staging. See ``msa_plan_cache.set_global_msa_geometry``. + from .msa_plan_cache import MsaPlanCacheGeometry, set_global_msa_geometry + + set_global_msa_geometry( + MsaPlanCacheGeometry( + num_q_heads=int(self.m3_config.num_q_heads), + num_kv_heads=int(self.m3_config.num_kv_heads), + num_index_heads=int(self.m3_config.num_index_heads), + head_dim=int(self.m3_config.head_dim), + block_size=int(self.m3_config.block_size), + topk=int(self.m3_config.topk), + init_blocks=int(self.m3_config.init_blocks), + local_blocks=int(self.m3_config.local_blocks), + ) + ) def forward_sparse( self, diff --git a/tensorrt_llm/_torch/attention_backend/sparse/minimax_m3/msa_plan_cache.py b/tensorrt_llm/_torch/attention_backend/sparse/minimax_m3/msa_plan_cache.py new file mode 100644 index 000000000000..213b868c674f --- /dev/null +++ b/tensorrt_llm/_torch/attention_backend/sparse/minimax_m3/msa_plan_cache.py @@ -0,0 +1,284 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Persistent paged-KV table staging for the MiniMax-M3 sparse path. + +Provides the CUDA-graph-stable device buffers the M3 sparse attention +reads every step: + + * :func:`build_stable_kv_indices` -- computes the flat paged-KV page + table (global page ids) into a persistent int32 buffer using + vectorized ops (avoiding the Python-per-batch loop the reference + helper uses). The buffer's ``data_ptr()`` is stable across calls, + so a captured CUDA graph keeps reading current values on every + replay. + * :class:`MsaPlanCache` -- owns those buffers plus the per-rank + sparse geometry; refreshed once per scheduler step from + :meth:`MiniMaxM3AttentionMetadata.prepare` (outside any capture + window). + * :func:`set_global_msa_geometry` -- process-wide geometry + registration from the backend constructor, so CUDA-graph metadata + instances (created separately by the graph runner) can stage + before any forward has run. + +Historical note: this module used to also stage ``fmha_sm100_plan`` +outputs for the legacy MSA host-driver decode. That driver was +replaced by the in-tree graph-safe driver +(``decode_wrapper/``); prefill plans in-forward (it always +runs eagerly) and decode assembles launch arguments from device +tensors, so no plan staging remains. +""" + +from __future__ import annotations + +from dataclasses import dataclass +from typing import Optional, Tuple + +import torch + +# --------------------------------------------------------------------------- +# kv_indices: vectorized, in-place into persistent buffer +# --------------------------------------------------------------------------- + + +def build_stable_kv_indices( + *, + req_to_token: torch.Tensor, + slot_ids: torch.Tensor, + seq_lens: torch.Tensor, + seq_lens_cpu: torch.Tensor, + page_size: int, + dst: torch.Tensor, + kv_page_indptr_dst: torch.Tensor, +) -> Tuple[torch.Tensor, torch.Tensor]: + """Compute the flat paged-KV page table into ``dst``. + + Vectorized replacement for + :func:`msa_backend._build_kv_indices_and_lens`. The old helper + looped in Python (``for b in range(batch)``) and allocated a fresh + ``kv_indices`` per call via ``torch.cat``, which is fine in eager + mode but has no stable ``data_ptr()`` under CUDA graph replay. + + Parameters + ---------- + req_to_token : ``[max_reqs, max_kv_len]`` int32 on cache device. + slot_ids : ``[batch]`` int32 on cache device. ``req_to_token`` + row indices for the current batch. + seq_lens : ``[batch]`` int32 on cache device. Per-request effective + KV length. + seq_lens_cpu : ``[batch]`` int32 or int64 on CPU. Same values, used + for the CPU-side sizing math so we do not force a D2H sync. + page_size : int, must equal ``req_to_token.stride`` block width + (the M3 KV cache manager enforces this). + dst : preallocated int32 buffer sized ``>= max_batch * max_pages``. + The output ``kv_indices`` view is ``dst[:total_pages]``. + kv_page_indptr_dst : preallocated int32 buffer sized ``>= batch + 1``. + The output ``kv_page_indptr`` view is ``kv_page_indptr_dst[:batch+1]``. + + Returns + ------- + (kv_indices, kv_page_indptr) as views into ``dst`` / ``kv_page_indptr_dst``. + The two views' ``data_ptr()`` is stable across calls because it + aliases the destination buffers. + """ + device = req_to_token.device + batch = int(seq_lens_cpu.shape[0]) + max_kv_len = int(req_to_token.shape[1]) + max_pages_per_seq = max_kv_len // page_size + + if batch == 0: + return dst[:0], kv_page_indptr_dst[:1].zero_() + + # Number of pages per request: ceil(seq_len / page_size). Do it on + # CPU so kv_page_indptr can be prepared as ints for the row/col + # gather without triggering a D2H sync on seq_lens. + seq_lens_cpu_long = seq_lens_cpu.to(torch.long).cpu() + num_pages_cpu = (seq_lens_cpu_long + page_size - 1) // page_size + num_pages_cpu = num_pages_cpu.clamp_min(0) + total_pages = int(num_pages_cpu.sum().item()) + if total_pages > int(dst.shape[0]): + raise RuntimeError( + f"MSA kv_indices persistent buffer too small: capacity {int(dst.shape[0])} " + f"< total pages {total_pages}. Increase max_kv_indices on allocate()." + ) + + # Build kv_page_indptr on CPU then copy the prefix into the + # persistent CPU-shadow-less buffer. Values are per-batch cumulative + # page counts (starts at 0). + kv_page_indptr_cpu = torch.empty(batch + 1, dtype=torch.int32) + kv_page_indptr_cpu[0] = 0 + kv_page_indptr_cpu[1:].copy_(num_pages_cpu.to(torch.int32).cumsum(0)) + kv_page_indptr_dst[: batch + 1].copy_( + kv_page_indptr_cpu.to(device=device, non_blocking=True), non_blocking=True + ) + + # Vectorized page-index gather: + # req_rows = req_to_token[slot_ids] -> [batch, max_kv_len] slot ids + # For each request b, valid pages are indices 0..num_pages[b]-1; + # the p-th page's first-slot column is p * page_size; its page + # id is req_rows[b, p*page_size] // page_size (mirroring the + # in-tree helper). We build a max-sized (batch, max_pages_per_seq) + # grid, gather with clamped column indices, then mask trailing + # invalid pages before scatter into ``dst``. + slot_ids_long = slot_ids.to(torch.long) + req_rows = req_to_token.index_select(0, slot_ids_long).to(torch.long) + + max_valid_pages = max(1, max_pages_per_seq) + pages_grid = torch.arange(max_valid_pages, device=device, dtype=torch.long) + # Column index of the first slot of page p: p * page_size, clamped + # to ``max_kv_len - 1`` for out-of-range pages so the gather does + # not fault. Out-of-range page ids are trimmed by the batch mask + # below. + col_idx = (pages_grid * page_size).clamp_max(max(0, max_kv_len - 1)) + # Broadcast to [batch, max_pages_per_seq] + col_idx_b = col_idx.unsqueeze(0).expand(batch, -1) + # The gathered values are *global* page ids into the paged cache; + # ``req_rows`` holds valid slot ids by construction so no value + # clamp is applied. (An earlier version clamped to + # ``max_pages_per_seq - 1`` — the per-request page count — which + # collapsed the page table for every request whose global page ids + # exceed that count and corrupted all requests after the first; + # same defect as msa_backend._build_kv_indices_and_lens, fixed in + # brb/feat/minimax_m3_mxfp8_msa commit 677bcb45e5.) + gathered = (req_rows.gather(1, col_idx_b) // page_size).to(torch.int32) + + # Build a valid-page mask per request: + # mask[b, p] = p < num_pages[b] + num_pages_dev = num_pages_cpu.to(device=device, dtype=torch.long, non_blocking=True) + mask = pages_grid.unsqueeze(0) < num_pages_dev.unsqueeze(1) # [batch, max_pages_per_seq] + + # Compact into the flat ``dst`` prefix using boolean indexing. + # torch.masked_select preserves row-major (batch, page) order which + # matches the ``concat([pages_of(seq_0), pages_of(seq_1), ...])`` + # layout kv_page_indptr encodes. + packed = torch.masked_select(gathered, mask) + dst[:total_pages].copy_(packed, non_blocking=True) + + return dst[:total_pages], kv_page_indptr_dst[: batch + 1] + + +# --------------------------------------------------------------------------- +# Cross-run cache +# --------------------------------------------------------------------------- + + +_GLOBAL_MSA_GEOMETRY: Optional["MsaPlanCacheGeometry"] = None + + +def set_global_msa_geometry(geometry: "MsaPlanCacheGeometry") -> None: + """Register the per-rank M3 sparse geometry process-wide. + + Called from ``MiniMaxM3MSARuntimeBackend.__init__`` — i.e. at layer + construction, before any forward and therefore before any CUDA + graph capture. ``MiniMaxM3AttentionMetadata.prepare()`` reads this + so the MSA plan / kv-indices staging runs for *every* metadata + instance, including the separate instances the CUDA graph runner + creates. Publishing only from the first sparse forward (the old + scheme) left graph-capture metadata without a geometry: their + ``prepare()`` skipped the plan pre-build and the captured forward + fell back to in-forward planning, freezing capture-time host values + into every replay. + + All sparse layers on a rank share one geometry; the first writer + wins and later identical writes are no-ops. + """ + global _GLOBAL_MSA_GEOMETRY + if _GLOBAL_MSA_GEOMETRY is None: + _GLOBAL_MSA_GEOMETRY = geometry + + +def get_global_msa_geometry() -> Optional["MsaPlanCacheGeometry"]: + return _GLOBAL_MSA_GEOMETRY + + +@dataclass +class MsaPlanCacheGeometry: + """Per-rank M3 model geometry needed to allocate the plan cache. + + Populated by the model layer (``MiniMaxM3Attention._sparse_forward``) + the first time it dispatches through the MSA backend. All sparse + layers share the same geometry so the value written by the first + layer is authoritative for the rest. + """ + + num_q_heads: int + num_kv_heads: int + num_index_heads: int + head_dim: int + block_size: int + topk: int + init_blocks: int + local_blocks: int + + +class MsaPlanCache: + """Persistent paged-KV table staging for the M3 sparse decode path. + + Historically this cache also staged ``fmha_sm100_plan`` outputs for + the legacy MSA host-driver decode; that path was removed once the + in-tree graph-safe driver (``decode_wrapper``) reached + bit-parity and passed GSM8K under CUDA graphs. What remains is the + piece both prefill and decode still need: ``kv_indices`` / + ``kv_page_indptr`` computed each step into persistent device + buffers whose ``data_ptr()`` is stable across CUDA graph replays. + + Lifecycle + --------- + 1. Allocated lazily on the first ``build_from_metadata`` call once + the geometry / device / capacity are known. + 2. Every subsequent ``build_from_metadata`` call rewrites the + buffer contents in-place for the current scheduler step. + """ + + def __init__( + self, + *, + device: torch.device, + geometry: MsaPlanCacheGeometry, + max_batch: int, + max_kv_indices: int, + ): + self.device = device + self.geometry = geometry + self.max_batch = int(max_batch) + self.max_kv_indices = int(max_kv_indices) + self.kv_indices_buf = torch.zeros(self.max_kv_indices, dtype=torch.int32, device=device) + self.kv_page_indptr_buf = torch.zeros(self.max_batch + 1, dtype=torch.int32, device=device) + # Populated on each build_from_metadata call. + self.kv_indices: Optional[torch.Tensor] = None + self.kv_page_indptr: Optional[torch.Tensor] = None + + def build_from_metadata( + self, + *, + req_to_token: torch.Tensor, + slot_ids: torch.Tensor, + seq_lens: torch.Tensor, + seq_lens_cpu: torch.Tensor, + page_size: int, + ) -> None: + """Refresh ``kv_indices`` / ``kv_page_indptr`` for one step. + + Runs entirely outside CUDA graph capture (called from + :meth:`MiniMaxM3AttentionMetadata.prepare`); the captured + forward reads the persistent buffers on every replay. + """ + kv_indices, kv_page_indptr = build_stable_kv_indices( + req_to_token=req_to_token, + slot_ids=slot_ids, + seq_lens=seq_lens, + seq_lens_cpu=seq_lens_cpu, + page_size=page_size, + dst=self.kv_indices_buf, + kv_page_indptr_dst=self.kv_page_indptr_buf, + ) + self.kv_indices = kv_indices + self.kv_page_indptr = kv_page_indptr + + +__all__ = [ + "MsaPlanCache", + "MsaPlanCacheGeometry", + "build_stable_kv_indices", + "get_global_msa_geometry", + "set_global_msa_geometry", +] diff --git a/tensorrt_llm/_torch/models/modeling_minimaxm3.py b/tensorrt_llm/_torch/models/modeling_minimaxm3.py index 0b8b9f3e00ad..54db2fe7ac83 100644 --- a/tensorrt_llm/_torch/models/modeling_minimaxm3.py +++ b/tensorrt_llm/_torch/models/modeling_minimaxm3.py @@ -1148,6 +1148,44 @@ def _sparse_forward( "ModelConfig so the standard attention-backend dispatch selects " "the M3 sparse runtime." ) + + # Publish the per-rank sparse geometry on the outer + # AttentionMetadata the first time any sparse layer dispatches. + # ``MiniMaxM3AttentionMetadata.prepare`` reads this on subsequent + # steps to pre-build the MSA (fmha_sm100) plan into + # CUDA-graph-stable buffers -- see + # :mod:`tensorrt_llm._torch.attention_backend.sparse.minimax_m3.msa_plan_cache`. + # Doing this attach here (not in ``__init__``) is what avoids a + # chicken-and-egg dependency: the metadata instance is created + # by the pyexecutor before any layer runs, and the FIRST sparse + # forward is always the eager warmup pass (no capture in + # flight), so writing to ``_msa_geometry`` at this point is + # CUDA-graph-safe. All sparse layers share the same m3_config + # on a given rank, so any layer's write is authoritative for + # the rest. + # + # IMPORTANT: publish on the metadata *class*, not the instance. + # CUDA graph capture uses separate metadata instances (created + # via ``create_cuda_graph_metadata``); an instance attribute + # written during eager warmup would be invisible to them, their + # ``prepare()`` would skip the plan pre-build, and the captured + # decode would fall back to in-forward planning — freezing + # capture-time host values into every replay (the original + # NaN-after-layer-3 CUDA graph bug). + if getattr(attn_metadata, "_msa_geometry", None) is None: + from ..attention_backend.sparse.minimax_m3.msa_plan_cache import MsaPlanCacheGeometry + + m3_config = self.attn.m3_config + type(attn_metadata)._msa_geometry = MsaPlanCacheGeometry( + num_q_heads=int(m3_config.num_q_heads), + num_kv_heads=int(m3_config.num_kv_heads), + num_index_heads=int(m3_config.num_index_heads), + head_dim=int(m3_config.head_dim), + block_size=int(m3_config.block_size), + topk=int(m3_config.topk), + init_blocks=int(m3_config.init_blocks), + local_blocks=int(m3_config.local_blocks), + ) o = self.attn.forward( q, k, diff --git a/tests/unittest/_torch/attention/sparse/test_minimax_m3_decode_driver_vs_msa.py b/tests/unittest/_torch/attention/sparse/test_minimax_m3_decode_driver_vs_msa.py new file mode 100644 index 000000000000..f07861713795 --- /dev/null +++ b/tests/unittest/_torch/attention/sparse/test_minimax_m3_decode_driver_vs_msa.py @@ -0,0 +1,467 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""M3/M4 validation: in-tree decode driver vs MSA api path, bit-exact. + +Runs the same JIT-compiled SM100 kernel binaries through (a) MSA's +host-centric ``fmha_sm100_plan`` / ``fmha_sm100`` driver and (b) the +in-tree graph-safe ``dispatch.M3DecodeKernelDriver``, on identical +inputs, and asserts bit-equality: + +* proxy MQA max-score pass — uniform and heterogeneous KV lens; +* top-k block selection — bit-diff vs ``sparse_topk_select`` on + uniform lens (global == per-row there), reference-diff vs a pure + Python implementation on heterogeneous lens; +* sparse block-GQA pass — same ``kv_block_indexes`` fed to both; +* full pipeline under CUDA graph capture/replay with mutated lengths + and contents between replays (the property MSA's driver lacks). + +Requires SM100 + the ``fmha_sm100`` package (see requirements.txt). +""" + +import math + +import pytest +import torch + +# Full-model per-rank M3 geometry. +NUM_Q_HEADS = 64 +NUM_KV_HEADS = 4 +NUM_INDEX_HEADS = 4 +HEAD_DIM = 128 +PAGE_SIZE = 128 +TOPK = 16 +INIT_BLOCKS = 0 +LOCAL_BLOCKS = 1 +MAX_KV_LEN = 2048 # 16 pages -> max_k_tiles rounds to 128, same as MSA's + +SM_SCALE = HEAD_DIM**-0.5 +IDX_SM_SCALE = HEAD_DIM**-0.5 + + +def _require_env(): + if not torch.cuda.is_available(): + pytest.skip("CUDA required") + major, _ = torch.cuda.get_device_capability() + if major != 10: + pytest.skip("SM100 (Blackwell) required") + try: + import fmha_sm100 # noqa: F401 + except ImportError: + pytest.skip("fmha_sm100 (MSA) not importable") + + +def _geometry(max_batch): + from tensorrt_llm._torch.attention_backend.sparse.minimax_m3.decode_wrapper.dispatch import ( # noqa: E501 + M3DecodeGeometry, + ) + + return M3DecodeGeometry( + num_q_heads=NUM_Q_HEADS, + num_kv_heads=NUM_KV_HEADS, + num_index_heads=NUM_INDEX_HEADS, + head_dim=HEAD_DIM, + page_size=PAGE_SIZE, + topk=TOPK, + init_blocks=INIT_BLOCKS, + local_blocks=LOCAL_BLOCKS, + max_batch=max_batch, + max_kv_len=MAX_KV_LEN, + ) + + +def _make_inputs(kv_lens, seed=0, pool_pages=None): + """Build synthetic paged caches + decode Q for the given KV lengths. + + ``pool_pages`` fixes the physical page-pool size so tensors keep + identical shapes across calls (required by the CUDA graph test's + in-place refreshes). + """ + device = torch.device("cuda") + gen = torch.Generator(device="cuda").manual_seed(seed) + batch = len(kv_lens) + kv_lens_t = torch.tensor(kv_lens, dtype=torch.int32) + num_pages = [(kv_len + PAGE_SIZE - 1) // PAGE_SIZE for kv_len in kv_lens] + total_pages = sum(num_pages) + if pool_pages is None: + pool_pages = max(total_pages, 1) + assert pool_pages >= total_pages, "pool too small for requested kv_lens" + + def r(*shape, dtype=torch.bfloat16, scale=1.0): + return (torch.randn(*shape, generator=gen, device=device, dtype=torch.float32) * scale).to( + dtype + ) + + q = r(batch, NUM_Q_HEADS, HEAD_DIM, scale=0.5) + idx_q = r(batch, NUM_INDEX_HEADS, HEAD_DIM, scale=0.5) + k_paged = r(pool_pages, NUM_KV_HEADS, PAGE_SIZE, HEAD_DIM, scale=0.5) + v_paged = r(pool_pages, NUM_KV_HEADS, PAGE_SIZE, HEAD_DIM, scale=0.5) + idx_k_paged = r(pool_pages, 1, PAGE_SIZE, HEAD_DIM, scale=0.5) + + # Shuffled physical page assignment (exercises kv_indices gather). + perm = torch.randperm(pool_pages, generator=gen, device=device)[:total_pages] + kv_indices = perm.to(torch.int32) + kv_page_indptr = torch.zeros(batch + 1, dtype=torch.int32, device=device) + kv_page_indptr[1:] = torch.cumsum(torch.tensor(num_pages, dtype=torch.int32, device=device), 0) + + return { + "batch": batch, + "q": q, + "idx_q": idx_q, + "k_paged": k_paged, + "v_paged": v_paged, + "idx_k_paged": idx_k_paged, + "kv_indices": kv_indices, + "kv_page_indptr": kv_page_indptr, + "kv_lens_cpu": kv_lens_t, + "seq_lens_dev": kv_lens_t.to(device), + } + + +# --------------------------------------------------------------------------- +# MSA reference path (mirrors msa_backend.py decode exactly) +# --------------------------------------------------------------------------- + + +def _msa_proxy_max_score(inp): + import fmha_sm100 + + batch = inp["batch"] + qo_lens_cpu = torch.ones(batch, dtype=torch.int32) + qo_offset_cpu = (inp["kv_lens_cpu"] - 1).to(torch.int32) + plan = fmha_sm100.fmha_sm100_plan( + qo_lens_cpu, + inp["kv_lens_cpu"], + NUM_INDEX_HEADS, + num_kv_heads=1, + qo_offset=qo_offset_cpu, + page_size=PAGE_SIZE, + output_maxscore=True, + causal=False, + num_kv_splits=1, + ) + _, max_score = fmha_sm100.fmha_sm100( + inp["idx_q"], + inp["idx_k_paged"], + inp["idx_k_paged"], + plan, + kv_indices=inp["kv_indices"], + output_o=False, + output_maxscore=True, + sm_scale=IDX_SM_SCALE, + ) + return max_score + + +def _msa_topk(max_score, kv_lens_cpu): + import fmha_sm100 + + max_valid = int(((kv_lens_cpu + PAGE_SIZE - 1) // PAGE_SIZE).max().item()) + return fmha_sm100.sparse_topk_select( + max_score.contiguous(), + TOPK, + num_valid_pages=max_valid, + force_begin_blocks=INIT_BLOCKS, + force_end_blocks=LOCAL_BLOCKS, + ) + + +def _msa_sparse(inp, kv_block_indexes, causal=False): + """MSA reference sparse pass. + + ``causal=False`` is the production decode configuration: MSA then + overwrites ``qo_offset`` with the *global* ``max_kv_len``, which in + sparse mode (no secondary seqlen clip) lets short requests attend + stale positions inside forced/OOB blocks — a real MSA hetero-batch + defect. ``causal=True`` keeps the per-request ``kv_len - 1`` + offsets and is the exact-semantics reference the in-tree driver + implements; both agree on uniform batches. + """ + import fmha_sm100 + + batch = inp["batch"] + qo_lens_cpu = torch.ones(batch, dtype=torch.int32) + qo_offset_cpu = (inp["kv_lens_cpu"] - 1).to(torch.int32) + plan = fmha_sm100.fmha_sm100_plan( + qo_lens_cpu, + inp["kv_lens_cpu"], + NUM_Q_HEADS, + num_kv_heads=NUM_KV_HEADS, + qo_offset=qo_offset_cpu, + page_size=PAGE_SIZE, + kv_block_num=TOPK, + causal=causal, + num_kv_splits=1, + ) + out, _ = fmha_sm100.fmha_sm100( + inp["q"], + inp["k_paged"], + inp["v_paged"], + plan, + kv_indices=inp["kv_indices"], + kv_block_indexes=kv_block_indexes, + sm_scale=SM_SCALE, + output_maxscore=False, + ) + return out + + +# --------------------------------------------------------------------------- +# In-tree driver path +# --------------------------------------------------------------------------- + + +def _driver(max_batch): + from tensorrt_llm._torch.attention_backend.sparse.minimax_m3.decode_wrapper.dispatch import ( # noqa: E501 + get_decode_driver, + ) + + return get_decode_driver(_geometry(max_batch), torch.device("cuda")) + + +def _intree_proxy(driver, inp): + return driver.proxy_max_score( + inp["idx_q"], + inp["idx_k_paged"], + seq_lens=inp["seq_lens_dev"], + kv_page_indptr=inp["kv_page_indptr"], + kv_indices=inp["kv_indices"], + sm_scale=IDX_SM_SCALE, + ) + + +def _intree_sparse(driver, inp, kv_block_indexes): + return driver.sparse_attention( + inp["q"], + inp["k_paged"], + inp["v_paged"], + kv_block_indexes, + seq_lens=inp["seq_lens_dev"], + kv_page_indptr=inp["kv_page_indptr"], + kv_indices=inp["kv_indices"], + sm_scale=SM_SCALE, + ) + + +# --------------------------------------------------------------------------- +# Pure-Python top-k reference (per-row semantics) +# --------------------------------------------------------------------------- + + +def _reference_topk(max_score_kv, kv_lens): + """Per-(token, kv-head) reference: force init/local, top-k, ascending.""" + num_kv_heads, max_k_tiles, total_q = max_score_kv.shape + scores = max_score_kv.float().cpu().numpy() + out = torch.full((total_q, num_kv_heads, TOPK), -1, dtype=torch.int32) + for t in range(total_q): + valid = (int(kv_lens[t]) + PAGE_SIZE - 1) // PAGE_SIZE + for h in range(num_kv_heads): + row = scores[h, :, t].copy() + row[valid:] = -math.inf + for k in range(min(INIT_BLOCKS, valid)): + row[k] = math.inf + for k in range(max(valid - LOCAL_BLOCKS, 0), valid): + row[k] = math.inf + order = sorted(range(max_k_tiles), key=lambda i: (-row[i], i))[:TOPK] + picked = sorted(i for i in order if row[i] != -math.inf) + for j, blk in enumerate(picked): + out[t, h, j] = blk + return out + + +# --------------------------------------------------------------------------- +# Tests +# --------------------------------------------------------------------------- + +UNIFORM_LENS = [512] * 8 +HETERO_LENS = [1, 130, 257, 128, 511, 1024, 33, 900] + + +@pytest.mark.parametrize("kv_lens", [UNIFORM_LENS, HETERO_LENS], ids=["uniform", "hetero"]) +def test_proxy_max_score_bitdiff(kv_lens): + _require_env() + inp = _make_inputs(kv_lens) + driver = _driver(max_batch=len(kv_lens)) + + ms_ref = _msa_proxy_max_score(inp) + ms_new = _intree_proxy(driver, inp) + torch.cuda.synchronize() + + assert ms_ref.shape == ms_new.shape, f"{ms_ref.shape} vs {ms_new.shape}" + same = ms_ref == ms_new + both_neginf = torch.isinf(ms_ref) & torch.isinf(ms_new) & (ms_ref == ms_new) + mismatch = (~same & ~both_neginf).sum().item() + assert mismatch == 0, f"proxy max_score mismatches: {mismatch}/{ms_ref.numel()}" + + +def test_topk_bitdiff_uniform(): + _require_env() + inp = _make_inputs(UNIFORM_LENS) + driver = _driver(max_batch=len(UNIFORM_LENS)) + + ms = _msa_proxy_max_score(inp) + blocks_ref = _msa_topk(ms, inp["kv_lens_cpu"]) + blocks_new = driver.select_blocks(ms, seq_lens=inp["seq_lens_dev"]) + torch.cuda.synchronize() + + assert torch.equal(blocks_ref, blocks_new), ( + f"topk mismatch rows: {(blocks_ref != blocks_new).any(dim=-1).nonzero()[:8].tolist()}" + ) + + +def test_topk_reference_hetero(): + _require_env() + inp = _make_inputs(HETERO_LENS) + driver = _driver(max_batch=len(HETERO_LENS)) + + ms = _intree_proxy(driver, inp) + blocks_new = driver.select_blocks(ms, seq_lens=inp["seq_lens_dev"]).cpu() + torch.cuda.synchronize() + blocks_ref = _reference_topk(ms, HETERO_LENS) + + assert torch.equal(blocks_ref, blocks_new), ( + f"per-row topk mismatch, first rows: ref={blocks_ref[:2].tolist()} " + f"new={blocks_new[:2].tolist()}" + ) + + +@pytest.mark.parametrize( + "kv_lens", + [UNIFORM_LENS, [130] * 8, HETERO_LENS], + ids=["uniform", "uniform_partial_page", "hetero"], +) +def test_sparse_gqa_bitdiff(kv_lens): + _require_env() + inp = _make_inputs(kv_lens) + driver = _driver(max_batch=len(kv_lens)) + + # Use MSA's own block selection for both sides to isolate the + # sparse kernel + driver comparison. Heterogeneous batches need + # the causal=True reference (per-request offsets — see _msa_sparse + # docstring); uniform batches match the production causal=False + # path bit-exactly as well. + ms = _msa_proxy_max_score(inp) + blocks = _msa_topk(ms, inp["kv_lens_cpu"]) + + # causal=True reference whenever MSA's causal=False global-offset + # shortcut would attend stale positions (hetero lens, or a + # partially filled last page). + needs_exact_ref = len(set(kv_lens)) > 1 or any(kv_len % PAGE_SIZE for kv_len in kv_lens) + out_ref = _msa_sparse(inp, blocks, causal=needs_exact_ref) + out_new = _intree_sparse(driver, inp, blocks) + torch.cuda.synchronize() + + assert out_ref.shape == out_new.shape + assert torch.equal(out_ref, out_new), ( + f"sparse out mismatch: max abs diff " + f"{(out_ref.float() - out_new.float()).abs().max().item()}" + ) + + +def test_full_pipeline_bitdiff_uniform(): + _require_env() + inp = _make_inputs(UNIFORM_LENS) + driver = _driver(max_batch=len(UNIFORM_LENS)) + + ms_ref = _msa_proxy_max_score(inp) + blocks_ref = _msa_topk(ms_ref, inp["kv_lens_cpu"]) + out_ref = _msa_sparse(inp, blocks_ref) + + ms_new = _intree_proxy(driver, inp) + blocks_new = driver.select_blocks(ms_new, seq_lens=inp["seq_lens_dev"]) + out_new = _intree_sparse(driver, inp, blocks_new) + torch.cuda.synchronize() + + assert torch.equal(blocks_ref, blocks_new) + assert torch.equal(out_ref, out_new) + + +def test_cuda_graph_replay_tracks_device_state(): + """Capture once, mutate lengths + contents, replay: must match eager. + + This is exactly the failure mode of the MSA driver (frozen plan) — + the in-tree driver must produce bit-identical results to its own + eager execution for every replay. + """ + _require_env() + batch = 8 + driver = _driver(max_batch=batch) + driver.warmup_shapes(batch) + + # Persistent input buffers the graph will read. + pool_pages = batch * (MAX_KV_LEN // PAGE_SIZE) + inp0 = _make_inputs([256] * batch, seed=1, pool_pages=pool_pages) + seq_lens = inp0["seq_lens_dev"].clone() + kv_page_indptr = inp0["kv_page_indptr"].clone() + kv_indices_buf = torch.zeros( + batch * (MAX_KV_LEN // PAGE_SIZE), dtype=torch.int32, device="cuda" + ) + n0 = inp0["kv_indices"].shape[0] + kv_indices_buf[:n0] = inp0["kv_indices"] + q = inp0["q"].clone() + idx_q = inp0["idx_q"].clone() + k_paged = inp0["k_paged"].clone() + v_paged = inp0["v_paged"].clone() + idx_k_paged = inp0["idx_k_paged"].clone() + + def run_pipeline(): + ms = driver.proxy_max_score( + idx_q, + idx_k_paged, + seq_lens=seq_lens, + kv_page_indptr=kv_page_indptr, + kv_indices=kv_indices_buf, + sm_scale=IDX_SM_SCALE, + ) + blocks = driver.select_blocks(ms, seq_lens=seq_lens) + return driver.sparse_attention( + q, + k_paged, + v_paged, + blocks, + seq_lens=seq_lens, + kv_page_indptr=kv_page_indptr, + kv_indices=kv_indices_buf, + sm_scale=SM_SCALE, + ) + + # Warm up (JIT + shape caches + allocator) then capture. + out_view = run_pipeline() + torch.cuda.synchronize() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + out_view = run_pipeline() + + for step, lens in enumerate( + [[256] * batch, [384] * batch, [128, 512, 1920, 256, 640, 64, 2048, 300]] + ): + # Refresh device state in place (as prepare() would). + inp = _make_inputs(lens, seed=10 + step, pool_pages=pool_pages) + seq_lens.copy_(inp["seq_lens_dev"]) + kv_page_indptr.copy_(inp["kv_page_indptr"]) + n = inp["kv_indices"].shape[0] + kv_indices_buf.zero_() + kv_indices_buf[:n] = inp["kv_indices"] + q.copy_(inp["q"]) + idx_q.copy_(inp["idx_q"]) + k_paged.copy_(inp["k_paged"]) + v_paged.copy_(inp["v_paged"]) + idx_k_paged.copy_(inp["idx_k_paged"]) + + graph.replay() + torch.cuda.synchronize() + replay_out = out_view.clone() + + eager_out = run_pipeline() + torch.cuda.synchronize() + + assert torch.equal(replay_out, eager_out), ( + f"replay step {step} (lens={lens[:4]}...) diverges from eager: " + f"max abs diff " + f"{(replay_out.float() - eager_out.float()).abs().max().item()}" + ) + + +if __name__ == "__main__": + import sys + + sys.exit(pytest.main([__file__, "-v", "-x"])) From 077ccb187e2977afcebbf26d8c2ae5891f08725e Mon Sep 17 00:00:00 2001 From: WeiHaocheng <20514172+WeiHaocheng@users.noreply.github.com> Date: Mon, 6 Jul 2026 05:43:23 -0700 Subject: [PATCH 4/4] [None][fix] Fix M3 attention semantic issue Signed-off-by: WeiHaocheng <20514172+WeiHaocheng@users.noreply.github.com> --- .../attention_backend/fmha/msa_proxy_mqa.py | 15 ++++-- .../attention_backend/fmha/msa_sparse_gqa.py | 10 ++-- .../sparse/minimax_m3/metadata.py | 29 ++++++++++- .../_torch/models/modeling_minimaxm3.py | 50 +++++++++++++++++-- tensorrt_llm/llmapi/llm_args.py | 6 +++ 5 files changed, 99 insertions(+), 11 deletions(-) diff --git a/tensorrt_llm/_torch/attention_backend/fmha/msa_proxy_mqa.py b/tensorrt_llm/_torch/attention_backend/fmha/msa_proxy_mqa.py index 70c3d6162d7b..86370aa0ceb8 100644 --- a/tensorrt_llm/_torch/attention_backend/fmha/msa_proxy_mqa.py +++ b/tensorrt_llm/_torch/attention_backend/fmha/msa_proxy_mqa.py @@ -28,6 +28,7 @@ from __future__ import annotations +import importlib.util from typing import Optional import torch @@ -56,9 +57,17 @@ class MsaProxyMqaFmha(IndexerProxyFmha): @classmethod def is_available(cls, attn=None) -> bool: - try: - import fmha_sm100 # noqa: F401 - except ImportError: + # Probe with find_spec instead of importing. fmha_sm100's import + # has module-level side effects: it imports tvm_ffi — the FFI + # runtime flashinfer owns in this process — and registers global + # functions into it. is_available() runs at attention-layer + # construction for every layer, and pulling tvm_ffi up that + # early (outside flashinfer's own initialization order) + # intermittently corrupts the flashinfer dense-attention path + # (~1/3 of processes emit garbage logits mid-decode). The real + # import is deferred to first kernel use, by which point + # flashinfer has initialized tvm_ffi itself. + if importlib.util.find_spec("fmha_sm100") is None: logger.debug("MsaProxyMqaFmha is unavailable: fmha_sm100 package not installed.") return False if not torch.cuda.is_available(): diff --git a/tensorrt_llm/_torch/attention_backend/fmha/msa_sparse_gqa.py b/tensorrt_llm/_torch/attention_backend/fmha/msa_sparse_gqa.py index b4ab275536f7..0d8f4af4f394 100644 --- a/tensorrt_llm/_torch/attention_backend/fmha/msa_sparse_gqa.py +++ b/tensorrt_llm/_torch/attention_backend/fmha/msa_sparse_gqa.py @@ -30,6 +30,7 @@ from __future__ import annotations +import importlib.util from typing import Optional import torch @@ -60,9 +61,12 @@ class MsaSparseGqaFmha(BlockSparseFmha): @classmethod def is_available(cls, attn=None) -> bool: - try: - import fmha_sm100 # noqa: F401 - except ImportError: + # Probe with find_spec instead of importing — fmha_sm100's import + # side effects (early tvm_ffi import + global-func registration) + # intermittently corrupt the flashinfer dense-attention path. + # See MsaProxyMqaFmha.is_available for the full story; the real + # import happens at first kernel use in forward_block_sparse. + if importlib.util.find_spec("fmha_sm100") is None: logger.debug("MsaSparseGqaFmha is unavailable: fmha_sm100 package not installed.") return False if not torch.cuda.is_available(): diff --git a/tensorrt_llm/_torch/attention_backend/sparse/minimax_m3/metadata.py b/tensorrt_llm/_torch/attention_backend/sparse/minimax_m3/metadata.py index 6311fceff978..592260391740 100644 --- a/tensorrt_llm/_torch/attention_backend/sparse/minimax_m3/metadata.py +++ b/tensorrt_llm/_torch/attention_backend/sparse/minimax_m3/metadata.py @@ -40,6 +40,14 @@ class MiniMaxM3SparseParams(SparseParams): algorithm: Literal["minimax_m3"] = field(init=False, default="minimax_m3") num_index_heads: int = 4 + # Global (pre-TP-shard) KV head count of the model. Needed to + # localize ``num_index_heads`` per rank: index head ``i`` pairs 1:1 + # with KV head ``i`` (SGLang/HF reference semantics), so a rank + # holding KV heads ``[s, e)`` must score with index heads + # ``[s*g, e*g)`` only, where ``g = num_index_heads_global // + # num_kv_heads_global``. ``None`` means "assume the per-rank KV + # head count is the global one" (single-GPU / tests). + num_kv_heads_global: Optional[int] = None sparse_index_dim: int = 128 block_size: int = 128 topk: int = 16 @@ -120,12 +128,31 @@ def from_sparse_params( ) -> "MiniMaxM3SparseConfig": """Build a kernel param bundle from lowered ``MiniMaxM3SparseParams`` and the per-rank model geometry. + + ``sparse_params.num_index_heads`` is the *global* index-head + count; the per-rank count is derived here from the per-rank KV + head count so that index head ``i`` stays paired 1:1 with KV + head ``i`` under TP (the model layer slices ``idx_q`` with the + matching offsets — see ``modeling_minimaxm3.py``). Selecting + blocks from all index heads' scores (the pre-fix behaviour) + gave every KV head the union/max over all index heads instead + of its own head's top-k. """ + num_kv_heads_global = int(sparse_params.num_kv_heads_global or num_kv_heads) + if int(sparse_params.num_index_heads) % num_kv_heads_global != 0: + raise ValueError( + f"num_index_heads ({sparse_params.num_index_heads}) must be divisible " + f"by the global num_kv_heads ({num_kv_heads_global})" + ) + index_group = int(sparse_params.num_index_heads) // num_kv_heads_global + # min() covers tp_size > num_kv_heads_global (KV heads duplicated + # across ranks: one KV head and its paired index head per rank). + num_index_heads_local = index_group * min(int(num_kv_heads), num_kv_heads_global) return cls( num_q_heads=int(num_q_heads), num_kv_heads=int(num_kv_heads), head_dim=int(head_dim), - num_index_heads=int(sparse_params.num_index_heads), + num_index_heads=num_index_heads_local, sparse_index_dim=int(sparse_params.sparse_index_dim), block_size=int(sparse_params.block_size), topk=int(sparse_params.topk), diff --git a/tensorrt_llm/_torch/models/modeling_minimaxm3.py b/tensorrt_llm/_torch/models/modeling_minimaxm3.py index 54db2fe7ac83..473f5744b696 100644 --- a/tensorrt_llm/_torch/models/modeling_minimaxm3.py +++ b/tensorrt_llm/_torch/models/modeling_minimaxm3.py @@ -629,10 +629,41 @@ def __init__( self.sparse_local_block = int(sparse_cfg.get("sparse_local_block", 1)) self.sparse_score_type = str(sparse_cfg.get("sparse_score_type", "max")) - # index_q_proj is **replicated** across TP ranks. The sparse - # forward reshapes idx_q to - # ``[num_tokens, sparse_num_index_heads, sparse_index_dim]``, - # which requires the rank-local idx_q to carry all heads. + # Index head i pairs 1:1 with KV head i in the reference + # model (SGLang/HF): KV head i's GQA group attends only to + # the blocks selected by index head i. Under TP the KV + # heads are sharded, so each rank must score with the + # index heads paired to its local KV heads — the sparse + # forward slices ``idx_q`` down to + # ``[num_tokens, num_local_index_heads, sparse_index_dim]`` + # with the offsets computed here. (Scoring with all index + # heads and reducing — the pre-fix behaviour — gave every + # KV head the union/max over all index heads' selections.) + num_kv_heads_global = int(config.num_key_value_heads) + if self.sparse_num_index_heads % num_kv_heads_global != 0: + raise ValueError( + f"sparse_num_index_heads ({self.sparse_num_index_heads}) must be " + f"divisible by num_key_value_heads ({num_kv_heads_global})" + ) + index_group = self.sparse_num_index_heads // num_kv_heads_global + if self.tp_size <= num_kv_heads_global: + # Contiguous KV-head shard: rank r holds KV heads + # [r * local, (r + 1) * local), mirroring the qkv_proj + # column split. + kv_head_start = self.tp_rank * self.num_key_value_heads + else: + # KV heads duplicated across ranks (tp > global KV): + # rank r holds the single KV head r * kv // tp. + kv_head_start = (self.tp_rank * num_kv_heads_global) // self.tp_size + self.index_head_start = kv_head_start * index_group + self.num_local_index_heads = ( + min(self.num_key_value_heads, num_kv_heads_global) * index_group + ) + + # index_q_proj stays **replicated** across TP ranks (the + # 4-head GEMM is negligible); the forward slices the local + # heads from its output, keeping checkpoint loading + # unchanged. index_q_total = self.sparse_num_index_heads * self.sparse_index_dim self.index_q_proj = Linear( config.hidden_size, @@ -1062,6 +1093,17 @@ def _sparse_forward( q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1) idx_q = self.index_q_proj(hidden_states) idx_k = self.index_k_proj(hidden_states) + if self.num_local_index_heads != self.sparse_num_index_heads: + # Keep only the index heads paired with this rank's KV + # heads (1:1 pairing — see __init__). Slicing before the + # per-head norm and RoPE is equivalent and cheaper. + idx_q = ( + idx_q.view(-1, self.sparse_num_index_heads, self.sparse_index_dim)[ + :, self.index_head_start : self.index_head_start + self.num_local_index_heads + ] + .reshape(idx_q.shape[0], self.num_local_index_heads * self.sparse_index_dim) + .contiguous() + ) # 2. Per-head Gemma RMSNorm on both branches. q, k = self.apply_qk_norm(q, k) diff --git a/tensorrt_llm/llmapi/llm_args.py b/tensorrt_llm/llmapi/llm_args.py index 10d6155f4fcc..6ebb061c6515 100644 --- a/tensorrt_llm/llmapi/llm_args.py +++ b/tensorrt_llm/llmapi/llm_args.py @@ -663,8 +663,14 @@ def to_sparse_params(self, **kwargs): from tensorrt_llm._torch.attention_backend.sparse.minimax_m3.metadata import \ MiniMaxM3SparseParams + # Global KV head count, used by the backend to localize the + # index heads per TP rank (index head i pairs 1:1 with KV head + # i in the reference model). + num_kv_heads_global = getattr(kwargs.get("pretrained_config"), + "num_key_value_heads", None) return MiniMaxM3SparseParams( num_index_heads=self.sparse_num_index_heads, + num_kv_heads_global=num_kv_heads_global, sparse_index_dim=self.sparse_index_dim, block_size=self.sparse_block_size, topk=self.sparse_topk_blocks,