[https://nvbugs/6221483][fix] Revert auto_deploy _mamba_ssm_prepare_metadata to pre-#13566 state

tensorrt_llm/_torch/auto_deploy/custom_ops/mamba/mamba_backend_common.py

@@ -19,10 +19,7 @@
 from torch._ops import OpOverloadPacket
 from torch.fx import Node
 
-from tensorrt_llm._torch.modules.mamba.mamba2_metadata import (
-    compute_extra_chunks_cpu,
-    cu_seqlens_to_chunk_indices_offsets_triton,
-)
+from tensorrt_llm._torch.modules.mamba.mamba2_metadata import cu_seqlens_to_chunk_indices_offsets
 from tensorrt_llm._torch.modules.mamba.ssd_combined import mamba_chunk_scan_combined
 
 from ..._compat import KvCacheConfig
@@ -44,39 +41,25 @@ def _mamba_ssm_prepare_metadata(
     position_ids: torch.Tensor,
     batch_info_host: torch.Tensor,
     seq_len: torch.Tensor,
-    seq_len_host: torch.Tensor,
     cu_seqlen: torch.Tensor,
     # EXTRA METADATA PROVIDED BY THE DESCRIPTOR
     chunk_size: int,
 ) -> List[torch.Tensor]:
     """Prepare metadata for cached SSM transform.
 
     Returns a tuple of (chunk_indices, chunk_offsets, seq_idx_prefill).
-
-    Uses seq_len_host (CPU tensor) and batch_info_host to derive total_seqlens
-    and extra_chunks without GPU->CPU synchronization, preventing deadlocks
-    when NCCL collectives from MoE expert-parallel layers are pending.
     """
     device = cu_seqlen.device
     batch_info = BatchInfo(batch_info_host)
 
     num_prefill, _, _ = batch_info.get_num_sequences()
 
     if num_prefill > 0:
-        num_prefill_tokens, _, _ = batch_info.get_num_tokens()
-
-        _extra = compute_extra_chunks_cpu(seq_len_host, num_prefill, chunk_size)
-
-        chunk_indices, chunk_offsets = cu_seqlens_to_chunk_indices_offsets_triton(
-            cu_seqlen[: num_prefill + 1],
-            chunk_size,
-            total_seqlens=num_prefill_tokens,
-            extra_chunks=_extra,
+        chunk_indices, chunk_offsets = cu_seqlens_to_chunk_indices_offsets(
+            cu_seqlen[: num_prefill + 1], chunk_size
         )
         seq_idx_prefill = torch.repeat_interleave(
-            torch.arange(num_prefill, device=device, dtype=torch.int32),
-            seq_len[:num_prefill],
-            output_size=num_prefill_tokens,
+            torch.arange(num_prefill, device=device, dtype=torch.int32), seq_len[:num_prefill]
         ).view(1, -1)
     else:
         chunk_indices = torch.empty(0, dtype=torch.int32, device=device)
@@ -92,7 +75,6 @@ def _mamba_ssm_prepare_metadata_fake(
     position_ids: torch.Tensor,
     batch_info_host: torch.Tensor,
     seq_len: torch.Tensor,
-    seq_len_host: torch.Tensor,
     cu_seqlen: torch.Tensor,
     # EXTRA METADATA PROVIDED BY THE DESCRIPTOR
     chunk_size: int,

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_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_autodeploy.py::TestQwen3_5_397B_MoE::test_nvfp4[8] SKIP (https://nvbugs/6211441)
 accuracy/test_llm_api_pytorch.py::TestDeepSeekR1::test_nvfp4_multi_gpus[throughput] SKIP (https://nvbugs/6084775)