[dev] [DeepSeek-v4] Packed Sequence (THD) support for DSv4 Hybrid Attention

megatron/core/datasets/data_schedule.py

@@ -778,12 +778,14 @@ def get_batch_on_this_rank_for_sequence_packing(
         else None
     )
 
-    # Use padded cumulative lengths for THD partitioning so token slices follow
-    # the padded sequence boundaries consumed by attention kernels.
+    # cu_seqlens_q/kv hold the original (unpadded) boundaries so downstream
+    # loss paths (e.g. CSA indexer KL) can identify padding rows.
+    # cu_seqlens_q/kv_padded hold the padded boundaries consumed by attention
+    # kernels and THD partitioning.
     packed_seq_params = PackedSeqParams(
         qkv_format="thd",
-        cu_seqlens_q=cu_seqlens_padded,
-        cu_seqlens_kv=cu_seqlens_padded,
+        cu_seqlens_q=cu_seqlens,
+        cu_seqlens_kv=cu_seqlens,
         cu_seqlens_q_padded=cu_seqlens_padded,
         cu_seqlens_kv_padded=cu_seqlens_padded,
         max_seqlen_q=max_seqlen,

megatron/core/fusions/fused_mla_yarn_rope_apply.py

@@ -29,17 +29,28 @@
 
 @triton.jit
 def _get_thd_token_idx(cu_seqlens, pid_m, seq_num, cp_rank, cp_size):
-    token_idx = -1
-    this_seq_len = 0
+    # Cast ``pid_m`` and ``cu_seqlens`` loads to a single shared dtype so
+    # the loop-body reassignments don't surface as
+    # "initial value is int32 but redefined as int64" in newer Triton
+    # versions (which promote ``// Python_int`` to int64).
+    pid_m = pid_m.to(tl.int64)
+    token_idx = tl.full((), -1, dtype=tl.int64)
+    this_seq_len = tl.full((), 0, dtype=tl.int64)
     seq_idx = 0
-    last_cum_seqlen = tl.load(cu_seqlens) // cp_size
+    last_cum_seqlen = tl.load(cu_seqlens).to(tl.int64) // cp_size
     while seq_idx < seq_num:
-        cur_cum_seqlen = tl.load(cu_seqlens + seq_idx + 1) // cp_size
+        cur_cum_seqlen = tl.load(cu_seqlens + seq_idx + 1).to(tl.int64) // cp_size
         if token_idx == -1 and cur_cum_seqlen > pid_m:
             token_idx = pid_m - last_cum_seqlen
             this_seq_len = cur_cum_seqlen - last_cum_seqlen
         last_cum_seqlen = cur_cum_seqlen
         seq_idx += 1
+    # Padding tokens beyond cu_seqlens[-1] (from THD CUDA-graph padding)
+    # never match any sequence, leaving token_idx == -1.  Clamp to 0 so
+    # the cos/sin table loads stay in-bounds; the wrong RoPE result is
+    # harmless because padding positions are excluded by loss_mask.
+    if token_idx == -1:
+        token_idx = tl.full((), 0, dtype=tl.int64)
     if cp_size > 1:
         if token_idx < this_seq_len // 2:
             token_idx = token_idx + cp_rank * this_seq_len // 2

megatron/core/models/hybrid/hybrid_model.py

@@ -557,6 +557,7 @@ def forward(
                 rotary_pos_emb=rotary_pos_emb,
                 packed_seq_params=packed_seq_params,
                 embedding=self.embedding,
+                padding_mask=padding_mask,
             )
 
         if not self.post_process: