fix: pass total_tokens for SSM seq_idx in packed sequences

src/megatron/bridge/training/gpt_step.py

@@ -253,6 +253,7 @@ def _forward_step_common(
             "max_seqlen": max_seqlen,
             "cu_seqlens_unpadded": cu_seqlens_unpadded,
             "cu_seqlens_unpadded_argmin": cu_seqlens_unpadded_argmin,
+            "total_tokens": tokens.size(1) if tokens is not None else labels.size(1),
         }
         forward_args["packed_seq_params"] = get_packed_seq_params(packed_seq_params)
 

src/megatron/bridge/training/llava_step.py

@@ -198,6 +198,7 @@ def forward_step(
             "cu_seqlens": cu_seqlens,
             "cu_seqlens_argmin": cu_seqlens_argmin,
             "max_seqlen": max_seqlen,
+            "total_tokens": input_ids.size(1) if input_ids is not None else labels.size(1),
         }
         forward_args["packed_seq_params"] = get_packed_seq_params(packed_seq_params)
 

src/megatron/bridge/training/utils/packed_seq_utils.py

@@ -27,8 +27,9 @@ def get_packed_seq_params(batch: dict[str, torch.Tensor]) -> PackedSeqParams:
 
     Args:
         batch: A dictionary containing packed-sequence metadata. Expected keys:
-            `cu_seqlens`, optional `cu_seqlens_unpadded`, optional argmins, and
-            optional `max_seqlen`.
+            `cu_seqlens`, optional `cu_seqlens_unpadded`, optional argmins,
+            optional `max_seqlen`, and optional `total_tokens` (required for
+            hybrid SSM/Mamba models to generate ``seq_idx``).
 
     Returns:
         PackedSeqParams with identical q/kv parameters and `qkv_format` set to
@@ -57,6 +58,7 @@ def get_packed_seq_params(batch: dict[str, torch.Tensor]) -> PackedSeqParams:
             cu_seqlens_unpadded = cu_seqlens_unpadded[: torch.argmin(cu_seqlens_unpadded)]
 
     max_seqlen = batch["max_seqlen"].squeeze() if "max_seqlen" in batch else None
+    total_tokens = batch.get("total_tokens")
 
     # When cu_seqlens_unpadded is present (pad_seq_to_mult > 1), pass both unpadded and padded
     # for proper THD CP support. Otherwise, just use cu_seqlens_padded to avoid slower TE kernel.
@@ -68,6 +70,7 @@ def get_packed_seq_params(batch: dict[str, torch.Tensor]) -> PackedSeqParams:
             cu_seqlens_kv_padded=cu_seqlens_padded,
             max_seqlen_q=max_seqlen,
             max_seqlen_kv=max_seqlen,
+            total_tokens=total_tokens,
             qkv_format="thd",
         )
     else:
@@ -76,5 +79,6 @@ def get_packed_seq_params(batch: dict[str, torch.Tensor]) -> PackedSeqParams:
             cu_seqlens_kv=cu_seqlens_padded,
             max_seqlen_q=max_seqlen,
             max_seqlen_kv=max_seqlen,
+            total_tokens=total_tokens,
             qkv_format="thd",
         )

src/megatron/bridge/training/vlm_step.py

@@ -448,6 +448,7 @@ def forward_step(
             "cu_seqlens": cu_seqlens,
             "max_seqlen": max_seqlen,
             "cu_seqlens_argmin": cu_seqlens_argmin,
+            "total_tokens": tokens.size(1) if tokens is not None else labels.size(1),
         }
         forward_args["packed_seq_params"] = get_packed_seq_params(packed_seq_params)
 

tests/unit_tests/training/utils/test_packed_seq_utils.py

@@ -153,6 +153,60 @@ def test_single_sequence(self):
         torch.testing.assert_close(result.cu_seqlens_q, expected)
         assert result.cu_seqlens_q_padded is None  # No unpadded, so no padded variants
 
+    def test_total_tokens_generates_seq_idx(self):
+        """Test that passing total_tokens causes PackedSeqParams to generate seq_idx.
+
+        This is critical for hybrid SSM/Mamba models in varlen (packed sequence)
+        settings. Without total_tokens, seq_idx remains None and SSM state bleeds
+        across sequence boundaries.
+        """
+        batch = {
+            "cu_seqlens": torch.IntTensor([0, 5, 7, 11, -1]),
+            "cu_seqlens_argmin": torch.tensor(4),
+            "max_seqlen": torch.tensor(6),
+            "total_tokens": 16,
+        }
+
+        result = get_packed_seq_params(batch)
+
+        assert result.total_tokens == 16
+        assert result.seq_idx is not None
+        # seq_idx maps each token position to its sequence index:
+        # seq 0: tokens 0-4 (len 5), seq 1: tokens 5-6 (len 2),
+        # seq 2: tokens 7-10 (len 4), seq 3: tokens 11-15 (len 5)
+        expected_seq_idx = torch.IntTensor([[0, 0, 0, 0, 0, 1, 1, 2, 2, 2, 2, 3, 3, 3, 3, 3]])
+        torch.testing.assert_close(result.seq_idx, expected_seq_idx)
+
+    def test_without_total_tokens_seq_idx_is_none(self):
+        """Test that omitting total_tokens leaves seq_idx as None (backward compat)."""
+        batch = {
+            "cu_seqlens": torch.IntTensor([0, 128, 256, -1]),
+            "cu_seqlens_argmin": torch.tensor(3),
+            "max_seqlen": torch.tensor(128),
+        }
+
+        result = get_packed_seq_params(batch)
+
+        assert result.total_tokens is None
+        assert result.seq_idx is None
+
+    def test_total_tokens_with_cu_seqlens_unpadded(self):
+        """Test total_tokens flows through when cu_seqlens_unpadded is present."""
+        batch = {
+            "cu_seqlens": torch.IntTensor([0, 128, 256, 384, -1]),
+            "cu_seqlens_argmin": torch.tensor(4),
+            "cu_seqlens_unpadded": torch.IntTensor([0, 120, 245, 370, -1]),
+            "cu_seqlens_unpadded_argmin": torch.tensor(4),
+            "max_seqlen": torch.tensor(128),
+            "total_tokens": 384,
+        }
+
+        result = get_packed_seq_params(batch)
+
+        assert result.total_tokens == 384
+        # seq_idx should be generated from cu_seqlens_q_padded (the padded variant)
+        assert result.seq_idx is not None
+
     def test_performance_no_unnecessary_padded_variants(self):
         """Verify that when unpadded is not provided, padded variants are None.