From a4ffe09968720b85cf5fb7cdb055f677a87c1e6a Mon Sep 17 00:00:00 2001 From: Gal Hubara Agam <96368689+galagam@users.noreply.github.com> Date: Tue, 26 May 2026 06:18:17 +0000 Subject: [PATCH 1/3] [https://nvbugs/6215690][fix] AutoDeploy: Fix alignment issue in FI SSM The decode path requires contiguous aligned to 128 bytes. Add _fi_align() before passing x/B/C to FlashInfer SSM kernel in the decode path. data.contiguous() is the real fix: with TP >= 8, the hidden-state tensor is a strided view (non-contiguous), which FlashInfer rejects. This regression was invisible at world_size=4 where the same sharding happened to produce a contiguous layout. clone() is purely defensive: for the current Mamba2 configs, row sizes (num_heads * head_dim * dtype_size and n_groups * ssm_state_size * dtype_size) are always multiples of 128 bytes, so a contiguous slice can never be misaligned. The clone() path guards against future configs where that assumption breaks. The extend path does not require contiguous memory and remains as-is. Signed-off-by: Gal Hubara Agam <96368689+galagam@users.noreply.github.com> --- .../custom_ops/mamba/flashinfer_backend_mamba.py | 15 +++++++++++++++ tests/integration/test_lists/waives.txt | 1 - 2 files changed, 15 insertions(+), 1 deletion(-) diff --git a/tensorrt_llm/_torch/auto_deploy/custom_ops/mamba/flashinfer_backend_mamba.py b/tensorrt_llm/_torch/auto_deploy/custom_ops/mamba/flashinfer_backend_mamba.py index 67ad432a2231..563cfe34c9b9 100644 --- a/tensorrt_llm/_torch/auto_deploy/custom_ops/mamba/flashinfer_backend_mamba.py +++ b/tensorrt_llm/_torch/auto_deploy/custom_ops/mamba/flashinfer_backend_mamba.py @@ -36,6 +36,17 @@ ) +def _fi_align(t: torch.Tensor) -> torch.Tensor: + """Ensure 128-byte alignment required by FlashInfer kernels. + + - Contiguous + aligned: contiguous() is a no-op, returns t unchanged. + - Non-contiguous: contiguous() allocates fresh aligned storage, returns it. + - Contiguous + misaligned: contiguous() is a no-op, clone() forces a new aligned allocation. + """ + t = t.contiguous() + return t if t.data_ptr() % 128 == 0 else t.clone() + + @torch.library.custom_op( "auto_deploy::flashinfer_cached_ssm", mutates_args=("ssm_state_cache", "intermediate_ssm_state_cache"), @@ -205,6 +216,10 @@ def _flashinfer_cached_ssm( D_full, ) = decode_inputs + x_decode = _fi_align(x_decode) + B_decode = _fi_align(B_decode) + C_decode = _fi_align(C_decode) + slot_idx_decode_i32 = slot_idx_decode.to(torch.int32) y_decode = _flashinfer_ssm_update( ssm_state_cache, diff --git a/tests/integration/test_lists/waives.txt b/tests/integration/test_lists/waives.txt index a53606807422..3e1f556fe9d9 100644 --- a/tests/integration/test_lists/waives.txt +++ b/tests/integration/test_lists/waives.txt @@ -17,7 +17,6 @@ accuracy/test_llm_api_autodeploy.py::TestMiniMaxM2::test_finegrained_fp8 SKIP (h accuracy/test_llm_api_autodeploy.py::TestNemotronNanoV3::test_accuracy[nvfp4-1-trtllm] SKIP (https://nvbugs/6200112) accuracy/test_llm_api_autodeploy.py::TestNemotronSuperV3::test_functional_small[bf16] SKIP (https://nvbugs/6162114) accuracy/test_llm_api_autodeploy.py::TestNemotronSuperV3::test_functional_small[fp8] SKIP (https://nvbugs/6162114) -accuracy/test_llm_api_autodeploy.py::TestNemotronUltraV3::test_accuracy[nvfp4-8] SKIP (https://nvbugs/6215690) accuracy/test_llm_api_autodeploy.py::TestQwen3_5_397B_MoE::test_bf16_small[4] SKIP (https://nvbugs/6158397) accuracy/test_llm_api_pytorch.py::TestDeepSeekR1::test_fp8_blockscale[throughput_mtp] SKIP (https://nvbugs/6215736) accuracy/test_llm_api_pytorch.py::TestDeepSeekR1::test_nvfp4_multi_gpus[throughput] SKIP (https://nvbugs/6084775) From 1e79f7640478d108caafb5f55ba4b1977b4e72ae Mon Sep 17 00:00:00 2001 From: Gal Hubara Agam <96368689+galagam@users.noreply.github.com> Date: Tue, 26 May 2026 06:27:18 +0000 Subject: [PATCH 2/3] also unwaive TestNemotronSuperV3::test_functional_small Signed-off-by: Gal Hubara Agam <96368689+galagam@users.noreply.github.com> --- tests/integration/test_lists/waives.txt | 2 -- 1 file changed, 2 deletions(-) diff --git a/tests/integration/test_lists/waives.txt b/tests/integration/test_lists/waives.txt index 3e1f556fe9d9..44bbd74d0ee1 100644 --- a/tests/integration/test_lists/waives.txt +++ b/tests/integration/test_lists/waives.txt @@ -15,8 +15,6 @@ accuracy/test_llm_api_autodeploy.py::TestGemma4MoE::test_bf16 SKIP (https://nvbu accuracy/test_llm_api_autodeploy.py::TestGemmaE2B::test_gemma4_e2b_it SKIP (https://nvbugs/6194934) accuracy/test_llm_api_autodeploy.py::TestMiniMaxM2::test_finegrained_fp8 SKIP (https://nvbugs/6158397) accuracy/test_llm_api_autodeploy.py::TestNemotronNanoV3::test_accuracy[nvfp4-1-trtllm] SKIP (https://nvbugs/6200112) -accuracy/test_llm_api_autodeploy.py::TestNemotronSuperV3::test_functional_small[bf16] SKIP (https://nvbugs/6162114) -accuracy/test_llm_api_autodeploy.py::TestNemotronSuperV3::test_functional_small[fp8] SKIP (https://nvbugs/6162114) accuracy/test_llm_api_autodeploy.py::TestQwen3_5_397B_MoE::test_bf16_small[4] SKIP (https://nvbugs/6158397) accuracy/test_llm_api_pytorch.py::TestDeepSeekR1::test_fp8_blockscale[throughput_mtp] SKIP (https://nvbugs/6215736) accuracy/test_llm_api_pytorch.py::TestDeepSeekR1::test_nvfp4_multi_gpus[throughput] SKIP (https://nvbugs/6084775) From 1292c5a32200e21871d931bd62ce520f9625739b Mon Sep 17 00:00:00 2001 From: Gal Hubara Agam <96368689+galagam@users.noreply.github.com> Date: Tue, 26 May 2026 12:33:30 -0700 Subject: [PATCH 3/3] waive https://nvbugs/6221483 Signed-off-by: Gal Hubara Agam <96368689+galagam@users.noreply.github.com> --- tests/integration/test_lists/waives.txt | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/tests/integration/test_lists/waives.txt b/tests/integration/test_lists/waives.txt index 44bbd74d0ee1..cca261193925 100644 --- a/tests/integration/test_lists/waives.txt +++ b/tests/integration/test_lists/waives.txt @@ -15,6 +15,10 @@ accuracy/test_llm_api_autodeploy.py::TestGemma4MoE::test_bf16 SKIP (https://nvbu accuracy/test_llm_api_autodeploy.py::TestGemmaE2B::test_gemma4_e2b_it SKIP (https://nvbugs/6194934) accuracy/test_llm_api_autodeploy.py::TestMiniMaxM2::test_finegrained_fp8 SKIP (https://nvbugs/6158397) accuracy/test_llm_api_autodeploy.py::TestNemotronNanoV3::test_accuracy[nvfp4-1-trtllm] SKIP (https://nvbugs/6200112) +accuracy/test_llm_api_autodeploy.py::TestNemotronSuperV3::test_mtp[bf16_ws4_180gb-flashinfer] SKIP (https://nvbugs/6221483) +accuracy/test_llm_api_autodeploy.py::TestNemotronSuperV3::test_mtp[bf16_ws4_180gb-trtllm] SKIP (https://nvbugs/6221483) +accuracy/test_llm_api_autodeploy.py::TestNemotronSuperV3::test_mtp[fp8_ws4_80gb-trtllm] SKIP (https://nvbugs/6221483) +accuracy/test_llm_api_autodeploy.py::TestNemotronSuperV3::test_mtp[nvfp4_ws4_80gb-trtllm] SKIP (https://nvbugs/6221483) accuracy/test_llm_api_autodeploy.py::TestQwen3_5_397B_MoE::test_bf16_small[4] SKIP (https://nvbugs/6158397) accuracy/test_llm_api_pytorch.py::TestDeepSeekR1::test_fp8_blockscale[throughput_mtp] SKIP (https://nvbugs/6215736) accuracy/test_llm_api_pytorch.py::TestDeepSeekR1::test_nvfp4_multi_gpus[throughput] SKIP (https://nvbugs/6084775)