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)