[None][fix] Apply MSA patch in place so it lands regardless of wheel build

scripts/build_wheel.py

@@ -80,8 +80,14 @@ 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."""
+def apply_msa_patch(project_dir: Path) -> None:
+    """Apply TensorRT-LLM's downstream MSA patch in place, idempotently.
+
+    The patch is applied directly to the pinned 3rdparty/MSA submodule working
+    tree so every consumer sees the patched fmha_sm100 sources: the runtime
+    importer, editable installs, and wheel packaging. Re-running is a no-op, so
+    it is safe to call on every build.
+    """
     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"
@@ -90,25 +96,37 @@ def stage_msa_package(project_dir: Path, build_dir: Path) -> Path:
             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())
+    git_env["GIT_CEILING_DIRECTORIES"] = str(msa_source_dir.parent.resolve())
+
+    # A clean reverse-apply means the patch is already present, so a forward
+    # apply would fail. Skip in that case.
+    already_applied = run(
+        ["git", "apply", "--reverse", "--check",
+         str(msa_patch)],
+        cwd=msa_source_dir,
+        env=git_env,
+        stdout=DEVNULL,
+        stderr=DEVNULL,
+    ).returncode == 0
+    if already_applied:
+        print(f"-- MSA patch already applied at {msa_package_dir}; skipping.")
+        return
+
+    # Verify a clean forward apply before mutating the working tree.
     run(
         ["git", "apply", "--check", str(msa_patch)],
-        cwd=staging_dir,
+        cwd=msa_source_dir,
         env=git_env,
         check=True,
     )
     run(
         ["git", "apply", str(msa_patch)],
-        cwd=staging_dir,
+        cwd=msa_source_dir,
         env=git_env,
         check=True,
     )
-    return staging_dir / "python" / "fmha_sm100"
+    print(f"-- Applied MSA patch to {msa_package_dir}.")
 
 
 def clear_folder(folder_path):
@@ -564,6 +582,10 @@ def main(*,
     if any(not (project_dir / submodule / ".git").exists()
            for submodule in submodules):
         build_run('git submodule update --init --recursive')
+
+    # Apply the downstream MSA patch before any packaging so every consumer
+    # loads the patched sources.
+    apply_msa_patch(project_dir)
     on_windows = platform.system() == "Windows"
     requirements_filename = "requirements-dev-windows.txt" if on_windows else "requirements-dev.txt"
 
@@ -1204,14 +1226,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:

setup.py

@@ -406,10 +406,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")
-}
+# scripts/build_wheel.py patches this submodule tree in place, so this path
+# always points at the patched fmha_sm100 sources.
+msa_package_dir = {"fmha_sm100": "3rdparty/MSA/python/fmha_sm100"}
 packages += ["fmha_sm100"]
 
 # https://setuptools.pypa.io/en/latest/references/keywords.html

tensorrt_llm/_torch/attention_backend/sparse/minimax_m3/msa_utils.py

@@ -51,6 +51,33 @@ def msa_package_available() -> bool:
     return importlib.util.find_spec("fmha_sm100") is not None
 
 
+# Symbol added by 3rdparty/patches/msa_strided_paged_kv.patch. Its presence in
+# the imported module confirms the downstream MSA patch was applied to the
+# loaded fmha_sm100 sources.
+_MSA_PATCH_MARKER = "_prepare_paged_hnd_input"
+
+
+def _require_msa_patch() -> None:
+    """Fail fast if the loaded fmha_sm100 is missing the downstream patch.
+
+    The patch avoids paged K/V materialization during prefill and is applied in
+    place on the 3rdparty/MSA submodule by scripts/build_wheel.py. An unpatched
+    copy would silently regress prefill, so raise an actionable error that names
+    the loaded module file.
+    """
+    from fmha_sm100.cute import interface
+
+    if not hasattr(interface, _MSA_PATCH_MARKER):
+        raise RuntimeError(
+            "The imported fmha_sm100 is missing the TensorRT-LLM MSA patch "
+            f"(expected symbol '{_MSA_PATCH_MARKER}' in "
+            f"'{getattr(interface, '__file__', '<unknown>')}'). Rebuild with "
+            "scripts/build_wheel.py, which applies "
+            "3rdparty/patches/msa_strided_paged_kv.patch in place, or apply the "
+            "patch to 3rdparty/MSA manually."
+        )
+
+
 def require_msa_module():
     """Import fmha_sm100 from the MSA submodule or raise a clear error.
 
@@ -69,6 +96,7 @@ def require_msa_module():
             "MSA git submodule at 3rdparty/MSA. Initialize it with "
             "'git submodule update --init --recursive', or install fmha_sm100."
         ) from exc
+    _require_msa_patch()
     return fmha_sm100
 
 

tests/integration/test_lists/test-db/l0_dgx_b200_m3.yml

@@ -13,6 +13,6 @@ l0_dgx_b200_m3:
       stage: pre_merge
       backend: pytorch
   tests:
-  - accuracy/test_llm_api_pytorch.py::TestMiniMaxM3::test_nvfp4[use_msa=True]
-  - accuracy/test_llm_api_pytorch.py::TestMiniMaxM3::test_nvfp4_eagle3[tp_size=4-ep_size=4-attention_dp=False-overlap_scheduler=True-use_msa=True-cuda_graph=True]
-  - accuracy/test_llm_api_pytorch.py::TestMiniMaxM3::test_nvfp4_eagle3[tp_size=4-ep_size=4-attention_dp=True-overlap_scheduler=True-use_msa=True-cuda_graph=True]
+  - accuracy/test_llm_api_pytorch.py::TestMiniMaxM3::test_nvfp4[use_msa=True-eval_mode=inferencemax]
+  - accuracy/test_llm_api_pytorch.py::TestMiniMaxM3::test_nvfp4_eagle3[tp_size=4-ep_size=4-attention_dp=False-overlap_scheduler=True-use_msa=True-cuda_graph=True-eval_mode=inferencemax]
+  - accuracy/test_llm_api_pytorch.py::TestMiniMaxM3::test_nvfp4_eagle3[tp_size=4-ep_size=4-attention_dp=True-overlap_scheduler=True-use_msa=True-cuda_graph=True-eval_mode=inferencemax]

tests/unittest/scripts/test_build_wheel.py

@@ -15,38 +15,58 @@
 # limitations under the License.
 
 import importlib.util
+import shutil
 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")
+MSA_PATCH = Path("3rdparty/patches/msa_strided_paged_kv.patch")
 
 _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
+apply_msa_patch = _BUILD_WHEEL.apply_msa_patch
 
 
-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():
+def _stage_project(tmp_path: Path) -> Path:
+    """Copy the MSA submodule and patch into a throwaway project tree.
+
+    The real submodule tree is left untouched; the copy drops .git so the apply
+    runs against a plain working tree, as it does after a fresh checkout.
+    """
+    source_msa = REPO_ROOT / "3rdparty" / "MSA"
+    if not (source_msa / MSA_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"
+    project_dir = tmp_path / "project"
+    (project_dir / "3rdparty" / "patches").mkdir(parents=True)
+    shutil.copytree(
+        source_msa, project_dir / "3rdparty" / "MSA", ignore=shutil.ignore_patterns(".git")
+    )
+    shutil.copy(REPO_ROOT / MSA_PATCH, project_dir / MSA_PATCH)
+    return project_dir
+
+
+def test_apply_msa_patch_is_idempotent_in_place(tmp_path):
+    project_dir = _stage_project(tmp_path)
+    patched_interface = project_dir / "3rdparty" / "MSA" / MSA_INTERFACE
+
+    apply_msa_patch(project_dir)
+    assert "def _prepare_paged_hnd_input" in patched_interface.read_text()
 
-    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
+    # A second call must short-circuit via the reverse-check guard rather than
+    # raise, leaving the patched content in place.
+    apply_msa_patch(project_dir)
+    assert "def _prepare_paged_hnd_input" in patched_interface.read_text()
 
 
-def test_stage_msa_package_requires_initialized_submodule(tmp_path):
+def test_apply_msa_patch_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")
+        apply_msa_patch(project_dir)