[None][perf] Avoid paged MSA K/V materialization during prefill

.gitignore

@@ -57,6 +57,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

3rdparty/CMakeLists.txt

@@ -101,11 +101,27 @@ foreach(DEP_IDX RANGE ${DEP_COUNT_MINUS_ONE})
   # Extract optional fields with error handling
   string(JSON DEP_GIT_SHALLOW            ERROR_VARIABLE _err GET "${DEP_OBJECT}" "git_shallow")
   string(JSON DEP_SOURCE_SUBDIR          ERROR_VARIABLE _err GET "${DEP_OBJECT}" "source_subdir")
+  string(JSON DEP_GIT_SUBMODULES         ERROR_VARIABLE _git_submodules_error GET "${DEP_OBJECT}" "git_submodules")
   string(JSON DEP_GIT_SUBMODULES_RECURSE ERROR_VARIABLE _err GET "${DEP_OBJECT}" "git_submodules_recurse")
   string(JSON DEP_USE_URL                ERROR_VARIABLE _err GET "${DEP_OBJECT}" "use_url")
   string(JSON DEP_PATCH_FILE            ERROR_VARIABLE _err GET "${DEP_OBJECT}" "patch_file")
   # cmake-format: on
 
+  if(DEP_NAME STREQUAL "cutlass")
+    set(TRTLLM_CUTLASS_TAG "${DEP_GIT_TAG}")
+  elseif(DEP_NAME STREQUAL "msa")
+    # MSA JIT-compiles against the CUTLASS headers staged by TRT-LLM. Force
+    # revalidation when TRT-LLM changes its CUTLASS pin.
+    set(MSA_VALIDATED_CUTLASS_TAG "v4.4.2")
+    if(NOT MSA_VALIDATED_CUTLASS_TAG STREQUAL TRTLLM_CUTLASS_TAG)
+      message(
+        FATAL_ERROR
+          "MSA ${DEP_GIT_TAG} was validated with CUTLASS ${MSA_VALIDATED_CUTLASS_TAG}, "
+          "but fetch_content.json selects ${TRTLLM_CUTLASS_TAG}. Revalidate MSA and "
+          "update MSA_VALIDATED_CUTLASS_TAG.")
+    endif()
+  endif()
+
   # Build FetchContent_Declare arguments
   set(FETCH_ARGS "${DEP_NAME}")
 
@@ -135,11 +151,15 @@ 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}' || (echo 'Patch state for ${DEP_NAME} is inconsistent. Remove ${CMAKE_BINARY_DIR}/_deps/${DEP_NAME}-src and reconfigure.' >&2 && false)"
     )
   endif()
 
-  FetchContent_Declare(${FETCH_ARGS})
+  if(_git_submodules_error STREQUAL "NOTFOUND")
+    FetchContent_Declare(${FETCH_ARGS} GIT_SUBMODULES "${DEP_GIT_SUBMODULES}")
+  else()
+    FetchContent_Declare(${FETCH_ARGS})
+  endif()
 
   # Special handling: Export deep_ep commit to global property
   if(DEP_NAME STREQUAL "deep_ep_download")

3rdparty/fetch_content.json

@@ -68,6 +68,14 @@
       "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",
+      "git_submodules": "",
+      "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",

3rdparty/patches/msa_strided_paged_kv.patch

@@ -0,0 +1,225 @@
+diff --git a/python/fmha_sm100/cute/interface.py b/python/fmha_sm100/cute/interface.py
+index d72b17a..eca5b8c 100644
+--- a/python/fmha_sm100/cute/interface.py
++++ b/python/fmha_sm100/cute/interface.py
+@@ -136,6 +136,35 @@ 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:
++        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])
++    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) >= 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 +765,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..b5f078b 100644
+--- a/python/fmha_sm100/cute/test_sparse_atten.py
++++ b/python/fmha_sm100/cute/test_sparse_atten.py
+@@ -61,6 +61,81 @@ 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)
++
++
++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 +1861,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])

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)

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

jenkins/UpdateTestDurations.groovy

@@ -130,7 +130,8 @@ 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)
+                        // Initialize submodules for older arbitrary refs; this is a no-op after their removal.
+                        trtllm_utils.checkoutSource(sourceRepo, params.TARGET_BRANCH, LLM_ROOT, true, false)
                     }
                 }
             }

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