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