[TRTLLM-14019][feat] Add MiniMax-M3 MSA sparse attention backend [revised]

.gitmodules

@@ -0,0 +1,3 @@
+[submodule "3rdparty/MSA"]
+	path = 3rdparty/MSA
+	url = https://gitlab.com/nvidia/tensorrt-llm/oss-components/msa.git

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
 --------------------------------------------------------------------------------

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"

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

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)
 

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")

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) {

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 \

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

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

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={

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",
 ]