diff --git a/.gitmodules b/.gitmodules index e69de29bb2d1..627760b34da2 100644 --- a/.gitmodules +++ b/.gitmodules @@ -0,0 +1,3 @@ +[submodule "3rdparty/MSA"] + path = 3rdparty/MSA + url = https://gitlab.com/nvidia/tensorrt-llm/oss-components/msa.git diff --git a/3rdparty/MSA b/3rdparty/MSA new file mode 160000 index 000000000000..e2ebe7656649 --- /dev/null +++ b/3rdparty/MSA @@ -0,0 +1 @@ +Subproject commit e2ebe7656649f619af0ad1d457b534283034655e diff --git a/LICENSE b/LICENSE index 8ba867f30567..ed5ca3b261ba 100644 --- a/LICENSE +++ b/LICENSE @@ -19,6 +19,13 @@ Original Source: https://github.com/Dao-AILab/causal-conv1d Copyright (c) 2024, Tri Dao. Licensed under the BSD 3-Clause License +-------------------------------------------------------------------------------- +CUTLASS +-------------------------------------------------------------------------------- +Original Source: https://github.com/NVIDIA/cutlass +Copyright (c) 2017 - 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +Licensed under the BSD 3-Clause License + -------------------------------------------------------------------------------- flash-attention -------------------------------------------------------------------------------- @@ -34,6 +41,14 @@ Original Source: https://github.com/fla-org/flash-linear-attention Copyright (c) 2023-2025 Songlin Yang Licensed under the MIT License +-------------------------------------------------------------------------------- +FlashInfer +-------------------------------------------------------------------------------- +Original Source: https://github.com/flashinfer-ai/flashinfer +Copyright 2025-2026 NVIDIA +Copyright 2023-2026 FlashInfer community (https://flashinfer.ai/) +Licensed under the Apache License 2.0 + -------------------------------------------------------------------------------- InstructEval -------------------------------------------------------------------------------- @@ -59,6 +74,14 @@ Original Source: https://github.com/state-spaces/mamba Copyright 2023 Tri Dao, Albert Gu Licensed under the Apache License 2.0 +-------------------------------------------------------------------------------- +MSA (MiniMax Sparse Attention) +-------------------------------------------------------------------------------- +Original Source: https://github.com/MiniMax-AI/MSA +Copyright (c) 2026 MiniMax +Licensed under the MIT License + + -------------------------------------------------------------------------------- Quack -------------------------------------------------------------------------------- diff --git a/jenkins/Build.groovy b/jenkins/Build.groovy index da963453cd52..b1d541c04b13 100644 --- a/jenkins/Build.groovy +++ b/jenkins/Build.groovy @@ -384,7 +384,7 @@ def runLLMBuild(pipeline, buildFlags, tarName, is_linux_x86_64) sh "ccache -sv" sh "rm -rf **/*.xml *.tar.gz" - trtllm_utils.checkoutSource(LLM_REPO, env.gitlabCommit, LLM_ROOT, false, true) + trtllm_utils.checkoutSource(LLM_REPO, env.gitlabCommit, LLM_ROOT, true, true) if (env.alternativeTRT) { sh "cd ${LLM_ROOT} && sed -i 's#tensorrt~=.*\$#tensorrt#g' requirements.txt && cat requirements.txt" } @@ -459,7 +459,7 @@ def buildWheelInContainer(pipeline, libraries=[], triple=X86_64_TRIPLE, clean=fa sh "cat ${CCACHE_DIR}/ccache.conf" // Step 1: cloning tekit source code - trtllm_utils.checkoutSource(LLM_REPO, env.gitlabCommit, LLM_ROOT, false, true) + trtllm_utils.checkoutSource(LLM_REPO, env.gitlabCommit, LLM_ROOT, true, true) if (env.alternativeTRT) { trtllm_utils.replaceWithAlternativeTRT(env.alternativeTRT, cpver) sh "cd ${LLM_ROOT} && sed -i 's#tensorrt~=.*\$#tensorrt#g' requirements.txt && cat requirements.txt" diff --git a/jenkins/BuildDockerImage.groovy b/jenkins/BuildDockerImage.groovy index ce94eb1ece1c..d22c9dc23c69 100644 --- a/jenkins/BuildDockerImage.groovy +++ b/jenkins/BuildDockerImage.groovy @@ -300,7 +300,7 @@ def buildImage(config, imageKeyToTag) stage (config.stageName) { // Step 1: Clone TRT-LLM source codes // If using a forked repo, svc_tensorrt needs to have the access to the forked repo. - trtllm_utils.checkoutSource(LLM_REPO, LLM_COMMIT_OR_BRANCH, LLM_ROOT, false, true) + trtllm_utils.checkoutSource(LLM_REPO, LLM_COMMIT_OR_BRANCH, LLM_ROOT, true, true) } // Step 2: Build the images diff --git a/jenkins/L0_MergeRequest.groovy b/jenkins/L0_MergeRequest.groovy index 89aead5f4fc0..3f276522ee54 100644 --- a/jenkins/L0_MergeRequest.groovy +++ b/jenkins/L0_MergeRequest.groovy @@ -325,10 +325,10 @@ def setupPipelineEnvironment(pipeline, testFilter, globalVars) // NB: getContainerURIs reads files in ${LLM_ROOT}/jenkins/ if (env.gitlabMergeRequestLastCommit) { env.gitlabCommit = env.gitlabMergeRequestLastCommit - trtllm_utils.checkoutSource(LLM_REPO, env.gitlabCommit, LLM_ROOT, false, true) + trtllm_utils.checkoutSource(LLM_REPO, env.gitlabCommit, LLM_ROOT, true, true) } else { branch = env.gitlabBranch ? env.gitlabBranch : "main" - trtllm_utils.checkoutSource(LLM_REPO, branch, LLM_ROOT, false, true) + trtllm_utils.checkoutSource(LLM_REPO, branch, LLM_ROOT, true, true) checkoutCommit = sh (script: "cd ${LLM_ROOT} && git rev-parse HEAD",returnStdout: true).trim() env.gitlabCommit = checkoutCommit } @@ -451,7 +451,7 @@ def launchReleaseCheck(pipeline, globalVars) sh "pip3 config set global.break-system-packages true" sh "git config --global --add safe.directory \"*\"" // Step 1: Clone TRT-LLM source codes - trtllm_utils.checkoutSource(LLM_REPO, env.gitlabCommit, LLM_ROOT, false, true) + trtllm_utils.checkoutSource(LLM_REPO, env.gitlabCommit, LLM_ROOT, true, true) sh "cd ${LLM_ROOT} && git config --unset-all core.hooksPath" // Step 2: Run guardwords scan @@ -1213,7 +1213,7 @@ def collectTestResults(pipeline, testFilter, globalVars) echo "Result File Number: ${resultFileNumber}, Downloaded: ${resultFileDownloadedNumber}" sh "find . -name results-\\*.tar.gz -type f -exec tar -zxvf {} \\; || true" - trtllm_utils.checkoutSource(LLM_REPO, env.gitlabCommit, LLM_ROOT, false, true) + trtllm_utils.checkoutSource(LLM_REPO, env.gitlabCommit, LLM_ROOT, true, true) junit(testResults: '**/results*.xml', allowEmptyResults : true) diff --git a/jenkins/L0_Test.groovy b/jenkins/L0_Test.groovy index e131587db787..7b2bca500f54 100644 --- a/jenkins/L0_Test.groovy +++ b/jenkins/L0_Test.groovy @@ -2889,7 +2889,7 @@ def runLLMDocBuild(pipeline, config) sh "pwd && ls -alh" sh "env | sort" // allow to checkout from forked repo, svc_tensorrt needs to have access to the repo, otherwise clone will fail - trtllm_utils.checkoutSource(LLM_REPO, env.gitlabCommit, LLM_ROOT, false, true) + trtllm_utils.checkoutSource(LLM_REPO, env.gitlabCommit, LLM_ROOT, true, true) sh "mkdir TensorRT-LLM" sh "cp -r ${LLM_ROOT}/ TensorRT-LLM/src/" trtllm_utils.llmExecStepWithRetry(pipeline, script: "git config --global --add safe.directory \"*\"") @@ -4316,7 +4316,7 @@ def runLLMBuild( sh "env | sort" sh "ccache -sv" - trtllm_utils.checkoutSource(LLM_REPO, env.gitlabCommit, "tensorrt_llm", false, true) + trtllm_utils.checkoutSource(LLM_REPO, env.gitlabCommit, "tensorrt_llm", true, true) if (env.alternativeTRT) { sh "cd tensorrt_llm/ && sed -i 's#tensorrt~=.*\$#tensorrt#g' requirements.txt && cat requirements.txt" } @@ -5442,7 +5442,7 @@ def launchTestJobs(pipeline, testFilter) trtllm_utils.llmExecStepWithRetry(pipeline, script: 'rm -rf $(python3 -c "import site; print(site.getsitepackages()[0])")/nvidia_cutlass_dsl*') } trtllm_utils.llmExecStepWithRetry(pipeline, script: "apt-get update && apt-get install -y python3-pip git rsync curl wget") - trtllm_utils.checkoutSource(LLM_REPO, env.gitlabCommit, LLM_ROOT, false, true) + trtllm_utils.checkoutSource(LLM_REPO, env.gitlabCommit, LLM_ROOT, true, true) trtllm_utils.llmExecStepWithRetry(pipeline, script: "pip3 config set global.break-system-packages true") trtllm_utils.llmExecStepWithRetry(pipeline, script: "pip3 install requests") trtllm_utils.llmExecStepWithRetry(pipeline, script: "pip3 uninstall -y tensorrt") diff --git a/jenkins/TensorRT_LLM_PLC.groovy b/jenkins/TensorRT_LLM_PLC.groovy index dc3c1735e2d1..651e431ea96f 100644 --- a/jenkins/TensorRT_LLM_PLC.groovy +++ b/jenkins/TensorRT_LLM_PLC.groovy @@ -140,7 +140,7 @@ def checkoutSource () def LLM_REPO = getLLMRepo() sh "git config --global --add safe.directory ${env.WORKSPACE}" def ref = params.ref - trtllm_utils.checkoutSource(LLM_REPO, ref, env.WORKSPACE, false, true) + trtllm_utils.checkoutSource(LLM_REPO, ref, env.WORKSPACE, true, true) } def getPulseToken(serviceId, scopes) { diff --git a/jenkins/runPerfSanityTriage.groovy b/jenkins/runPerfSanityTriage.groovy index 86591da57cc7..bc31ad53e318 100644 --- a/jenkins/runPerfSanityTriage.groovy +++ b/jenkins/runPerfSanityTriage.groovy @@ -89,7 +89,7 @@ pipeline { container("trt-llm") { script { sh "pwd && ls -alh" - trtllm_utils.checkoutSource(LLM_REPO, params.BRANCH, LLM_ROOT, false, false) + trtllm_utils.checkoutSource(LLM_REPO, params.BRANCH, LLM_ROOT, true, false) def commandsBase64 = params.COMMANDS.bytes.encodeBase64().toString() sh """ cd ${LLM_ROOT}/jenkins/scripts/perf && python3 perf_sanity_triage.py \ diff --git a/requirements.txt b/requirements.txt index 5ccea71003f1..663b5829012a 100644 --- a/requirements.txt +++ b/requirements.txt @@ -73,6 +73,8 @@ blobfile openai-harmony==0.0.4 nvidia-cutlass-dsl[cu13]==4.5.0; python_version >= "3.10" nvidia-matmul-heuristics==0.1.0.27; python_version >= "3.10" # analytic GEMM heuristics for CuTe DSL autotuner tactic pruning +quack-kernels>=0.2.10; python_version >= "3.10" # required for MinimaxM3 MSA +jinja2 # required for MinimaxM3 MSA plotly numexpr partial_json_parser diff --git a/scripts/attribution/scan/metadata/msa.yml b/scripts/attribution/scan/metadata/msa.yml new file mode 100644 index 000000000000..ab212fc1b664 --- /dev/null +++ b/scripts/attribution/scan/metadata/msa.yml @@ -0,0 +1,5 @@ +name: msa +description: MiniMax Sparse Attention (fmha_sm100) kernels for SM100 sparse attention +source: submodule +directory_matches: +- 3rdparty/MSA diff --git a/setup.py b/setup.py index 13781a22f2c9..97ded73d1f8a 100644 --- a/setup.py +++ b/setup.py @@ -423,6 +423,9 @@ def extract_from_precompiled(precompiled_location: str, package_data: list[str], # internal absolute imports (e.g., "from triton_kernels.foo import bar") work. packages += find_packages(include=["triton_kernels", "triton_kernels.*"]) +msa_package_dir = {"fmha_sm100": "3rdparty/MSA/python/fmha_sm100"} +packages += ["fmha_sm100"] + # https://setuptools.pypa.io/en/latest/references/keywords.html setup( name='tensorrt_llm', @@ -437,6 +440,7 @@ def extract_from_precompiled(precompiled_location: str, package_data: list[str], url="https://github.com/NVIDIA/TensorRT-LLM", download_url="https://github.com/NVIDIA/TensorRT-LLM/tags", packages=packages, + package_dir=msa_package_dir, exclude_package_data=exclude_package_data, # TODO Add windows support for python bindings. classifiers=[ @@ -449,8 +453,17 @@ def extract_from_precompiled(precompiled_location: str, package_data: list[str], license="Apache License 2.0", keywords="nvidia tensorrt deeplearning inference", package_data={ - 'tensorrt_llm': package_data, + 'tensorrt_llm': + package_data, 'triton_kernels': ['LICENSE', 'VERSION', 'README.md'], + 'fmha_sm100': [ + '*.py', + 'csrc/**/*', + 'cute/**/*', + 'cutlass/include/**/*', + 'cutlass/tools/util/include/**/*', + 'cutlass/LICENSE.txt', + ], }, license_files=get_license(), entry_points={ diff --git a/tensorrt_llm/_torch/attention_backend/fmha/__init__.py b/tensorrt_llm/_torch/attention_backend/fmha/__init__.py index 1c3981abcf91..b2cd1e75ec7e 100644 --- a/tensorrt_llm/_torch/attention_backend/fmha/__init__.py +++ b/tensorrt_llm/_torch/attention_backend/fmha/__init__.py @@ -16,6 +16,7 @@ from .fallback import FallbackFmha from .flashinfer_trtllm_gen import FlashInferTrtllmGenFmha from .interface import Fmha +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 @@ -27,6 +28,7 @@ "Fmha", "FmhaCls", "FmhaParams", + "MsaSparseGqaFmha", "PhasedFmha", "get_enabled_fmha_lib_classes", ] 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..0bbb1ad1509a --- /dev/null +++ b/tensorrt_llm/_torch/attention_backend/fmha/msa_sparse_gqa.py @@ -0,0 +1,237 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Block-sparse GQA FMHA backed by MSA's fmha_sm100 kernel. + +MsaSparseGqaFmha wraps the fmha_sm100 paged sparse GQA kernel and +participates in the standard TrtllmAttention.forward dispatch loop. The +owning MiniMax-M3 MSA attention layer runs an MsaIndexer to select the +per-query KV blocks and publishes them on forward_args.sparse_prediction; +this class attends over them. +""" + +from __future__ import annotations + +from typing import TYPE_CHECKING, Optional + +import torch + +from tensorrt_llm._utils import is_sm_100f + +from .interface import Fmha + +if TYPE_CHECKING: + from tensorrt_llm._torch.attention_backend.interface import AttentionForwardArgs + from tensorrt_llm._torch.attention_backend.trtllm import ( + TrtllmAttention, + TrtllmAttentionMetadata, + ) + + +def run_msa_sparse_gqa( + q: torch.Tensor, + k_paged: torch.Tensor, + v_paged: torch.Tensor, + kv_block_indexes: Optional[torch.Tensor] = None, + *, + kv_indices: torch.Tensor, + sm_scale: float, + qo_lens_cpu: Optional[torch.Tensor] = None, + kv_lens_cpu: Optional[torch.Tensor] = None, + qo_offset_cpu: Optional[torch.Tensor] = None, + causal: bool = True, + head_dim: int = 128, + plan: Optional[tuple] = None, + out: Optional[torch.Tensor] = None, + use_fp8: bool = False, +) -> None: + """Run fmha_sm100 paged GQA (plan/run split). + + `kv_block_indexes`: if set, sparse top-k mode (fixed `kv_block_num=topk`); + if None, dense mode attending all pages in `kv_indices`. + `plan`: prebuilt execution plan; if None, built inline from the CPU length + tensors (eager prefill/tests vs. CUDA-graph decode). + `out`: destination buffer the kernel writes in place. + `use_fp8`: FP8 KV cache. The caller must pass FP8 `q` to match the FP8 paged + K/V, since the kernel variant shares one dtype across q/k/v. Also selects the + FP8 AOT kernels for an inline sparse-prefill plan; no-op for the decode planner. + """ + from tensorrt_llm._torch.attention_backend.sparse.minimax_m3.msa_utils import require_msa_module + + fmha_sm100 = require_msa_module() + + if q.dim() != 3: + raise ValueError( + f"MsaSparseGqaFmha expects q [total_q, num_qo_heads, head_dim]; got {tuple(q.shape)}." + ) + if q.shape[-1] != head_dim: + raise NotImplementedError( + f"MsaSparseGqaFmha supports head_dim={head_dim}; got {q.shape[-1]}." + ) + if k_paged.dim() != 4 or v_paged.dim() != 4: + raise ValueError( + "MsaSparseGqaFmha expects paged KV [num_pages, num_kv_heads, page_size, head_dim]; " + f"got k={tuple(k_paged.shape)}, v={tuple(v_paged.shape)}." + ) + if k_paged.shape != v_paged.shape: + raise ValueError( + f"MsaSparseGqaFmha requires k and v to share shape; " + f"got k={tuple(k_paged.shape)}, v={tuple(v_paged.shape)}." + ) + + if plan is None: + # kv_block_num is planned only for the sparse (block-indexed) path; + # dense paged GQA leaves it unset and attends the full page table. + kv_block_num = int(kv_block_indexes.shape[-1]) if kv_block_indexes is not None else -1 + plan = fmha_sm100.fmha_sm100_plan( + qo_lens_cpu, + kv_lens_cpu, + int(q.shape[1]), # num query heads. + num_kv_heads=int(k_paged.shape[1]), + qo_offset=qo_offset_cpu, + page_size=int(k_paged.shape[2]), + kv_block_num=kv_block_num, + causal=causal, + num_kv_splits=1, + use_fp8_kvcache=use_fp8, + ) + fmha_sm100.fmha_sm100( + q, + k_paged, + v_paged, + plan, + kv_indices=kv_indices, + kv_block_indexes=kv_block_indexes, + out=out, + sm_scale=sm_scale, + output_maxscore=False, + ) + + +def run_msa_paged_gqa( + attn: "TrtllmAttention", + q: torch.Tensor, + k: Optional[torch.Tensor], + v: Optional[torch.Tensor], + metadata: "TrtllmAttentionMetadata", + output: torch.Tensor, + *, + kv_block_indexes: Optional[torch.Tensor], + plan: Optional[tuple], +) -> None: + """Write the new-token main K/V, then run paged GQA into output in place. + + Shared by the sparse layers (kv_block_indexes is the per-query top-k table, + with the sparse plan) and the dense layers (kv_block_indexes None, with the + dense plan, attending the full page table). fmha_sm100 reads the paged cache + directly, so the new-token K/V must be resident before the run. + """ + from tensorrt_llm._torch.attention_backend.sparse.minimax_m3.msa_utils import ( + msa_paged_kv, + write_msa_main_kv, + ) + + layer_idx = attn.layer_idx + head_dim = attn.head_dim + kv_cache_manager = metadata.kv_cache_manager + num_tokens = int(q.shape[0]) + if k is not None and v is not None: + write_msa_main_kv( + kv_cache_manager, layer_idx, metadata.msa_out_cache_loc[:num_tokens], k, v + ) + + q_view = q.view(num_tokens, attn.num_heads, head_dim) + out_view = output.view(num_tokens, attn.num_heads, head_dim) + k_paged, v_paged = msa_paged_kv(kv_cache_manager, layer_idx) + sm_scale = (head_dim**-0.5) / float(attn.q_scaling) + + # The fmha_sm100 variant is chosen from q.dtype and shares one dtype across + # q/k/v, so q must be FP8 to match an FP8 paged K/V. MiniMax-M3 has no + # KV-cache scales, so the scale is 1.0 and this is a plain E4M3 cast. + use_fp8 = k_paged.dtype == torch.float8_e4m3fn + if use_fp8: + q_view = q_view.to(torch.float8_e4m3fn) + + run_msa_sparse_gqa( + q_view, + k_paged, + v_paged, + kv_block_indexes, + kv_indices=metadata.msa_kv_indices, + sm_scale=sm_scale, + qo_lens_cpu=metadata.msa_qo_lens_cpu, + kv_lens_cpu=metadata.msa_kv_lens_cpu, + qo_offset_cpu=metadata.msa_qo_offset_cpu, + causal=True, + head_dim=head_dim, + plan=plan, + out=out_view, + use_fp8=use_fp8, + ) + + +class MsaSparseGqaFmha(Fmha): + """SM100 paged GQA FMHA powered by MSA's fmha_sm100 kernel. + + Handles every MiniMax-M3 MSA layer. Sparse layers pass the indexer's + selected KV block indices on forward_args.sparse_prediction.sparse_attn_indices + and attend those blocks; dense layers leave the indices None and attend the + full page table. + + Inherits Fmha rather than PhasedFmha: fmha_sm100 takes a single plan and + the selected block indices span the whole batch, so it handles a mixed + context and generation batch in one call and there is no + context/generation split from PhasedFmha to reuse. Requires head_dim 128 + and 4-D HND paged K/V. + """ + + @classmethod + def is_available(cls, attn: Optional["TrtllmAttention"] = None) -> bool: + # fmha_sm100 runs only on the SM100 family and ships in the MSA git + # submodule, so it is unavailable off SM100 or without the package. + # Imported lazily because the minimax_m3 package init imports the trtllm + # attention classes, which a module-scope import here would cycle with. + from tensorrt_llm._torch.attention_backend.sparse.minimax_m3.msa_utils import ( + msa_package_available, + ) + + if not is_sm_100f() or not msa_package_available(): + return False + # Only the MiniMax-M3 MSA layer uses this library. Matching the lowered + # sparse algorithm lets the base create_fmha_libs add it to that layer + # alone, so no create_fmha_libs override is needed. + return attn.sparse_params is not None and attn.sparse_params.algorithm == "minimax_m3" + + def forward( + self, + q: torch.Tensor, + k: Optional[torch.Tensor], + v: Optional[torch.Tensor], + metadata: "TrtllmAttentionMetadata", + forward_args: "AttentionForwardArgs", + ) -> None: + output = forward_args.output + if output is None: + raise RuntimeError(f"{type(self).__name__} requires an output buffer.") + + # Sparse layers attend the per-query top-k blocks with the sparse plan; + # dense layers leave the indices None and attend the full page table + # with the dense plan. + kv_block_indexes = forward_args.sparse_prediction.sparse_attn_indices + plan = ( + metadata.msa_decode_gqa_plan + if kv_block_indexes is not None + else metadata.msa_decode_dense_plan + ) + run_msa_paged_gqa( + self.attn, + q, + k, + v, + metadata, + output, + kv_block_indexes=kv_block_indexes, + plan=plan, + ) + + +__all__ = ["MsaSparseGqaFmha"] diff --git a/tensorrt_llm/_torch/attention_backend/fmha/registry.py b/tensorrt_llm/_torch/attention_backend/fmha/registry.py index 97467657e9b5..2b3ef3fc7ff3 100644 --- a/tensorrt_llm/_torch/attention_backend/fmha/registry.py +++ b/tensorrt_llm/_torch/attention_backend/fmha/registry.py @@ -22,10 +22,24 @@ FmhaCls: TypeAlias = type[Fmha] -FMHA_LIBS: dict[str, FmhaCls] = { - "flashinfer_trtllm_gen": FlashInferTrtllmGenFmha, - "fallback": FallbackFmha, -} + +def init_fmha_libs() -> dict[str, "FmhaCls"]: + """Build the ordered FMHA library registry. + + Backend classes are imported inside this factory rather than at module + scope, so backends can import trtllm attention classes at module scope + without an import cycle. + """ + from .msa_sparse_gqa import MsaSparseGqaFmha + + return { + "msa_sparse_gqa": MsaSparseGqaFmha, + "flashinfer_trtllm_gen": FlashInferTrtllmGenFmha, + "fallback": FallbackFmha, + } + + +FMHA_LIBS: dict[str, FmhaCls] = init_fmha_libs() DEFAULT_FMHA_LIBS: tuple[str, ...] = tuple(FMHA_LIBS) @@ -78,4 +92,5 @@ def get_enabled_fmha_lib_classes() -> list[FmhaCls]: "FMHA_LIBS", "FmhaCls", "get_enabled_fmha_lib_classes", + "init_fmha_libs", ] diff --git a/tensorrt_llm/_torch/attention_backend/sparse/__init__.py b/tensorrt_llm/_torch/attention_backend/sparse/__init__.py index 0b2994441bf0..f293f9547506 100644 --- a/tensorrt_llm/_torch/attention_backend/sparse/__init__.py +++ b/tensorrt_llm/_torch/attention_backend/sparse/__init__.py @@ -1,13 +1,3 @@ -# yapf: disable -from .minimax_m3 import (MiniMaxM3SparseAttention, - MiniMaxM3SparseAttentionMetadata, - MiniMaxM3SparseConfig, MiniMaxM3SparseIndexCache, - allocate_minimax_m3_static_buffers, - build_runtime_metadata_from_kv_manager, - get_minimax_m3_attention_backend_cls, - get_minimax_m3_kv_cache_manager_cls, - minimax_m3_sparse_decode, minimax_m3_sparse_prefill) -# yapf: enable from .utils import (get_flashinfer_sparse_attn_attention_backend, get_sparse_attn_kv_cache_manager, get_trtllm_sparse_attn_attention_backend, @@ -18,14 +8,4 @@ "get_vanilla_sparse_attn_attention_backend", "get_trtllm_sparse_attn_attention_backend", "get_flashinfer_sparse_attn_attention_backend", - "MiniMaxM3SparseAttention", - "MiniMaxM3SparseAttentionMetadata", - "MiniMaxM3SparseConfig", - "MiniMaxM3SparseIndexCache", - "allocate_minimax_m3_static_buffers", - "build_runtime_metadata_from_kv_manager", - "get_minimax_m3_attention_backend_cls", - "get_minimax_m3_kv_cache_manager_cls", - "minimax_m3_sparse_decode", - "minimax_m3_sparse_prefill", ] 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..522bf2d5b9ee 100644 --- a/tensorrt_llm/_torch/attention_backend/sparse/minimax_m3/__init__.py +++ b/tensorrt_llm/_torch/attention_backend/sparse/minimax_m3/__init__.py @@ -4,24 +4,37 @@ Layered as: - * :mod:`.kernels` -- OpenAI Triton kernels (per-block max - score, masked softmax for sparse GQA). - * :mod:`.metadata` -- ``MiniMaxM3SparseConfig`` / - ``MiniMaxM3SparseAttentionMetadata`` - dataclasses, CUDA-graph-stable buffer - allocator + builder, and the - :class:`AttentionMetadata` subclass - factory. - * :mod:`.cache_manager` -- standalone side index cache used by tests - and the :class:`KVCacheManagerV2` - subclass factory. - * :mod:`.backend` -- the algorithm itself (vectorized - paged-cache helpers, prefill / decode - entry points, the thin - :class:`MiniMaxM3SparseAttention` - orchestrator) and the - :class:`AttentionBackend` subclass - factory. + * :mod:`.triton_metadata` -- ``MiniMaxM3TritonSparseAttentionMetadata`` + dataclass, CUDA-graph-stable buffer + allocator + builder, and the + :class:`AttentionMetadata` subclass + factory for the Triton reference path. + * :mod:`.cache_manager` -- standalone side index cache used by tests + and the :class:`KVCacheManagerV2` + subclass factory. Shared by both backends. + * :mod:`.common` -- backend-neutral config bundles, the paged + KV-slot writer, block-priority sentinels, and + the paged-cache slot mapping builder shared by + both backends. + * :mod:`.msa_utils` -- MSA-only (fmha_sm100) helpers: import guard, + kernel precondition constants, HND paged-cache + adapters, main-KV writer, page-table builder, + valid-block counting, and top-k selection. + * :mod:`.triton_kernels` -- OpenAI Triton kernels (per-block max + score, masked softmax for sparse GQA). + * :mod:`.triton_backend` -- the Triton reference algorithm (vectorized + paged-cache helpers, prefill / decode + entry points, the thin + :class:`MiniMaxM3TritonSparseAttention` + orchestrator) and its + :class:`AttentionBackend` subclass + factory. + * :mod:`.msa_backend` -- the MSA (fmha_sm100) backend, its flat + metadata, and the backend factory. + * :mod:`.msa_indexer` -- the MSA proxy scoring + top-k block + selection submodule. + * :mod:`.msa_availability`-- SM100 and fmha_sm100 gating for the MSA + path. This package's public surface re-exports the names callers historically imported from ``...sparse.minimax_m3`` so external @@ -29,48 +42,20 @@ working unchanged. """ -# Re-export the algorithm-internal helpers focused unit tests reach -# into so the package preserves the surface the monolithic module -# exposed. These are not part of ``__all__`` (still package-private) -# but stay importable as ``from ...minimax_m3 import _write_main_kv_slots``. -from .backend import ( # noqa: F401 - MiniMaxM3SparseAttention, - _compute_index_attn_chunk_q, - _compute_sparse_gqa_chunk_q, +# The dense Triton oracle in the model imports these paged-cache helpers, so +# they stay importable from the package. They are package-private and are not +# part of __all__. Every other backend/metadata/config symbol is imported +# directly from its defining submodule by the code that needs it. +from .cache_manager import MiniMaxM3KVCacheManagerV2 +from .msa_backend import MiniMaxM3MsaSparseAttention +from .triton_backend import ( # noqa: F401 + MiniMaxM3SparseRuntimeBackend, _gather_paged_batched, - _index_attention_and_select, - _write_main_kv_slots, _write_main_kv_slots_to_pool, - get_minimax_m3_attention_backend_cls, - minimax_m3_sparse_decode, - minimax_m3_sparse_prefill, -) -from .cache_manager import ( - MiniMaxM3KVCacheManagerV2, - MiniMaxM3SparseIndexCache, - get_minimax_m3_kv_cache_manager_cls, -) -from .metadata import ( - MiniMaxM3SparseAttentionMetadata, - MiniMaxM3SparseConfig, - allocate_minimax_m3_static_buffers, - build_runtime_metadata_from_kv_manager, - get_minimax_m3_attention_metadata_cls, - replace_metadata, ) __all__ = [ "MiniMaxM3KVCacheManagerV2", - "MiniMaxM3SparseAttention", - "MiniMaxM3SparseAttentionMetadata", - "MiniMaxM3SparseConfig", - "MiniMaxM3SparseIndexCache", - "allocate_minimax_m3_static_buffers", - "build_runtime_metadata_from_kv_manager", - "get_minimax_m3_attention_backend_cls", - "get_minimax_m3_attention_metadata_cls", - "get_minimax_m3_kv_cache_manager_cls", - "minimax_m3_sparse_decode", - "minimax_m3_sparse_prefill", - "replace_metadata", + "MiniMaxM3MsaSparseAttention", + "MiniMaxM3SparseRuntimeBackend", ] 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 a46f812d7a63..cd2b98e6bcb1 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 @@ -13,7 +13,7 @@ from __future__ import annotations -from typing import List, Optional +from typing import List, Optional, Sequence import torch @@ -417,6 +417,8 @@ def _get_batch_cache_indices_by_pool_id( *, pool_id: int = 0, is_kv_aggregate: bool = True, + num_blocks_per_seq: Optional[Sequence[int]] = None, + index_scale: Optional[int] = None, ): """Return per-request slot ids in ``[0, num_slots)`` directly. @@ -427,15 +429,27 @@ def _get_batch_cache_indices_by_pool_id( breaks the V1 conversion and produces out-of-bounds block ids during V2 warmup. + ``index_scale`` (per-layer scale supplied by the base + :meth:`get_batch_cache_indices`) is accepted for + signature compatibility but intentionally ignored: this + override bypasses the scale conversion entirely. + Bypass the conversion: the M3 forward path indexes paged views (built by :meth:`get_buffers` / :meth:`get_index_k_buffer`) directly by slot id. ``BAD_PAGE_INDEX`` slots stay as 0 to match the legacy padding contract. + + ``num_blocks_per_seq``, when provided by the base + :meth:`get_batch_cache_indices`, truncates each request's slot + ids to the blocks it actually owns (matching the base method's + contract); the padded tail is discarded by attention callers. """ res = [] - for req_id in request_ids: + for req_idx, req_id in enumerate(request_ids): idx_tensor = torch.as_tensor(self.kv_cache_map[req_id].get_base_page_indices(pool_id)) + if num_blocks_per_seq is not None: + idx_tensor = idx_tensor[: num_blocks_per_seq[req_idx]] res.append( ( torch.where( diff --git a/tensorrt_llm/_torch/attention_backend/sparse/minimax_m3/common.py b/tensorrt_llm/_torch/attention_backend/sparse/minimax_m3/common.py new file mode 100644 index 000000000000..9f2863846df2 --- /dev/null +++ b/tensorrt_llm/_torch/attention_backend/sparse/minimax_m3/common.py @@ -0,0 +1,252 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Shared building blocks for the MiniMax-M3 sparse attention backends. + +Both the Triton reference and the MSA (fmha_sm100) path share these +backend-neutral pieces: the lowered parameter and per-rank kernel config +bundles, block-priority sentinels, KV-slot writers, and the paged-cache +slot mapping builder. MSA-only helpers live in :mod:`.msa_utils`. +""" + +from __future__ import annotations + +from dataclasses import dataclass, field +from typing import TYPE_CHECKING, List, Literal, Optional, Tuple + +import torch + +from ..params import SparseMetadataParams, SparseParams + +if TYPE_CHECKING: + from tensorrt_llm.mapping import Mapping + +# Sentinel scores that force init and local blocks into the top-k regardless +# of their computed score. Init outranks local. +_INIT_SCORE = 1e30 +_LOCAL_SCORE = 1e29 + + +@dataclass(frozen=True) +class MiniMaxM3SparseParams(SparseParams): + """Lowered runtime parameters for the MiniMax-M3 sparse backend.""" + + algorithm: Literal["minimax_m3"] = field(init=False, default="minimax_m3") + num_index_heads: int = 4 + sparse_index_dim: int = 128 + block_size: int = 128 + topk: int = 16 + init_blocks: int = 0 + local_blocks: int = 1 + score_type: str = "max" + disable_index_value: bool = True + implementation: Literal["triton", "msa"] = "triton" + + @property + def indices_block_size(self) -> int: + """Block granularity of the selected sparse indices. + + Read by the shared TrtllmAttention forward when publishing the + sparse prediction. It equals the per-block scoring size. + """ + return self.block_size + + +@dataclass(frozen=True) +class MiniMaxM3SparseMetadataParams(SparseMetadataParams): + """Metadata-facing MiniMax-M3 sparse geometry.""" + + global_num_q_heads: int = 0 + global_num_kv_heads: int = 0 + num_index_heads: int = 4 + topk: int = 16 + + def sharded_head_counts(self, mapping: Optional["Mapping"] = None) -> Tuple[int, int]: + """Return per-rank (num_q_heads, num_kv_heads) for mapping. + + Matches the model's attention sharding: no split under attention data + parallelism, otherwise split by tp_size. + """ + if mapping is not None and not getattr(mapping, "enable_attention_dp", False): + tp_size = int(getattr(mapping, "tp_size", 1) or 1) + else: + tp_size = 1 + + def _shard(num_heads: int) -> int: + return (int(num_heads) + tp_size - 1) // tp_size + + return _shard(self.global_num_q_heads), _shard(self.global_num_kv_heads) + + +@dataclass(frozen=True) +class MiniMaxM3SparseConfig: + """Per-rank kernel parameter bundle for MiniMax-M3 sparse attention. + + This is **not** a user-facing config (use + :class:`tensorrt_llm.llmapi.llm_args.MiniMaxM3SparseAttentionConfig` + for that). It is the layer-invariant, post-TP-shard parameter bundle + that backend kernels and reference helpers consume. The user knobs + come from :class:`MiniMaxM3SparseParams`; ``num_q_heads`` / + ``num_kv_heads`` / ``head_dim`` come from the per-rank model + geometry and must be supplied by the caller (typically via + :meth:`from_sparse_params`). + """ + + num_q_heads: int + num_kv_heads: int + head_dim: int + num_index_heads: int + sparse_index_dim: int + block_size: int + topk: int + init_blocks: int = 0 + local_blocks: int = 1 + score_type: str = "max" + + def __post_init__(self) -> None: + if self.num_q_heads % self.num_kv_heads != 0: + raise ValueError( + f"num_q_heads ({self.num_q_heads}) must be divisible by " + f"num_kv_heads ({self.num_kv_heads})" + ) + if self.num_index_heads % self.num_kv_heads != 0: + raise ValueError( + f"num_index_heads ({self.num_index_heads}) must be divisible " + f"by num_kv_heads ({self.num_kv_heads})" + ) + if self.block_size <= 0: + raise ValueError(f"block_size must be > 0, got {self.block_size}") + if self.topk <= 0: + raise ValueError(f"topk must be > 0, got {self.topk}") + if self.init_blocks < 0: + raise ValueError(f"init_blocks must be >= 0, got {self.init_blocks}") + if self.local_blocks < 0: + raise ValueError(f"local_blocks must be >= 0, got {self.local_blocks}") + if self.score_type != "max": + # SGLang exposes only "max" today and that is what the MiniMax-M3 + # checkpoint config specifies. Reject anything else explicitly so + # a config drift surfaces immediately. + raise ValueError( + f"score_type={self.score_type!r} is not supported " + "(only 'max' matches the SGLang reference)" + ) + + @classmethod + def from_sparse_params( + cls, + sparse_params: "MiniMaxM3SparseParams", + *, + num_q_heads: int, + num_kv_heads: int, + head_dim: int, + ) -> "MiniMaxM3SparseConfig": + """Build a kernel param bundle from lowered ``MiniMaxM3SparseParams`` + and the per-rank model geometry. + """ + 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), + sparse_index_dim=int(sparse_params.sparse_index_dim), + block_size=int(sparse_params.block_size), + topk=int(sparse_params.topk), + init_blocks=int(sparse_params.init_blocks), + local_blocks=int(sparse_params.local_blocks), + score_type=str(sparse_params.score_type), + ) + + +def write_kv_slots( + cache: torch.Tensor, + out_cache_loc: torch.Tensor, + values: torch.Tensor, + *, + layout: Literal["NHD", "HND"] = "NHD", +) -> None: + """Write per-token values into a K, V, or index-K cache at given slots. + + Handles a 3-D flat-slot cache and a 4-D paged view. `layout` sets the paged + axis order: "NHD" is [num_pages, tokens_per_block, num_heads, channel], + "HND" is [num_pages, num_heads, tokens_per_block, channel]. The paged view + is non-contiguous, so the slot id is split into (page, within) and written + by multi-dim assignment. `values` is always [num_tokens, num_heads, channel]. + """ + with torch.no_grad(): + if cache.ndim >= 4: + token_axis = 2 if layout == "HND" else 1 + tokens_per_block = int(cache.shape[token_axis]) + out_long = out_cache_loc.to(torch.long) + page = out_long // tokens_per_block + within = out_long % tokens_per_block + if layout == "HND": + # Advanced indices on dims 0 and 2 broadcast to [num_tokens] and + # move front, giving a [num_tokens, num_heads, channel] target. + cache[page, :, within, :] = values.to(cache.dtype) + else: + cache[page, within] = values.to(cache.dtype) + else: + cache.index_copy_(0, out_cache_loc.to(torch.long), values.to(cache.dtype)) + + +def build_paged_kv_slot_mapping( + *, + kv_cache_manager, + request_ids, + qo_lens_cpu: torch.Tensor, + qo_offset_cpu: torch.Tensor, + device: torch.device, +) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """Build the backend-neutral paged-cache slot mapping. + + Returns (req_to_token, slot_ids, out_cache_loc), derived only from the paged + KV cache manager and the per-request query geometry, with no dependency on + any backend-specific metadata. + + req_to_token is the [batch, max_kv_len] int32 map from (request, position) + to a global slot id, expanded from get_block_ids_per_seq with + tokens_per_block as block_id * tokens_per_block + offset_within_block. + slot_ids is the [batch] identity row index into req_to_token. out_cache_loc + lists the per-new-token slot ids in flattened query order: request b + contributes positions qo_offset[b] through qo_offset[b] + qo_lens[b] - 1. + That one formula covers prefill (qo_offset is the prefix length) and decode + (qo_offset is kv_len - 1 with qo_len 1). + + The req_to_token reads that build out_cache_loc sync the host, so call this + only from prepare(), never from the forward path. + """ + tokens_per_block = int(kv_cache_manager.tokens_per_block) + # block_ids_per_seq is a [batch, max_blocks_per_seq] tensor; row b holds the + # block ids assigned to request_ids[b] in order. + block_ids = kv_cache_manager.get_block_ids_per_seq(list(request_ids)) + batch = int(qo_lens_cpu.shape[0]) + max_blocks = int(block_ids.shape[1]) + max_kv_len = max_blocks * tokens_per_block + + # Expand block ids -> per-token slot ids. + block_ids_dev = block_ids.to(device).to(torch.int64) + within_block = torch.arange(tokens_per_block, device=device, dtype=torch.int64) + # Outer product per batch entry: [batch, max_blocks, tokens_per_block] + slot_grid = block_ids_dev.unsqueeze(-1) * tokens_per_block + within_block + req_to_token = slot_grid.reshape(batch, max_kv_len).to(torch.int32) + slot_ids = torch.arange(batch, device=device, dtype=torch.int32) + + # out_cache_loc: per-new-token slot ids, in flattened query-token order. + req_to_token_cpu = req_to_token.to("cpu") + qo_lens_list = qo_lens_cpu.to(torch.long).tolist() + qo_offset_list = qo_offset_cpu.to(torch.long).tolist() + out_cache_loc_list: List[int] = [] + for b in range(batch): + start = int(qo_offset_list[b]) + for offset in range(int(qo_lens_list[b])): + out_cache_loc_list.append(int(req_to_token_cpu[b, start + offset].item())) + out_cache_loc = torch.tensor(out_cache_loc_list, dtype=torch.int32, device=device) + return req_to_token, slot_ids, out_cache_loc + + +__all__ = [ + "MiniMaxM3SparseConfig", + "MiniMaxM3SparseMetadataParams", + "MiniMaxM3SparseParams", + "build_paged_kv_slot_mapping", + "write_kv_slots", +] diff --git a/tensorrt_llm/_torch/attention_backend/sparse/minimax_m3/msa_availability.py b/tensorrt_llm/_torch/attention_backend/sparse/minimax_m3/msa_availability.py new file mode 100644 index 000000000000..c4751cdec2e4 --- /dev/null +++ b/tensorrt_llm/_torch/attention_backend/sparse/minimax_m3/msa_availability.py @@ -0,0 +1,41 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Availability checks for the MiniMax-M3 MSA sparse attention kernels. + +The MSA kernels are provided by the fmha_sm100 package from the MSA git +submodule at 3rdparty/MSA and run only on the SM100 architecture family +(SM100 and SM103). These helpers gate backend selection so a request for the +MSA path fails early with a clear message on unsupported systems. +""" + +from __future__ import annotations + +from tensorrt_llm._utils import get_sm_version, is_sm_100f + +from .msa_utils import msa_package_available + +# fmha_sm100 runs on the SM100 architecture family (SM100 and SM103). Other +# architectures, including SM120, are not supported. +MSA_PACKAGE = "fmha_sm100" + + +def ensure_msa_available() -> None: + """Raise RuntimeError if the MSA sparse attention path cannot run here.""" + if not msa_package_available(): + raise RuntimeError( + f"MiniMax-M3 MSA sparse attention requires the {MSA_PACKAGE} kernels " + "from the MSA git submodule at 3rdparty/MSA. Initialize it with " + "'git submodule update --init --recursive'." + ) + if not is_sm_100f(): + sm_version = get_sm_version() + raise RuntimeError( + "MiniMax-M3 MSA sparse attention requires an SM100 or SM103 device, " + f"but the current device reports SM version {sm_version}." + ) + + +__all__ = [ + "MSA_PACKAGE", + "ensure_msa_available", +] 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..86cd67fa7272 --- /dev/null +++ b/tensorrt_llm/_torch/attention_backend/sparse/minimax_m3/msa_backend.py @@ -0,0 +1,789 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""MSA-backed MiniMax-M3 sparse attention on the TrtllmAttention stack. + + * MiniMaxM3MsaSparseAttention subclasses TrtllmAttention and reuses its + inherited forward, overriding only the sparse hooks and owning an + MsaIndexer. + * The main sparse GQA runs through the registered MsaSparseGqaFmha. + * The indexer calls fmha_sm100 directly to produce the per-query selected + block indices, which the model layer threads through + forward_args.topk_indices. + * MiniMaxM3MsaSparseAttentionMetadata subclasses TrtllmAttentionMetadata and + stores its per-forward MSA tensors in CUDA-graph-stable buffers. + The buffers are allocated once in __post_init__ via + get_empty(capture_graph=...), and prepare() copies the per-step values + into them. The standard CUDAGraphRunner clones one metadata per graph + batch size (create_cuda_graph_metadata), so no per-batch-size cache is + needed here. + +The classes subclass TrtllmAttention and TrtllmAttentionMetadata, imported at +module scope. This is cycle-free because the fmha registry defers its +MsaSparseGqaFmha import (see fmha/registry.py), so trtllm's import chain does +not reach this module. +""" + +from __future__ import annotations + +from dataclasses import dataclass +from typing import Optional, Tuple + +import torch + +from tensorrt_llm._torch.attention_backend.interface import AttentionForwardArgs +from tensorrt_llm._torch.attention_backend.trtllm import TrtllmAttention, TrtllmAttentionMetadata + +from .common import ( + MiniMaxM3SparseConfig, + MiniMaxM3SparseMetadataParams, + build_paged_kv_slot_mapping, + write_kv_slots, +) +from .msa_indexer import MsaIndexer +from .msa_utils import ( + MSA_REQUIRED_HEAD_DIM, + MSA_REQUIRED_TOPK, + build_kv_page_indices, + per_token_valid_blocks, + require_msa_module, +) + + +def _cache_device(meta) -> torch.device: + """Device hosting the paged KV buffers, else the current CUDA device.""" + kv_cache_manager = meta.kv_cache_manager + if kv_cache_manager is not None: + try: + return kv_cache_manager.get_buffers(0).device + except Exception: + pass + return torch.device(f"cuda:{torch.cuda.current_device()}") + + +def _worst_case_proxy_max_k_tiles( + fmha_sm100, + *, + num_index_heads: int, + kv_cache_manager, + max_batch: int, +) -> int: + """Return max_k_tiles for a proxy plan at the manager's max KV length.""" + page_size = int(kv_cache_manager.tokens_per_block) + max_kv_len = int(kv_cache_manager.max_blocks_per_seq) * page_size + qo_lens = torch.ones(max_batch, dtype=torch.int32) + kv_lens = torch.full((max_batch,), max_kv_len, dtype=torch.int32) + qo_offset = kv_lens - qo_lens + proxy_plan = fmha_sm100.fmha_sm100_plan( + qo_lens, + kv_lens, + num_index_heads, + num_kv_heads=1, + qo_offset=qo_offset, + page_size=page_size, + output_maxscore=True, + num_kv_splits=1, + causal=True, + ) + return int(proxy_plan[3]["max_k_tiles"]) + + +# Per-step fmha_sm100 plan tensors that must live in CUDA-graph-stable buffers. +# At num_kv_splits=1 the plan carries no split-KV workspaces, and +# cute_workspace_buffer is the vendor's cached scratch (kept by reference, not +# copied). +_MSA_PLAN_STABLE_KEYS = ( + "packed_work_range", + "packed_work_info", + "qo_segment_offsets", + "kv_segment_offsets", + "kv_page_indptr", + "qo_segment_lens", + "kv_segment_lens", + "qo_offset", +) +_MSA_PLAN_INT64_KEYS = ("packed_work_range", "packed_work_info") +# fmha_sm100 sizes packed_work_info at 131072 * max(num_kv_splits, 1); forcing +# num_kv_splits=1 pins this worklist width. +_MSA_PACKED_WORK_INFO_LEN = 131072 +_MSA_SPLIT_KV_KEYS = ( + "kv_tile_begin_indices", + "kv_tile_end_indices", + "kv_split_indices", + "num_kv_splits_per_row", + "workspace_o", + "workspace_lse", +) + + +class _MsaGraphSafePlan: + """CUDA-graph-stable mirror of one fmha_sm100 decode plan. + + Owns fixed device buffers for the per-step plan worklists. refresh() copies + a freshly built plan into them and returns a plan tuple pointing at the + stable buffers, so the captured fmha_sm100 run reads addresses that do not + change across replays. Mirrors FlashInfer's fixed indptr/indices buffers. + + Only valid at num_kv_splits=1: the plan then has no split-KV workspaces + (refresh() asserts this), and cute_workspace_buffer and the scalar fields + pass through unchanged. + """ + + def __init__(self, metadata, name: str, *, max_batch: int, num_ctas: int, capture_graph: bool): + buffers = metadata.cuda_graph_buffers + self._buf = {} + # Set by refresh(), read through the plan property. + self._plan: Optional[tuple] = None + # cute_workspace_buffer must keep a fixed address across steps for the + # captured graph to replay correctly. Pin it on first use and fail if + # it moves. + self._ws_ptr: Optional[int] = None + for key in _MSA_PLAN_STABLE_KEYS: + if key == "packed_work_range": + shape = (num_ctas,) + elif key == "packed_work_info": + shape = (_MSA_PACKED_WORK_INFO_LEN,) + elif key in ("qo_segment_offsets", "kv_segment_offsets", "kv_page_indptr"): + shape = (max_batch + 1,) + else: + shape = (max_batch,) + dtype = torch.int64 if key in _MSA_PLAN_INT64_KEYS else torch.int32 + self._buf[key] = metadata.get_empty( + buffers, + shape, + cache_name=f"{name}_{key}", + dtype=dtype, + capture_graph=capture_graph, + ) + + @property + def plan(self) -> Optional[tuple]: + """The current graph-safe plan tuple, or None if no decode plan is live.""" + return self._plan + + def reset(self) -> None: + """Drop the live plan tuple (e.g. for a prefill/mixed or captured step).""" + self._plan = None + + def refresh(self, plan_tuple) -> tuple: + has_mixed, split, batch, decode, prefill = plan_tuple + if has_mixed: + raise RuntimeError( + "MSA decode expects a single (non-mixed) fmha_sm100 plan; a decode " + "batch must be pure decode." + ) + for key in _MSA_SPLIT_KV_KEYS: + if decode.get(key) is not None: + raise RuntimeError( + f"MSA decode plan used split-KV workspace {key!r}; num_kv_splits=1 " + "is required for graph-safe decode." + ) + ws = decode.get("cute_workspace_buffer") + if ws is not None: + if self._ws_ptr is None: + self._ws_ptr = ws.data_ptr() + elif ws.data_ptr() != self._ws_ptr: + raise RuntimeError( + "cute_workspace_buffer moved across steps; the fmha_sm100 plan " + "is not CUDA-graph safe." + ) + rebuilt = dict(decode) + for key in _MSA_PLAN_STABLE_KEYS: + src = decode.get(key) + if src is None: + continue + n = int(src.shape[0]) + dst = self._buf[key] + if n > dst.shape[0]: + raise ValueError( + f"MSA plan buffer {key} ({dst.shape[0]}) is smaller than the plan tensor ({n})." + ) + dst[:n].copy_(src, non_blocking=True) + rebuilt[key] = dst[:n] + self._plan = (has_mixed, split, batch, rebuilt, prefill) + return self._plan + + +@dataclass(init=False) +class MiniMaxM3MsaSparseAttentionMetadata(TrtllmAttentionMetadata): + """TrtllmAttentionMetadata for MiniMax-M3 MSA sparse layers. + + Tensors read inside the captured forward are CUDA-graph-stable: the + cache slots (msa_out_cache_loc), page table (msa_kv_indices), and proxy + scratch (msa_max_score, msa_n_valid_blocks) are allocated once from the + manager's worst-case geometry. msa_out_cache_loc, msa_kv_indices, and + msa_n_valid_blocks are refreshed via copy_, while the fmha_sm100 proxy pass + writes msa_max_score directly (see msa_proxy_max_score_view). Decode-plan + worklists live on _MsaGraphSafePlan owners, surfaced via msa_decode_*_plan. + + Length inputs to fmha_sm100_plan (msa_qo_lens_cpu, msa_kv_lens_cpu, + msa_qo_offset_cpu) are host properties of the base seq_lens/kv_lens, + read only while building plans in prepare() (outside capture), so they + need no graph-stable storage. Plans are built in _build_decode_plans and + are absent for prefill/mixed batches, which run eagerly. + """ + + # Graph-stable buffers; consumers slice to the live count at the call + # site. Filled once the current step's cache write is prepared. + msa_out_cache_loc: Optional[torch.Tensor] = None + msa_kv_indices: Optional[torch.Tensor] = None + msa_max_score: Optional[torch.Tensor] = None + msa_n_valid_blocks: Optional[torch.Tensor] = None + + # _msa_buffers_ready gates the once-only device buffers; + # _msa_fields_ready marks that the current step's buffers are populated. + _msa_buffers_ready: bool = False + _msa_fields_ready: bool = False + # Sparse geometry the decode plans need. + _msa_params: Optional[MiniMaxM3SparseMetadataParams] = None + # Plan owners, created lazily when the decode plans are first built and + # reused across steps. Each owns its graph-safe plan buffers and the + # current refreshed plan tuple. + _msa_proxy_plan: Optional["_MsaGraphSafePlan"] = None + _msa_gqa_plan: Optional["_MsaGraphSafePlan"] = None + _msa_dense_plan: Optional["_MsaGraphSafePlan"] = None + + def __post_init__(self) -> None: + super().__post_init__() + params = self.sparse_metadata_params + self._msa_params = params if isinstance(params, MiniMaxM3SparseMetadataParams) else None + self._create_msa_buffers() + + @property + def msa_qo_lens_cpu(self) -> Optional[torch.Tensor]: + """Per-request query length (host int32), from the base seq_lens.""" + seq_lens = self.seq_lens + if seq_lens is None: + return None + out = seq_lens[: self.num_seqs] + return out if out.dtype == torch.int32 else out.to(torch.int32) + + @property + def msa_kv_lens_cpu(self) -> Optional[torch.Tensor]: + """Per-request KV length, cached plus new tokens (host int32).""" + kv_lens = getattr(self, "kv_lens", None) + if self.seq_lens is None or kv_lens is None: + return None + out = kv_lens[: self.num_seqs] + return out if out.dtype == torch.int32 else out.to(torch.int32) + + @property + def msa_qo_offset_cpu(self) -> Optional[torch.Tensor]: + """Per-request causal offset (kv_len - qo_len), the cached prefix length.""" + qo = self.msa_qo_lens_cpu + kv = self.msa_kv_lens_cpu + if qo is None or kv is None: + return None + return kv - qo + + @property + def msa_decode_proxy_plan(self) -> Optional[tuple]: + """Proxy (max-score) plan tuple, or None outside decode.""" + plan = self._msa_proxy_plan + return plan.plan if plan is not None else None + + @property + def msa_decode_gqa_plan(self) -> Optional[tuple]: + """Sparse GQA plan tuple, or None outside decode.""" + plan = self._msa_gqa_plan + return plan.plan if plan is not None else None + + @property + def msa_decode_dense_plan(self) -> Optional[tuple]: + """Dense GQA plan tuple, shared by dense layers 0 to 2.""" + plan = self._msa_dense_plan + return plan.plan if plan is not None else None + + def _create_msa_buffers(self) -> None: + """Allocate the CUDA-graph-stable MSA device buffers. + + Buffers come from the shared graph buffer pool so they are reserved + under capture. Sizing follows the worst-case graph geometry: + max_num_tokens for cache slots, max_num_sequences * max_blocks_per_seq + for the page table, and worst-case max_k_tiles for proxy scratch. + """ + kv_cache_manager = self.kv_cache_manager + self._msa_buffers_ready = False + if kv_cache_manager is None or not hasattr(kv_cache_manager, "get_index_k_buffer"): + return + capture_graph = self.is_cuda_graph + buffers = self.cuda_graph_buffers + max_num_sequences = int(self.max_num_sequences) + max_blocks_per_seq = int(kv_cache_manager.max_blocks_per_seq) + max_total_pages = max_num_sequences * max_blocks_per_seq + max_num_tokens = int(self.max_num_tokens) + + self.msa_out_cache_loc = self.get_empty( + buffers, + (max_num_tokens,), + cache_name="msa_out_cache_loc", + dtype=torch.int32, + capture_graph=capture_graph, + ) + self.msa_kv_indices = self.get_empty( + buffers, + (max_total_pages,), + cache_name="msa_kv_indices", + dtype=torch.int32, + capture_graph=capture_graph, + ) + # The proxy scratch needs the fmha_sm100 plan geometry. This metadata + # exists only for the MSA backend, whose selection already required the + # kernels, so a failed import here is a hard error rather than a reason + # to skip allocation. + params = self._msa_params + if params is not None: + fmha_sm100 = require_msa_module() + max_k_tiles = _worst_case_proxy_max_k_tiles( + fmha_sm100, + num_index_heads=params.num_index_heads, + kv_cache_manager=kv_cache_manager, + max_batch=max_num_sequences, + ) + self._alloc_msa_proxy_scratch( + num_index_heads=params.num_index_heads, + max_batch=max_num_sequences, + max_k_tiles=max_k_tiles, + capture_graph=capture_graph, + ) + self._msa_buffers_ready = True + + def _alloc_msa_proxy_scratch( + self, + *, + num_index_heads: int, + max_batch: int, + max_k_tiles: int, + capture_graph: bool, + ) -> None: + """Allocate the flat proxy max-score store and the valid-block scratch. + + The store is sized for the worst-case max_k_tiles so one allocation + serves every decode step. msa_proxy_max_score_view slices the per-step + shape out of it. + """ + buffers = self.cuda_graph_buffers + self.msa_max_score = self.get_empty( + buffers, + (num_index_heads * max_k_tiles * max_batch,), + cache_name="msa_max_score", + dtype=torch.float32, + capture_graph=capture_graph, + ) + self.msa_n_valid_blocks = self.get_empty( + buffers, + (max_batch,), + cache_name="msa_n_valid_blocks", + dtype=torch.int32, + capture_graph=capture_graph, + ) + + def _ensure_msa_decode_scratch_buffers( + self, + *, + num_index_heads: int, + max_batch: int, + capture_graph: bool, + required_max_k_tiles: int, + ) -> None: + """Ensure proxy scratch buffers exist and cover the current plan.""" + required_numel = num_index_heads * required_max_k_tiles * max_batch + if self.msa_max_score is not None: + if self.msa_max_score.numel() < required_numel: + raise ValueError( + f"msa_max_score backing store ({self.msa_max_score.numel()} " + f"elements) is smaller than the decode plan needs " + f"({required_numel} = {num_index_heads} heads * " + f"{required_max_k_tiles} k-tiles * {max_batch} batch)." + ) + return + + kv_cache_manager = self.kv_cache_manager + if kv_cache_manager is None: + return + + fmha_sm100 = require_msa_module() + max_k_tiles = _worst_case_proxy_max_k_tiles( + fmha_sm100, + num_index_heads=num_index_heads, + kv_cache_manager=kv_cache_manager, + max_batch=max_batch, + ) + if max_k_tiles < required_max_k_tiles: + raise ValueError( + f"Worst-case max_k_tiles ({max_k_tiles}) is less than the " + f"decode plan ({required_max_k_tiles})." + ) + self._alloc_msa_proxy_scratch( + num_index_heads=num_index_heads, + max_batch=max_batch, + max_k_tiles=max_k_tiles, + capture_graph=capture_graph, + ) + + def prepare(self) -> None: + super().prepare() + self._build_msa_fields() + self._build_decode_plans() + + def _build_decode_plans(self) -> None: + """Build the graph-safe decode plans and buffers for this step. + + Runs in prepare(), outside CUDA graph capture. The plans are + layer-invariant for MiniMax-M3, so they are built once per step from + the shared sparse geometry, mirrored into CUDA-graph-stable buffers, + and reused by every layer. Mirrors FlashInfer's plan() split. + Prefill/mixed batches leave the plans cleared and run eagerly. + """ + # Drop any plan tuples from the previous step; the msa_decode_*_plan + # properties then report None until they are rebuilt below. + for plan in (self._msa_proxy_plan, self._msa_gqa_plan, self._msa_dense_plan): + if plan is not None: + plan.reset() + if not self._msa_fields_ready: + return + # A decode batch is pure generation (no context requests). + if int(self.num_contexts or 0) > 0: + return + # Geometry is captured in __post_init__; skip when it is unavailable. + params = self._msa_params + if params is None: + return + num_index_heads = params.num_index_heads + num_q_heads, num_kv_heads = params.sharded_head_counts(self.mapping) + topk = params.topk + + fmha_sm100 = require_msa_module() + qo_lens_cpu = self.msa_qo_lens_cpu + kv_lens_cpu = self.msa_kv_lens_cpu + qo_offset_cpu = self.msa_qo_offset_cpu + if qo_lens_cpu is None or kv_lens_cpu is None or qo_offset_cpu is None: + return + batch = int(qo_lens_cpu.shape[0]) + device = _cache_device(self) + page_size = int(self.kv_cache_manager.tokens_per_block) + capture_graph = self.is_cuda_graph + max_batch = int(self.max_num_sequences) + + # Proxy plan: MQA (num_kv_heads=1) max-score pass over the index + # branch; output_maxscore feeds the indexer's top-k block selection. + proxy_plan = fmha_sm100.fmha_sm100_plan( + qo_lens_cpu, + kv_lens_cpu, + num_index_heads, + num_kv_heads=1, + qo_offset=qo_offset_cpu, + page_size=page_size, + output_maxscore=True, + num_kv_splits=1, + causal=True, + ) + # Sparse-layer plan: kv_block_num=topk limits attention to top-k blocks. + gqa_plan = fmha_sm100.fmha_sm100_plan( + qo_lens_cpu, + kv_lens_cpu, + num_q_heads, + num_kv_heads=num_kv_heads, + qo_offset=qo_offset_cpu, + page_size=page_size, + kv_block_num=topk, + num_kv_splits=1, + causal=True, + ) + # Dense-layer plan: no kv_block_num, so it attends the full page table. + dense_plan = fmha_sm100.fmha_sm100_plan( + qo_lens_cpu, + kv_lens_cpu, + num_q_heads, + num_kv_heads=num_kv_heads, + qo_offset=qo_offset_cpu, + page_size=page_size, + num_kv_splits=1, + causal=True, + ) + + required_max_k_tiles = int(proxy_plan[3]["max_k_tiles"]) + self._ensure_msa_decode_scratch_buffers( + num_index_heads=num_index_heads, + max_batch=max_batch, + capture_graph=capture_graph, + required_max_k_tiles=required_max_k_tiles, + ) + + # Allocate the graph-safe plan owners once per metadata; later steps + # only refresh their contents below. + if self._msa_proxy_plan is None: + num_ctas = torch.cuda.get_device_properties(device).multi_processor_count + self._msa_proxy_plan = _MsaGraphSafePlan( + self, + "msa_proxy_plan", + max_batch=max_batch, + num_ctas=num_ctas, + capture_graph=capture_graph, + ) + self._msa_gqa_plan = _MsaGraphSafePlan( + self, + "msa_gqa_plan", + max_batch=max_batch, + num_ctas=num_ctas, + capture_graph=capture_graph, + ) + self._msa_dense_plan = _MsaGraphSafePlan( + self, + "msa_dense_plan", + max_batch=max_batch, + num_ctas=num_ctas, + capture_graph=capture_graph, + ) + + # refresh() stores each plan tuple on its owner, surfaced by the + # msa_decode_*_plan properties. + self._msa_proxy_plan.refresh(proxy_plan) + self._msa_gqa_plan.refresh(gqa_plan) + self._msa_dense_plan.refresh(dense_plan) + + n_valid = per_token_valid_blocks( + qo_lens_cpu, kv_lens_cpu, qo_offset_cpu, causal=True, block_size=page_size + ) + self.msa_n_valid_blocks[:batch].copy_(n_valid.to(torch.int32), non_blocking=True) + + def _build_msa_fields(self) -> None: + """Populate the MSA cache-write buffers for this step. + + The page table and per-new-token cache slots are derived via the + build_paged_kv_slot_mapping helper, then copied into the persistent + buffers. The transient builder tensors are discarded. + """ + self._msa_fields_ready = False + if not self._msa_buffers_ready: + return + request_ids = self.request_ids + qo_lens_cpu = self.msa_qo_lens_cpu + kv_lens_cpu = self.msa_kv_lens_cpu + qo_offset_cpu = self.msa_qo_offset_cpu + if request_ids is None or qo_lens_cpu is None: + return + batch_size = int(qo_lens_cpu.shape[0]) + if batch_size == 0: + return + + kv_cache_manager = self.kv_cache_manager + cache_device = _cache_device(self) + page_size = int(kv_cache_manager.tokens_per_block) + + is_prefill = int(self.num_contexts or 0) > 0 + if not is_prefill and int(qo_lens_cpu.max().item()) > 1: + raise NotImplementedError( + "MiniMax-M3 MSA attention does not support speculative decoding " + "(multiple query tokens per decode step). Disable speculative " + "decoding or use the non-MSA MiniMax-M3 backend." + ) + + # Built in prepare() (outside capture), so these transients are + # fine: forwards read only the persistent buffers filled below. + # qo_offset is the prefix length, so one build covers prefill + # (num_cached) and decode (kv_len - 1 with qo_len 1). + req_to_token, slot_ids, out_cache_loc = build_paged_kv_slot_mapping( + kv_cache_manager=kv_cache_manager, + request_ids=request_ids, + qo_lens_cpu=qo_lens_cpu, + qo_offset_cpu=qo_offset_cpu, + device=cache_device, + ) + kv_indices = build_kv_page_indices(req_to_token, slot_ids, kv_lens_cpu, page_size) + + total_new_tokens = int(out_cache_loc.shape[0]) + total_pages = int(kv_indices.shape[0]) + if total_new_tokens > self.msa_out_cache_loc.shape[0]: + raise ValueError( + f"MSA out_cache_loc buffer ({self.msa_out_cache_loc.shape[0]}) is " + f"smaller than the step's new-token count ({total_new_tokens})." + ) + if total_pages > self.msa_kv_indices.shape[0]: + raise ValueError( + f"MSA kv_indices buffer ({self.msa_kv_indices.shape[0]}) is " + f"smaller than the step's page count ({total_pages})." + ) + + self.msa_out_cache_loc[:total_new_tokens].copy_(out_cache_loc, non_blocking=True) + self.msa_kv_indices[:total_pages].copy_(kv_indices, non_blocking=True) + self._msa_fields_ready = True + + def msa_idx_k_cache(self, layer_idx: int) -> torch.Tensor: + """Paged index-K view for the indexer; HND conversion is done there.""" + return self.kv_cache_manager.get_index_k_buffer(layer_idx) + + def msa_write_idx_k(self, layer_idx: int, idx_k: torch.Tensor) -> None: + """Write the new-token index-K into the side cache at out_cache_loc.""" + cache = self.msa_idx_k_cache(layer_idx) + sparse_index_dim = int(cache.shape[-1]) + num_tokens = int(idx_k.shape[0]) + write_kv_slots( + cache, + self.msa_out_cache_loc[:num_tokens], + idx_k.reshape(num_tokens, 1, sparse_index_dim), + ) + + def msa_proxy_max_score_view( + self, num_index_heads: int, plan_max_k_tiles: int, num_tokens: int + ) -> torch.Tensor: + """Return a contiguous [num_index_heads, plan_max_k_tiles, num_tokens] view. + + fmha_sm100 ignores the passed tensor's strides and writes a contiguous + [num_index_heads, plan_max_k_tiles, total_q] block sized by the current + decode plan, so it must receive a tensor contiguous in exactly that + shape. The view is taken from the flat store's prefix starting at offset + 0, so its data_ptr is stable for CUDA graph replay. Capture builds the + decode plan at the worst-case max_k_tiles, so replays only shrink it. + """ + store = self.msa_max_score + numel = num_index_heads * plan_max_k_tiles * num_tokens + if numel > store.numel(): + raise ValueError( + f"msa_max_score backing store ({store.numel()} elements) is " + f"smaller than the proxy view needs ({numel} = {num_index_heads} " + f"heads * {plan_max_k_tiles} k-tiles * {num_tokens} tokens)." + ) + return store[:numel].view(num_index_heads, plan_max_k_tiles, num_tokens) + + +class MiniMaxM3MsaSparseAttention(TrtllmAttention): + """MSA-backed MiniMax-M3 sparse attention.""" + + Metadata = MiniMaxM3MsaSparseAttentionMetadata + + def __init__( + self, + layer_idx: int, + num_heads: int, + head_dim: int, + num_kv_heads: Optional[int] = None, + quant_config=None, + *, + sparse_params, + **kwargs, + ): + TrtllmAttention.__init__( + self, + layer_idx, + num_heads, + head_dim, + num_kv_heads=num_kv_heads, + quant_config=quant_config, + sparse_params=sparse_params, + **kwargs, + ) + self.m3_config = MiniMaxM3SparseConfig.from_sparse_params( + sparse_params, + num_q_heads=num_heads, + num_kv_heads=num_kv_heads or num_heads, + head_dim=head_dim, + ) + self.disable_index_value = bool(sparse_params.disable_index_value) + self._validate_msa_preconditions() + self.indexer = MsaIndexer(self.m3_config) + + def _validate_msa_preconditions(self) -> None: + config = self.m3_config + if not self.disable_index_value: + raise NotImplementedError( + "MSA backend requires disable_index_value=True; the proxy pass " + "consumes only the max score and has no index-V path." + ) + if config.head_dim != MSA_REQUIRED_HEAD_DIM: + raise NotImplementedError( + f"MSA backend requires head_dim={MSA_REQUIRED_HEAD_DIM}, 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 requires topk={MSA_REQUIRED_TOPK}, got {config.topk}." + ) + + @classmethod + def support_fused_rope(cls) -> bool: + # The MiniMax-M3 model layer applies partial RoPE to the main and + # index branches explicitly. + return False + + def run_indexer( + self, + idx_q: torch.Tensor, + idx_k: torch.Tensor, + metadata, + *, + idx_sm_scale: Optional[float] = None, + ) -> torch.Tensor: + """Write the index-K cache and return the selected block indices. + + The model layer runs this before forward and threads the result through + forward_args.topk_indices. Returns [total_q, num_kv_heads, topk]. + Decode uses the prebuilt graph-safe proxy plan; prefill plans + eagerly. + """ + config = self.m3_config + idx_sm_scale = idx_sm_scale if idx_sm_scale is not None else config.sparse_index_dim**-0.5 + num_tokens = int(idx_q.shape[0]) + idx_q_view = idx_q.view(num_tokens, config.num_index_heads, config.sparse_index_dim) + idx_k_view = idx_k.view(num_tokens, 1, config.sparse_index_dim) + + metadata.msa_write_idx_k(self.layer_idx, idx_k_view) + idx_k_cache = metadata.msa_idx_k_cache(self.layer_idx) + + # One selection path: decode passes the prebuilt graph-safe proxy + # plan plus the proxy scratch shaped to the live query count; prefill + # leaves them None and the proxy plan is built inline. + proxy_plan = metadata.msa_decode_proxy_plan + if proxy_plan is not None: + # proxy_plan is (has_mixed, split, batch, decode_dict, prefill); + # decode_dict carries max_k_tiles for the contiguous score view. + plan_max_k_tiles = int(proxy_plan[3]["max_k_tiles"]) + max_score = metadata.msa_proxy_max_score_view( + config.num_index_heads, plan_max_k_tiles, num_tokens + ) + n_valid_blocks = metadata.msa_n_valid_blocks[:num_tokens] + else: + max_score = None + n_valid_blocks = None + return self.indexer.select_blocks( + idx_q_view, + idx_k_cache, + idx_sm_scale=idx_sm_scale, + kv_indices=metadata.msa_kv_indices, + qo_lens_cpu=metadata.msa_qo_lens_cpu, + kv_lens_cpu=metadata.msa_kv_lens_cpu, + qo_offset_cpu=metadata.msa_qo_offset_cpu, + proxy_plan=proxy_plan, + max_score=max_score, + n_valid_blocks=n_valid_blocks, + ) + + def sparse_attn_predict( + self, + q: torch.Tensor, + k: Optional[torch.Tensor], + metadata, + forward_args: "AttentionForwardArgs", + ) -> Tuple[Optional[torch.Tensor], Optional[torch.Tensor]]: + # The model layer runs run_indexer and passes the selected block + # indices through forward_args.topk_indices. Publish them as the + # sparse attention indices MsaSparseGqaFmha reads. + return forward_args.topk_indices, None + + def sparse_kv_predict( + self, + q: torch.Tensor, + k: Optional[torch.Tensor], + metadata, + forward_args: "AttentionForwardArgs", + ) -> Tuple[Optional[torch.Tensor], Optional[torch.Tensor]]: + return None, None + + +__all__ = [ + "MiniMaxM3MsaSparseAttention", + "MiniMaxM3MsaSparseAttentionMetadata", +] diff --git a/tensorrt_llm/_torch/attention_backend/sparse/minimax_m3/msa_indexer.py b/tensorrt_llm/_torch/attention_backend/sparse/minimax_m3/msa_indexer.py new file mode 100644 index 000000000000..a6a2e10e0dc1 --- /dev/null +++ b/tensorrt_llm/_torch/attention_backend/sparse/minimax_m3/msa_indexer.py @@ -0,0 +1,201 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""MiniMax-M3 MSA sparse-attention indexer. + +Mirrors the DSA indexer pattern: a submodule owned by the sparse backend +that runs the predictor pass and returns the per-query selected KV block +indices the main attention consumes. It calls fmha_sm100 directly in +output_maxscore mode, reduces the per-index-head max score to KV-head +granularity, and selects the top-k blocks per query. + +Results are [total_q, num_kv_heads, topk] int32, ascending with -1 padding. +""" + +from __future__ import annotations + +from typing import TYPE_CHECKING, Optional + +import torch + +from .msa_utils import ( + MSA_REQUIRED_TOPK, + cache_view_to_msa_paged, + per_token_valid_blocks, + require_msa_module, + select_blocks_from_maxscore, +) + +if TYPE_CHECKING: + from .common import MiniMaxM3SparseConfig + + +def _proxy_max_score( + 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 fmha_sm100 MQA proxy pass and return the per-block max score. + + Follows MSA's two-call pattern: fmha_sm100_plan builds the plan with + output_maxscore and num_kv_heads 1, then fmha_sm100 runs with output_o + disabled so only the per-block max score is produced. Returns + [num_index_heads, max_k_tiles, total_q] float32. + """ + fmha_sm100 = require_msa_module() + + if idx_q.dim() != 3: + raise ValueError( + "MsaIndexer expects idx_q [total_q, num_index_heads, head_dim]; " + f"got {tuple(idx_q.shape)}." + ) + if idx_k_paged.dim() != 4 or idx_k_paged.shape[1] != 1: + raise ValueError( + "MsaIndexer expects MQA paged index-K [num_pages, 1, page_size, head_dim]; " + f"got {tuple(idx_k_paged.shape)}." + ) + + 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_kv_heads=1, + qo_offset=qo_offset_cpu, + page_size=page_size, + output_maxscore=True, + causal=causal, + num_kv_splits=1, + ) + _, max_score = fmha_sm100.fmha_sm100( + idx_q, + idx_k_paged, + idx_k_paged, + proxy_plan, + kv_indices=kv_indices, + output_o=False, + output_maxscore=True, + sm_scale=sm_scale, + ) + return max_score + + +def _group_max_reduce( + max_score: torch.Tensor, + config: "MiniMaxM3SparseConfig", +) -> torch.Tensor: + """Reduce per-index-head max score to per-KV-head granularity by amax. + + Index heads are assumed to be grouped contiguously per KV head, so head h + maps to KV group h // group. + """ + group, rem = divmod(config.num_index_heads, config.num_kv_heads) + if rem != 0: + raise ValueError( + "num_index_heads must be divisible by num_kv_heads for group max " + f"reduce; got num_index_heads={config.num_index_heads}, " + f"num_kv_heads={config.num_kv_heads}." + ) + if group > 1: + return max_score.view( + config.num_kv_heads, group, max_score.shape[1], max_score.shape[2] + ).amax(dim=1) + return max_score + + +class MsaIndexer: + """Predictor submodule: proxy MQA scoring and top-k block selection. + + Owned by the MSA attention layer. Stateless in eager mode: it reads the + per-forward page table and lengths from the attention metadata and calls + the kernel directly. + """ + + def __init__(self, config: "MiniMaxM3SparseConfig"): + self.config = config + + def select_blocks( + self, + idx_q: torch.Tensor, + idx_k_cache: torch.Tensor, + *, + idx_sm_scale: float, + kv_indices: torch.Tensor, + qo_lens_cpu: Optional[torch.Tensor] = None, + kv_lens_cpu: Optional[torch.Tensor] = None, + qo_offset_cpu: Optional[torch.Tensor] = None, + proxy_plan: Optional[tuple] = None, + max_score: Optional[torch.Tensor] = None, + n_valid_blocks: Optional[torch.Tensor] = None, + ) -> torch.Tensor: + """Return [total_q, num_kv_heads, topk] selected block indices. + + Plan/run split, mirroring the sparse GQA. When `proxy_plan` is None + (prefill and focused tests) the proxy plan is built inline and the + per-query valid-block count is derived here; when provided (CUDA-graph + decode) the proxy runs from the prebuilt plan into the preallocated + `max_score` buffer with a precomputed `n_valid_blocks`, so there is + no host sync inside the captured region. The same top-k selection serves + both, and generation is the one-query-token-per-request special case. + """ + config = self.config + idx_k_paged = cache_view_to_msa_paged(idx_k_cache) + + if proxy_plan is None: + max_score = _proxy_max_score( + idx_q, + idx_k_paged, + qo_lens_cpu=qo_lens_cpu, + kv_lens_cpu=kv_lens_cpu, + qo_offset_cpu=qo_offset_cpu, + kv_indices=kv_indices, + sm_scale=idx_sm_scale, + causal=True, + ) + else: + fmha_sm100 = require_msa_module() + _, max_score = fmha_sm100.fmha_sm100( + idx_q, + idx_k_paged, + idx_k_paged, + proxy_plan, + kv_indices=kv_indices, + output_o=False, + output_maxscore=True, + max_score=max_score, + sm_scale=idx_sm_scale, + ) + max_score_kv = _group_max_reduce(max_score, config) + + if n_valid_blocks is None: + n_valid_blocks = per_token_valid_blocks( + qo_lens_cpu, + kv_lens_cpu, + qo_offset_cpu, + causal=True, + block_size=int(idx_k_paged.shape[2]), + ) + # The empty-selection guard uses a host sync, so it only runs on the + # eager path; a decode batch always has valid blocks. + if n_valid_blocks.numel() == 0 or int(n_valid_blocks.max().item()) <= 0: + return torch.full( + (idx_q.shape[0], config.num_kv_heads, MSA_REQUIRED_TOPK), + -1, + dtype=torch.int32, + device=idx_q.device, + ) + return select_blocks_from_maxscore( + max_score_kv, + topk=MSA_REQUIRED_TOPK, + n_valid_blocks=n_valid_blocks, + init_blocks=config.init_blocks, + local_blocks=config.local_blocks, + ) + + +__all__ = ["MsaIndexer"] diff --git a/tensorrt_llm/_torch/attention_backend/sparse/minimax_m3/msa_utils.py b/tensorrt_llm/_torch/attention_backend/sparse/minimax_m3/msa_utils.py new file mode 100644 index 000000000000..83380bc730c0 --- /dev/null +++ b/tensorrt_llm/_torch/attention_backend/sparse/minimax_m3/msa_utils.py @@ -0,0 +1,257 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""MSA (fmha_sm100) specific helpers for the MiniMax-M3 sparse backend.""" + +from __future__ import annotations + +import importlib.util +import sys +from pathlib import Path +from typing import Optional, Tuple + +import torch + +from .common import _INIT_SCORE, _LOCAL_SCORE, write_kv_slots + +# fmha_sm100 ships only head_dim 128 variants and the MiniMax-M3 checkpoint +# selects topk 16. Callers enforce these early so a misconfiguration fails +# with a clear message rather than a cryptic shape error inside the kernel. +MSA_REQUIRED_TOPK = 16 +MSA_REQUIRED_HEAD_DIM = 128 + +# Path of the fmha_sm100 package inside the MSA git submodule relative to the +# repository root (see 3rdparty/MSA/LICENSE and 3rdparty/MSA/NOTICE). +_MSA_PYTHON_RELPATH = Path("3rdparty") / "MSA" / "python" + + +def _find_msa_python_dir() -> Optional[Path]: + """Locate the fmha_sm100 package dir by walking up from this file. + + Returns None in installed layouts where the 3rdparty submodule is not + shipped. Walking up avoids hardcoding this module's depth below the + repository root. + """ + for parent in Path(__file__).resolve().parents: + candidate = parent / _MSA_PYTHON_RELPATH + if candidate.is_dir(): + return candidate + return None + + +def _ensure_msa_on_path() -> None: + """Prepend the MSA python package directory to sys.path if present.""" + msa_python = _find_msa_python_dir() + if msa_python is not None and str(msa_python) not in sys.path: + sys.path.insert(0, str(msa_python)) + + +def msa_package_available() -> bool: + """True if fmha_sm100 can be imported (submodule checkout or installed).""" + _ensure_msa_on_path() + return importlib.util.find_spec("fmha_sm100") is not None + + +def require_msa_module(): + """Import fmha_sm100 from the MSA submodule or raise a clear error. + + The import is deferred to first kernel use so the MSA backend can be + advertised in the config schema on systems where the kernels cannot load. + The 3rdparty/MSA/python directory is added to sys.path first, so a source + checkout with the submodule initialized resolves without a separate install. + A missing package is a hard error, never a silent fallback to another backend. + """ + _ensure_msa_on_path() + try: + import fmha_sm100 + except ImportError as exc: + raise RuntimeError( + "MiniMax-M3 MSA attention requires the fmha_sm100 kernels from the " + "MSA git submodule at 3rdparty/MSA. Initialize it with " + "'git submodule update --init --recursive', or install fmha_sm100." + ) from exc + return fmha_sm100 + + +def cache_view_to_msa_paged(cache_view: torch.Tensor) -> torch.Tensor: + """Convert a KV cache view to the fmha_sm100 HND paged layout. + + A 4-D paged view [num_pages, page_size, num_heads, head_dim] permutes to + [num_pages, num_heads, page_size, head_dim]. A 3-D flat-slot cache + [num_slots, num_heads, head_dim] is treated as one virtual page, giving + [1, num_heads, num_slots, head_dim]. + """ + if cache_view.dim() == 4: + return cache_view.permute(0, 2, 1, 3).contiguous() + if cache_view.dim() == 3: + return cache_view.permute(1, 0, 2).unsqueeze(0).contiguous() + raise ValueError( + f"Unsupported cache view rank {cache_view.dim()} for MSA paged conversion; " + "expected 3 (flat-slot) or 4 (paged)." + ) + + +def msa_paged_kv(kv_cache_manager, layer_idx: int) -> Tuple[torch.Tensor, torch.Tensor]: + """Return per-layer paged K and V in fmha_sm100 HND layout, zero-copy. + + The cache is stored head-major (see `write_msa_main_kv`), so the "HND" + buffer view is already the [num_slots, num_kv_heads, page_size, head_dim] + layout fmha_sm100 expects. The kernel reads the page and head strides at + runtime and needs only each page's [page_size, head_dim] block to be + contiguous, which this view satisfies, so no copy is required. + """ + buffers = kv_cache_manager.get_buffers(layer_idx, kv_layout="HND") + return buffers[:, 0], buffers[:, 1] + + +def write_msa_main_kv( + kv_cache_manager, + layer_idx: int, + out_cache_loc: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, +) -> None: + """Write new-token K and V into the paged main cache at out_cache_loc. + + fmha_sm100 reads the paged cache directly, so the new-token K and V must be + resident before the sparse GQA runs. The write uses the head-major HND view + so `msa_paged_kv` can return a zero-copy view. + """ + buffers = kv_cache_manager.get_buffers(layer_idx, kv_layout="HND") + k_view, v_view = buffers[:, 0], buffers[:, 1] + num_kv_heads = int(k_view.shape[1]) + head_dim = int(k_view.shape[3]) + num_tokens = int(k.shape[0]) + write_kv_slots( + k_view, out_cache_loc, k.reshape(num_tokens, num_kv_heads, head_dim), layout="HND" + ) + write_kv_slots( + v_view, out_cache_loc, v.reshape(num_tokens, num_kv_heads, head_dim), layout="HND" + ) + + +def build_kv_page_indices( + req_to_token: torch.Tensor, + slot_ids: torch.Tensor, + kv_lens_cpu: torch.Tensor, + page_size: int, +) -> torch.Tensor: + """Build the flattened per-request page table fmha_sm100 consumes. + + Returns int32 global page ids concatenated per request. A request's + pages come from the first slot of each page in its req_to_token row. + Page ids are global and non-contiguous in production, so they are not + clamped to a per-request bound. + """ + device = req_to_token.device + req_rows = req_to_token.index_select(0, slot_ids.to(torch.long)).to(torch.long) + batch = int(req_rows.shape[0]) + kv_lens_list = kv_lens_cpu.to(torch.long).tolist() + + page_lists = [] + for b in range(batch): + kv_len = int(kv_lens_list[b]) + if kv_len <= 0: + continue + num_pages = (kv_len + page_size - 1) // page_size + page_starts = torch.arange(num_pages, device=device, dtype=torch.long) * page_size + page_ids = req_rows[b].gather(0, page_starts) // page_size + page_lists.append(page_ids.to(torch.int32)) + + if page_lists: + return torch.cat(page_lists, dim=0) + return torch.empty(0, dtype=torch.int32, device=device) + + +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: + """Return the per-query number of valid KV blocks, on CPU. + + Expands per-request lengths and offsets to a per-token vector so block + selection can honour each query token's own causal extent. + """ + 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: + """Select per-query top-k blocks from per-KV-head block scores. + + Applies init and local forced blocks and per-query valid-block masking + on the amax-reduced scores [num_kv_heads, n_blocks, total_q]. Returns + [total_q, num_kv_heads, topk] int32 ascending block ids with -1 tail + padding. + """ + 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() + block_ids = torch.arange(n_blocks, device=device, dtype=torch.long) + nvb = n_valid_blocks.to(device=device, dtype=torch.long) + + 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) + local_mask = (block_ids.view(1, -1) >= local_start.view(-1, 1)) & ( + block_ids.view(1, -1) < nvb.view(-1, 1) + ) + scores = torch.where(local_mask.unsqueeze(1), torch.full_like(scores, _LOCAL_SCORE), scores) + block_valid = block_ids.view(1, -1) < nvb.view(-1, 1) + scores = scores.masked_fill(~block_valid.unsqueeze(1), float("-inf")) + + k = min(topk, n_blocks) + vals, idx = scores.topk(k=k, dim=-1) + 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) + + +__all__ = [ + "MSA_REQUIRED_HEAD_DIM", + "MSA_REQUIRED_TOPK", + "build_kv_page_indices", + "cache_view_to_msa_paged", + "msa_package_available", + "msa_paged_kv", + "per_token_valid_blocks", + "require_msa_module", + "select_blocks_from_maxscore", + "write_msa_main_kv", +] diff --git a/tensorrt_llm/_torch/attention_backend/sparse/minimax_m3/backend.py b/tensorrt_llm/_torch/attention_backend/sparse/minimax_m3/triton_backend.py similarity index 63% rename from tensorrt_llm/_torch/attention_backend/sparse/minimax_m3/backend.py rename to tensorrt_llm/_torch/attention_backend/sparse/minimax_m3/triton_backend.py index 090c6d7956fd..844fd979469c 100644 --- a/tensorrt_llm/_torch/attention_backend/sparse/minimax_m3/backend.py +++ b/tensorrt_llm/_torch/attention_backend/sparse/minimax_m3/triton_backend.py @@ -46,7 +46,7 @@ ----------------- All scalar max lengths (``max_seqlen_q``, ``max_seqlen_k``) are -pre-computed CPU-side in :meth:`MiniMaxM3SparseAttentionMetadata.prepare` +pre-computed CPU-side in :meth:`MiniMaxM3TritonSparseAttentionMetadata.prepare` and stored as plain Python ints. The hot path uses only batched tensor ops with static shapes derived from those CPU-side scalars; no ``.item()`` or other GPU-CPU sync runs inside the forward @@ -56,29 +56,29 @@ from __future__ import annotations -import functools import os from dataclasses import dataclass from typing import TYPE_CHECKING, Optional, Tuple import torch -from .kernels import triton_block_max_score, triton_sparse_softmax -from .metadata import ( - MiniMaxM3SparseAttentionMetadata, - MiniMaxM3SparseConfig, +from ...interface import ( + AttentionBackend, + AttentionForwardArgs, + AttentionMetadata, + merge_attention_forward_args, +) +from .common import _INIT_SCORE, _LOCAL_SCORE, MiniMaxM3SparseConfig, write_kv_slots +from .triton_kernels import triton_block_max_score, triton_sparse_softmax +from .triton_metadata import ( + MiniMaxM3AttentionMetadata, + MiniMaxM3TritonSparseAttentionMetadata, ensure_metadata_on_device, - get_minimax_m3_attention_metadata_cls, ) if TYPE_CHECKING: from .cache_manager import MiniMaxM3SparseIndexCache - from .metadata import MiniMaxM3SparseParams - -# Sentinel block score for blocks that init / local priority forces into -# the top-k regardless of their numerical score. -_INIT_SCORE = 1e30 -_LOCAL_SCORE = 1e29 + from .common import MiniMaxM3SparseParams # --------------------------------------------------------------------------- @@ -143,53 +143,6 @@ def _gather_paged_batched( return flat.view(batch, max_k, *cache.shape[1:]) -def _assert_paged_write_in_bounds( - name: str, - cache: torch.Tensor, - page: torch.Tensor, - within: torch.Tensor, -) -> None: - """Optional CPU-side bounds check for paged-cache writes. - - The runtime computes per-token slot ids from - ``KVCacheManagerV2``'s block ids; if the runtime ever produces a - block id that does not fit the per-layer view's dim-0 the write - falls into another layer's coalesced memory and corrupts the - cache, or fires the CUDA ``IndexKernel.cu`` device-side assert - during fancy indexing. Both are far enough away from the root - cause to be hard to triage. - - When ``TRTLLM_MINIMAX_M3_DEBUG_BOUNDS`` is set this check runs a - CPU-side max/min comparison against the cache's dim-0 and dim-2 - bounds, surfacing the misindex with the exact tensor names and - values instead of a device-side assert spam. It is opt-in - because the comparison forces a CPU sync. - """ - if not os.environ.get("TRTLLM_MINIMAX_M3_DEBUG_BOUNDS"): - return - num_pages = int(cache.shape[0]) - tokens_per_block = int(cache.shape[1]) if cache.ndim == 4 else int(cache.shape[2]) - page_max = int(page.max().item()) if page.numel() else -1 - page_min = int(page.min().item()) if page.numel() else 0 - within_max = int(within.max().item()) if within.numel() else -1 - within_min = int(within.min().item()) if within.numel() else 0 - assert 0 <= page_min and page_max < num_pages, ( - f"{name}: page index out of bounds — page.min={page_min} " - f"page.max={page_max} but cache.shape[0]={num_pages} " - f"(shape={tuple(cache.shape)}). This usually means the " - f"runtime's get_block_ids_per_seq produced a block id wider " - f"than the per-layer paged view's dim-0; check that the M3 " - f"override path returns slot ids in [0, num_slots)." - ) - assert 0 <= within_min and within_max < tokens_per_block, ( - f"{name}: within-page offset out of bounds — within.min=" - f"{within_min} within.max={within_max} but tokens_per_block=" - f"{tokens_per_block} (shape={tuple(cache.shape)}). This " - f"usually means out_cache_loc was computed with a different " - f"tokens_per_block than the cache was allocated with." - ) - - def _write_main_kv_slots_to_pool( pool: torch.Tensor, kv_index: int, @@ -201,40 +154,12 @@ def _write_main_kv_slots_to_pool( ``pool`` is the 5-D main K/V pool returned by :meth:`KVCacheManagerV2.get_buffers` with the NHD layout ``[num_pages, kv_factor, tokens_per_block, num_kv_heads, head_dim]``. - ``values`` has shape ``[num_new_tokens, num_kv_heads, head_dim]`` - and ``out_cache_loc`` is the 1-D ``[num_new_tokens]`` int tensor of - flat slot ids the caller wants to update. - - The write decomposes each flat slot id into - ``(page = s // tokens_per_block, within = s % tokens_per_block)`` - and uses multi-dim fancy-index assignment so the writes propagate - to the underlying pool storage. The previously used pattern - ``pool[:, kv_index].reshape(-1, num_kv_heads, head_dim) - .index_copy_(0, ...)`` instead wrote into a silent copy (see - :func:`_gather_paged_batched`), so the next forward call read - zeros for the prefilled positions. - - The optional CPU-side bounds assertion (enabled when the - ``TRTLLM_MINIMAX_M3_DEBUG_BOUNDS`` env var is set) catches - block_ids overflowing the pool's dim-0 before the device-side - ``IndexKernel.cu`` assert fires deep inside the kernel. The - assertion is a CPU sync, so the env var keeps it opt-in for - production runs that need a clean fast path. + ``values`` has shape ``[num_new_tokens, num_kv_heads, head_dim]`` and + ``out_cache_loc`` is the 1-D ``[num_new_tokens]`` int tensor of flat slot + ids to update. ``pool[:, kv_index]`` is a storage-sharing view, so the + shared :func:`common.write_kv_slots` propagates the write to the pool. """ - tokens_per_block = int(pool.shape[2]) - out_long = out_cache_loc.to(torch.long) - page = out_long // tokens_per_block - within = out_long % tokens_per_block - _assert_paged_write_in_bounds("pool", pool, page, within) - # KV-cache writes never need to participate in autograd. Wrap the - # fancy-index assignment in ``torch.no_grad()`` so callers that - # enter this path with an active grad context (e.g. unit tests - # exercising :class:`MiniMaxM3Attention` without ``inference_mode``) - # do not trip the "leaf Variable that requires grad is being used - # in an in-place operation" autograd guard on the view chain. - with torch.no_grad(): - # Multi-dim fancy assignment writes into the underlying pool buffer. - pool[page, kv_index, within] = values.to(pool.dtype) + write_kv_slots(pool[:, kv_index], out_cache_loc, values) def _write_main_kv_slots( @@ -242,39 +167,13 @@ def _write_main_kv_slots( out_cache_loc: torch.Tensor, values: torch.Tensor, ) -> None: - """Layout-aware writer for K (or V) caches used by the M3 backend. - - Supports two layouts, mirroring :func:`_gather_paged_batched`: - - * **3-D flat-slot** ``[num_slots, num_kv_heads, head_dim]``: used - by focused unit tests that allocate the cache as a contiguous - flat-slot tensor. ``index_copy_(0, ...)`` writes propagate - because the tensor IS the storage. - * **4-D multi-dim paged** ``[num_pages, tokens_per_block, - num_kv_heads, head_dim]``: used when the cache is a view of - ``kv_pool[:, 0]`` / ``kv_pool[:, 1]``. The view is - non-contiguous (its dim-0 stride is 2× the contiguous stride - because dim 1 separates K from V in the pool), so - ``index_copy_(0, ...)`` would silently fork a copy and the - write would be lost. Decompose the flat slot id into - ``(page, within)`` and use multi-dim fancy assignment so the - write propagates through the view to the underlying pool. + """Write per-new-token K or V into a cache view via the shared writer. + + Delegates to :func:`common.write_kv_slots`, which handles both the 3-D + flat-slot layout used by focused unit tests and the 4-D paged view of + ``kv_pool[:, 0]`` / ``kv_pool[:, 1]``. """ - # KV-cache writes never need to participate in autograd. Wrap both - # branches in ``torch.no_grad()`` so callers that enter this path - # with an active grad context (e.g. unit tests exercising - # :class:`MiniMaxM3Attention` without ``inference_mode``) do not - # trip the autograd in-place guard on the cache view chain. - with torch.no_grad(): - if cache.ndim >= 4: - tokens_per_block = int(cache.shape[1]) - out_long = out_cache_loc.to(torch.long) - page = out_long // tokens_per_block - within = out_long % tokens_per_block - _assert_paged_write_in_bounds("cache", cache, page, within) - cache[page, within] = values.to(cache.dtype) - else: - cache.index_copy_(0, out_cache_loc.to(torch.long), values.to(cache.dtype)) + write_kv_slots(cache, out_cache_loc, values) def _scatter_topk_to_block_mask( @@ -699,7 +598,7 @@ def minimax_m3_sparse_decode( v_cache: torch.Tensor, idx_k_cache: torch.Tensor, idx_v_cache: Optional[torch.Tensor], - metadata: MiniMaxM3SparseAttentionMetadata, + metadata: MiniMaxM3TritonSparseAttentionMetadata, config: MiniMaxM3SparseConfig, *, disable_index_value: bool, @@ -794,7 +693,7 @@ def minimax_m3_sparse_prefill( idx_q: torch.Tensor, idx_k_cache: torch.Tensor, idx_v_cache: Optional[torch.Tensor], - metadata: MiniMaxM3SparseAttentionMetadata, + metadata: MiniMaxM3TritonSparseAttentionMetadata, config: MiniMaxM3SparseConfig, *, disable_index_value: bool, @@ -890,18 +789,8 @@ def minimax_m3_sparse_prefill( # --------------------------------------------------------------------------- -# Lazy import alias to avoid a circular import at module load — the side -# cache class lives in ``cache_manager`` which imports from this module -# is fine (no cycle), but keeping the indirection explicit makes the -# dependency direction in this file clearer. -def _import_index_cache_cls(): - from .cache_manager import MiniMaxM3SparseIndexCache - - return MiniMaxM3SparseIndexCache - - @dataclass -class MiniMaxM3SparseAttention: +class MiniMaxM3TritonSparseAttention: """Thin orchestrator for :func:`minimax_m3_sparse_prefill` and :func:`minimax_m3_sparse_decode`. @@ -909,7 +798,7 @@ class MiniMaxM3SparseAttention: :class:`MiniMaxM3SparseIndexCache`. The caller is responsible for routing the projected Q, K, V, ``idx_q``, ``idx_k`` (and optional ``idx_v``) tensors plus the populated - :class:`MiniMaxM3SparseAttentionMetadata`. + :class:`MiniMaxM3TritonSparseAttentionMetadata`. """ config: MiniMaxM3SparseConfig @@ -946,7 +835,7 @@ def forward( idx_q: torch.Tensor, k_cache: torch.Tensor, v_cache: torch.Tensor, - metadata: MiniMaxM3SparseAttentionMetadata, + metadata: MiniMaxM3TritonSparseAttentionMetadata, *, disable_index_value: bool, sm_scale: Optional[float] = None, @@ -1001,322 +890,304 @@ def forward( # --------------------------------------------------------------------------- -@functools.lru_cache(maxsize=1) -def get_minimax_m3_attention_backend_cls(): - """Return :class:`MiniMaxM3SparseRuntimeBackend` (lazy import). +class MiniMaxM3SparseRuntimeBackend(AttentionBackend[AttentionMetadata]): + """:class:`AttentionBackend` for MiniMax-M3 sparse layers. - Deferring the :class:`AttentionBackend` import keeps the algorithm - module usable from test paths that do not need the runtime backend. - """ - from ...interface import ( - AttentionBackend, - AttentionForwardArgs, - AttentionMetadata, - merge_attention_forward_args, - ) + Constructed under the standard ``create_attention(...)`` dispatch + when ``SparseAttentionConfig(algorithm='minimax_m3', ...)`` is + configured. Drives the MiniMax-M3 sparse algorithm directly via + :func:`minimax_m3_sparse_prefill` and + :func:`minimax_m3_sparse_decode`. - metadata_cls = get_minimax_m3_attention_metadata_cls() - - class MiniMaxM3SparseRuntimeBackend(AttentionBackend[AttentionMetadata]): - """:class:`AttentionBackend` for MiniMax-M3 sparse layers. - - Constructed under the standard ``create_attention(...)`` dispatch - when ``SparseAttentionConfig(algorithm='minimax_m3', ...)`` is - configured. Drives the MiniMax-M3 sparse algorithm directly via - :func:`minimax_m3_sparse_prefill` and - :func:`minimax_m3_sparse_decode`. - - The standard :class:`AttentionForwardArgs` surface does not - carry ``idx_q`` / ``idx_k`` slots, so the model layer threads - the index branch through ``**kwargs`` of :meth:`forward`. When - ``forward`` is called without ``idx_q`` it raises - :class:`NotImplementedError` with a pointer at the model layer - — the backend's ``forward`` is **executable**, but it is not a - substitute for the MiniMax-specific projection / norm / RoPE - steps the model layer is responsible for. - - The backend exposes - :meth:`forward_sparse` for callers that want a name-explicit - entry point (the model layer calls it directly); :meth:`forward` - is the standard contract entry point and routes to - :meth:`forward_sparse` when ``idx_q`` is supplied. - """ + The standard :class:`AttentionForwardArgs` surface does not + carry ``idx_q`` / ``idx_k`` slots, so the model layer threads + the index branch through ``**kwargs`` of :meth:`forward`. When + ``forward`` is called without ``idx_q`` it raises + :class:`NotImplementedError` with a pointer at the model layer + — the backend's ``forward`` is **executable**, but it is not a + substitute for the MiniMax-specific projection / norm / RoPE + steps the model layer is responsible for. + + The backend exposes + :meth:`forward_sparse` for callers that want a name-explicit + entry point (the model layer calls it directly); :meth:`forward` + is the standard contract entry point and routes to + :meth:`forward_sparse` when ``idx_q`` is supplied. + """ - Metadata = metadata_cls + Metadata = MiniMaxM3AttentionMetadata - def __init__( - self, - layer_idx: int, - num_heads: int, - head_dim: int, - num_kv_heads: Optional[int] = None, - quant_config=None, - sparse_params: Optional["MiniMaxM3SparseParams"] = None, + def __init__( + self, + layer_idx: int, + num_heads: int, + head_dim: int, + num_kv_heads: Optional[int] = None, + quant_config=None, + sparse_params: Optional["MiniMaxM3SparseParams"] = None, + **kwargs, + ): + if sparse_params is None: + raise ValueError("sparse_params is required for MiniMaxM3SparseRuntimeBackend") + super().__init__( + layer_idx, + num_heads, + head_dim, + num_kv_heads=num_kv_heads, + quant_config=quant_config, + sparse_params=sparse_params, **kwargs, - ): - if sparse_params is None: - raise ValueError("sparse_params is required for MiniMaxM3SparseRuntimeBackend") - super().__init__( - layer_idx, - num_heads, - head_dim, - num_kv_heads=num_kv_heads, - quant_config=quant_config, - sparse_params=sparse_params, - **kwargs, + ) + self.m3_config = MiniMaxM3SparseConfig.from_sparse_params( + sparse_params, + num_q_heads=num_heads, + num_kv_heads=num_kv_heads or num_heads, + head_dim=head_dim, + ) + self.disable_index_value = bool(sparse_params.disable_index_value) + + @staticmethod + def support_fused_rope() -> bool: + # The MiniMax-M3 model layer applies RoPE explicitly because + # both the main and the index branches need partial RoPE, + # and the standard fused-RoPE attention op does not have a + # hook for the index branch. + return False + + 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: "MiniMaxM3TritonSparseAttentionMetadata", + sm_scale: Optional[float] = None, + idx_sm_scale: Optional[float] = None, + output: Optional[torch.Tensor] = None, + ) -> torch.Tensor: + """Execute the MiniMax-M3 sparse path end-to-end. + + Inputs: + ``q``, ``k``, ``v`` : new-token projections, already + per-head norm + RoPE applied. + ``idx_q``, ``idx_k`` : index-branch projections, already + per-head norm + RoPE applied. + ``idx_v`` : index-V projection (only when + ``disable_index_value=False``). + ``k_cache``, ``v_cache``: flat-slot view of the paged + main K/V cache, + ``[num_slots, num_kv_heads, head_dim]``. + ``idx_k_cache`` : side index-K cache, + ``[num_slots, 1, sparse_index_dim]``. + ``idx_v_cache`` : side index-V cache (or ``None``). + ``out_cache_loc`` : ``[num_new_tokens]`` int slot + indices to write the new + token's K/V/idx_K to. + ``m3_metadata`` : populated + :class:`MiniMaxM3TritonSparseAttentionMetadata`. + ``output`` : optional preallocated final output, + ``[num_tokens, num_q_heads * head_dim]``. + + Returns ``[num_tokens, num_q_heads * head_dim]``. + """ + 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) + + # Production paths build the M3 metadata on the cache + # device in ``MiniMaxM3AttentionMetadata.prepare`` (called + # outside any CUDA-graph capture window), and test paths + # construct it directly on the desired device. So all + # metadata tensors should already live on ``k_cache.device`` + # by this point. We keep a same-device pass for resilience + # against legacy test callers that produce metadata on a + # different device, but it must not introduce CPU->GPU + # copies inside the capture window. ``ensure_metadata_on_device`` + # is a no-op when each tensor is already on the target + # device; under capture that no-op path is the contract. + 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, ) - self.m3_config = MiniMaxM3SparseConfig.from_sparse_params( - sparse_params, - num_q_heads=num_heads, - num_kv_heads=num_kv_heads or num_heads, - head_dim=head_dim, + ): + m3_metadata = ensure_metadata_on_device(m3_metadata, cache_device) + + # Write new K/V/idx_K to the configured slots. + # ``out_cache_loc`` comes from the pre-built attachment so + # it already lives on the cache device. The write goes + # through :func:`_write_main_kv_slots`, which is layout- + # aware: + # + # * 4-D multi-dim paged caches (the production V2 path: + # main K/V is ``kv_pool[:, 0]`` / ``kv_pool[:, 1]``, and + # ``idx_k_cache`` is the V2 4-D paged view ``[num_pages, + # tokens_per_block, 1, sparse_index_dim]``): decomposes + # each slot id into + # ``(page, within)`` and uses multi-dim fancy + # assignment so the write propagates to the underlying + # pool. A plain ``index_copy_(0, ...)`` would either + # silently fork a copy (non-contiguous main K/V view) + # or raise a shape mismatch (4-D index-K view), but + # the layout-aware helper sidesteps both failure + # modes. + # * 3-D flat-slot caches (focused unit tests that + # allocate plain ``torch.zeros((num_slots, num_heads, + # channel))`` tensors): falls back to + # ``index_copy_(0, ...)`` because the tensor IS the + # storage. + _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 idx_v is not None and idx_v_cache is not None: + idx_v_view = idx_v.view(num_tokens, 1, sparse_index_dim) + _write_main_kv_slots(idx_v_cache, out_cache_loc, idx_v_view) + + if m3_metadata.is_prefill: + _, o = minimax_m3_sparse_prefill( + q_view, + k_cache, + v_cache, + idx_q_view, + idx_k_cache, + None if self.disable_index_value else idx_v_cache, + m3_metadata, + self.m3_config, + disable_index_value=self.disable_index_value, + sm_scale=sm_scale, + idx_sm_scale=idx_sm_scale, + output=output, ) - self.disable_index_value = bool(sparse_params.disable_index_value) - - @staticmethod - def support_fused_rope() -> bool: - # The MiniMax-M3 model layer applies RoPE explicitly because - # both the main and the index branches need partial RoPE, - # and the standard fused-RoPE attention op does not have a - # hook for the index branch. - return False - - 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, - output: Optional[torch.Tensor] = None, - ) -> torch.Tensor: - """Execute the MiniMax-M3 sparse path end-to-end. - - Inputs: - ``q``, ``k``, ``v`` : new-token projections, already - per-head norm + RoPE applied. - ``idx_q``, ``idx_k`` : index-branch projections, already - per-head norm + RoPE applied. - ``idx_v`` : index-V projection (only when - ``disable_index_value=False``). - ``k_cache``, ``v_cache``: flat-slot view of the paged - main K/V cache, - ``[num_slots, num_kv_heads, head_dim]``. - ``idx_k_cache`` : side index-K cache, - ``[num_slots, 1, sparse_index_dim]``. - ``idx_v_cache`` : side index-V cache (or ``None``). - ``out_cache_loc`` : ``[num_new_tokens]`` int slot - indices to write the new - token's K/V/idx_K to. - ``m3_metadata`` : populated - :class:`MiniMaxM3SparseAttentionMetadata`. - ``output`` : optional preallocated final output, - ``[num_tokens, num_q_heads * head_dim]``. - - Returns ``[num_tokens, num_q_heads * head_dim]``. - """ - 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) - - # Production paths build the M3 metadata on the cache - # device in ``MiniMaxM3AttentionMetadata.prepare`` (called - # outside any CUDA-graph capture window), and test paths - # construct it directly on the desired device. So all - # metadata tensors should already live on ``k_cache.device`` - # by this point. We keep a same-device pass for resilience - # against legacy test callers that produce metadata on a - # different device, but it must not introduce CPU->GPU - # copies inside the capture window. ``ensure_metadata_on_device`` - # is a no-op when each tensor is already on the target - # device; under capture that no-op path is the contract. - 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) - - # Write new K/V/idx_K to the configured slots. - # ``out_cache_loc`` comes from the pre-built attachment so - # it already lives on the cache device. The write goes - # through :func:`_write_main_kv_slots`, which is layout- - # aware: - # - # * 4-D multi-dim paged caches (the production V2 path: - # main K/V is ``kv_pool[:, 0]`` / ``kv_pool[:, 1]``, and - # ``idx_k_cache`` is the V2 4-D paged view ``[num_pages, - # tokens_per_block, 1, sparse_index_dim]``): decomposes - # each slot id into - # ``(page, within)`` and uses multi-dim fancy - # assignment so the write propagates to the underlying - # pool. A plain ``index_copy_(0, ...)`` would either - # silently fork a copy (non-contiguous main K/V view) - # or raise a shape mismatch (4-D index-K view), but - # the layout-aware helper sidesteps both failure - # modes. - # * 3-D flat-slot caches (focused unit tests that - # allocate plain ``torch.zeros((num_slots, num_heads, - # channel))`` tensors): falls back to - # ``index_copy_(0, ...)`` because the tensor IS the - # storage. - _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 idx_v is not None and idx_v_cache is not None: - idx_v_view = idx_v.view(num_tokens, 1, sparse_index_dim) - _write_main_kv_slots(idx_v_cache, out_cache_loc, idx_v_view) - - if m3_metadata.is_prefill: - _, o = minimax_m3_sparse_prefill( - q_view, - k_cache, - v_cache, - idx_q_view, - idx_k_cache, - None if self.disable_index_value else idx_v_cache, - m3_metadata, - self.m3_config, - disable_index_value=self.disable_index_value, - sm_scale=sm_scale, - idx_sm_scale=idx_sm_scale, - output=output, - ) - else: - _, o = minimax_m3_sparse_decode( - q_view, - idx_q_view, - k_cache, - v_cache, - idx_k_cache, - None if self.disable_index_value else idx_v_cache, - m3_metadata, - self.m3_config, - disable_index_value=self.disable_index_value, - sm_scale=sm_scale, - idx_sm_scale=idx_sm_scale, - output=output, - ) - return o - - def forward( - self, - q: torch.Tensor, - k: Optional[torch.Tensor], - v: Optional[torch.Tensor], - metadata=None, - forward_args: Optional[AttentionForwardArgs] = None, - *, - output: Optional[torch.Tensor] = None, - idx_q: Optional[torch.Tensor] = None, - idx_k: Optional[torch.Tensor] = None, - idx_v: Optional[torch.Tensor] = None, - k_cache: Optional[torch.Tensor] = None, - v_cache: Optional[torch.Tensor] = None, - idx_k_cache: Optional[torch.Tensor] = None, - idx_v_cache: Optional[torch.Tensor] = None, - out_cache_loc: Optional[torch.Tensor] = None, - m3_metadata: Optional["MiniMaxM3SparseAttentionMetadata"] = None, - sm_scale: Optional[float] = None, - idx_sm_scale: Optional[float] = None, - **kwargs, - ) -> torch.Tensor: - """Standard ``AttentionBackend.forward`` entry point. - - The MiniMax-M3 sparse path needs the index branch projection - and the M3-shaped metadata; both arrive through keyword - arguments because the standard - :class:`AttentionForwardArgs` surface does not carry them. - - When ``idx_q`` is omitted, this method raises - :class:`NotImplementedError` to make the misuse loud — a - generic AttentionBackend dispatch site cannot drive this - backend without supplying the index branch. - """ - forward_args = merge_attention_forward_args(forward_args, kwargs) - if ( - output is not None - and forward_args.output is not None - and output is not forward_args.output - ): - raise ValueError("output was supplied both directly and through forward_args") - if output is None: - output = forward_args.output - if idx_q is None or idx_k is None or m3_metadata is None: - raise NotImplementedError( - f"MiniMaxM3SparseRuntimeBackend.forward (layer " - f"{self.layer_idx}) requires the M3 index branch and " - "metadata to be passed as keyword arguments " - "(`idx_q`, `idx_k`, `m3_metadata`, " - "`out_cache_loc`, `k_cache`, `v_cache`, " - "`idx_k_cache`). The standard AttentionForwardArgs " - "surface does not carry them; the model layer " - "(`MiniMaxM3Attention.forward`) supplies them when " - "calling this backend." - ) - if ( - k is None - or v is None - or k_cache is None - or v_cache is None - or idx_k_cache is None - or out_cache_loc is None - ): - raise ValueError( - "MiniMaxM3SparseRuntimeBackend.forward requires k, v, " - "k_cache, v_cache, idx_k_cache, and out_cache_loc to " - "be supplied alongside idx_q / idx_k / m3_metadata." - ) - return self.forward_sparse( - q=q, - k=k, - v=v, - idx_q=idx_q, - idx_k=idx_k, - idx_v=idx_v, - k_cache=k_cache, - v_cache=v_cache, - idx_k_cache=idx_k_cache, - idx_v_cache=idx_v_cache, - out_cache_loc=out_cache_loc, - m3_metadata=m3_metadata, + else: + _, o = minimax_m3_sparse_decode( + q_view, + idx_q_view, + k_cache, + v_cache, + idx_k_cache, + None if self.disable_index_value else idx_v_cache, + m3_metadata, + self.m3_config, + disable_index_value=self.disable_index_value, sm_scale=sm_scale, idx_sm_scale=idx_sm_scale, output=output, ) + return o - return MiniMaxM3SparseRuntimeBackend + def forward( + self, + q: torch.Tensor, + k: Optional[torch.Tensor], + v: Optional[torch.Tensor], + metadata=None, + forward_args: Optional[AttentionForwardArgs] = None, + *, + output: Optional[torch.Tensor] = None, + idx_q: Optional[torch.Tensor] = None, + idx_k: Optional[torch.Tensor] = None, + idx_v: Optional[torch.Tensor] = None, + k_cache: Optional[torch.Tensor] = None, + v_cache: Optional[torch.Tensor] = None, + idx_k_cache: Optional[torch.Tensor] = None, + idx_v_cache: Optional[torch.Tensor] = None, + out_cache_loc: Optional[torch.Tensor] = None, + m3_metadata: Optional["MiniMaxM3TritonSparseAttentionMetadata"] = None, + sm_scale: Optional[float] = None, + idx_sm_scale: Optional[float] = None, + **kwargs, + ) -> torch.Tensor: + """Standard ``AttentionBackend.forward`` entry point. + + The MiniMax-M3 sparse path needs the index branch projection + and the M3-shaped metadata; both arrive through keyword + arguments because the standard + :class:`AttentionForwardArgs` surface does not carry them. + + When ``idx_q`` is omitted, this method raises + :class:`NotImplementedError` to make the misuse loud — a + generic AttentionBackend dispatch site cannot drive this + backend without supplying the index branch. + """ + forward_args = merge_attention_forward_args(forward_args, kwargs) + if ( + output is not None + and forward_args.output is not None + and output is not forward_args.output + ): + raise ValueError("output was supplied both directly and through forward_args") + if output is None: + output = forward_args.output + if idx_q is None or idx_k is None or m3_metadata is None: + raise NotImplementedError( + f"MiniMaxM3SparseRuntimeBackend.forward (layer " + f"{self.layer_idx}) requires the M3 index branch and " + "metadata to be passed as keyword arguments " + "(`idx_q`, `idx_k`, `m3_metadata`, " + "`out_cache_loc`, `k_cache`, `v_cache`, " + "`idx_k_cache`). The standard AttentionForwardArgs " + "surface does not carry them; the model layer " + "(`MiniMaxM3Attention.forward`) supplies them when " + "calling this backend." + ) + if ( + k is None + or v is None + or k_cache is None + or v_cache is None + or idx_k_cache is None + or out_cache_loc is None + ): + raise ValueError( + "MiniMaxM3SparseRuntimeBackend.forward requires k, v, " + "k_cache, v_cache, idx_k_cache, and out_cache_loc to " + "be supplied alongside idx_q / idx_k / m3_metadata." + ) + return self.forward_sparse( + q=q, + k=k, + v=v, + idx_q=idx_q, + idx_k=idx_k, + idx_v=idx_v, + k_cache=k_cache, + v_cache=v_cache, + idx_k_cache=idx_k_cache, + idx_v_cache=idx_v_cache, + out_cache_loc=out_cache_loc, + m3_metadata=m3_metadata, + sm_scale=sm_scale, + idx_sm_scale=idx_sm_scale, + output=output, + ) __all__ = [ - "MiniMaxM3SparseAttention", - "get_minimax_m3_attention_backend_cls", + "MiniMaxM3SparseRuntimeBackend", + "MiniMaxM3TritonSparseAttention", "minimax_m3_sparse_decode", "minimax_m3_sparse_prefill", ] diff --git a/tensorrt_llm/_torch/attention_backend/sparse/minimax_m3/kernels.py b/tensorrt_llm/_torch/attention_backend/sparse/minimax_m3/triton_kernels.py similarity index 100% rename from tensorrt_llm/_torch/attention_backend/sparse/minimax_m3/kernels.py rename to tensorrt_llm/_torch/attention_backend/sparse/minimax_m3/triton_kernels.py diff --git a/tensorrt_llm/_torch/attention_backend/sparse/minimax_m3/metadata.py b/tensorrt_llm/_torch/attention_backend/sparse/minimax_m3/triton_metadata.py similarity index 61% rename from tensorrt_llm/_torch/attention_backend/sparse/minimax_m3/metadata.py rename to tensorrt_llm/_torch/attention_backend/sparse/minimax_m3/triton_metadata.py index 925ca6c0189b..be4b8ae24143 100644 --- a/tensorrt_llm/_torch/attention_backend/sparse/minimax_m3/metadata.py +++ b/tensorrt_llm/_torch/attention_backend/sparse/minimax_m3/triton_metadata.py @@ -1,131 +1,31 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 -"""MiniMax-M3 sparse attention configuration + per-forward metadata. +"""MiniMax-M3 Triton reference per-forward metadata. Contains: - * :class:`MiniMaxM3SparseConfig` -- post-TP-shard kernel - parameter bundle. - * :class:`MiniMaxM3SparseAttentionMetadata` -- per-forward metadata - dataclass with a - CUDA-graph-safe - :meth:`prepare`. + * :class:`MiniMaxM3TritonSparseAttentionMetadata` -- per-forward metadata + dataclass with a CUDA-graph-safe :meth:`prepare`. * Helpers to migrate metadata across devices, build it from a real - :class:`KVCacheManagerV2`, and pre-allocate CUDA-graph-stable - buffers. - * :func:`get_minimax_m3_attention_metadata_cls` -- lazy factory for - the :class:`AttentionMetadata` subclass the pyexecutor wires into - the M3 sparse layer's forward path. + :class:`KVCacheManagerV2`, and pre-allocate CUDA-graph-stable buffers. + * :class:`MiniMaxM3AttentionMetadata` -- the :class:`AttentionMetadata` + subclass the pyexecutor wires into the M3 sparse layer's forward path. """ from __future__ import annotations import dataclasses -import functools from dataclasses import dataclass, field -from typing import List, Literal, Optional, Tuple +from typing import List, Optional, Tuple import torch -from ..params import SparseParams - - -@dataclass(frozen=True) -class MiniMaxM3SparseParams(SparseParams): - """Lowered runtime parameters for the MiniMax-M3 sparse backend.""" - - algorithm: Literal["minimax_m3"] = field(init=False, default="minimax_m3") - num_index_heads: int = 4 - sparse_index_dim: int = 128 - block_size: int = 128 - topk: int = 16 - init_blocks: int = 0 - local_blocks: int = 1 - score_type: str = "max" - disable_index_value: bool = True - - -@dataclass(frozen=True) -class MiniMaxM3SparseConfig: - """Per-rank kernel parameter bundle for MiniMax-M3 sparse attention. - - This is **not** a user-facing config (use - :class:`tensorrt_llm.llmapi.llm_args.MiniMaxM3SparseAttentionConfig` - for that). It is the layer-invariant, post-TP-shard parameter bundle - that backend kernels and reference helpers consume. The user knobs - come from :class:`MiniMaxM3SparseParams`; ``num_q_heads`` / - ``num_kv_heads`` / ``head_dim`` come from the per-rank model - geometry and must be supplied by the caller (typically via - :meth:`from_sparse_params`). - """ - - num_q_heads: int - num_kv_heads: int - head_dim: int - num_index_heads: int - sparse_index_dim: int - block_size: int - topk: int - init_blocks: int = 0 - local_blocks: int = 1 - score_type: str = "max" - - def __post_init__(self) -> None: - if self.num_q_heads % self.num_kv_heads != 0: - raise ValueError( - f"num_q_heads ({self.num_q_heads}) must be divisible by " - f"num_kv_heads ({self.num_kv_heads})" - ) - if self.num_index_heads % self.num_kv_heads != 0: - raise ValueError( - f"num_index_heads ({self.num_index_heads}) must be divisible " - f"by num_kv_heads ({self.num_kv_heads})" - ) - if self.block_size <= 0: - raise ValueError(f"block_size must be > 0, got {self.block_size}") - if self.topk <= 0: - raise ValueError(f"topk must be > 0, got {self.topk}") - if self.init_blocks < 0: - raise ValueError(f"init_blocks must be >= 0, got {self.init_blocks}") - if self.local_blocks < 0: - raise ValueError(f"local_blocks must be >= 0, got {self.local_blocks}") - if self.score_type != "max": - # SGLang exposes only "max" today and that is what the MiniMax-M3 - # checkpoint config specifies. Reject anything else explicitly so - # a config drift surfaces immediately. - raise ValueError( - f"score_type={self.score_type!r} is not supported " - "(only 'max' matches the SGLang reference)" - ) - - @classmethod - def from_sparse_params( - cls, - sparse_params: "MiniMaxM3SparseParams", - *, - num_q_heads: int, - num_kv_heads: int, - head_dim: int, - ) -> "MiniMaxM3SparseConfig": - """Build a kernel param bundle from lowered ``MiniMaxM3SparseParams`` - and the per-rank model geometry. - """ - 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), - sparse_index_dim=int(sparse_params.sparse_index_dim), - block_size=int(sparse_params.block_size), - topk=int(sparse_params.topk), - init_blocks=int(sparse_params.init_blocks), - local_blocks=int(sparse_params.local_blocks), - score_type=str(sparse_params.score_type), - ) +from ...interface import AttentionMetadata +from .common import build_paged_kv_slot_mapping @dataclass -class MiniMaxM3SparseAttentionMetadata: +class MiniMaxM3TritonSparseAttentionMetadata: """Per-forward metadata for MiniMax-M3 sparse attention. Mirrors the shape of SGLang's @@ -269,12 +169,12 @@ def prepare(self) -> None: def ensure_metadata_on_device( - metadata: "MiniMaxM3SparseAttentionMetadata", + metadata: "MiniMaxM3TritonSparseAttentionMetadata", device: torch.device, -) -> "MiniMaxM3SparseAttentionMetadata": +) -> "MiniMaxM3TritonSparseAttentionMetadata": """Return ``metadata`` with every GPU-consumed tensor on ``device``. - Constructs a new :class:`MiniMaxM3SparseAttentionMetadata` whose + Constructs a new :class:`MiniMaxM3TritonSparseAttentionMetadata` whose tensor fields are migrated to ``device`` when they live elsewhere. The CPU-side mirror ``seq_lens_cpu`` is preserved because the algorithm reads it only when explicitly noted (e.g. for @@ -306,18 +206,6 @@ def _move(t: Optional[torch.Tensor]) -> Optional[torch.Tensor]: ) -def replace_metadata( - metadata: MiniMaxM3SparseAttentionMetadata, - **changes, -) -> MiniMaxM3SparseAttentionMetadata: - """Helper around :func:`dataclasses.replace` for ``metadata``. - - Provided so callers can build a decode metadata from a prefill - metadata without manually re-typing every field. - """ - return dataclasses.replace(metadata, **changes) - - def allocate_minimax_m3_static_buffers( *, max_num_sequences: int, @@ -405,6 +293,82 @@ def allocate_minimax_m3_static_buffers( } +def _build_runtime_metadata_fresh( + *, + kv_cache_manager, + request_ids, + seq_lens: torch.Tensor, + seq_lens_cpu: torch.Tensor, + is_prefill: bool, + prefix_lens: Optional[torch.Tensor], + extend_seq_lens_cpu: Optional[List[int]], + device: torch.device, +) -> Tuple[MiniMaxM3TritonSparseAttentionMetadata, torch.Tensor]: + """Fresh-allocation build of the Triton reference metadata. + + Delegates the backend-neutral req_to_token, slot_ids and out_cache_loc + derivation to common.build_paged_kv_slot_mapping and adds the Triton-only + fields: cu_seqlens_q, prefix_lens on device, and the scalar max_seqlen + values and per-query-token tensors that prepare() computes. Used when no + graph-stable static_buffers are supplied; the static_buffers path in + build_runtime_metadata_from_kv_manager keeps its own in-place buffer writes. + """ + batch = int(seq_lens.shape[0]) + seq_lens_dev = seq_lens.to(device) if seq_lens.device != device else seq_lens + + if is_prefill: + if extend_seq_lens_cpu is None: + raise ValueError("prefill metadata requires extend_seq_lens_cpu") + if prefix_lens is None: + raise ValueError("prefill metadata requires prefix_lens") + prefix_lens_dev = prefix_lens.to(device) if prefix_lens.device != device else prefix_lens + qo_lens_cpu = torch.tensor([int(x) for x in extend_seq_lens_cpu], dtype=torch.int32) + qo_offset_cpu = prefix_lens.detach().to(device="cpu", dtype=torch.int32) + req_to_token, slot_ids, out_cache_loc = build_paged_kv_slot_mapping( + kv_cache_manager=kv_cache_manager, + request_ids=request_ids, + qo_lens_cpu=qo_lens_cpu, + qo_offset_cpu=qo_offset_cpu, + device=device, + ) + cu_q: List[int] = [0] + for ext in extend_seq_lens_cpu: + cu_q.append(cu_q[-1] + int(ext)) + cu_seqlens_q = torch.tensor(cu_q, dtype=torch.int32, device=device) + meta = MiniMaxM3TritonSparseAttentionMetadata( + is_prefill=True, + req_to_token=req_to_token, + slot_ids=slot_ids, + seq_lens=seq_lens_dev, + seq_lens_cpu=seq_lens_cpu, + prefix_lens=prefix_lens_dev, + cu_seqlens_q=cu_seqlens_q, + extend_seq_lens_cpu=list(extend_seq_lens_cpu), + q_batch_row=None, + q_positions=None, + ) + else: + # Decode: the new token sits at position seq_lens[b] - 1. + qo_lens_cpu = torch.ones(batch, dtype=torch.int32) + qo_offset_cpu = seq_lens_cpu.detach().to(device="cpu", dtype=torch.int32) - 1 + req_to_token, slot_ids, out_cache_loc = build_paged_kv_slot_mapping( + kv_cache_manager=kv_cache_manager, + request_ids=request_ids, + qo_lens_cpu=qo_lens_cpu, + qo_offset_cpu=qo_offset_cpu, + device=device, + ) + meta = MiniMaxM3TritonSparseAttentionMetadata( + is_prefill=False, + req_to_token=req_to_token, + slot_ids=slot_ids, + seq_lens=seq_lens_dev, + seq_lens_cpu=seq_lens_cpu, + ) + meta.prepare() + return meta, out_cache_loc + + def build_runtime_metadata_from_kv_manager( *, kv_cache_manager, @@ -416,8 +380,8 @@ def build_runtime_metadata_from_kv_manager( extend_seq_lens_cpu: Optional[List[int]] = None, device: Optional[torch.device] = None, static_buffers: Optional[dict] = None, -) -> Tuple[MiniMaxM3SparseAttentionMetadata, torch.Tensor]: - """Build a :class:`MiniMaxM3SparseAttentionMetadata` from a real +) -> Tuple[MiniMaxM3TritonSparseAttentionMetadata, torch.Tensor]: + """Build a :class:`MiniMaxM3TritonSparseAttentionMetadata` from a real :class:`MiniMaxM3KVCacheManagerV2`. Returns the populated metadata plus an ``out_cache_loc`` tensor @@ -452,6 +416,21 @@ def build_runtime_metadata_from_kv_manager( the end-to-end runtime path without going through the full LLM forward. """ + if device is None: + device = seq_lens.device + + if static_buffers is None: + return _build_runtime_metadata_fresh( + kv_cache_manager=kv_cache_manager, + request_ids=request_ids, + seq_lens=seq_lens, + seq_lens_cpu=seq_lens_cpu, + is_prefill=is_prefill, + prefix_lens=prefix_lens, + extend_seq_lens_cpu=extend_seq_lens_cpu, + device=device, + ) + tokens_per_block = int(kv_cache_manager.tokens_per_block) # block_ids_per_seq is a [batch_size, max_blocks_per_seq] tensor; row b # holds the block ids assigned to request_ids[b] in order. @@ -459,8 +438,6 @@ def build_runtime_metadata_from_kv_manager( batch = int(seq_lens.shape[0]) max_blocks = int(block_ids.shape[1]) max_kv_len = max_blocks * tokens_per_block - if device is None: - device = seq_lens.device if static_buffers is not None: if static_buffers.get("device") != device: @@ -631,7 +608,7 @@ def build_runtime_metadata_from_kv_manager( cu_seqlens_q = torch.tensor(cu_q, dtype=torch.int32, device=device) q_batch_row = None q_positions = None - meta = MiniMaxM3SparseAttentionMetadata( + meta = MiniMaxM3TritonSparseAttentionMetadata( is_prefill=True, req_to_token=req_to_token, slot_ids=slot_ids, @@ -663,7 +640,7 @@ def build_runtime_metadata_from_kv_manager( out_cache_loc = out_cache_loc_buf[:batch] else: out_cache_loc = torch.tensor(out_cache_loc_list, dtype=torch.int32, device=device) - meta = MiniMaxM3SparseAttentionMetadata( + meta = MiniMaxM3TritonSparseAttentionMetadata( is_prefill=False, req_to_token=req_to_token, slot_ids=slot_ids, @@ -674,242 +651,222 @@ def build_runtime_metadata_from_kv_manager( return meta, out_cache_loc -@functools.lru_cache(maxsize=1) -def get_minimax_m3_attention_metadata_cls(): - """Return :class:`MiniMaxM3AttentionMetadata` (lazy import). +class MiniMaxM3AttentionMetadata(AttentionMetadata): + """:class:`AttentionMetadata` that pre-builds MiniMax-M3 metadata. + + Overrides :meth:`prepare` so the M3-sparse + :class:`MiniMaxM3TritonSparseAttentionMetadata` and the per-new-token + ``out_cache_loc`` are built **once per scheduler step**, on the + cache device, before the model forward runs. The result is + stored as ``self.minimax_m3 = {"metadata": m3_meta, + "out_cache_loc": out_cache_loc}`` so the model layer's + ``_dense_forward`` and ``_sparse_forward`` can read it without + any device migration. - The class extends :class:`AttentionMetadata` so the pyexecutor's - metadata-creation/prepare hooks (model_engine.py) drive M3 metadata - construction outside the CUDA-graph capture window. Building the - M3-sparse ``req_to_token`` / ``slot_ids`` / ``out_cache_loc`` - tensors during ``prepare()`` lands them on the GPU **before** the - forward call; the forward path then reads from the pre-built - attachment and performs no CPU->GPU copies, which is required for - CUDA-graph capture safety (``cudaErrorStreamCaptureUnsupported`` - fires for CPU->GPU ``memcpyAsync`` calls inside a captured stream). + Test paths that build their own metadata can short-circuit by + attaching ``attn_metadata.minimax_m3`` directly before calling + the forward; those paths do not go through :meth:`prepare`. """ - from ...interface import AttentionMetadata - - class MiniMaxM3AttentionMetadata(AttentionMetadata): - """:class:`AttentionMetadata` that pre-builds MiniMax-M3 metadata. - - Overrides :meth:`prepare` so the M3-sparse - :class:`MiniMaxM3SparseAttentionMetadata` and the per-new-token - ``out_cache_loc`` are built **once per scheduler step**, on the - cache device, before the model forward runs. The result is - stored as ``self.minimax_m3 = {"metadata": m3_meta, - "out_cache_loc": out_cache_loc}`` so the model layer's - ``_dense_forward`` and ``_sparse_forward`` can read it without - any device migration. - - Test paths that build their own metadata can short-circuit by - attaching ``attn_metadata.minimax_m3`` directly before calling - the forward; those paths do not go through :meth:`prepare`. + + minimax_m3: Optional[dict] = None + # Lazily allocated dict of persistent device buffers used to keep + # ``MiniMaxM3TritonSparseAttentionMetadata`` tensor addresses stable + # across CUDA-graph capture/replay. None until the first + # ``prepare()`` call decides to use them (``is_cuda_graph`` / + # graph-stable mode). + _m3_static_buffers: Optional[dict] = None + + def _maybe_get_m3_static_buffers( + self, cache_device: torch.device, kv_cache_manager + ) -> Optional[dict]: + """Return persistent M3 buffers when graph stability is + required. + + Allocates the persistent buffer dict the first time it is + needed and caches it on ``self._m3_static_buffers``. We + allocate the buffers under two conditions: + + * ``self.is_cuda_graph`` is True -- the captured graph + requires stable ``data_ptr()`` across replays; OR + * the previous prepare() already allocated buffers -- + we keep using them so the algorithm sees the same + addresses even between non-graph and graph-mode calls + (which can happen when the model engine alternates + between eager warmup and graph replay). + + Returns ``None`` when no static buffers should be used (e.g. + eager-only test paths that rely on per-call allocations). """ + need_static = ( + bool(getattr(self, "is_cuda_graph", False)) or self._m3_static_buffers is not None + ) + if not need_static: + return None + if self._m3_static_buffers is not None: + bufs = self._m3_static_buffers + if bufs.get("device") == cache_device: + return bufs + + # First-time use: return an empty placeholder dict. + # ``build_runtime_metadata_from_kv_manager`` performs the + # actual allocation lazily on the first call where the + # current scheduler step's geometry (max_kv_len from the + # manager's block-id table, total_q from extend_seq_lens, + # actual batch size after CUDA-graph padding) is known. + # That removes the need to predict the warmup geometry up + # front. The first allocation pins the buffer addresses + # for the rest of this metadata instance's lifetime, so all + # subsequent prepare() calls reuse the same ``data_ptr()``s + # and CUDA graph capture/replay stays valid. + placeholder: dict = { + "device": cache_device, + # Caller-provided hints used by the lazy allocator below + # when it sizes the persistent buffers on the first real + # prepare() call. + "max_num_sequences_hint": int( + getattr(self, "max_num_sequences", None) or self.max_num_requests + ), + "max_num_tokens_hint": int( + getattr(self, "max_num_tokens", None) + or (int(getattr(self, "max_num_sequences", None) or self.max_num_requests)) + ), + } + self._m3_static_buffers = placeholder + return placeholder - minimax_m3: Optional[dict] = None - # Lazily allocated dict of persistent device buffers used to keep - # ``MiniMaxM3SparseAttentionMetadata`` tensor addresses stable - # across CUDA-graph capture/replay. None until the first - # ``prepare()`` call decides to use them (``is_cuda_graph`` / - # graph-stable mode). - _m3_static_buffers: Optional[dict] = None - - def _maybe_get_m3_static_buffers( - self, cache_device: torch.device, kv_cache_manager - ) -> Optional[dict]: - """Return persistent M3 buffers when graph stability is - required. - - Allocates the persistent buffer dict the first time it is - needed and caches it on ``self._m3_static_buffers``. We - allocate the buffers under two conditions: - - * ``self.is_cuda_graph`` is True -- the captured graph - requires stable ``data_ptr()`` across replays; OR - * the previous prepare() already allocated buffers -- - we keep using them so the algorithm sees the same - addresses even between non-graph and graph-mode calls - (which can happen when the model engine alternates - between eager warmup and graph replay). - - Returns ``None`` when no static buffers should be used (e.g. - eager-only test paths that rely on per-call allocations). - """ - need_static = ( - bool(getattr(self, "is_cuda_graph", False)) or self._m3_static_buffers is not None + def prepare(self) -> None: + super().prepare() + + # Always rebuild the M3 metadata block on each prepare() + # call so it reflects the current scheduler step's seq_lens + # / request_ids / num_cached_tokens. Production + # ``model_engine`` invokes ``prepare()`` outside any CUDA + # graph capture window, so the (potentially expensive) build + # is safe to perform here. + # + # When CUDA graph is enabled the inner ``build_runtime_metadata_from_kv_manager`` + # call writes into the persistent ``_m3_static_buffers`` so + # the captured graph keeps reading from stable ``data_ptr()``s + # across replays. Without this the captured ``index_select`` + # over ``req_to_token``/``slot_ids`` reads from freed warmup + # memory and either produces wrong tokens or fires + # ``Indexing.cu:1515`` ``srcIndex < srcSelectDimSize``. + self.minimax_m3 = None + + # Production path: build the M3 metadata from the standard + # AttentionMetadata fields. Requires kv_cache_manager + the + # M3 sparse-cache contract. + kv_cache_manager = getattr(self, "kv_cache_manager", None) + if kv_cache_manager is None or not hasattr(kv_cache_manager, "get_index_k_buffer"): + # Not an M3 KV cache manager: nothing to build. The + # forward path will raise a clear error if the M3 + # backend ends up dispatched without the M3 cache. + return + request_ids = getattr(self, "request_ids", None) + seq_lens = self.seq_lens + if request_ids is None or seq_lens is None: + return + num_contexts = int(getattr(self, "num_contexts", 0) or 0) + batch_size = int(seq_lens.shape[0]) + if batch_size == 0: + return + + # The cache device hosts every paged buffer; this is the + # device the forward path consumes. + try: + layer_buf = kv_cache_manager.get_buffers(0) + cache_device = layer_buf.device + except Exception: + cache_device = torch.device(f"cuda:{torch.cuda.current_device()}") + + seq_lens_cpu = ( + getattr(self, "seq_lens_cpu", None) + if hasattr(self, "seq_lens_cpu") + else seq_lens.detach().to("cpu") + ) + if seq_lens_cpu is None: + seq_lens_cpu = seq_lens.detach().to("cpu") + + kv_cache_params = getattr(self, "kv_cache_params", None) + num_cached_per_seq = ( + kv_cache_params.num_cached_tokens_per_seq + if kv_cache_params is not None + else [0] * batch_size + ) + + # ``attn_metadata.seq_lens`` from the PyExecutor is the + # per-step new-token count. The M3 sparse-attention algorithm + # consumes a *cumulative* kv length: ``minimax_m3_sparse_*`` + # masks reads against ``metadata.seq_lens`` as the per-request + # K-side extent. Compute that cumulative kv length per request + # and feed it into the algorithm metadata builder. + kv_lens_cpu_list = [ + int(num_cached_per_seq[b]) + int(seq_lens_cpu[b].item()) for b in range(batch_size) + ] + kv_lens_cpu = torch.tensor(kv_lens_cpu_list, dtype=torch.int32) + kv_lens_dev = kv_lens_cpu.to(device=cache_device, non_blocking=True) + + static_buffers = self._maybe_get_m3_static_buffers(cache_device, kv_cache_manager) + + # Any batch containing a context (prefill or chunked extend) + # request takes the extend path. For prefill rows + # ``num_cached_per_seq`` is ``prefix_lens`` and the full new + # chunk is ``extend_seq_len``; for decode rows + # ``num_cached`` is ``kv_len - 1`` and ``extend_seq_len`` is + # 1, so the same builder produces the correct one-slot + # entry. Pure-decode batches (``num_contexts == 0``) still + # take the decode optimization for CUDA-graph warmup + # geometry. + # + # Mixed prefill+decode batches always take the extend path: + # the prefill kernel handles decode rows as 1-slot extends. + # The decode branch below is a pure-decode-only perf + # specialization. (iter-131 regression: previously a wrong + # predicate routed mixed batches into the decode branch and + # crashed in index_copy_.) + is_extend = num_contexts > 0 + if is_extend: + prefix_lens_list = [int(num_cached_per_seq[b]) for b in range(batch_size)] + extend_seq_lens_cpu = [ + kv_lens_cpu_list[b] - prefix_lens_list[b] for b in range(batch_size) + ] + prefix_lens = torch.tensor( + prefix_lens_list, + dtype=torch.int32, + device=cache_device, ) - if not need_static: - return None - if self._m3_static_buffers is not None: - bufs = self._m3_static_buffers - if bufs.get("device") == cache_device: - return bufs - - # First-time use: return an empty placeholder dict. - # ``build_runtime_metadata_from_kv_manager`` performs the - # actual allocation lazily on the first call where the - # current scheduler step's geometry (max_kv_len from the - # manager's block-id table, total_q from extend_seq_lens, - # actual batch size after CUDA-graph padding) is known. - # That removes the need to predict the warmup geometry up - # front. The first allocation pins the buffer addresses - # for the rest of this metadata instance's lifetime, so all - # subsequent prepare() calls reuse the same ``data_ptr()``s - # and CUDA graph capture/replay stays valid. - placeholder: dict = { - "device": cache_device, - # Caller-provided hints used by the lazy allocator below - # when it sizes the persistent buffers on the first real - # prepare() call. - "max_num_sequences_hint": int( - getattr(self, "max_num_sequences", None) or self.max_num_requests - ), - "max_num_tokens_hint": int( - getattr(self, "max_num_tokens", None) - or (int(getattr(self, "max_num_sequences", None) or self.max_num_requests)) - ), - } - self._m3_static_buffers = placeholder - return placeholder - - def prepare(self) -> None: - super().prepare() - - # Always rebuild the M3 metadata block on each prepare() - # call so it reflects the current scheduler step's seq_lens - # / request_ids / num_cached_tokens. Production - # ``model_engine`` invokes ``prepare()`` outside any CUDA - # graph capture window, so the (potentially expensive) build - # is safe to perform here. - # - # When CUDA graph is enabled the inner ``build_runtime_metadata_from_kv_manager`` - # call writes into the persistent ``_m3_static_buffers`` so - # the captured graph keeps reading from stable ``data_ptr()``s - # across replays. Without this the captured ``index_select`` - # over ``req_to_token``/``slot_ids`` reads from freed warmup - # memory and either produces wrong tokens or fires - # ``Indexing.cu:1515`` ``srcIndex < srcSelectDimSize``. - self.minimax_m3 = None - - # Production path: build the M3 metadata from the standard - # AttentionMetadata fields. Requires kv_cache_manager + the - # M3 sparse-cache contract. - kv_cache_manager = getattr(self, "kv_cache_manager", None) - if kv_cache_manager is None or not hasattr(kv_cache_manager, "get_index_k_buffer"): - # Not an M3 KV cache manager: nothing to build. The - # forward path will raise a clear error if the M3 - # backend ends up dispatched without the M3 cache. - return - request_ids = getattr(self, "request_ids", None) - seq_lens = self.seq_lens - if request_ids is None or seq_lens is None: - return - num_contexts = int(getattr(self, "num_contexts", 0) or 0) - batch_size = int(seq_lens.shape[0]) - if batch_size == 0: - return - - # The cache device hosts every paged buffer; this is the - # device the forward path consumes. - try: - layer_buf = kv_cache_manager.get_buffers(0) - cache_device = layer_buf.device - except Exception: - cache_device = torch.device(f"cuda:{torch.cuda.current_device()}") - - seq_lens_cpu = ( - getattr(self, "seq_lens_cpu", None) - if hasattr(self, "seq_lens_cpu") - else seq_lens.detach().to("cpu") + m3_meta, out_cache_loc = build_runtime_metadata_from_kv_manager( + kv_cache_manager=kv_cache_manager, + request_ids=request_ids, + seq_lens=kv_lens_dev, + seq_lens_cpu=kv_lens_cpu, + is_prefill=True, + prefix_lens=prefix_lens, + extend_seq_lens_cpu=extend_seq_lens_cpu, + device=cache_device, + static_buffers=static_buffers, ) - if seq_lens_cpu is None: - seq_lens_cpu = seq_lens.detach().to("cpu") - - kv_cache_params = getattr(self, "kv_cache_params", None) - num_cached_per_seq = ( - kv_cache_params.num_cached_tokens_per_seq - if kv_cache_params is not None - else [0] * batch_size + else: + m3_meta, out_cache_loc = build_runtime_metadata_from_kv_manager( + kv_cache_manager=kv_cache_manager, + request_ids=request_ids, + seq_lens=kv_lens_dev, + seq_lens_cpu=kv_lens_cpu, + is_prefill=False, + device=cache_device, + static_buffers=static_buffers, ) - # ``attn_metadata.seq_lens`` from the PyExecutor is the - # per-step new-token count. The M3 sparse-attention algorithm - # consumes a *cumulative* kv length: ``minimax_m3_sparse_*`` - # masks reads against ``metadata.seq_lens`` as the per-request - # K-side extent. Compute that cumulative kv length per request - # and feed it into the algorithm metadata builder. - kv_lens_cpu_list = [ - int(num_cached_per_seq[b]) + int(seq_lens_cpu[b].item()) for b in range(batch_size) - ] - kv_lens_cpu = torch.tensor(kv_lens_cpu_list, dtype=torch.int32) - kv_lens_dev = kv_lens_cpu.to(device=cache_device, non_blocking=True) - - static_buffers = self._maybe_get_m3_static_buffers(cache_device, kv_cache_manager) - - # Any batch containing a context (prefill or chunked extend) - # request takes the extend path. For prefill rows - # ``num_cached_per_seq`` is ``prefix_lens`` and the full new - # chunk is ``extend_seq_len``; for decode rows - # ``num_cached`` is ``kv_len - 1`` and ``extend_seq_len`` is - # 1, so the same builder produces the correct one-slot - # entry. Pure-decode batches (``num_contexts == 0``) still - # take the decode optimization for CUDA-graph warmup - # geometry. - # - # Mixed prefill+decode batches always take the extend path: - # the prefill kernel handles decode rows as 1-slot extends. - # The decode branch below is a pure-decode-only perf - # specialization. (iter-131 regression: previously a wrong - # predicate routed mixed batches into the decode branch and - # crashed in index_copy_.) - is_extend = num_contexts > 0 - if is_extend: - prefix_lens_list = [int(num_cached_per_seq[b]) for b in range(batch_size)] - extend_seq_lens_cpu = [ - kv_lens_cpu_list[b] - prefix_lens_list[b] for b in range(batch_size) - ] - prefix_lens = torch.tensor( - prefix_lens_list, - dtype=torch.int32, - device=cache_device, - ) - m3_meta, out_cache_loc = build_runtime_metadata_from_kv_manager( - kv_cache_manager=kv_cache_manager, - request_ids=request_ids, - seq_lens=kv_lens_dev, - seq_lens_cpu=kv_lens_cpu, - is_prefill=True, - prefix_lens=prefix_lens, - extend_seq_lens_cpu=extend_seq_lens_cpu, - device=cache_device, - static_buffers=static_buffers, - ) - else: - m3_meta, out_cache_loc = build_runtime_metadata_from_kv_manager( - kv_cache_manager=kv_cache_manager, - request_ids=request_ids, - seq_lens=kv_lens_dev, - seq_lens_cpu=kv_lens_cpu, - is_prefill=False, - device=cache_device, - static_buffers=static_buffers, - ) - - self.minimax_m3 = { - "metadata": m3_meta, - "out_cache_loc": out_cache_loc, - } - - return MiniMaxM3AttentionMetadata + self.minimax_m3 = { + "metadata": m3_meta, + "out_cache_loc": out_cache_loc, + } __all__ = [ - "MiniMaxM3SparseConfig", - "MiniMaxM3SparseAttentionMetadata", + "MiniMaxM3AttentionMetadata", + "MiniMaxM3TritonSparseAttentionMetadata", "allocate_minimax_m3_static_buffers", "build_runtime_metadata_from_kv_manager", "ensure_metadata_on_device", - "get_minimax_m3_attention_metadata_cls", - "replace_metadata", ] diff --git a/tensorrt_llm/_torch/attention_backend/sparse/utils.py b/tensorrt_llm/_torch/attention_backend/sparse/utils.py index f438d998939f..1762268bf242 100644 --- a/tensorrt_llm/_torch/attention_backend/sparse/utils.py +++ b/tensorrt_llm/_torch/attention_backend/sparse/utils.py @@ -42,14 +42,30 @@ def get_sparse_attn_kv_cache_manager( ) +def _resolve_minimax_m3_backend_cls( + sparse_params: "SparseParams") -> Type["AttentionBackend"]: + """Select the MiniMax-M3 sparse backend from the lowered params. + + The Triton reference is the default. When implementation is 'msa' the MSA + (fmha_sm100) backend is used instead, gated on SM100 availability so an + unsupported system fails early rather than at kernel launch. + """ + from .minimax_m3 import MiniMaxM3SparseRuntimeBackend + if getattr(sparse_params, "implementation", "triton") == "msa": + from .minimax_m3 import MiniMaxM3MsaSparseAttention + from .minimax_m3.msa_availability import ensure_msa_available + ensure_msa_available() + return MiniMaxM3MsaSparseAttention + return MiniMaxM3SparseRuntimeBackend + + def get_vanilla_sparse_attn_attention_backend( sparse_params: "SparseParams") -> Type["AttentionBackend"]: - from .minimax_m3 import get_minimax_m3_attention_backend_cls from .rocket import RocketVanillaAttention if sparse_params.algorithm == "rocket": return RocketVanillaAttention elif sparse_params.algorithm == "minimax_m3": - return get_minimax_m3_attention_backend_cls() + return _resolve_minimax_m3_backend_cls(sparse_params) else: raise ValueError( f"Unsupported sparse attention algorithm in vanilla attention backend: {sparse_params.algorithm}" @@ -62,7 +78,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_params.algorithm == "rocket": return RocketTrtllmAttention @@ -78,7 +93,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_params) else: raise ValueError( f"Unsupported sparse attention algorithm in trtllm attention backend: {sparse_params.algorithm}" @@ -87,9 +102,8 @@ def get_trtllm_sparse_attn_attention_backend( def get_flashinfer_sparse_attn_attention_backend( sparse_params: "SparseParams") -> Type["AttentionBackend"]: - from .minimax_m3 import get_minimax_m3_attention_backend_cls if sparse_params.algorithm == "minimax_m3": - return get_minimax_m3_attention_backend_cls() + return _resolve_minimax_m3_backend_cls(sparse_params) raise ValueError( f"Unsupported sparse attention algorithm in flashinfer attention backend: {sparse_params.algorithm}" ) diff --git a/tensorrt_llm/_torch/models/modeling_minimaxm3.py b/tensorrt_llm/_torch/models/modeling_minimaxm3.py index e38b9f8be115..2416c13665df 100644 --- a/tensorrt_llm/_torch/models/modeling_minimaxm3.py +++ b/tensorrt_llm/_torch/models/modeling_minimaxm3.py @@ -35,7 +35,17 @@ from tensorrt_llm.models.modeling_utils import QuantConfig from ..attention_backend import AttentionMetadata -from ..attention_backend.interface import PositionalEmbeddingParams, RopeParams +from ..attention_backend.interface import ( + AttentionForwardArgs, + PositionalEmbeddingParams, + RopeParams, +) +from ..attention_backend.sparse.minimax_m3 import ( + MiniMaxM3MsaSparseAttention, + MiniMaxM3SparseRuntimeBackend, + _gather_paged_batched, + _write_main_kv_slots_to_pool, +) from ..distributed import AllReduce, AllReduceParams, MiniMaxAllReduceRMS from ..modules.attention import Attention from ..modules.decoder_layer import DecoderLayer @@ -911,11 +921,6 @@ def _dense_attention_core( output: torch.Tensor, ) -> torch.Tensor: """Run dense cache updates and attention into ``output``.""" - from ..attention_backend.sparse.minimax_m3 import ( - _gather_paged_batched, - _write_main_kv_slots_to_pool, - ) - kv_cache_manager = getattr(attn_metadata, "kv_cache_manager", None) if kv_cache_manager is None: raise RuntimeError( @@ -1117,11 +1122,27 @@ def _attention_core( attn_metadata: AttentionMetadata, output: torch.Tensor, ) -> torch.Tensor: - if self.is_sparse_attention_layer: - assert idx_q is not None and idx_k is not None - return self._sparse_attention_core(q, k, v, idx_q, idx_k, attn_metadata, output) - assert idx_q is None and idx_k is None - return self._dense_attention_core(q, k, v, attn_metadata, output) + # The MSA backend runs the sparse GQA or dense paged GQA through its + # inherited forward; this layer selects the top-k blocks (sparse only) + # and builds the forward_args the FMHA reads. + if isinstance(self.attn, MiniMaxM3MsaSparseAttention): + if self.is_sparse_attention_layer: + assert idx_q is not None and idx_k is not None + # Publish the selected blocks so the FMHA runs the sparse path. + kv_block_indexes = self.attn.run_indexer(idx_q, idx_k, attn_metadata) + forward_args = AttentionForwardArgs(output=output, topk_indices=kv_block_indexes) + else: + assert idx_q is None and idx_k is None + # No top-k selection means the FMHA attends the full page table. + forward_args = AttentionForwardArgs(output=output) + self.attn.forward(q, k, v, attn_metadata, forward_args=forward_args) + return output + else: + if self.is_sparse_attention_layer: + assert idx_q is not None and idx_k is not None + return self._sparse_attention_core(q, k, v, idx_q, idx_k, attn_metadata, output) + assert idx_q is None and idx_k is None + return self._dense_attention_core(q, k, v, attn_metadata, output) def _sparse_forward( self, @@ -1141,7 +1162,7 @@ def _sparse_forward( 4. Pull paged main K/V cache (reshaped to flat-slot view) and paged side index-K cache from the :class:`MiniMaxM3KVCacheManagerV2`. - 5. Build a :class:`MiniMaxM3SparseAttentionMetadata` from the + 5. Build a :class:`MiniMaxM3TritonSparseAttentionMetadata` from the standard :class:`AttentionMetadata` (using ``request_ids`` + ``seq_lens`` + ``num_cached_tokens_per_seq``). 6. Write the new token's K/V/idx_K to the slots named by the @@ -1154,7 +1175,7 @@ def _sparse_forward( Production callers (the LLM API path) drive :meth:`MiniMaxM3AttentionMetadata.prepare` outside any CUDA-graph capture window; that method attaches a pre-built - :class:`MiniMaxM3SparseAttentionMetadata` and an + :class:`MiniMaxM3TritonSparseAttentionMetadata` and an ``out_cache_loc`` tensor as ``attn_metadata.minimax_m3``. Test callers attach the same dict manually. This forward path always reads the pre-built attachment and never builds metadata @@ -1198,7 +1219,7 @@ def _sparse_attention_core( attn_metadata: AttentionMetadata, output: torch.Tensor, ) -> torch.Tensor: - """Run sparse cache updates and attention into ``output``.""" + """Run sparse cache updates and attention into ``output`` (Triton path).""" kv_cache_manager = getattr(attn_metadata, "kv_cache_manager", None) if kv_cache_manager is None: raise RuntimeError( @@ -1268,10 +1289,7 @@ def _sparse_attention_core( # registers :class:`MiniMaxM3SparseRuntimeBackend` as # ``self.attn``; any other backend on a sparse layer is a # configuration error. - from ..attention_backend.sparse.minimax_m3 import get_minimax_m3_attention_backend_cls - - m3_backend_cls = get_minimax_m3_attention_backend_cls() - if not isinstance(self.attn, m3_backend_cls): + if not isinstance(self.attn, MiniMaxM3SparseRuntimeBackend): raise RuntimeError( f"MiniMax-M3 sparse forward (layer {self.layer_idx}) requires " f"self.attn to be a MiniMaxM3SparseRuntimeBackend, got " diff --git a/tensorrt_llm/_torch/pyexecutor/config_utils.py b/tensorrt_llm/_torch/pyexecutor/config_utils.py index 5f1e0cbc0a26..601ab716dfcb 100644 --- a/tensorrt_llm/_torch/pyexecutor/config_utils.py +++ b/tensorrt_llm/_torch/pyexecutor/config_utils.py @@ -96,6 +96,11 @@ def is_mla(config): return False +def is_minimax_m3(sparse_attention_config): + """True when the sparse attention config selects the MiniMax-M3 algorithm.""" + return sparse_attention_config is not None and sparse_attention_config.algorithm == "minimax_m3" + + def is_qwen3_next(config): return hasattr( config, 'architectures' diff --git a/tensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.py b/tensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.py index de16d5c5000c..0bf7db14928e 100644 --- a/tensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.py +++ b/tensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.py @@ -1155,7 +1155,12 @@ def append_to_kv_heads_per_layer( self._log_kv_cache_pool_lifecycle_mapping() - def _prepare_page_table_tensor(self, index_mapper_capacity: int) -> None: + def _build_pool_mapping_tensors(self): + """Build the (kv_cache_pool_pointers, kv_cache_pool_mapping) tensors. + + An overridable hook for subclasses whose pools coalesce extra + per-layer buffers alongside K/V. + """ kv_cache_pool_pointers_list = [] kv_cache_pool_mapping_list = [] block_scale_pool_pointers_list = [] @@ -1253,18 +1258,22 @@ def _prepare_page_table_tensor(self, index_mapper_capacity: int) -> None: [pool_pointers[1], block_scale_pool_pointers[1]], ] - self.kv_cache_pool_pointers = torch.tensor( + kv_cache_pool_pointers = torch.tensor( kv_cache_pool_pointers_list, dtype=torch.int64, device="cpu", pin_memory=prefer_pinned(), ) - self.kv_cache_pool_mapping = torch.tensor( + kv_cache_pool_mapping = torch.tensor( kv_cache_pool_mapping_list, dtype=torch.int32, device="cpu", pin_memory=prefer_pinned(), ) + return kv_cache_pool_pointers, kv_cache_pool_mapping + + def _prepare_page_table_tensor(self, index_mapper_capacity: int) -> None: + self.kv_cache_pool_pointers, self.kv_cache_pool_mapping = self._build_pool_mapping_tensors() self.index_scales = torch.empty( self.num_pools, dtype=torch.int32, pin_memory=prefer_pinned(), device="cpu" ) diff --git a/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py b/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py index f82207881c8f..7920c8942e71 100644 --- a/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py +++ b/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py @@ -40,7 +40,7 @@ from ._util import (KvCacheCreator, _adjust_torch_mem_fraction, create_py_executor_instance, instantiate_sampler, is_mla, validate_feature_combination) -from .config_utils import is_hybrid_linear +from .config_utils import is_hybrid_linear, is_minimax_m3 from .connectors.kv_cache_connector import KvCacheConnectorManager from .dwdp import DwdpManager from .guided_decoder import CapturableGuidedDecoder, GuidedDecoder @@ -420,6 +420,13 @@ def create_py_executor( if llm_args.attn_backend == "VANILLA": tokens_per_block = max_num_tokens + # The MSA kernels require a page size of 128; the Triton reference uses TRT-LLM's default + # of 32. + m3_sparse_config = llm_args.sparse_attention_config + if is_minimax_m3(m3_sparse_config): + tokens_per_block = 128 if m3_sparse_config.implementation == "msa" else 32 + kv_cache_config.tokens_per_block = tokens_per_block + if llm_args.attn_backend in ["FLASHINFER", "FLASHINFER_STAR_ATTENTION"]: # Workaround for flashinfer and star attention if kv_cache_config.enable_block_reuse: diff --git a/tensorrt_llm/llmapi/llm_args.py b/tensorrt_llm/llmapi/llm_args.py index 220f3b540163..d06faa6e9c0d 100644 --- a/tensorrt_llm/llmapi/llm_args.py +++ b/tensorrt_llm/llmapi/llm_args.py @@ -680,20 +680,18 @@ class MiniMaxM3SparseAttentionConfig(BaseSparseAttentionConfig): Drives the two-step sparse attention used by MiniMax-M3 layers 3..N: - 1. An index attention branch projects a per-head Q vector and a - **single replicated** K vector, scores main K/V cache blocks, - and selects the top-``topk`` blocks per ``(num_kv_heads, q_token)`` - pair (with ``init_blocks`` forced at the head and ``local_blocks`` - forced at the tail). + 1. An index attention branch projects a per-head Q vector and a single + replicated K vector, scores main K/V cache blocks, and selects the + top-k blocks per (num_kv_heads, q_token) pair, with init_blocks forced + at the head and local_blocks forced at the tail. 2. A sparse GQA attention runs only over the selected blocks. - The selected backend at runtime uses - :class:`tensorrt_llm._torch.attention_backend.sparse.minimax_m3.MiniMaxM3SparseAttention` - on top of a :class:`MiniMaxM3KVCacheManagerV2` that allocates a - paged side index-K cache (``[num_slots, 1, sparse_index_dim]``) - parallel to the main K/V cache. The M3 checkpoint sets - ``disable_index_value=True`` on every sparse layer so no index V - cache is allocated for the bring-up. + At runtime one of the MiniMax-M3 sparse attention backends under + tensorrt_llm._torch.attention_backend.sparse.minimax_m3 is selected. The + chosen backend runs on top of a MiniMaxM3KVCacheManagerV2 that allocates a + paged side index-K cache of shape [num_slots, 1, sparse_index_dim] parallel + to the main K/V cache. The M3 checkpoint sets disable_index_value=True on + every sparse layer, so no index V cache is allocated. """ algorithm: Literal["minimax_m3"] = "minimax_m3" @@ -730,6 +728,34 @@ class MiniMaxM3SparseAttentionConfig(BaseSparseAttentionConfig): default=True, description="If True, skip the index V branch (M3 checkpoint default).", ) + num_attention_heads: Optional[int] = Field( + default=None, + description= + "Global number of attention (query) heads. When unset, it falls back " + "to pretrained_config.num_attention_heads.", + ) + num_key_value_heads: Optional[int] = Field( + default=None, + description= + "Global number of key/value heads. When unset, it falls back to " + "pretrained_config.num_key_value_heads, then to num_attention_heads.", + ) + implementation: Literal["triton", "msa"] = Field( + default="triton", + description= + "Sparse attention implementation: 'triton' reference (default) or 'msa' " + "(fmha_sm100 kernels). The 'msa' implementation requires an SM100 GPU, " + "the fmha_sm100 package, and sparse_block_size == 128.", + status="prototype", + ) + + @model_validator(mode="after") + def _validate_msa_block_size(self): + if self.implementation == "msa" and self.sparse_block_size != 128: + raise ValueError( + "MiniMax-M3 'msa' implementation requires sparse_block_size == " + f"128, got {self.sparse_block_size}.") + return self def supports_backend(self, backend: str) -> bool: return backend == "pytorch" @@ -738,7 +764,7 @@ def get_indices_block_size(self) -> int: return self.sparse_block_size def to_sparse_params(self, **kwargs): - from tensorrt_llm._torch.attention_backend.sparse.minimax_m3.metadata import \ + from tensorrt_llm._torch.attention_backend.sparse.minimax_m3.common import \ MiniMaxM3SparseParams return MiniMaxM3SparseParams( @@ -750,6 +776,36 @@ 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, + implementation=self.implementation, + ) + + def to_sparse_metadata_params(self, **kwargs): + """Lower into MiniMaxM3SparseMetadataParams for the attention metadata. + + Head counts resolve as this config, then pretrained_config, then a + default; num_key_value_heads falls back to num_attention_heads. Setting + them on the config lets tests skip building a pretrained_config. + """ + from tensorrt_llm._torch.attention_backend.sparse.minimax_m3.common import \ + MiniMaxM3SparseMetadataParams + + pretrained_config = kwargs.get("pretrained_config", None) + + def _value(name: str, default=None): + value = getattr(self, name) + if value is not None: + return value + if pretrained_config is not None: + return getattr(pretrained_config, name, default) + return default + + num_attention_heads = int(_value("num_attention_heads", 0)) + num_kv_heads = int(_value("num_key_value_heads", num_attention_heads)) + return MiniMaxM3SparseMetadataParams( + global_num_q_heads=num_attention_heads, + global_num_kv_heads=num_kv_heads, + num_index_heads=self.sparse_num_index_heads, + topk=self.sparse_topk_blocks, ) diff --git a/tensorrt_llm/usage/llm_args_golden_manifest.json b/tensorrt_llm/usage/llm_args_golden_manifest.json index 6f5b5385a66d..f5d0dd8925d2 100644 --- a/tensorrt_llm/usage/llm_args_golden_manifest.json +++ b/tensorrt_llm/usage/llm_args_golden_manifest.json @@ -1397,6 +1397,16 @@ "kind": "value", "path": "sparse_attention_config.enable_heuristic_topk" }, + { + "allowed_values": [ + "triton", + "msa" + ], + "annotation": "Literal['triton', 'msa']", + "converter": "", + "kind": "categorical", + "path": "sparse_attention_config.implementation" + }, { "allowed_values": [], "annotation": "Optional[int]", @@ -1459,6 +1469,20 @@ "kind": "categorical", "path": "sparse_attention_config.kt_cache_dtype" }, + { + "allowed_values": [], + "annotation": "Optional[int]", + "converter": "", + "kind": "value", + "path": "sparse_attention_config.num_attention_heads" + }, + { + "allowed_values": [], + "annotation": "Optional[int]", + "converter": "", + "kind": "value", + "path": "sparse_attention_config.num_key_value_heads" + }, { "allowed_values": [], "annotation": "Optional[int]", diff --git a/tests/integration/defs/accuracy/references/gsm8k.yaml b/tests/integration/defs/accuracy/references/gsm8k.yaml index afcd0c8136d8..ef32f74f1ba0 100644 --- a/tests/integration/defs/accuracy/references/gsm8k.yaml +++ b/tests/integration/defs/accuracy/references/gsm8k.yaml @@ -492,6 +492,9 @@ MiniMaxAI/MiniMax-M3-MXFP8: nvidia/MiniMax-M3-NVFP4: - quant_algo: MIXED_PRECISION accuracy: 88 + - quant_algo: MIXED_PRECISION + kv_cache_quant_algo: FP8 + accuracy: 86 nvidia/NVIDIA-Nemotron-Nano-9B-v2: - accuracy: 85.027 - quant_algo: FP8 diff --git a/tests/integration/defs/accuracy/references/mmlu.yaml b/tests/integration/defs/accuracy/references/mmlu.yaml index cbce62591221..a870fbba6402 100644 --- a/tests/integration/defs/accuracy/references/mmlu.yaml +++ b/tests/integration/defs/accuracy/references/mmlu.yaml @@ -284,6 +284,9 @@ MiniMaxAI/MiniMax-M3-MXFP8: nvidia/MiniMax-M3-NVFP4: - quant_algo: MIXED_PRECISION accuracy: 83 + - quant_algo: MIXED_PRECISION + kv_cache_quant_algo: FP8 + accuracy: 81 moonshotai/Kimi-K2-Instruct: - quant_algo: FP8_BLOCK_SCALES accuracy: 87.65 diff --git a/tests/integration/defs/accuracy/test_llm_api_pytorch.py b/tests/integration/defs/accuracy/test_llm_api_pytorch.py index 0bf4298efaad..77305811dd98 100644 --- a/tests/integration/defs/accuracy/test_llm_api_pytorch.py +++ b/tests/integration/defs/accuracy/test_llm_api_pytorch.py @@ -7683,11 +7683,12 @@ def test_auto_dtype(self, tp_size, ep_size): @pytest.mark.skip_less_device(4) @pytest.mark.skip_less_device_memory(140000) - @parametrize_with_ids("tp_size,ep_size", [(4, 4)]) - def test_mxfp8(self, tp_size, ep_size): + @parametrize_with_ids("use_msa", [False, True]) + def test_mxfp8(self, use_msa): # 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. + # BF16. + tp_size = ep_size = 4 model_name = "MiniMaxAI/MiniMax-M3-MXFP8" model_path = f"{llm_models_root()}/MiniMax-M3-MXFP8" # Halving TP from the BF16 reference (TP=8) doubles per-rank @@ -7696,7 +7697,8 @@ def test_mxfp8(self, tp_size, ep_size): # under the PyTorch cap. kv_cache_config = KvCacheConfig(free_gpu_memory_fraction=0.4, enable_block_reuse=False) - sparse_attention_config = MiniMaxM3SparseAttentionConfig() + sparse_attention_config = MiniMaxM3SparseAttentionConfig( + implementation="msa" if use_msa else "triton") with LLM(model_path, tensor_parallel_size=tp_size, moe_expert_parallel_size=ep_size, @@ -7712,15 +7714,18 @@ def test_mxfp8(self, tp_size, ep_size): task = GSM8K(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)]) - def test_mxfp8_piecewise_cuda_graph(self, tp_size, ep_size): + @parametrize_with_ids("use_msa", [False, True]) + def test_mxfp8_piecewise_cuda_graph(self, use_msa): + tp_size = ep_size = 4 model_name = "MiniMaxAI/MiniMax-M3-MXFP8" model_path = f"{llm_models_root()}/MiniMax-M3-MXFP8" kv_cache_config = KvCacheConfig(free_gpu_memory_fraction=0.5, - enable_block_reuse=False) - sparse_attention_config = MiniMaxM3SparseAttentionConfig() + enable_block_reuse=False, + dtype="fp8" if use_msa else "auto") + sparse_attention_config = MiniMaxM3SparseAttentionConfig( + implementation="msa" if use_msa else "triton") cuda_graph_config = CudaGraphConfig( enable_padding=True, batch_sizes=[1, 2, 4, 8, 12, 16, 24, 32]) torch_compile_config = TorchCompileConfig( @@ -7747,16 +7752,19 @@ def test_mxfp8_piecewise_cuda_graph(self, tp_size, ep_size): @pytest.mark.skip_less_device(4) @pytest.mark.skip_less_device_memory(140000) - @parametrize_with_ids("tp_size,ep_size", [(4, 4)]) - def test_nvfp4(self, tp_size, ep_size): + @parametrize_with_ids("use_msa", [False, True]) + def test_nvfp4(self, use_msa): # NVFP4 checkpoint: MXFP8 base layers with NVFP4 routed experts - # (MIXED_PRECISION checkpoint); the KV cache stays in BF16 and the - # sparse attention path is unchanged from BF16. + # (MIXED_PRECISION checkpoint). The MSA path runs an FP8 KV cache; the + # Triton path keeps the KV cache in BF16. + tp_size = ep_size = 4 model_name = "nvidia/MiniMax-M3-NVFP4" model_path = f"{llm_models_root()}/MiniMax-M3-NVFP4" kv_cache_config = KvCacheConfig(free_gpu_memory_fraction=0.6, - enable_block_reuse=False) - sparse_attention_config = MiniMaxM3SparseAttentionConfig() + enable_block_reuse=False, + dtype="fp8" if use_msa else "auto") + sparse_attention_config = MiniMaxM3SparseAttentionConfig( + implementation="msa" if use_msa else "triton") moe_config = MoeConfig(backend="CUTLASS") with LLM(model_path, tensor_parallel_size=tp_size, diff --git a/tests/integration/test_lists/qa/llm_function_core.txt b/tests/integration/test_lists/qa/llm_function_core.txt index 02c43b89f8dd..2a81bff62940 100644 --- a/tests/integration/test_lists/qa/llm_function_core.txt +++ b/tests/integration/test_lists/qa/llm_function_core.txt @@ -659,9 +659,10 @@ 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=4-ep_size=4] TIMEOUT (180) -accuracy/test_llm_api_pytorch.py::TestMiniMaxM3::test_mxfp8_piecewise_cuda_graph[tp_size=8-ep_size=8] TIMEOUT (180) -accuracy/test_llm_api_pytorch.py::TestMiniMaxM3::test_nvfp4[tp_size=4-ep_size=4] TIMEOUT (180) +accuracy/test_llm_api_pytorch.py::TestMiniMaxM3::test_mxfp8[use_msa=False] TIMEOUT (180) +accuracy/test_llm_api_pytorch.py::TestMiniMaxM3::test_mxfp8_piecewise_cuda_graph[use_msa=False] TIMEOUT (180) +accuracy/test_llm_api_pytorch.py::TestMiniMaxM3::test_nvfp4[use_msa=False] TIMEOUT (180) +accuracy/test_llm_api_pytorch.py::TestMiniMaxM3::test_nvfp4[use_msa=True] TIMEOUT (60) accuracy/test_llm_api_pytorch.py::TestMinistral8BInstruct::test_auto_dtype accuracy/test_llm_api_pytorch.py::TestMinistral8BInstruct::test_fp8 accuracy/test_llm_api_pytorch.py::TestMistralLarge3_675B::test_fp8[latency_moe_deepgemm] diff --git a/tests/integration/test_lists/test-db/l0_dgx_b200.yml b/tests/integration/test_lists/test-db/l0_dgx_b200.yml index cc96800d204a..0a8988c17edc 100644 --- a/tests/integration/test_lists/test-db/l0_dgx_b200.yml +++ b/tests/integration/test_lists/test-db/l0_dgx_b200.yml @@ -51,7 +51,7 @@ l0_dgx_b200: - disaggregated/test_disaggregated.py::test_disaggregated_deepseek_v3_lite_fp8_nixl[DeepSeek-V3-Lite-fp8] - disaggregated/test_disaggregated.py::test_disaggregated_gpt_oss_120b_harmony[gpt_oss/gpt-oss-120b] - accuracy/test_llm_api_pytorch.py::TestDeepSeekR1::test_nvfp4_multi_gpus[latency_adp_lmtp_tp4] - - accuracy/test_llm_api_pytorch.py::TestMiniMaxM2::test_4gpus[attention_dp=False-cuda_graph=True-overlap_scheduler=True-tp_size=4-ep_size=4] TIMEOUT (60) + - accuracy/test_llm_api_pytorch.py::TestMiniMaxM3::test_nvfp4[use_msa=True] TIMEOUT (60) - accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_bfloat16_4gpus[pp4-mtp_nextn=0-attention_dp=False-cuda_graph=False-overlap_scheduler=False-torch_compile=False] TIMEOUT (60) - unittest/_torch/modeling/test_modeling_deepseekv4.py - accuracy/test_llm_api_pytorch.py::TestDeepSeekV4Flash::test_auto_dtype TIMEOUT (60) @@ -236,7 +236,7 @@ l0_dgx_b200: - accuracy/test_llm_api_pytorch.py::TestDeepSeekV32::test_nvfp4_multi_gpus[fp4_indexer_dsl_mtp3] TIMEOUT (60) - accuracy/test_llm_api_pytorch.py::TestDeepSeekV32::test_nvfp4_multi_gpus[baseline_pp4_mtp1] TIMEOUT (60) - accuracy/test_llm_api_pytorch.py::TestDeepSeekV32::test_nvfp4_multi_gpus_chunked_prefill[baseline_fp8kv] TIMEOUT (60) - - accuracy/test_llm_api_pytorch.py::TestMiniMaxM3::test_mxfp8_piecewise_cuda_graph[tp_size=8-ep_size=8] TIMEOUT (180) + - accuracy/test_llm_api_pytorch.py::TestMiniMaxM3::test_mxfp8_piecewise_cuda_graph[use_msa=False] TIMEOUT (180) - accuracy/test_llm_api_pytorch.py::TestKimiK25::test_nvfp4[tp8_attn_dp] TIMEOUT (60) - accuracy/test_llm_api_pytorch.py::TestKimiK25::test_nvfp4[ep8] TIMEOUT (60) - accuracy/test_llm_api_pytorch.py::TestKimiK25::test_nvfp4[dep8] TIMEOUT (60) @@ -263,6 +263,7 @@ l0_dgx_b200: backend: pytorch orchestrator: mpi tests: + - accuracy/test_llm_api_pytorch.py::TestMiniMaxM2::test_4gpus[attention_dp=False-cuda_graph=True-overlap_scheduler=True-tp_size=4-ep_size=4] TIMEOUT (60) - accuracy/test_llm_api_pytorch.py::TestQwen3_5_35B_A3B::test_bf16[tp2-CUTLASS] - accuracy/test_llm_api_pytorch.py::TestLlama3_1_8BInstruct::test_bfloat16_4gpus[tp4-attn_backend=FLASHINFER-torch_compile=False] - accuracy/test_llm_api_pytorch.py::TestLlama3_1_8BInstruct::test_fp8_4gpus[tp4-fp8kv=False-attn_backend=FLASHINFER-torch_compile=False] diff --git a/tests/integration/test_lists/waives.txt b/tests/integration/test_lists/waives.txt index 7e4d467787cf..e4ab50f31a82 100644 --- a/tests/integration/test_lists/waives.txt +++ b/tests/integration/test_lists/waives.txt @@ -180,8 +180,8 @@ full:B200/accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_bfloat16[mt full:B200/accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_bfloat16[mtp_nextn=2-attention_dp=True-cuda_graph=True-overlap_scheduler=True-torch_compile=True-enable_chunked_prefill=True-v2_kv_cache=True] SKIP (https://nvbugs/6422343) full:B200/accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_cute_dsl_nvfp4_4gpus[tp4-fp8kv=False-attention_dp=False-cuda_graph=False-overlap_scheduler=False-torch_compile=True] SKIP (https://nvbugs/6474888) full:B200/accuracy/test_llm_api_pytorch.py::TestMiniMaxM3::test_auto_dtype[tp_size=8-ep_size=8] SKIP (https://nvbugs/6384747) -full:B200/accuracy/test_llm_api_pytorch.py::TestMiniMaxM3::test_mxfp8[tp_size=4-ep_size=4] SKIP (https://nvbugs/6424188) -full:B200/accuracy/test_llm_api_pytorch.py::TestMiniMaxM3::test_nvfp4[tp_size=4-ep_size=4] SKIP (https://nvbugs/6424188) +full:B200/accuracy/test_llm_api_pytorch.py::TestMiniMaxM3::test_mxfp8[use_msa=False] SKIP (https://nvbugs/6424188) +full:B200/accuracy/test_llm_api_pytorch.py::TestMiniMaxM3::test_nvfp4[use_msa=False] SKIP (https://nvbugs/6424188) full:B200/disaggregated/test_disaggregated.py::test_disaggregated_stress_test[input8k-output1k-conc512-deepseek_r1_v2_fp4_stress] SKIP (https://nvbugs/6472256) full:B200/disaggregated/test_disaggregated.py::test_disaggregated_stress_test[input8k-output1k-conc512-qwen3_32b_fp8_stress] SKIP (https://nvbugs/6472256) full:B200/llmapi/test_llm_api_pytorch_moe_lora.py::test_qwen_moe_routed_expert_multi_lora_varying_ranks[cudagraph] SKIP (https://nvbugs/6475623) @@ -206,9 +206,8 @@ full:B300/accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_bfloat16[mt full:B300/accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_cute_dsl_nvfp4_4gpus[tp4-fp8kv=False-attention_dp=False-cuda_graph=False-overlap_scheduler=False-torch_compile=True] SKIP (https://nvbugs/6474888) full:B300/accuracy/test_llm_api_pytorch.py::TestGemma3_1BInstruct::test_fp8_prequantized[torch_compile=True] SKIP (https://nvbugs/6475346) full:B300/accuracy/test_llm_api_pytorch.py::TestMiniMaxM3::test_auto_dtype[tp_size=8-ep_size=8] SKIP (https://nvbugs/6445375) -full:B300/accuracy/test_llm_api_pytorch.py::TestMiniMaxM3::test_mxfp8[tp_size=4-ep_size=4] SKIP (https://nvbugs/6424188) -full:B300/accuracy/test_llm_api_pytorch.py::TestMiniMaxM3::test_mxfp8_piecewise_cuda_graph[tp_size=8-ep_size=8] SKIP (https://nvbugs/6442594) -full:B300/accuracy/test_llm_api_pytorch.py::TestMiniMaxM3::test_nvfp4[tp_size=4-ep_size=4] SKIP (https://nvbugs/6445375) +full:B300/accuracy/test_llm_api_pytorch.py::TestMiniMaxM3::test_mxfp8[use_msa=False] SKIP (https://nvbugs/6424188) +full:B300/accuracy/test_llm_api_pytorch.py::TestMiniMaxM3::test_nvfp4[use_msa=False] SKIP (https://nvbugs/6445375) full:B300/accuracy/test_llm_api_pytorch.py::TestNemotronV3Ultra::test_nvfp4_8gpus[attention_dp_off-trtllm] SKIP (https://nvbugs/6474894) full:B300/disaggregated/test_auto_scaling.py::test_service_discovery[http-kv_cache_aware] SKIP (https://nvbugs/6474892) full:B300/disaggregated/test_disaggregated.py::test_disaggregated_ctxpp2_genpp2[TinyLlama-1.1B-Chat-v1.0] SKIP (https://nvbugs/6322073) @@ -221,8 +220,8 @@ full:GB200/accuracy/test_dwdp_disaggregated_serving.py::TestDwdpDeepSeekV3Lite:: full:GB200/accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_bfloat16[mtp_nextn=2-attention_dp=True-cuda_graph=True-overlap_scheduler=True-torch_compile=True-enable_chunked_prefill=True-v2_kv_cache=False] SKIP (https://nvbugs/6400067) full:GB200/accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_bfloat16[mtp_nextn=2-attention_dp=True-cuda_graph=True-overlap_scheduler=True-torch_compile=True-enable_chunked_prefill=True-v2_kv_cache=True] SKIP (https://nvbugs/6400067) full:GB200/accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_cute_dsl_nvfp4_4gpus[tp4-fp8kv=False-attention_dp=False-cuda_graph=False-overlap_scheduler=False-torch_compile=True] SKIP (https://nvbugs/6474888) -full:GB200/accuracy/test_llm_api_pytorch.py::TestMiniMaxM3::test_mxfp8[tp_size=4-ep_size=4] SKIP (https://nvbugs/6479471) -full:GB200/accuracy/test_llm_api_pytorch.py::TestMiniMaxM3::test_nvfp4[tp_size=4-ep_size=4] SKIP (https://nvbugs/6479471) +full:GB200/accuracy/test_llm_api_pytorch.py::TestMiniMaxM3::test_mxfp8[use_msa=False] SKIP (https://nvbugs/6479471) +full:GB200/accuracy/test_llm_api_pytorch.py::TestMiniMaxM3::test_nvfp4[use_msa=False] SKIP (https://nvbugs/6479471) full:GB200/accuracy/test_llm_api_pytorch.py::TestNemotronV3Ultra::test_nvfp4_4gpus_block_reuse[ADP4] SKIP (https://nvbugs/6474894) full:GB200/accuracy/test_llm_api_pytorch.py::TestQwen3_5_35B_A3B::test_fp8_moe_dflash SKIP (https://nvbugs/6316985) full:GB200/accuracy/test_llm_api_pytorch.py::TestQwen3_5_397B_A17B::test_nvfp4_4gpus_online_eplb[moe_backend=TRTLLM] SKIP (https://nvbugs/6462928) @@ -244,8 +243,8 @@ full:GB300/accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_nvfp4_4gpu full:GB300/accuracy/test_llm_api_pytorch.py::TestGemma3_1BInstruct::test_fp8_prequantized[torch_compile=True] SKIP (https://nvbugs/6475346) full:GB300/accuracy/test_llm_api_pytorch.py::TestLlama3_1_8BInstruct::test_bfloat16_4gpus[pp4-attn_backend=FLASHINFER-torch_compile=False] SKIP (https://nvbugs/6385771) full:GB300/accuracy/test_llm_api_pytorch.py::TestLlama3_1_8BInstruct::test_fp8_4gpus[pp4-fp8kv=True-attn_backend=FLASHINFER-torch_compile=False] SKIP (https://nvbugs/6385771) -full:GB300/accuracy/test_llm_api_pytorch.py::TestMiniMaxM3::test_mxfp8[tp_size=4-ep_size=4] SKIP (https://nvbugs/6422502) -full:GB300/accuracy/test_llm_api_pytorch.py::TestMiniMaxM3::test_nvfp4[tp_size=4-ep_size=4] SKIP (https://nvbugs/6422502) +full:GB300/accuracy/test_llm_api_pytorch.py::TestMiniMaxM3::test_mxfp8[use_msa=False] SKIP (https://nvbugs/6422502) +full:GB300/accuracy/test_llm_api_pytorch.py::TestMiniMaxM3::test_nvfp4[use_msa=False] SKIP (https://nvbugs/6422502) full:GB300/accuracy/test_llm_api_pytorch.py::TestQwen3_5_35B_A3B::test_fp8_moe_dflash SKIP (https://nvbugs/6316985) full:GB300/accuracy/test_llm_api_pytorch_multimodal.py::TestGemma3_27BInstruct::test_fp8_prequantized SKIP (https://nvbugs/6479708) full:GB300/test_e2e.py::test_qwen_e2e_cpprunner_large_new_tokens[DeepSeek-R1-Distill-Qwen-1.5B-DeepSeek-R1-Distill-Qwen-1.5B] SKIP (https://nvbugs/6414760) @@ -422,6 +421,8 @@ unittest/_torch/visual_gen/multi_gpu/test_ulysses_async.py::test_capture_smoke S unittest/_torch/visual_gen/test_attention_integration.py::test_sage_attention_self_attention[fp8-2-1560] SKIP (https://nvbugs/6198760) unittest/_torch/visual_gen/test_attention_integration.py::test_sage_attention_self_attention[int8-2-1560] SKIP (https://nvbugs/6198760) unittest/auto_deploy/multigpu/custom_ops SKIP (https://nvbugs/6403920) +unittest/disaggregated/test_kv_transfer.py SKIP (https://nvbugs/6403793) +unittest/disaggregated/test_kv_transfer.py::test_transfer_worker_v2[tp1_pp4_to_tp2_pp2] SKIP (https://nvbugs/6445316) unittest/disaggregated/test_kv_transfer.py::test_transfer_worker_v2[tp4_pp1_to_tp2_pp2] SKIP (https://nvbugs/6426834) unittest/executor/test_rpc.py::TestRpcCorrectness::test_incremental_task_async SKIP (https://nvbugs/5741476) unittest/executor/test_rpc_proxy.py SKIP (https://nvbugs/5605741) 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..58d651631608 --- /dev/null +++ b/tests/unittest/_torch/attention/sparse/test_minimax_m3_msa_backend.py @@ -0,0 +1,76 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Structural tests for the MiniMax-M3 MSA sparse attention backend. + +These validate backend selection and decode scratch-buffer sizing without +launching kernels. Numerical parity against the Triton reference is covered +by the SM100 integration accuracy test. +""" + +import pytest +import torch + +from tensorrt_llm._torch.attention_backend.sparse.minimax_m3 import MiniMaxM3MsaSparseAttention +from tensorrt_llm._torch.attention_backend.sparse.utils import _resolve_minimax_m3_backend_cls +from tensorrt_llm.llmapi.llm_args import MiniMaxM3SparseAttentionConfig + + +def test_resolver_selects_msa_backend_when_available(monkeypatch): + import tensorrt_llm._torch.attention_backend.sparse.minimax_m3.msa_availability as avail + + monkeypatch.setattr(avail, "ensure_msa_available", lambda: None) + params = MiniMaxM3SparseAttentionConfig(implementation="msa").to_sparse_params() + assert _resolve_minimax_m3_backend_cls(params) is MiniMaxM3MsaSparseAttention + + +def test_msa_requires_block_size_128(): + # The MSA implementation is fixed to a 128-token page size; a mismatched + # sparse_block_size must fail loudly at config construction rather than being + # silently overridden at runtime. + with pytest.raises(ValueError, match=r"sparse_block_size == 128"): + MiniMaxM3SparseAttentionConfig(implementation="msa", sparse_block_size=64) + + # The Triton reference is unaffected by the constraint. + cfg = MiniMaxM3SparseAttentionConfig(implementation="triton", sparse_block_size=64) + assert cfg.sparse_block_size == 64 + + +def test_msa_metadata_rejects_undersized_max_score_buffer(): + metadata_cls = MiniMaxM3MsaSparseAttention.Metadata + metadata = metadata_cls.__new__(metadata_cls) + # Flat backing store sized for 4 heads * 8 k-tiles * 2 batch = 64 elements, + # too small for the plan's required 4 * 16 * 2 = 128. + metadata.msa_max_score = torch.zeros(4 * 8 * 2) + metadata.kv_cache_manager = None + + with pytest.raises(ValueError, match=r"msa_max_score backing store"): + metadata._ensure_msa_decode_scratch_buffers( + num_index_heads=4, + max_batch=2, + capture_graph=False, + required_max_k_tiles=16, + ) + + +def test_msa_proxy_max_score_view_is_contiguous_over_stable_store(): + """The proxy view fed to fmha_sm100 must be contiguous in the exact + [num_index_heads, plan_max_k_tiles, num_tokens] shape the kernel writes, + backed by a stable store so its data_ptr survives CUDA graph replay. + """ + metadata_cls = MiniMaxM3MsaSparseAttention.Metadata + metadata = metadata_cls.__new__(metadata_cls) + # Worst-case store: 4 heads * 16 k-tiles * 8 batch. + num_index_heads, worst_k, max_batch = 4, 16, 8 + metadata.msa_max_score = torch.zeros(num_index_heads * worst_k * max_batch) + store_ptr = metadata.msa_max_score.data_ptr() + + # A smaller live step still yields a contiguous view sized to that step, + # which is what the kernel's stride-agnostic write requires. + view = metadata.msa_proxy_max_score_view(num_index_heads, 5, 3) + assert view.shape == (num_index_heads, 5, 3) + assert view.is_contiguous() + assert view.data_ptr() == store_ptr + + # Oversized requests are rejected rather than silently corrupting memory. + with pytest.raises(ValueError, match=r"msa_max_score backing store"): + metadata.msa_proxy_max_score_view(num_index_heads, worst_k, max_batch + 1) diff --git a/tests/unittest/_torch/executor/test_py_executor_creator_mla_cache_reuse_sync.py b/tests/unittest/_torch/executor/test_py_executor_creator_mla_cache_reuse_sync.py index 15cab5d360c0..c2a9a56a87bb 100644 --- a/tests/unittest/_torch/executor/test_py_executor_creator_mla_cache_reuse_sync.py +++ b/tests/unittest/_torch/executor/test_py_executor_creator_mla_cache_reuse_sync.py @@ -182,6 +182,7 @@ def _make_llm_args(): trust_remote_code=False, mm_encoder_only=False, enable_chunked_prefill=False, + sparse_attention_config=None, attn_backend="TRTLLM", speculative_config=None, disable_overlap_scheduler=True,