From 334db275ddce00773c5d9231c080170242558787 Mon Sep 17 00:00:00 2001 From: peihengh <259410613+peihu-nv@users.noreply.github.com> Date: Fri, 17 Jul 2026 00:04:56 -0700 Subject: [PATCH 1/5] [None][perf] Avoid paged MSA K/V materialization during prefill Signed-off-by: peihengh <259410613+peihu-nv@users.noreply.github.com> --- 3rdparty/patches/msa_strided_paged_kv.patch | 184 ++++++++++++++++++++ scripts/build_wheel.py | 41 ++++- setup.py | 5 +- tests/unittest/scripts/test_build_wheel.py | 52 ++++++ 4 files changed, 278 insertions(+), 4 deletions(-) create mode 100644 3rdparty/patches/msa_strided_paged_kv.patch create mode 100644 tests/unittest/scripts/test_build_wheel.py diff --git a/3rdparty/patches/msa_strided_paged_kv.patch b/3rdparty/patches/msa_strided_paged_kv.patch new file mode 100644 index 000000000000..743587c5608b --- /dev/null +++ b/3rdparty/patches/msa_strided_paged_kv.patch @@ -0,0 +1,184 @@ +diff --git a/python/fmha_sm100/cute/interface.py b/python/fmha_sm100/cute/interface.py +index d72b17a..e27a74f 100644 +--- a/python/fmha_sm100/cute/interface.py ++++ b/python/fmha_sm100/cute/interface.py +@@ -136,6 +136,32 @@ def _prepare_paged_kv_for_tma(k, v, blk_kv: int): + return k, v + + ++def _prepare_paged_hnd_input(tensor: torch.Tensor, blk_kv: int) -> torch.Tensor: ++ """Keep TMA-compatible paged HND views strided across physical pages. ++ ++ The sparse prefill kernel imports the runtime tensor strides through ++ DLPack. It requires each ``[page_size, head_dim]`` head plane to be packed, ++ but the outer physical-page stride may include other coalesced cache roles. ++ Materialize every other layout to preserve the public API's old contract. ++ """ ++ if tensor.ndim != 4 or int(tensor.shape[2]) != int(blk_kv): ++ return tensor.contiguous() ++ ++ head_dim = int(tensor.shape[3]) ++ page_size = int(tensor.shape[2]) ++ packed_within_page = ( ++ tensor.stride(3) == 1 ++ and tensor.stride(2) == head_dim ++ and tensor.stride(1) == page_size * head_dim ++ ) ++ alignment_bytes = 16 ++ aligned_for_tma = ( ++ tensor.data_ptr() % alignment_bytes == 0 ++ and tensor.stride(0) * tensor.element_size() % alignment_bytes == 0 ++ ) ++ return tensor if packed_within_page and aligned_for_tma else tensor.contiguous() ++ ++ + def _validate_cu_seqlens( + cu_seqlens: torch.Tensor, + *, +@@ -736,10 +762,21 @@ def sparse_atten_func( + max_seqlen_q = int(max_seqlen_q) + max_seqlen_k = int(max_seqlen_k) + ++ k_input = ( ++ k.contiguous() ++ if page_table is None ++ else _prepare_paged_hnd_input(k, blk_kv) ++ ) ++ v_input = ( ++ v.contiguous() ++ if page_table is None ++ else _prepare_paged_hnd_input(v, blk_kv) ++ ) ++ + return _sparse_atten_csr_varlen_forward( + q.contiguous(), +- k.contiguous(), +- v.contiguous(), ++ k_input, ++ v_input, + k2q_row_ptr.contiguous(), + k2q_q_indices.contiguous(), + int(topK), +diff --git a/python/fmha_sm100/cute/test_sparse_atten.py b/python/fmha_sm100/cute/test_sparse_atten.py +index 21c777e..c22beef 100644 +--- a/python/fmha_sm100/cute/test_sparse_atten.py ++++ b/python/fmha_sm100/cute/test_sparse_atten.py +@@ -61,6 +61,43 @@ DECODE_DIM = 128 + DECODE_KV_TOKEN_SWEEP = tuple(2**exp for exp in range(3, 21)) + + ++def test_prepare_paged_hnd_input_keeps_aligned_outer_page_stride() -> None: ++ pages, roles, heads, page_size, head_dim = 5, 3, 2, 128, 128 ++ pool = torch.empty( ++ pages, ++ roles, ++ heads, ++ page_size, ++ head_dim, ++ dtype=torch.float8_e4m3fn, ++ device="cuda", ++ ) ++ view = pool[:, 1] ++ ++ assert not view.is_contiguous() ++ prepared = sparse_interface._prepare_paged_hnd_input(view, page_size) ++ assert prepared.data_ptr() == view.data_ptr() ++ assert prepared.stride() == view.stride() ++ ++ ++def test_prepare_paged_hnd_input_materializes_unpacked_tokens() -> None: ++ pages, heads, page_size, head_dim = 5, 2, 128, 128 ++ storage = torch.empty( ++ pages, ++ heads, ++ page_size * 2, ++ head_dim, ++ dtype=torch.float8_e4m3fn, ++ device="cuda", ++ ) ++ view = storage[:, :, ::2, :] ++ ++ prepared = sparse_interface._prepare_paged_hnd_input(view, page_size) ++ assert prepared.is_contiguous() ++ assert prepared.data_ptr() != view.data_ptr() ++ torch.testing.assert_close(prepared, view, rtol=0, atol=0) ++ ++ + @contextmanager + def _nvtx_range(message: str): + torch.cuda.nvtx.range_push(message) +@@ -1786,6 +1823,74 @@ def test_sparse_page_atten( + + _assert_forward_close(out, out_ref, out_pt.float(), lse, lse_ref) + ++ ++def test_sparse_page_atten_strided_outer_page_matches_packed() -> None: ++ inputs = _build_paged_inputs( ++ batch=1, ++ seqlen_q=2048, ++ seqlen_kv=2048, ++ head_kv=2, ++ qhead_per_kv=16, ++ dim=128, ++ topk=16, ++ blk_kv=128, ++ causal=True, ++ page_size=128, ++ seqused_trim=0, ++ dtype=torch.float8_e4m3fn, ++ ) ++ k_packed = inputs["k_paged"].detach().clone() ++ v_packed = inputs["v_paged"].detach().clone() ++ pool = torch.empty( ++ k_packed.shape[0], ++ 4, ++ *k_packed.shape[1:], ++ dtype=k_packed.dtype, ++ device=k_packed.device, ++ ) ++ k_strided = pool[:, 1] ++ v_strided = pool[:, 3] ++ k_strided.copy_(k_packed) ++ v_strided.copy_(v_packed) ++ ++ assert not k_strided.is_contiguous() ++ assert not v_strided.is_contiguous() ++ assert ( ++ sparse_interface._prepare_paged_hnd_input(k_strided, inputs["blk_kv"]).data_ptr() ++ == k_strided.data_ptr() ++ ) ++ assert ( ++ sparse_interface._prepare_paged_hnd_input(v_strided, inputs["blk_kv"]).data_ptr() ++ == v_strided.data_ptr() ++ ) ++ ++ def run(k: torch.Tensor, v: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: ++ return sparse_atten_func( ++ inputs["q"], ++ k, ++ v, ++ inputs["k2q_row_ptr"], ++ inputs["k2q_q_indices"], ++ 16, ++ blk_kv=inputs["blk_kv"], ++ causal=True, ++ softmax_scale=inputs["softmax_scale"], ++ return_softmax_lse=True, ++ cu_seqlens_q=inputs["cu_seqlens_q"], ++ cu_seqlens_k=inputs["cu_seqlens_k"], ++ max_seqlen_q=inputs["max_seqlen_q"], ++ max_seqlen_k=inputs["max_seqlen_k"], ++ page_table=inputs["page_table"], ++ seqused_k=inputs["seqused_k"], ++ schedule=inputs["schedule"], ++ ) ++ ++ packed_out, packed_lse = run(k_packed, v_packed) ++ strided_out, strided_lse = run(k_strided, v_strided) ++ torch.testing.assert_close(strided_out, packed_out, rtol=0, atol=0) ++ torch.testing.assert_close(strided_lse, packed_lse, rtol=0, atol=0) ++ ++ + @pytest.mark.parametrize("paged", [False, True]) + @pytest.mark.parametrize("causal", [True]) + @pytest.mark.parametrize("batch", [3]) diff --git a/scripts/build_wheel.py b/scripts/build_wheel.py index 462b74d8e511..4e7f8a9d4407 100755 --- a/scripts/build_wheel.py +++ b/scripts/build_wheel.py @@ -80,6 +80,37 @@ def get_build_dir(build_dir, build_type): return build_dir +def stage_msa_package(project_dir: Path, build_dir: Path) -> Path: + """Copy pinned MSA sources and apply TensorRT-LLM's downstream patch.""" + msa_source_dir = project_dir / "3rdparty" / "MSA" + msa_package_dir = msa_source_dir / "python" / "fmha_sm100" + msa_patch = project_dir / "3rdparty" / "patches" / "msa_strided_paged_kv.patch" + if not msa_package_dir.is_dir(): + raise FileNotFoundError( + f"MSA sources are missing at {msa_package_dir}; initialize 3rdparty/MSA" + ) + + staging_dir = build_dir / "msa_patched" + if staging_dir.exists(): + rmtree(staging_dir) + copytree(msa_source_dir, staging_dir, ignore=shutil.ignore_patterns(".git")) + git_env = os.environ.copy() + git_env["GIT_CEILING_DIRECTORIES"] = str(staging_dir.parent.resolve()) + run( + ["git", "apply", "--check", str(msa_patch)], + cwd=staging_dir, + env=git_env, + check=True, + ) + run( + ["git", "apply", str(msa_patch)], + cwd=staging_dir, + env=git_env, + check=True, + ) + return staging_dir / "python" / "fmha_sm100" + + def clear_folder(folder_path): for item in os.listdir(folder_path): item_path = os.path.join(folder_path, item) @@ -1143,10 +1174,14 @@ def get_binding_lib(subdirectory, name): f"Copied auto-generated attributions to {project_dir / 'ATTRIBUTIONS.md'}" ) + msa_package_dir = stage_msa_package(project_dir, build_dir) + wheel_env = os.environ.copy() + wheel_env["TRTLLM_MSA_PACKAGE_DIR"] = str(msa_package_dir) + build_run( - f'\"{venv_python}\" -m build {project_dir} --skip-dependency-check {extra_wheel_build_args} --no-isolation --wheel --outdir "{dist_dir}"' - ) - env = os.environ.copy() + f'\"{venv_python}\" -m build {project_dir} --skip-dependency-check {extra_wheel_build_args} --no-isolation --wheel --outdir "{dist_dir}"', + env=wheel_env) + env = wheel_env.copy() if mypyc: env["TRTLLM_ENABLE_MYPYC"] = "1" else: diff --git a/setup.py b/setup.py index e71819c66fb9..8865078b3598 100644 --- a/setup.py +++ b/setup.py @@ -421,7 +421,10 @@ 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"} +msa_package_dir = { + "fmha_sm100": + os.environ.get("TRTLLM_MSA_PACKAGE_DIR", "3rdparty/MSA/python/fmha_sm100") +} packages += ["fmha_sm100"] # https://setuptools.pypa.io/en/latest/references/keywords.html diff --git a/tests/unittest/scripts/test_build_wheel.py b/tests/unittest/scripts/test_build_wheel.py new file mode 100644 index 000000000000..0b9061efd34d --- /dev/null +++ b/tests/unittest/scripts/test_build_wheel.py @@ -0,0 +1,52 @@ +#!/usr/bin/env python3 +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import importlib.util +from pathlib import Path + +import pytest + +REPO_ROOT = Path(__file__).resolve().parent.parent.parent.parent +SCRIPT_PATH = REPO_ROOT / "scripts" / "build_wheel.py" +MSA_INTERFACE = Path("python/fmha_sm100/cute/interface.py") + +_SPEC = importlib.util.spec_from_file_location("build_wheel", SCRIPT_PATH) +assert _SPEC is not None and _SPEC.loader is not None +_BUILD_WHEEL = importlib.util.module_from_spec(_SPEC) +_SPEC.loader.exec_module(_BUILD_WHEEL) +stage_msa_package = _BUILD_WHEEL.stage_msa_package + + +def test_stage_msa_package_applies_patch_without_modifying_submodule(tmp_path): + source_interface = REPO_ROOT / "3rdparty" / "MSA" / MSA_INTERFACE + if not source_interface.is_file(): + pytest.skip("3rdparty/MSA is not initialized") + + source_before = source_interface.read_bytes() + staged_package = stage_msa_package(REPO_ROOT, tmp_path) + staged_interface = staged_package / "cute" / "interface.py" + + assert b"def _prepare_paged_hnd_input" not in source_before + assert "def _prepare_paged_hnd_input" in staged_interface.read_text() + assert source_interface.read_bytes() == source_before + + +def test_stage_msa_package_requires_initialized_submodule(tmp_path): + project_dir = tmp_path / "project" + project_dir.mkdir() + + with pytest.raises(FileNotFoundError, match="initialize 3rdparty/MSA"): + stage_msa_package(project_dir, tmp_path / "build") From 5a38f2c867647417698a1bb8e686a995da9c86ae Mon Sep 17 00:00:00 2001 From: Tyler Burt <195370667+tburt-nv@users.noreply.github.com> Date: Tue, 28 Jul 2026 09:52:02 -0700 Subject: [PATCH 2/5] use fetch_content for patching Signed-off-by: Tyler Burt <195370667+tburt-nv@users.noreply.github.com> --- .gitignore | 1 + .gitmodules | 3 - 3rdparty/CMakeLists.txt | 2 +- 3rdparty/MSA | 1 - 3rdparty/fetch_content.json | 7 +++ cpp/CMakeLists.txt | 2 +- jenkins/Build.groovy | 4 +- jenkins/BuildDockerImage.groovy | 2 +- jenkins/L0_MergeRequest.groovy | 8 +-- jenkins/L0_Test.groovy | 6 +- jenkins/TensorRT_LLM_PLC.groovy | 2 +- jenkins/runPerfSanityTriage.groovy | 2 +- scripts/attribution/scan/metadata/msa.yml | 4 +- scripts/build_wheel.py | 62 +++++-------------- setup.py | 32 ++++++++-- .../attention_backend/fmha/msa_sparse_gqa.py | 4 +- .../sparse/minimax_m3/msa_availability.py | 11 ++-- .../sparse/minimax_m3/msa_utils.py | 40 ++---------- tests/unittest/scripts/test_build_wheel.py | 52 ---------------- 19 files changed, 77 insertions(+), 168 deletions(-) delete mode 100644 .gitmodules delete mode 160000 3rdparty/MSA delete mode 100644 tests/unittest/scripts/test_build_wheel.py diff --git a/.gitignore b/.gitignore index 9a9332335ca5..bfef25aae0d7 100644 --- a/.gitignore +++ b/.gitignore @@ -56,6 +56,7 @@ tensorrt_llm/pg_utils_bindings.*.so tensorrt_llm/flash_mla/ tensorrt_llm/flash_mla_cpp_tllm.*.so tensorrt_llm/flash_mla_cpp_tllm.pyi +/3rdparty/fmha_sm100/ tensorrt_llm/runtime/kv_cache_manager_v2/**/*.so **/*__mypyc*.so tensorrt_llm/scripts diff --git a/.gitmodules b/.gitmodules deleted file mode 100644 index 627760b34da2..000000000000 --- a/.gitmodules +++ /dev/null @@ -1,3 +0,0 @@ -[submodule "3rdparty/MSA"] - path = 3rdparty/MSA - url = https://gitlab.com/nvidia/tensorrt-llm/oss-components/msa.git diff --git a/3rdparty/CMakeLists.txt b/3rdparty/CMakeLists.txt index b9982a166b52..02a265608d88 100644 --- a/3rdparty/CMakeLists.txt +++ b/3rdparty/CMakeLists.txt @@ -135,7 +135,7 @@ foreach(DEP_IDX RANGE ${DEP_COUNT_MINUS_ONE}) PATCH_COMMAND bash -c - "patch -p1 --forward --batch --dry-run -i '${_patch_file}' && patch -p1 --forward --batch -i '${_patch_file}' || echo 'Patch already applied, skipping.'" + "patch -p1 --forward --batch --dry-run -i '${_patch_file}' && patch -p1 --forward --batch -i '${_patch_file}' || patch -p1 --reverse --batch --dry-run -i '${_patch_file}'" ) endif() diff --git a/3rdparty/MSA b/3rdparty/MSA deleted file mode 160000 index e2ebe7656649..000000000000 --- a/3rdparty/MSA +++ /dev/null @@ -1 +0,0 @@ -Subproject commit e2ebe7656649f619af0ad1d457b534283034655e diff --git a/3rdparty/fetch_content.json b/3rdparty/fetch_content.json index d86ebb57fc46..f1e2c1622f2e 100644 --- a/3rdparty/fetch_content.json +++ b/3rdparty/fetch_content.json @@ -68,6 +68,13 @@ "git_shallow": true, "source_subdir": "dont-add-this-project-with-add-subdirectory" }, + { + "name": "msa", + "git_repository": "https://gitlab.com/nvidia/tensorrt-llm/oss-components/msa.git", + "git_tag": "e2ebe7656649f619af0ad1d457b534283034655e", + "source_subdir": "dont-add-this-project-with-add-subdirectory", + "patch_file": "patches/msa_strided_paged_kv.patch" + }, { "name": "nanobind", "git_repository": "https://github.com/wjakob/nanobind", diff --git a/cpp/CMakeLists.txt b/cpp/CMakeLists.txt index 5dec92599ae6..316edb9c766b 100644 --- a/cpp/CMakeLists.txt +++ b/cpp/CMakeLists.txt @@ -258,7 +258,7 @@ endif() FetchContent_MakeAvailable(nanobind) include_directories(${CMAKE_BINARY_DIR}/_deps/nanobind-src/include) -FetchContent_MakeAvailable(cutlass cxxopts flashmla json xgrammar) +FetchContent_MakeAvailable(cutlass cxxopts flashmla json msa xgrammar) if(ENABLE_UCX) FetchContent_MakeAvailable(cppzmq ucxx) diff --git a/jenkins/Build.groovy b/jenkins/Build.groovy index b1d541c04b13..da963453cd52 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, true, true) + trtllm_utils.checkoutSource(LLM_REPO, env.gitlabCommit, LLM_ROOT, false, 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, true, true) + trtllm_utils.checkoutSource(LLM_REPO, env.gitlabCommit, LLM_ROOT, false, 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 d22c9dc23c69..ce94eb1ece1c 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, true, true) + trtllm_utils.checkoutSource(LLM_REPO, LLM_COMMIT_OR_BRANCH, LLM_ROOT, false, true) } // Step 2: Build the images diff --git a/jenkins/L0_MergeRequest.groovy b/jenkins/L0_MergeRequest.groovy index 3f276522ee54..89aead5f4fc0 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, true, true) + trtllm_utils.checkoutSource(LLM_REPO, env.gitlabCommit, LLM_ROOT, false, true) } else { branch = env.gitlabBranch ? env.gitlabBranch : "main" - trtllm_utils.checkoutSource(LLM_REPO, branch, LLM_ROOT, true, true) + trtllm_utils.checkoutSource(LLM_REPO, branch, LLM_ROOT, false, 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, true, true) + trtllm_utils.checkoutSource(LLM_REPO, env.gitlabCommit, LLM_ROOT, false, 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, true, true) + trtllm_utils.checkoutSource(LLM_REPO, env.gitlabCommit, LLM_ROOT, false, true) junit(testResults: '**/results*.xml', allowEmptyResults : true) diff --git a/jenkins/L0_Test.groovy b/jenkins/L0_Test.groovy index ed2e0a88f113..256c47570716 100644 --- a/jenkins/L0_Test.groovy +++ b/jenkins/L0_Test.groovy @@ -3104,7 +3104,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, true, true) + trtllm_utils.checkoutSource(LLM_REPO, env.gitlabCommit, LLM_ROOT, false, true) sh "mkdir TensorRT-LLM" sh "cp -r ${LLM_ROOT}/ TensorRT-LLM/src/" trtllm_utils.llmExecStepWithRetry(pipeline, script: "git config --global --add safe.directory \"*\"") @@ -4527,7 +4527,7 @@ def runLLMBuild( sh "env | sort" sh "ccache -sv" - trtllm_utils.checkoutSource(LLM_REPO, env.gitlabCommit, "tensorrt_llm", true, true) + trtllm_utils.checkoutSource(LLM_REPO, env.gitlabCommit, "tensorrt_llm", false, true) if (env.alternativeTRT) { sh "cd tensorrt_llm/ && sed -i 's#tensorrt~=.*\$#tensorrt#g' requirements.txt && cat requirements.txt" } @@ -5642,7 +5642,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, true, true) + trtllm_utils.checkoutSource(LLM_REPO, env.gitlabCommit, LLM_ROOT, false, 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 651e431ea96f..dc3c1735e2d1 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, true, true) + trtllm_utils.checkoutSource(LLM_REPO, ref, env.WORKSPACE, false, true) } def getPulseToken(serviceId, scopes) { diff --git a/jenkins/runPerfSanityTriage.groovy b/jenkins/runPerfSanityTriage.groovy index bc31ad53e318..86591da57cc7 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, true, false) + trtllm_utils.checkoutSource(LLM_REPO, params.BRANCH, LLM_ROOT, false, false) def commandsBase64 = params.COMMANDS.bytes.encodeBase64().toString() sh """ cd ${LLM_ROOT}/jenkins/scripts/perf && python3 perf_sanity_triage.py \ diff --git a/scripts/attribution/scan/metadata/msa.yml b/scripts/attribution/scan/metadata/msa.yml index ab212fc1b664..3248e2ea9020 100644 --- a/scripts/attribution/scan/metadata/msa.yml +++ b/scripts/attribution/scan/metadata/msa.yml @@ -1,5 +1,5 @@ name: msa description: MiniMax Sparse Attention (fmha_sm100) kernels for SM100 sparse attention -source: submodule +source: fetched directory_matches: -- 3rdparty/MSA +- fmha_sm100 diff --git a/scripts/build_wheel.py b/scripts/build_wheel.py index 4e7f8a9d4407..b5d29cd130a5 100755 --- a/scripts/build_wheel.py +++ b/scripts/build_wheel.py @@ -80,37 +80,6 @@ def get_build_dir(build_dir, build_type): return build_dir -def stage_msa_package(project_dir: Path, build_dir: Path) -> Path: - """Copy pinned MSA sources and apply TensorRT-LLM's downstream patch.""" - msa_source_dir = project_dir / "3rdparty" / "MSA" - msa_package_dir = msa_source_dir / "python" / "fmha_sm100" - msa_patch = project_dir / "3rdparty" / "patches" / "msa_strided_paged_kv.patch" - if not msa_package_dir.is_dir(): - raise FileNotFoundError( - f"MSA sources are missing at {msa_package_dir}; initialize 3rdparty/MSA" - ) - - staging_dir = build_dir / "msa_patched" - if staging_dir.exists(): - rmtree(staging_dir) - copytree(msa_source_dir, staging_dir, ignore=shutil.ignore_patterns(".git")) - git_env = os.environ.copy() - git_env["GIT_CEILING_DIRECTORIES"] = str(staging_dir.parent.resolve()) - run( - ["git", "apply", "--check", str(msa_patch)], - cwd=staging_dir, - env=git_env, - check=True, - ) - run( - ["git", "apply", str(msa_patch)], - cwd=staging_dir, - env=git_env, - check=True, - ) - return staging_dir / "python" / "fmha_sm100" - - def clear_folder(folder_path): for item in os.listdir(folder_path): item_path = os.path.join(folder_path, item) @@ -568,16 +537,6 @@ def main(*, project_dir = get_project_dir() os.chdir(project_dir) - # Get all submodules and check their folder exists. If not, - # invoke git submodule update - with open(project_dir / ".gitmodules", "r") as submodules_f: - submodules = [ - l.split("=")[1].strip() for l in submodules_f.readlines() - if "path = " in l - ] - if any(not (project_dir / submodule / ".git").exists() - for submodule in submodules): - build_run('git submodule update --init --recursive') on_windows = platform.system() == "Windows" requirements_filename = "requirements-dev-windows.txt" if on_windows else "requirements-dev.txt" @@ -1106,6 +1065,17 @@ def get_binding_lib(subdirectory, name): pkg_dir / "flash_mla", dirs_exist_ok=True) + # Stage the FetchContent-patched MSA package for setup.py packaging. + msa_src = build_dir / "_deps" / "msa-src" / "python" / "fmha_sm100" + msa_dst = project_dir / "3rdparty" / "fmha_sm100" + if not (msa_src / "cute" / "interface.py").is_file(): + raise FileNotFoundError( + f"MSA package missing at {msa_src}; CMake FetchContent for msa " + "did not populate the expected sources.") + if msa_dst.exists(): + rmtree(msa_dst) + install_tree(msa_src, msa_dst, dirs_exist_ok=True) + if not skip_stubs: with working_directory(pkg_dir): if on_windows: @@ -1174,14 +1144,10 @@ def get_binding_lib(subdirectory, name): f"Copied auto-generated attributions to {project_dir / 'ATTRIBUTIONS.md'}" ) - msa_package_dir = stage_msa_package(project_dir, build_dir) - wheel_env = os.environ.copy() - wheel_env["TRTLLM_MSA_PACKAGE_DIR"] = str(msa_package_dir) - build_run( - f'\"{venv_python}\" -m build {project_dir} --skip-dependency-check {extra_wheel_build_args} --no-isolation --wheel --outdir "{dist_dir}"', - env=wheel_env) - env = wheel_env.copy() + f'\"{venv_python}\" -m build {project_dir} --skip-dependency-check {extra_wheel_build_args} --no-isolation --wheel --outdir "{dist_dir}"' + ) + env = os.environ.copy() if mypyc: env["TRTLLM_ENABLE_MYPYC"] = "1" else: diff --git a/setup.py b/setup.py index 8865078b3598..50058675e674 100644 --- a/setup.py +++ b/setup.py @@ -58,6 +58,13 @@ def sanity_check(): 'If you are attempting to use the pip development mode (editable installation), ' 'please execute `scripts/build_wheel.py` first, and then run `pip install -e .`.' ) + if not (Path(__file__).resolve().parent / "3rdparty" / + "fmha_sm100").is_dir(): + raise ImportError( + 'The `fmha_sm100` package is missing. Please execute ' + '`scripts/build_wheel.py` first (CMake FetchContent stages it under ' + '3rdparty/fmha_sm100), or use TRTLLM_USE_PRECOMPILED to extract it ' + 'from a published wheel.') def get_version(): @@ -293,6 +300,15 @@ def extract_from_precompiled(precompiled_location: str, package_data: list[str], os.makedirs(dst_dir, exist_ok=True) print(f"Copying {rel_path} from local directory.") shutil.copy2(src_file, dst_file) + + source_fmha = os.path.join(precompiled_location, "3rdparty", + "fmha_sm100") + if os.path.isdir(source_fmha): + dst_fmha = os.path.join("3rdparty", "fmha_sm100") + print(f"Copying fmha_sm100 from local directory: {source_fmha}") + if os.path.isdir(dst_fmha): + shutil.rmtree(dst_fmha) + shutil.copytree(source_fmha, dst_fmha) return # Handle local file or remote URL @@ -338,6 +354,15 @@ def extract_from_precompiled(precompiled_location: str, package_data: list[str], if should_skip_precompiled_package_data(file.filename): continue + # Top-level MSA package in the wheel; stage under 3rdparty for + # setuptools package_dir, matching scripts/build_wheel.py. + if file.filename.startswith("fmha_sm100/"): + print( + f"Extracting and including {file.filename} from precompiled wheel." + ) + wheel.extract(file, path="3rdparty") + continue + # Skip .py files EXCEPT for generated C++ extension wrappers # (deep_gemm, deep_ep, flash_mla Python files are generated during build) if file.filename.endswith(".py"): @@ -421,10 +446,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": - os.environ.get("TRTLLM_MSA_PACKAGE_DIR", "3rdparty/MSA/python/fmha_sm100") -} +# fmha_sm100 is staged under 3rdparty/ by scripts/build_wheel.py from the +# CMake FetchContent tree (same packaging role as tensorrt_llm/deep_ep). +msa_package_dir = {"fmha_sm100": "3rdparty/fmha_sm100"} packages += ["fmha_sm100"] # https://setuptools.pypa.io/en/latest/references/keywords.html diff --git a/tensorrt_llm/_torch/attention_backend/fmha/msa_sparse_gqa.py b/tensorrt_llm/_torch/attention_backend/fmha/msa_sparse_gqa.py index 0bbb1ad1509a..982f77a08ac4 100644 --- a/tensorrt_llm/_torch/attention_backend/fmha/msa_sparse_gqa.py +++ b/tensorrt_llm/_torch/attention_backend/fmha/msa_sparse_gqa.py @@ -186,8 +186,8 @@ class MsaSparseGqaFmha(Fmha): @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. + # fmha_sm100 runs only on the SM100 family and is packaged in the + # wheel, so it is unavailable off SM100 or without the wheel. # Imported lazily because the minimax_m3 package init imports the trtllm # attention classes, which a module-scope import here would cycle with. from tensorrt_llm._torch.attention_backend.sparse.minimax_m3.msa_utils import ( 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 index c4751cdec2e4..6e423da224d2 100644 --- a/tensorrt_llm/_torch/attention_backend/sparse/minimax_m3/msa_availability.py +++ b/tensorrt_llm/_torch/attention_backend/sparse/minimax_m3/msa_availability.py @@ -2,10 +2,10 @@ # 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. +The MSA kernels are provided by the fmha_sm100 package bundled with +TensorRT-LLM and run only on the SM100 architecture family (SM100 and SM103). +These helpers gate backend selection so a request for the MSA path fails early +with a clear message on unsupported systems. """ from __future__ import annotations @@ -24,8 +24,7 @@ def ensure_msa_available() -> None: 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'." + "packaged with TensorRT-LLM. Reinstall TensorRT-LLM from a complete build." ) if not is_sm_100f(): sm_version = get_sm_version() diff --git a/tensorrt_llm/_torch/attention_backend/sparse/minimax_m3/msa_utils.py b/tensorrt_llm/_torch/attention_backend/sparse/minimax_m3/msa_utils.py index 83380bc730c0..8cf5b3de6b33 100644 --- a/tensorrt_llm/_torch/attention_backend/sparse/minimax_m3/msa_utils.py +++ b/tensorrt_llm/_torch/attention_backend/sparse/minimax_m3/msa_utils.py @@ -5,8 +5,6 @@ from __future__ import annotations import importlib.util -import sys -from pathlib import Path from typing import Optional, Tuple import torch @@ -19,55 +17,25 @@ 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 whether the packaged fmha_sm100 module can be imported.""" 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. + """Import the packaged fmha_sm100 module 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." + "MiniMax-M3 MSA attention requires the fmha_sm100 kernels packaged " + "with TensorRT-LLM. Reinstall TensorRT-LLM from a complete build." ) from exc return fmha_sm100 diff --git a/tests/unittest/scripts/test_build_wheel.py b/tests/unittest/scripts/test_build_wheel.py deleted file mode 100644 index 0b9061efd34d..000000000000 --- a/tests/unittest/scripts/test_build_wheel.py +++ /dev/null @@ -1,52 +0,0 @@ -#!/usr/bin/env python3 -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. - -import importlib.util -from pathlib import Path - -import pytest - -REPO_ROOT = Path(__file__).resolve().parent.parent.parent.parent -SCRIPT_PATH = REPO_ROOT / "scripts" / "build_wheel.py" -MSA_INTERFACE = Path("python/fmha_sm100/cute/interface.py") - -_SPEC = importlib.util.spec_from_file_location("build_wheel", SCRIPT_PATH) -assert _SPEC is not None and _SPEC.loader is not None -_BUILD_WHEEL = importlib.util.module_from_spec(_SPEC) -_SPEC.loader.exec_module(_BUILD_WHEEL) -stage_msa_package = _BUILD_WHEEL.stage_msa_package - - -def test_stage_msa_package_applies_patch_without_modifying_submodule(tmp_path): - source_interface = REPO_ROOT / "3rdparty" / "MSA" / MSA_INTERFACE - if not source_interface.is_file(): - pytest.skip("3rdparty/MSA is not initialized") - - source_before = source_interface.read_bytes() - staged_package = stage_msa_package(REPO_ROOT, tmp_path) - staged_interface = staged_package / "cute" / "interface.py" - - assert b"def _prepare_paged_hnd_input" not in source_before - assert "def _prepare_paged_hnd_input" in staged_interface.read_text() - assert source_interface.read_bytes() == source_before - - -def test_stage_msa_package_requires_initialized_submodule(tmp_path): - project_dir = tmp_path / "project" - project_dir.mkdir() - - with pytest.raises(FileNotFoundError, match="initialize 3rdparty/MSA"): - stage_msa_package(project_dir, tmp_path / "build") From 235f5528a34f59b7a4e02e3325a24bae73b1f3b9 Mon Sep 17 00:00:00 2001 From: peihengh <259410613+peihu-nv@users.noreply.github.com> Date: Tue, 28 Jul 2026 10:29:49 -0700 Subject: [PATCH 3/5] [None][fix] Consolidate MSA wheel packaging Signed-off-by: peihengh <259410613+peihu-nv@users.noreply.github.com> --- scripts/build_wheel.py | 5 +---- .../_torch/attention_backend/sparse/minimax_m3/msa_utils.py | 3 ++- 2 files changed, 3 insertions(+), 5 deletions(-) diff --git a/scripts/build_wheel.py b/scripts/build_wheel.py index b5d29cd130a5..9a0dbefabd74 100755 --- a/scripts/build_wheel.py +++ b/scripts/build_wheel.py @@ -1144,9 +1144,6 @@ def get_binding_lib(subdirectory, name): f"Copied auto-generated attributions to {project_dir / 'ATTRIBUTIONS.md'}" ) - build_run( - f'\"{venv_python}\" -m build {project_dir} --skip-dependency-check {extra_wheel_build_args} --no-isolation --wheel --outdir "{dist_dir}"' - ) env = os.environ.copy() if mypyc: env["TRTLLM_ENABLE_MYPYC"] = "1" @@ -1154,7 +1151,7 @@ def get_binding_lib(subdirectory, name): env["TRTLLM_ENABLE_MYPYC"] = "0" build_run( - f'\"{venv_python}\" -m build {project_dir} --skip-dependency-check {plat_name_arg} --no-isolation --wheel --outdir "{dist_dir}"', + f'\"{venv_python}\" -m build {project_dir} --skip-dependency-check {extra_wheel_build_args} --no-isolation --wheel --outdir "{dist_dir}"', env=env) if install: diff --git a/tensorrt_llm/_torch/attention_backend/sparse/minimax_m3/msa_utils.py b/tensorrt_llm/_torch/attention_backend/sparse/minimax_m3/msa_utils.py index 8cf5b3de6b33..ff6aa734b856 100644 --- a/tensorrt_llm/_torch/attention_backend/sparse/minimax_m3/msa_utils.py +++ b/tensorrt_llm/_torch/attention_backend/sparse/minimax_m3/msa_utils.py @@ -5,6 +5,7 @@ from __future__ import annotations import importlib.util +import types from typing import Optional, Tuple import torch @@ -23,7 +24,7 @@ def msa_package_available() -> bool: return importlib.util.find_spec("fmha_sm100") is not None -def require_msa_module(): +def require_msa_module() -> types.ModuleType: """Import the packaged fmha_sm100 module or raise a clear error. The import is deferred to first kernel use so the MSA backend can be From 722bfd6ac6af119e430251d637fcc4107666fd16 Mon Sep 17 00:00:00 2001 From: peihengh <259410613+peihu-nv@users.noreply.github.com> Date: Tue, 4 Aug 2026 20:16:55 -0700 Subject: [PATCH 4/5] [None][fix] Address MSA packaging review feedback Signed-off-by: peihengh <259410613+peihu-nv@users.noreply.github.com> --- 3rdparty/CMakeLists.txt | 2 +- 3rdparty/patches/msa_strided_paged_kv.patch | 55 +++++++++++++--- docs/source/installation/build-from-source.md | 1 - jenkins/Build.groovy | 4 +- jenkins/BuildDockerImage.groovy | 2 +- jenkins/L0_MergeRequest.groovy | 8 +-- jenkins/L0_Test.groovy | 6 +- jenkins/TensorRT_LLM_PLC.groovy | 2 +- jenkins/UpdateTestDurations.groovy | 2 +- jenkins/runPerfSanityTriage.groovy | 2 +- scripts/build_wheel.py | 16 ++++- setup.py | 5 ++ .../sparse/test_minimax_m3_msa_backend.py | 63 ++++++++++++++++++- 13 files changed, 142 insertions(+), 26 deletions(-) diff --git a/3rdparty/CMakeLists.txt b/3rdparty/CMakeLists.txt index 02a265608d88..01bb76553d8b 100644 --- a/3rdparty/CMakeLists.txt +++ b/3rdparty/CMakeLists.txt @@ -135,7 +135,7 @@ foreach(DEP_IDX RANGE ${DEP_COUNT_MINUS_ONE}) PATCH_COMMAND bash -c - "patch -p1 --forward --batch --dry-run -i '${_patch_file}' && patch -p1 --forward --batch -i '${_patch_file}' || patch -p1 --reverse --batch --dry-run -i '${_patch_file}'" + "patch -p1 --forward --batch --dry-run -i '${_patch_file}' && patch -p1 --forward --batch -i '${_patch_file}' || patch -p1 --reverse --batch --dry-run -i '${_patch_file}' || (echo 'Patch state for ${DEP_NAME} is inconsistent. Remove ${CMAKE_BINARY_DIR}/_deps/${DEP_NAME}-src and reconfigure.' >&2 && false)" ) endif() diff --git a/3rdparty/patches/msa_strided_paged_kv.patch b/3rdparty/patches/msa_strided_paged_kv.patch index 743587c5608b..c2d4617cc7a5 100644 --- a/3rdparty/patches/msa_strided_paged_kv.patch +++ b/3rdparty/patches/msa_strided_paged_kv.patch @@ -1,8 +1,8 @@ diff --git a/python/fmha_sm100/cute/interface.py b/python/fmha_sm100/cute/interface.py -index d72b17a..e27a74f 100644 +index d72b17a..eca5b8c 100644 --- a/python/fmha_sm100/cute/interface.py +++ b/python/fmha_sm100/cute/interface.py -@@ -136,6 +136,32 @@ def _prepare_paged_kv_for_tma(k, v, blk_kv: int): +@@ -136,6 +136,35 @@ def _prepare_paged_kv_for_tma(k, v, blk_kv: int): return k, v @@ -14,8 +14,10 @@ index d72b17a..e27a74f 100644 + but the outer physical-page stride may include other coalesced cache roles. + Materialize every other layout to preserve the public API's old contract. + """ -+ if tensor.ndim != 4 or int(tensor.shape[2]) != int(blk_kv): ++ if tensor.ndim != 4: + return tensor.contiguous() ++ if int(tensor.shape[2]) != int(blk_kv): ++ return tensor + + head_dim = int(tensor.shape[3]) + page_size = int(tensor.shape[2]) @@ -27,6 +29,7 @@ index d72b17a..e27a74f 100644 + alignment_bytes = 16 + aligned_for_tma = ( + tensor.data_ptr() % alignment_bytes == 0 ++ and tensor.stride(0) >= 0 + and tensor.stride(0) * tensor.element_size() % alignment_bytes == 0 + ) + return tensor if packed_within_page and aligned_for_tma else tensor.contiguous() @@ -35,7 +38,7 @@ index d72b17a..e27a74f 100644 def _validate_cu_seqlens( cu_seqlens: torch.Tensor, *, -@@ -736,10 +762,21 @@ def sparse_atten_func( +@@ -736,10 +765,21 @@ def sparse_atten_func( max_seqlen_q = int(max_seqlen_q) max_seqlen_k = int(max_seqlen_k) @@ -60,10 +63,10 @@ index d72b17a..e27a74f 100644 k2q_q_indices.contiguous(), int(topK), diff --git a/python/fmha_sm100/cute/test_sparse_atten.py b/python/fmha_sm100/cute/test_sparse_atten.py -index 21c777e..c22beef 100644 +index 21c777e..b5f078b 100644 --- a/python/fmha_sm100/cute/test_sparse_atten.py +++ b/python/fmha_sm100/cute/test_sparse_atten.py -@@ -61,6 +61,43 @@ DECODE_DIM = 128 +@@ -61,6 +61,81 @@ DECODE_DIM = 128 DECODE_KV_TOKEN_SWEEP = tuple(2**exp for exp in range(3, 21)) @@ -103,11 +106,49 @@ index 21c777e..c22beef 100644 + assert prepared.data_ptr() != view.data_ptr() + torch.testing.assert_close(prepared, view, rtol=0, atol=0) + ++ ++def test_prepare_paged_hnd_input_materializes_unaligned_outer_stride() -> None: ++ pages, heads, page_size, head_dim = 5, 2, 128, 128 ++ outer_stride = heads * page_size * head_dim + 1 ++ storage = torch.empty( ++ pages * outer_stride, ++ dtype=torch.float8_e4m3fn, ++ device="cuda", ++ ) ++ view = storage.as_strided( ++ (pages, heads, page_size, head_dim), ++ (outer_stride, page_size * head_dim, head_dim, 1), ++ ) ++ ++ prepared = sparse_interface._prepare_paged_hnd_input(view, page_size) ++ assert prepared.is_contiguous() ++ assert prepared.data_ptr() != view.data_ptr() ++ torch.testing.assert_close(prepared, view, rtol=0, atol=0) ++ ++ ++def test_prepare_paged_hnd_input_defers_page_size_validation() -> None: ++ pages, roles, heads, page_size, head_dim = 5, 2, 2, 128, 128 ++ pool = torch.empty( ++ pages, ++ roles, ++ heads, ++ page_size, ++ head_dim, ++ dtype=torch.float8_e4m3fn, ++ device="cuda", ++ ) ++ view = pool[:, 0] ++ ++ prepared = sparse_interface._prepare_paged_hnd_input(view, page_size // 2) ++ assert prepared.data_ptr() == view.data_ptr() ++ with pytest.raises(ValueError, match="page_size == blk_kv"): ++ sparse_interface._prepare_paged_kv_for_tma(prepared, prepared, page_size // 2) ++ + @contextmanager def _nvtx_range(message: str): torch.cuda.nvtx.range_push(message) -@@ -1786,6 +1823,74 @@ def test_sparse_page_atten( +@@ -1786,6 +1861,74 @@ def test_sparse_page_atten( _assert_forward_close(out, out_ref, out_pt.float(), lse, lse_ref) diff --git a/docs/source/installation/build-from-source.md b/docs/source/installation/build-from-source.md index 8ae5f551d920..d28151842094 100644 --- a/docs/source/installation/build-from-source.md +++ b/docs/source/installation/build-from-source.md @@ -21,7 +21,6 @@ git lfs install ```bash git clone https://github.com/NVIDIA/TensorRT-LLM.git cd TensorRT-LLM -git submodule update --init --recursive git lfs pull ``` 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 256c47570716..ed2e0a88f113 100644 --- a/jenkins/L0_Test.groovy +++ b/jenkins/L0_Test.groovy @@ -3104,7 +3104,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 \"*\"") @@ -4527,7 +4527,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" } @@ -5642,7 +5642,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/UpdateTestDurations.groovy b/jenkins/UpdateTestDurations.groovy index 78938d084e76..5b4dfeb12758 100644 --- a/jenkins/UpdateTestDurations.groovy +++ b/jenkins/UpdateTestDurations.groovy @@ -130,7 +130,7 @@ pipeline { container('trt-llm') { script { def sourceRepo = "https://github.com/${params.SOURCE_REPO}.git" - trtllm_utils.checkoutSource(sourceRepo, params.TARGET_BRANCH, LLM_ROOT, false, false) + trtllm_utils.checkoutSource(sourceRepo, params.TARGET_BRANCH, LLM_ROOT, true, false) } } } 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/scripts/build_wheel.py b/scripts/build_wheel.py index 9a0dbefabd74..9499492dbe0c 100755 --- a/scripts/build_wheel.py +++ b/scripts/build_wheel.py @@ -1074,7 +1074,21 @@ def get_binding_lib(subdirectory, name): "did not populate the expected sources.") if msa_dst.exists(): rmtree(msa_dst) - install_tree(msa_src, msa_dst, dirs_exist_ok=True) + msa_dst.mkdir(parents=True) + for python_source in msa_src.glob("*.py"): + install_file(python_source, msa_dst) + for relative_dir in ( + Path("csrc"), + Path("cute"), + Path("cutlass/include"), + Path("cutlass/tools/util/include"), + ): + install_tree( + msa_src / relative_dir, + msa_dst / relative_dir, + dirs_exist_ok=True, + ) + install_file(msa_src / "cutlass" / "LICENSE.txt", msa_dst / "cutlass") if not skip_stubs: with working_directory(pkg_dir): diff --git a/setup.py b/setup.py index 50058675e674..f6817c3a816a 100644 --- a/setup.py +++ b/setup.py @@ -344,6 +344,11 @@ def extract_from_precompiled(precompiled_location: str, package_data: list[str], wheel_path = precompiled_path with zipfile.ZipFile(wheel_path) as wheel: + dst_fmha = os.path.join("3rdparty", "fmha_sm100") + wheel_has_fmha = any( + file.filename.startswith("fmha_sm100/") for file in wheel.filelist) + if wheel_has_fmha and os.path.isdir(dst_fmha): + shutil.rmtree(dst_fmha) for file in wheel.filelist: # Skip yaml files if file.filename.endswith(".yaml"): diff --git a/tests/unittest/_torch/attention/sparse/test_minimax_m3_msa_backend.py b/tests/unittest/_torch/attention/sparse/test_minimax_m3_msa_backend.py index 58d651631608..fc49c8e8a8a0 100644 --- a/tests/unittest/_torch/attention/sparse/test_minimax_m3_msa_backend.py +++ b/tests/unittest/_torch/attention/sparse/test_minimax_m3_msa_backend.py @@ -2,15 +2,18 @@ # 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. +These validate backend selection, decode scratch-buffer sizing, and the paged +HND view contract passed to the packaged MSA kernel. Numerical parity against +the Triton reference is covered by the SM100 integration accuracy test. """ +from unittest.mock import Mock + import pytest import torch from tensorrt_llm._torch.attention_backend.sparse.minimax_m3 import MiniMaxM3MsaSparseAttention +from tensorrt_llm._torch.attention_backend.sparse.minimax_m3.msa_utils import msa_paged_kv from tensorrt_llm._torch.attention_backend.sparse.utils import _resolve_minimax_m3_backend_cls from tensorrt_llm.llmapi.llm_args import MiniMaxM3SparseAttentionConfig @@ -74,3 +77,57 @@ def test_msa_proxy_max_score_view_is_contiguous_over_stable_store(): # 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) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA") +def test_msa_paged_kv_preserves_tma_compatible_outer_stride() -> None: + from fmha_sm100.cute import interface as sparse_interface + + pages, roles, heads, page_size, head_dim = 5, 2, 2, 128, 128 + pool = torch.empty( + pages, + roles, + heads, + page_size, + head_dim, + dtype=torch.float8_e4m3fn, + device="cuda", + ) + kv_cache_manager = Mock() + kv_cache_manager.get_buffers.return_value = pool + + k_view, v_view = msa_paged_kv(kv_cache_manager, layer_idx=3) + + kv_cache_manager.get_buffers.assert_called_once_with(3, kv_layout="HND") + for view in (k_view, v_view): + assert not view.is_contiguous() + prepared = sparse_interface._prepare_paged_hnd_input(view, page_size) + assert prepared.data_ptr() == view.data_ptr() + assert prepared.stride() == view.stride() + + mismatched = sparse_interface._prepare_paged_hnd_input(k_view, page_size // 2) + assert mismatched.data_ptr() == k_view.data_ptr() + with pytest.raises(ValueError, match="page_size == blk_kv"): + sparse_interface._prepare_paged_kv_for_tma(mismatched, mismatched, page_size // 2) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA") +def test_msa_paged_hnd_input_materializes_unaligned_outer_stride() -> None: + from fmha_sm100.cute import interface as sparse_interface + + pages, heads, page_size, head_dim = 5, 2, 128, 128 + outer_stride = heads * page_size * head_dim + 1 + storage = torch.empty( + pages * outer_stride, + dtype=torch.float8_e4m3fn, + device="cuda", + ) + view = storage.as_strided( + (pages, heads, page_size, head_dim), + (outer_stride, page_size * head_dim, head_dim, 1), + ) + + prepared = sparse_interface._prepare_paged_hnd_input(view, page_size) + + assert prepared.is_contiguous() + assert prepared.data_ptr() != view.data_ptr() From a92acdf9b8504b3dac087861c2e87b878d1e662a Mon Sep 17 00:00:00 2001 From: peihengh <259410613+peihu-nv@users.noreply.github.com> Date: Tue, 4 Aug 2026 22:43:49 -0700 Subject: [PATCH 5/5] [None][fix] Honor MSA block-index strides for q_len <= 32 Signed-off-by: peihengh <259410613+peihu-nv@users.noreply.github.com> --- 3rdparty/patches/msa_strided_paged_kv.patch | 311 ++++++++++++++++++ .../sparse/test_minimax_m3_msa_backend.py | 88 ++++- 2 files changed, 397 insertions(+), 2 deletions(-) diff --git a/3rdparty/patches/msa_strided_paged_kv.patch b/3rdparty/patches/msa_strided_paged_kv.patch index c2d4617cc7a5..d08631a336ee 100644 --- a/3rdparty/patches/msa_strided_paged_kv.patch +++ b/3rdparty/patches/msa_strided_paged_kv.patch @@ -223,3 +223,314 @@ index 21c777e..b5f078b 100644 @pytest.mark.parametrize("paged", [False, True]) @pytest.mark.parametrize("causal", [True]) @pytest.mark.parametrize("batch", [3]) +diff --git a/python/fmha_sm100/csrc/fmha_sm100_inst.jinja b/python/fmha_sm100/csrc/fmha_sm100_inst.jinja +index 6949035..e59ecfb 100644 +--- a/python/fmha_sm100/csrc/fmha_sm100_inst.jinja ++++ b/python/fmha_sm100/csrc/fmha_sm100_inst.jinja +@@ -58,6 +58,8 @@ cudaError_t {{ func_name }}(const FMHACutlassSM100Params& p) { + p.kv_indices_ptr, p.kv_page_indptr_ptr, p.total_page_num, + p.max_score_ptr, p.max_k_tiles, + p.kv_block_indexes_ptr, p.kv_block_num, ++ p.kv_block_indexes_stride_q, p.kv_block_indexes_stride_h, ++ p.kv_block_indexes_stride_k, + p.pack_factor, p.q_stride_n_original, p.q_stride_h_original, p.h_r_original, + flashinfer::PackGQAUnpackParams{p.max_score_direct_ptr, p.total_qo_len_orig, + p.o_direct_ptr, p.num_qo_heads_orig, +diff --git a/python/fmha_sm100/csrc/fmha_sm100_params.h b/python/fmha_sm100/csrc/fmha_sm100_params.h +index 878a90c..f621511 100644 +--- a/python/fmha_sm100/csrc/fmha_sm100_params.h ++++ b/python/fmha_sm100/csrc/fmha_sm100_params.h +@@ -54,6 +54,9 @@ struct FMHACutlassSM100Params { + int max_k_tiles; + int* kv_block_indexes_ptr; + int kv_block_num; ++ int64_t kv_block_indexes_stride_q; ++ int64_t kv_block_indexes_stride_h; ++ int64_t kv_block_indexes_stride_k; + int pack_factor = 1; + int h_r_original = 0; + int q_stride_n_original = 0; +diff --git a/python/fmha_sm100/csrc/fmha_sm100_variant_run.cu.jinja b/python/fmha_sm100/csrc/fmha_sm100_variant_run.cu.jinja +index 7e21585..03e6c7a 100644 +--- a/python/fmha_sm100/csrc/fmha_sm100_variant_run.cu.jinja ++++ b/python/fmha_sm100/csrc/fmha_sm100_variant_run.cu.jinja +@@ -124,6 +124,12 @@ void FMHAVariantRun_{{ variant_name }}(ffi::TensorView workspace_buffer, ffi::Te + params.kv_block_num = maybe_kv_block_indexes.has_value() + ? static_cast(maybe_kv_block_indexes.value().size(2)) + : 0; ++ params.kv_block_indexes_stride_q = maybe_kv_block_indexes.has_value() ++ ? maybe_kv_block_indexes.value().stride(0) : 0; ++ params.kv_block_indexes_stride_h = maybe_kv_block_indexes.has_value() ++ ? maybe_kv_block_indexes.value().stride(1) : 0; ++ params.kv_block_indexes_stride_k = maybe_kv_block_indexes.has_value() ++ ? maybe_kv_block_indexes.value().stride(2) : 0; + + params.pack_factor = static_cast(pack_factor); + params.num_ctas = static_cast(packed_work_range.size(0)); +@@ -188,9 +194,12 @@ void FMHAVariantRun_{{ variant_name }}(ffi::TensorView workspace_buffer, ffi::Te + * maybe_max_score.value().size(1) + * maybe_max_score.value().size(2)) : 0; + params.gmem_bounds.kv_block_indexes_numel = maybe_kv_block_indexes.has_value() +- ? static_cast(maybe_kv_block_indexes.value().size(0) +- * maybe_kv_block_indexes.value().size(1) +- * maybe_kv_block_indexes.value().size(2)) : 0; ++ ? static_cast((maybe_kv_block_indexes.value().size(0) - 1) ++ * maybe_kv_block_indexes.value().stride(0) ++ + (maybe_kv_block_indexes.value().size(1) - 1) ++ * maybe_kv_block_indexes.value().stride(1) ++ + (maybe_kv_block_indexes.value().size(2) - 1) ++ * maybe_kv_block_indexes.value().stride(2) + 1) : 0; + params.gmem_bounds.split_kv_size = maybe_kv_tile_begin_indices.has_value() + ? static_cast(maybe_kv_tile_begin_indices.value().size(0)) : 0; + params.gmem_bounds.qo_offsets_size = maybe_qo_offsets.has_value() +diff --git a/python/fmha_sm100/csrc/include/fmha_cutlass_sm100.cuh b/python/fmha_sm100/csrc/include/fmha_cutlass_sm100.cuh +index b11abe6..f6f95f7 100644 +--- a/python/fmha_sm100/csrc/include/fmha_cutlass_sm100.cuh ++++ b/python/fmha_sm100/csrc/include/fmha_cutlass_sm100.cuh +@@ -126,6 +126,9 @@ struct FwdRunner { + int max_k_tiles = 0, + int* kv_block_indexes = nullptr, + int kv_block_num = 0, ++ int64_t kv_block_indexes_stride_q = 0, ++ int64_t kv_block_indexes_stride_h = 0, ++ int64_t kv_block_indexes_stride_k = 0, + int pack_factor = 1, + int q_stride_n_original = 0, + int q_stride_h_original = 0, +@@ -224,6 +227,7 @@ struct FwdRunner { + arguments = { + {problem_shape, + {{q, layout_Q, k, layout_K, v, layout_V, kv_indices, kv_page_indptr, kv_block_indexes, kv_block_num, ++ kv_block_indexes_stride_q, kv_block_indexes_stride_h, kv_block_indexes_stride_k, + #ifdef FMHA_GMEM_BOUNDS_CHECK + pack_gqa.gmem_bounds.kv_page_indptr_size, pack_gqa.gmem_bounds.kv_indices_size, pack_gqa.gmem_bounds.kv_block_indexes_numel, + #endif +@@ -245,6 +249,7 @@ struct FwdRunner { + arguments = { + problem_shape, + {{q, layout_Q, k, layout_K, v, layout_V, kv_indices, kv_page_indptr, kv_block_indexes, kv_block_num, ++ kv_block_indexes_stride_q, kv_block_indexes_stride_h, kv_block_indexes_stride_k, + #ifdef FMHA_GMEM_BOUNDS_CHECK + pack_gqa.gmem_bounds.kv_page_indptr_size, pack_gqa.gmem_bounds.kv_indices_size, pack_gqa.gmem_bounds.kv_block_indexes_numel, + #endif +@@ -436,6 +441,9 @@ cudaError_t run_fmha_fwd(void* workspace_buffer, DTypeIn* q, DTypeIn* k, DTypeIn + int max_k_tiles = 0, + int* kv_block_indexes = nullptr, + int kv_block_num = 0, ++ int64_t kv_block_indexes_stride_q = 0, ++ int64_t kv_block_indexes_stride_h = 0, ++ int64_t kv_block_indexes_stride_k = 0, + int pack_factor = 1, + int q_stride_n_original = 0, + int q_stride_h_original = 0, +@@ -455,6 +463,7 @@ cudaError_t run_fmha_fwd(void* workspace_buffer, DTypeIn* q, DTypeIn* k, DTypeIn + kv_indices, kv_page_indptr, total_page_num, + maybe_max_score, max_k_tiles, + kv_block_indexes, kv_block_num, ++ kv_block_indexes_stride_q, kv_block_indexes_stride_h, kv_block_indexes_stride_k, + pack_factor, q_stride_n_original, q_stride_h_original, h_r_original, + pack_gqa, num_ctas); + } +diff --git a/python/fmha_sm100/csrc/include/sm100_fmha_fwd_mainloop_tma_warpspecialized.hpp b/python/fmha_sm100/csrc/include/sm100_fmha_fwd_mainloop_tma_warpspecialized.hpp +index 5a5d0f3..843aae8 100644 +--- a/python/fmha_sm100/csrc/include/sm100_fmha_fwd_mainloop_tma_warpspecialized.hpp ++++ b/python/fmha_sm100/csrc/include/sm100_fmha_fwd_mainloop_tma_warpspecialized.hpp +@@ -1131,14 +1131,13 @@ struct Sm100FmhaFwdMainloopTmaWarpspecialized { + int skip_count, real_masked_count, unmasked_count; + int effective_end; + +- int kbi_off_s = 0; ++ int64_t kbi_off_s = 0; + if constexpr (kNeedSparse) { + int h_r_s = get<3, 0, 0>(problem_shape); +- int num_kv_heads_s = get<3, 0, 1>(problem_shape); + int kv_head_idx_s = get<2, 0>(blk_coord) / h_r_s; + int batch_idx_s = get<2, 1>(blk_coord); +- kbi_off_s = batch_idx_s * num_kv_heads_s * params.load.kv_block_num +- + kv_head_idx_s * params.load.kv_block_num; ++ kbi_off_s = static_cast(batch_idx_s) * params.load.kv_block_indexes_stride_q ++ + static_cast(kv_head_idx_s) * params.load.kv_block_indexes_stride_h; + } + + if constexpr (kNeedSparse) { +@@ -1154,7 +1153,10 @@ struct Sm100FmhaFwdMainloopTmaWarpspecialized { + + // Count padding blocks (backward scan for -1) + int valid_blks = params.load.kv_block_num; +- while (valid_blks > 0 && __ldg(¶ms.load.kv_block_indexes[kbi_off_s + valid_blks - 1]) < 0) ++ while (valid_blks > 0 ++ && __ldg(¶ms.load.kv_block_indexes[kbi_off_s ++ + static_cast(valid_blks - 1) ++ * params.load.kv_block_indexes_stride_k]) < 0) + valid_blks--; + + int valid_tiles = (valid_blks * KVPageSize + full_tile_kv - 1) / full_tile_kv; +@@ -1223,7 +1225,8 @@ struct Sm100FmhaFwdMainloopTmaWarpspecialized { + cS = domain_offset(make_coord(q_s, INT_MAX / 2), cS_base); + return false; + } +- int pos = __ldg(¶ms.load.kv_block_indexes[kbi_off_s + page_idx]); ++ int pos = __ldg(¶ms.load.kv_block_indexes[kbi_off_s ++ + static_cast(page_idx) * params.load.kv_block_indexes_stride_k]); + if (pos < 0) { + cS = domain_offset(make_coord(q_s, INT_MAX / 2), cS_base); + return false; +diff --git a/python/fmha_sm100/csrc/include/sm100_fmha_load_tma_warpspecialized.hpp b/python/fmha_sm100/csrc/include/sm100_fmha_load_tma_warpspecialized.hpp +index 0fdcbfa..c577f34 100644 +--- a/python/fmha_sm100/csrc/include/sm100_fmha_load_tma_warpspecialized.hpp ++++ b/python/fmha_sm100/csrc/include/sm100_fmha_load_tma_warpspecialized.hpp +@@ -112,6 +112,9 @@ struct Sm100FmhaLoadTmaWarpspecialized { + int* kv_page_indptr = nullptr; + int* kv_block_indexes = nullptr; + int kv_block_num = 0; ++ int64_t kv_block_indexes_stride_q = 0; ++ int64_t kv_block_indexes_stride_h = 0; ++ int64_t kv_block_indexes_stride_k = 0; + #ifdef FMHA_GMEM_BOUNDS_CHECK + int kv_page_indptr_size = 0; + int kv_indices_size = 0; +@@ -148,6 +151,9 @@ struct Sm100FmhaLoadTmaWarpspecialized { + int* kv_page_indptr = nullptr; + int* kv_block_indexes = nullptr; + int kv_block_num = 0; ++ int64_t kv_block_indexes_stride_q = 0; ++ int64_t kv_block_indexes_stride_h = 0; ++ int64_t kv_block_indexes_stride_k = 0; + #ifdef FMHA_GMEM_BOUNDS_CHECK + int kv_page_indptr_size = 0; + int kv_indices_size = 0; +@@ -218,6 +224,8 @@ struct Sm100FmhaLoadTmaWarpspecialized { + + Params p{tma_load_Q, layout_Q, tma_load_K, layout_K, tma_load_V, layout_V, + args.kv_indices, args.kv_page_indptr, args.kv_block_indexes, args.kv_block_num, ++ args.kv_block_indexes_stride_q, args.kv_block_indexes_stride_h, ++ args.kv_block_indexes_stride_k, + #ifdef FMHA_GMEM_BOUNDS_CHECK + args.kv_page_indptr_size, args.kv_indices_size, args.kv_block_indexes_numel, + #endif +@@ -525,11 +533,10 @@ struct Sm100FmhaLoadTmaWarpspecialized { + constexpr int effective_tile_kv = get<1>(TileShapeQK{}); + constexpr int tiles_per_page = KVPageSize / effective_tile_kv; + +- int kv_block_offset = 0; ++ int64_t kv_block_offset = 0; + if constexpr (kSparseAttnMode == SparseAttnMode::Sparse) { +- int num_kv_heads_val = get<3, 0, 1>(params_problem_shape); +- kv_block_offset = batch_idx * num_kv_heads_val * params.kv_block_num +- + kv_head_idx * params.kv_block_num; ++ kv_block_offset = static_cast(batch_idx) * params.kv_block_indexes_stride_q ++ + static_cast(kv_head_idx) * params.kv_block_indexes_stride_h; + } + + // Keep full 4D tensor (page_size, D, H_kv, P) — select head at copy time +@@ -549,7 +556,8 @@ struct Sm100FmhaLoadTmaWarpspecialized { + int page_for_lookup; + if constexpr (kSparseAttnMode == SparseAttnMode::Sparse) { + int sparse_idx = (logical_page < params.kv_block_num) +- ? __ldg(¶ms.kv_block_indexes[kv_block_offset + logical_page]) ++ ? __ldg(¶ms.kv_block_indexes[kv_block_offset ++ + static_cast(logical_page) * params.kv_block_indexes_stride_k]) + : -1; + page_for_lookup = (sparse_idx >= 0) ? sparse_idx : 0; + } else { +@@ -573,7 +581,8 @@ struct Sm100FmhaLoadTmaWarpspecialized { + int page_for_lookup; + if constexpr (kSparseAttnMode == SparseAttnMode::Sparse) { + int sparse_idx = (logical_page < params.kv_block_num) +- ? __ldg(¶ms.kv_block_indexes[kv_block_offset + logical_page]) ++ ? __ldg(¶ms.kv_block_indexes[kv_block_offset ++ + static_cast(logical_page) * params.kv_block_indexes_stride_k]) + : -1; + page_for_lookup = (sparse_idx >= 0) ? sparse_idx : 0; + } else { +diff --git a/tests/regression/test_noncontiguous_sparse_block_indexes.py b/tests/regression/test_noncontiguous_sparse_block_indexes.py +new file mode 100644 +index 0000000..6a162b0 +--- /dev/null ++++ b/tests/regression/test_noncontiguous_sparse_block_indexes.py +@@ -0,0 +1,85 @@ ++#!/usr/bin/env python3 ++# SPDX-FileCopyrightText: Copyright (c) 2026 MiniMax ++# SPDX-License-Identifier: MIT ++ ++"""Regression for sparse attention consuming a head-major selector view.""" ++ ++import torch ++ ++from fmha_sm100.api import fmha_sm100, fmha_sm100_plan ++ ++ ++def _run_attention(q, k, v, plan, kv_indices, selected): ++ out = torch.empty_like(q) ++ returned, _ = fmha_sm100( ++ q, ++ k, ++ v, ++ plan, ++ kv_indices=kv_indices, ++ kv_block_indexes=selected, ++ out=out, ++ sm_scale=128.0**-0.5, ++ output_maxscore=False, ++ ) ++ torch.cuda.synchronize() ++ assert returned.data_ptr() == out.data_ptr() ++ assert torch.isfinite(out).all() ++ return out ++ ++ ++def _check_query_length(query_len, num_kv_splits): ++ device = torch.device("cuda", 0) ++ total_selector_rows = query_len + 2 ++ q = torch.zeros((query_len, 32, 128), dtype=torch.bfloat16, device=device) ++ k = torch.zeros((8, 2, 128, 128), dtype=torch.bfloat16, device=device) ++ v = torch.empty_like(k) ++ for page in range(8): ++ v[page].fill_(page + 1) ++ kv_indices = torch.arange(8, dtype=torch.int32, device=device) ++ ++ intended = torch.tensor([0, 1, 2, 3], dtype=torch.int32, device=device) ++ poison = torch.tensor([4, 5, 6, 7], dtype=torch.int32, device=device) ++ logical = intended.expand(total_selector_rows, 2, 4).clone() ++ # These rows are outside the logical slice but occupy addresses reached by ++ # a pointer-only contiguous read for later query/head pairs. ++ logical[:2, 1, :] = poison ++ backing = logical.permute(1, 0, 2).contiguous() ++ selected_strided = backing.permute(1, 0, 2)[-query_len:] ++ selected_contiguous = logical[-query_len:].contiguous() ++ ++ assert torch.equal(selected_strided, selected_contiguous) ++ assert not selected_strided.is_contiguous() ++ ++ plan = fmha_sm100_plan( ++ torch.tensor([query_len], dtype=torch.int32), ++ torch.tensor([1024], dtype=torch.int32), ++ 32, ++ num_kv_heads=2, ++ qo_offset=torch.tensor([1024 - query_len], dtype=torch.int32), ++ num_kv_splits=num_kv_splits, ++ page_size=128, ++ output_maxscore=False, ++ kv_block_num=4, ++ causal=True, ++ device=device, ++ ) ++ assert plan[3]["MM-SA-Nv"] is False ++ ++ expected = _run_attention( ++ q, k, v, plan, kv_indices, selected_contiguous ++ ) ++ actual = _run_attention(q, k, v, plan, kv_indices, selected_strided) ++ torch.testing.assert_close(actual, expected, rtol=0, atol=0) ++ ++ ++def test_noncontiguous_sparse_block_indexes(): ++ torch.cuda.set_device(0) ++ for query_len in (1, 5): ++ for num_kv_splits in (1, 2): ++ _check_query_length(query_len, num_kv_splits) ++ ++ ++if __name__ == "__main__": ++ test_noncontiguous_sparse_block_indexes() ++ print("MSA_NONCONTIGUOUS_SPARSE_BLOCK_INDEX_REGRESSION=PASS") diff --git a/tests/unittest/_torch/attention/sparse/test_minimax_m3_msa_backend.py b/tests/unittest/_torch/attention/sparse/test_minimax_m3_msa_backend.py index 185f0d317baf..2e673c253aa8 100644 --- a/tests/unittest/_torch/attention/sparse/test_minimax_m3_msa_backend.py +++ b/tests/unittest/_torch/attention/sparse/test_minimax_m3_msa_backend.py @@ -3,8 +3,9 @@ """Structural tests for the MiniMax-M3 MSA sparse attention backend. These validate backend selection, decode scratch-buffer sizing, and the paged -HND view contract passed to the packaged MSA kernel. Numerical parity against -the Triton reference is covered by the SM100 integration accuracy test. +HND and sparse block-index stride contracts passed to the packaged MSA kernel. +Numerical parity against the Triton reference is covered by the SM100 +integration accuracy test. """ from unittest.mock import Mock @@ -280,3 +281,86 @@ def test_msa_proxy_max_score_strided_index_k_matches_packed(): assert not index_k_strided.is_contiguous() assert index_k_strided.stride(0) == coalescing_scale * page_size * head_dim assert torch.equal(strided_scores, packed_scores) + + +def _run_msa_sparse_attention( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + plan: tuple, + kv_indices: torch.Tensor, + selected: torch.Tensor, +) -> torch.Tensor: + from fmha_sm100.api import fmha_sm100 + + out = torch.empty_like(q) + returned, _ = fmha_sm100( + q, + k, + v, + plan, + kv_indices=kv_indices, + kv_block_indexes=selected, + out=out, + sm_scale=128.0**-0.5, + output_maxscore=False, + ) + torch.cuda.synchronize() + assert returned.data_ptr() == out.data_ptr() + assert torch.isfinite(out).all() + return out + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA") +@pytest.mark.parametrize("query_len", [1, 5]) +@pytest.mark.parametrize("num_kv_splits", [1, 2]) +def test_msa_sparse_attention_honors_noncontiguous_block_indexes( + query_len: int, + num_kv_splits: int, +) -> None: + from fmha_sm100.api import fmha_sm100_plan + + device = torch.device("cuda", 0) + torch.cuda.set_device(device) + total_selector_rows = query_len + 2 + q = torch.zeros((query_len, 32, 128), dtype=torch.bfloat16, device=device) + k = torch.zeros((8, 2, 128, 128), dtype=torch.bfloat16, device=device) + v = torch.empty_like(k) + for page in range(8): + v[page].fill_(page + 1) + kv_indices = torch.arange(8, dtype=torch.int32, device=device) + + intended = torch.tensor([0, 1, 2, 3], dtype=torch.int32, device=device) + poison = torch.tensor([4, 5, 6, 7], dtype=torch.int32, device=device) + logical = intended.expand(total_selector_rows, 2, 4).clone() + # These rows are outside the logical slice but occupy addresses reached by + # a pointer-only contiguous read for later query/head pairs. + logical[:2, 1, :] = poison + backing = logical.permute(1, 0, 2).contiguous() + selected_strided = backing.permute(1, 0, 2)[-query_len:] + selected_contiguous = logical[-query_len:].contiguous() + + assert selected_strided.stride() == (4, total_selector_rows * 4, 1) + assert not selected_strided.is_contiguous() + assert torch.equal(selected_strided, selected_contiguous) + + plan = fmha_sm100_plan( + torch.tensor([query_len], dtype=torch.int32), + torch.tensor([1024], dtype=torch.int32), + 32, + num_kv_heads=2, + qo_offset=torch.tensor([1024 - query_len], dtype=torch.int32), + num_kv_splits=num_kv_splits, + page_size=128, + output_maxscore=False, + kv_block_num=4, + causal=True, + device=device, + ) + short_plan = plan[3] + assert short_plan["MM-SA-Nv"] is False + assert short_plan["num_kv_splits"] == num_kv_splits + + expected = _run_msa_sparse_attention(q, k, v, plan, kv_indices, selected_contiguous) + actual = _run_msa_sparse_attention(q, k, v, plan, kv_indices, selected_strided) + torch.testing.assert_close(actual, expected, rtol=0, atol=0)