[https://nvbugs/6120981][fix] Switch to cu_seqlens_to_chunk_indices_offsets_triton with total_seqlens/extra_ch

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

@@ -19,7 +19,10 @@
 from torch._ops import OpOverloadPacket
 from torch.fx import Node
 
-from tensorrt_llm._torch.modules.mamba.mamba2_metadata import cu_seqlens_to_chunk_indices_offsets
+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.ssd_combined import mamba_chunk_scan_combined
 
 from ..._compat import KvCacheConfig
@@ -41,25 +44,39 @@ 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:
-        chunk_indices, chunk_offsets = cu_seqlens_to_chunk_indices_offsets(
-            cu_seqlen[: num_prefill + 1], chunk_size
+        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,
         )
         seq_idx_prefill = torch.repeat_interleave(
-            torch.arange(num_prefill, device=device, dtype=torch.int32), seq_len[:num_prefill]
+            torch.arange(num_prefill, device=device, dtype=torch.int32),
+            seq_len[:num_prefill],
+            output_size=num_prefill_tokens,
         ).view(1, -1)
     else:
         chunk_indices = torch.empty(0, dtype=torch.int32, device=device)
@@ -75,6 +92,7 @@ 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,

tensorrt_llm/_torch/modules/mamba/mamba2_metadata.py

@@ -63,6 +63,20 @@ def _cu_seqlens_triton_kernel(
     tl.store(chunk_offsets_ptr + offsets, chunk_offsets.to(tl.int32), mask=mask)
 
 
+def compute_extra_chunks_cpu(seq_lens, num_seqs: int, chunk_size: int) -> int:
+    """Count extra chunks caused by misaligned sequence boundaries.
+
+    Computes from CPU seq_lens to avoid GPU->CPU synchronization.
+    """
+    cumsum = 0
+    extra = 0
+    for i in range(num_seqs - 1):
+        cumsum += int(seq_lens[i])
+        if cumsum % chunk_size != 0:
+            extra += 1
+    return extra
+
+
 def cu_seqlens_to_chunk_indices_offsets_triton(
         cu_seqlens: torch.Tensor,
         chunk_size: int,
@@ -327,15 +341,8 @@ def prepare(self, attn_metadata: AttentionMetadata):
                 self.has_initial_states_cpu[:num_contexts].any())
 
             if self.use_initial_states:
-                # Compute extra_chunks using pure Python arithmetic on CPU
-                # seq_lens to avoid any GPU->CPU sync point.
-                _cs = self.chunk_size
-                _cumsum = 0
-                _extra = 0
-                for i in range(num_contexts - 1):
-                    _cumsum += int(attn_metadata.seq_lens[i])
-                    if _cumsum % _cs != 0:
-                        _extra += 1
+                _extra = compute_extra_chunks_cpu(attn_metadata.seq_lens,
+                                                  num_contexts, self.chunk_size)
 
                 self.chunk_indices, self.chunk_offsets = cu_seqlens_to_chunk_indices_offsets_triton(
                     self.cu_seqlens[:num_contexts + 1],

tests/integration/test_lists/waives.txt

@@ -15,7 +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::TestNemotronNanoV3::test_accuracy[nvfp4-4-trtllm] SKIP (https://nvbugs/6120981)
 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)