fix(dflash): remove incorrect padding shift and fix cross-block data leak

- Remove Eagle3-inherited padding (left-shift) calls in SGLang backend
  that misaligned input_ids/hidden_states with loss_mask. DFlash uses
  same-position prediction and does not need this shift.

- Fix cross-block data leak in random-anchor mode by changing context
  visibility from block-id comparison to original-position comparison,
  preventing overlapping blocks from leaking future hidden states.

Author: xiaomin-D

specforge/core/dflash.py

@@ -151,6 +151,8 @@ def _get_or_create_block_mask(
         kv_len: int,
         device: torch.device,
         block_ids: Optional[torch.Tensor] = None,
+        orig_positions: Optional[torch.Tensor] = None,
+        token_anchor_pos: Optional[torch.Tensor] = None,
     ) -> "BlockMask":
         """Get cached BlockMask or create a new one."""
         if block_ids is None:
@@ -165,17 +167,21 @@ def _get_or_create_block_mask(
 
         if block_ids is not None:
             _block_ids = block_ids
+            _orig_pos = orig_positions
+            _anchor_pos = token_anchor_pos
 
             def dflash_mask_fn(b, h, q_idx, kv_idx):
                 L = q_len
                 is_ctx = kv_idx < L
                 q_b = _block_ids[b, q_idx]
-                k_ctx = _block_ids[b, kv_idx.clamp(max=L - 1)]
+                k_ctx_id = _block_ids[b, kv_idx.clamp(max=L - 1)]
                 k_noise = _block_ids[b, (kv_idx - L).clamp(min=0, max=L - 1)]
                 q_valid = q_b >= 0
-                k_ctx_valid = k_ctx >= 0
+                k_ctx_valid = k_ctx_id >= 0
                 k_noise_valid = k_noise >= 0
-                ctx_visible = is_ctx & q_valid & k_ctx_valid & (k_ctx < q_b)
+                kv_orig = _orig_pos[b, kv_idx.clamp(max=L - 1)]
+                q_anchor = _anchor_pos[b, q_idx]
+                ctx_visible = is_ctx & q_valid & k_ctx_valid & (kv_orig < q_anchor)
                 noise_visible = (~is_ctx) & q_valid & k_noise_valid & (k_noise == q_b)
                 return ctx_visible | noise_visible
 
@@ -213,6 +219,8 @@ def _create_parallel_attention_mask(
         seq_len: int,
         device: torch.device,
         block_ids: Optional[torch.Tensor] = None,
+        orig_positions: Optional[torch.Tensor] = None,
+        token_anchor_pos: Optional[torch.Tensor] = None,
     ) -> torch.Tensor:
         """Create [bsz, L, 2L] attention mask for parallel training."""
         if block_ids is None:
@@ -226,11 +234,13 @@ def _create_parallel_attention_mask(
             full_mask.masked_fill_(~full_mask_bool, torch.finfo(torch.float32).min)
             return full_mask.unsqueeze(0).expand(bsz, -1, -1)
 
+        q_anchor = token_anchor_pos.unsqueeze(2)
+        k_orig = orig_positions.unsqueeze(1)
         q_ids = block_ids.unsqueeze(2)
         k_ids = block_ids.unsqueeze(1)
         q_valid = q_ids >= 0
         k_valid = k_ids >= 0
-        ctx_mask = q_valid & k_valid & (k_ids < q_ids)
+        ctx_mask = q_valid & k_valid & (k_orig < q_anchor)
         noise_mask = q_valid & k_valid & (k_ids == q_ids)
         full_mask_bool = torch.cat([ctx_mask, noise_mask], dim=2)
         full_mask = torch.zeros_like(full_mask_bool, dtype=torch.float32)
@@ -248,6 +258,8 @@ def forward(
         bsz, seq_len = input_ids.shape
         device = input_ids.device
         block_ids = None
+        orig_positions = None
+        token_anchor_pos = None
 
         if self.random_anchor and self.training:
             anchor_positions, block_keep_mask = self._sample_anchor_positions(
@@ -264,6 +276,10 @@ def forward(
             )
             effective_len = input_ids.shape[1]
             base_positions = block_positions
+            orig_positions = block_positions
+            token_anchor_pos = anchor_positions.repeat_interleave(
+                self.block_size, dim=1
+            )
         else:
             n_blocks = seq_len // self.block_size
             effective_len = n_blocks * self.block_size
@@ -291,10 +307,12 @@ def forward(
                 kv_len=effective_len * 2,
                 device=device,
                 block_ids=block_ids,
+                orig_positions=orig_positions,
+                token_anchor_pos=token_anchor_pos,
             )
         else:
             dflash_attn_mask = self._create_parallel_attention_mask(
-                bsz, effective_len, device, block_ids
+                bsz, effective_len, device, block_ids, orig_positions, token_anchor_pos
             )
             dflash_attn_mask = dflash_attn_mask.to(dtype=hidden_states.dtype)
             dflash_attn_mask = dflash_attn_mask.unsqueeze(1)

specforge/modeling/target/dflash_target_model.py

@@ -18,7 +18,6 @@
 from transformers import AutoModelForCausalLM
 
 from specforge.distributed import get_tp_group
-from specforge.utils import padding
 
 from .sglang_backend import SGLangRunner
 
@@ -235,10 +234,6 @@ def generate_dflash_data(
         attention_mask = torch.cat([d[1] for d in data_cache], dim=0)
         loss_mask = torch.cat([d[2] for d in data_cache], dim=0)
 
-        # Padding might be needed if batching varied lengths (but usually fixed length training)
-        hidden_states = padding(hidden_states, left=False)
-        input_ids = padding(input_ids, left=False)
-
         return DFlashTargetOutput(
             hidden_states=hidden_states,
             input_ids=input_ids,