From a783a2658dff8847949e8de92535b5d216e7d450 Mon Sep 17 00:00:00 2001 From: trtllm-agent <296075020+trtllm-agent@users.noreply.github.com> Date: Sat, 8 Aug 2026 18:02:17 -0700 Subject: [PATCH] [nvbugs/6402058][fix] Densify MLA bmm LHS on SM120/121 to avoid TMA MMU fault MLA._bmm_bf16_out passes the LHS as a head-major transpose of a [tokens, heads, dim] buffer, so its batch stride is the token count rather than the tile extent. On SM120/121 cuBLAS selects a TMA-based nvjet kernel for that layout, which cuTensorMapEncodeTiled cannot describe; the kernel faults with an MMU page fault that surfaces either as CUBLAS_STATUS_INTERNAL_ERROR or as a later illegal memory access at the next synchronize. Make the LHS contiguous on those architectures when it is not already, so a non-TMA kernel is selected. Only the LHS matters: bisection showed densifying the output buffer alone does not avoid the fault. This mirrors the existing SM120 accommodations in fp8_block_scaling_bmm_out and the CUTLASS TMA guard in fused_moe_cutlass.py. Also unwaive the test the fault was masking. Signed-off-by: trtllm-agent <296075020+trtllm-agent@users.noreply.github.com> --- tensorrt_llm/_torch/modules/mla.py | 8 ++++++++ tests/integration/test_lists/waives.txt | 1 - 2 files changed, 8 insertions(+), 1 deletion(-) diff --git a/tensorrt_llm/_torch/modules/mla.py b/tensorrt_llm/_torch/modules/mla.py index 6bdf977d131e..05eca4a26af3 100644 --- a/tensorrt_llm/_torch/modules/mla.py +++ b/tensorrt_llm/_torch/modules/mla.py @@ -1391,6 +1391,14 @@ def _bmm_bf16_out(self, a, b_no_transpose, b_transposed, output): if self.use_cute_dsl_bf16_bmm and is_sm_100f(): torch.ops.trtllm.cute_dsl_bf16_bmm_blackwell(a, b_no_transpose, output) else: + if get_sm_version() in (120, 121) and not a.is_contiguous(): + # `a` is a head-major transpose of a [tokens, heads, dim] buffer, so its + # batch stride is the token count rather than the tile extent. cuBLAS picks + # a TMA-based nvjet kernel for that layout on SM120/121, and + # cuTensorMapEncodeTiled cannot describe it -- the kernel then faults with + # an MMU page fault (surfaced as CUBLAS_STATUS_INTERNAL_ERROR or a later + # illegal memory access). Densify so a non-TMA kernel is selected. + a = a.contiguous() torch.ops.trtllm.bmm_out(a, b_transposed, output) def forward_absorption_generation( diff --git a/tests/integration/test_lists/waives.txt b/tests/integration/test_lists/waives.txt index 6a003bf64f0a..42bf19039154 100644 --- a/tests/integration/test_lists/waives.txt +++ b/tests/integration/test_lists/waives.txt @@ -28,7 +28,6 @@ accuracy/test_llm_api_pytorch.py::TestDeepSeekV32::test_dsa_host_cache_offload[h accuracy/test_llm_api_pytorch.py::TestDeepSeekV32::test_dsa_host_cache_offload[host_cache_offload_mtp3_no_adp] SKIP (https://nvbugs/6384357) accuracy/test_llm_api_pytorch.py::TestDeepSeekV32::test_nvfp4_multi_gpus_piecewise_cuda_graph[mtp3_fp8kv_chunked] SKIP (https://nvbugs/5989920) accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_bfloat16_4gpus[ep4-mtp_nextn=2-attention_dp=True-cuda_graph=True-overlap_scheduler=True-torch_compile=False] SKIP (https://nvbugs/6517844) -accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_bfloat16_4gpus[ep4-mtp_nextn=2-attention_dp=True-cuda_graph=True-overlap_scheduler=True-torch_compile=True] SKIP (https://nvbugs/6402058) accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_bfloat16_4gpus[pp4-mtp_nextn=0-attention_dp=False-cuda_graph=False-overlap_scheduler=False-torch_compile=False] SKIP (https://nvbugs/6278337) accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_bfloat16_4gpus[pp4-mtp_nextn=0-attention_dp=False-cuda_graph=False-overlap_scheduler=False-torch_compile=True] SKIP (https://nvbugs/6428057) accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_bfloat16_4gpus[pp4-mtp_nextn=0-attention_dp=False-cuda_graph=True-overlap_scheduler=False-torch_compile=False] SKIP (https://nvbugs/6278337)