Switch to producing document_ids in collate fn

Signed-off-by: Fynn Schmitt-Ulms <fschmitt@redhat.com>

Author: fynnsu

src/speculators/models/dflash/attention.py

@@ -5,7 +5,7 @@
 
 
 def create_anchor_block_mask_mod(
-    lengths: torch.Tensor,
+    document_ids: torch.Tensor,
     total_seq_len: int,
     anchor_positions: torch.Tensor,
     block_size: int,
@@ -29,7 +29,7 @@ def create_anchor_block_mask_mod(
         - may not attend to other synthetic blocks or later base tokens
 
     Args:
-        lengths: [num_docs] lengths of packed documents
+        document_ids: [total_seq_len] maps each position to its document index, -1 for padding
         total_seq_len: padded packed sequence width
         anchor_positions: [n_anchors] absolute positions into the packed base sequence
         block_size: number of query tokens per anchor block
@@ -42,7 +42,7 @@ def create_anchor_block_mask_mod(
     # Always use non_causal for full attn
     non_causal = sliding_window is None or sliding_window_non_causal
 
-    device = lengths.device
+    device = document_ids.device
     anchor_positions = anchor_positions.to(device=device, dtype=torch.long).contiguous()
 
     if anchor_positions.ndim != 1:
@@ -54,28 +54,6 @@ def create_anchor_block_mask_mod(
     q_len = n_anchors * block_size
     kv_len = total_seq_len + q_len
 
-    # Map each base-sequence position -> document id, padding -> -1
-    document_ids = torch.repeat_interleave(
-        torch.arange(lengths.shape[0], device=device, dtype=torch.long),
-        lengths,
-    )
-    if document_ids.numel() > total_seq_len:
-        raise ValueError(
-            f"sum(lengths)={document_ids.numel()} exceeds total_seq_len={total_seq_len}"
-        )
-    if document_ids.numel() < total_seq_len:
-        document_ids = torch.cat(
-            [
-                document_ids,
-                -1
-                * torch.ones(
-                    total_seq_len - document_ids.numel(),
-                    device=device,
-                    dtype=torch.long,
-                ),
-            ]
-        ).contiguous()
-
     if (oob := (anchor_positions < 0) | (anchor_positions >= total_seq_len)).any():
         raise ValueError(
             f"anchor_positions out of range: {anchor_positions[oob].tolist()}"

src/speculators/models/dflash/core.py

@@ -207,15 +207,15 @@ def mask_token_id(self) -> int:
     @torch.compiler.disable
     def _create_attention_mask(
         self,
-        lengths: torch.Tensor,
+        document_ids: torch.Tensor,
         total_seq_len: int,
         anchor_positions: torch.Tensor,
         device: torch.device,
         sliding_window: int | None = None,
         sliding_window_non_causal: bool = False,
     ):
         mask_mod, q_len, kv_len = create_anchor_block_mask_mod(
-            lengths=lengths.to(device),
+            document_ids=document_ids.squeeze(0).to(device),
             total_seq_len=total_seq_len,
             anchor_positions=anchor_positions,
             block_size=self.block_size,
@@ -232,7 +232,7 @@ def _create_attention_mask(
         )
 
     @torch.compiler.disable
-    def _build_attention_mask(self, loss_mask, lengths, device):
+    def _build_attention_mask(self, loss_mask, document_ids, device):
         total_seq_len = loss_mask.shape[1]
 
         anchor_positions, anchor_valid = select_anchors(
@@ -242,7 +242,7 @@ def _build_attention_mask(self, loss_mask, lengths, device):
         full_attn_mask = None
         if self.uses_full_attn:
             full_attn_mask = self._create_attention_mask(
-                lengths=lengths,
+                document_ids=document_ids,
                 total_seq_len=total_seq_len,
                 anchor_positions=anchor_positions,
                 device=device,
@@ -252,7 +252,7 @@ def _build_attention_mask(self, loss_mask, lengths, device):
         sliding_window_attn_mask = None
         if self.uses_sliding_window_attn:
             sliding_window_attn_mask = self._create_attention_mask(
-                lengths=lengths,
+                document_ids=document_ids,
                 total_seq_len=total_seq_len,
                 anchor_positions=anchor_positions,
                 device=device,
@@ -269,7 +269,7 @@ def forward(
         input_ids: torch.Tensor,  # shape: [1, total_seq_len]
         loss_mask: torch.Tensor,  # shape: [1, total_seq_len]
         verifier_last_hidden_states: torch.Tensor,  # shape: [1, total_seq_len, hidden_size] # noqa: E501
-        lengths: torch.Tensor | None = None,  # shape: [batch_size]
+        document_ids: torch.Tensor,  # shape: [1, total_seq_len]
         position_ids: torch.Tensor | None = None,  # shape: [1, total_seq_len]
         loss_fn=kl_div_loss,
         **kwargs,
@@ -278,15 +278,13 @@ def forward(
         total_seq_len = hidden_states.shape[1]
         num_anchors = self.config.max_anchors
 
-        if lengths is None:
-            lengths = torch.tensor([total_seq_len], dtype=torch.long, device=device)
         if position_ids is None:
             position_ids = 1 + torch.arange(
                 total_seq_len, dtype=torch.long, device=device
             ).unsqueeze(0)
 
         full_attn_mask, sliding_window_attn_mask, anchor_positions, anchor_valid = (
-            self._build_attention_mask(loss_mask, lengths, device)
+            self._build_attention_mask(loss_mask, document_ids, device)
         )
 
         mask_tokens_size = num_anchors * self.block_size

src/speculators/models/eagle3/attention.py

@@ -8,23 +8,7 @@
 from speculators.models.attention import ALL_ATTENTION_FUNCTIONS  # noqa: F401
 
 
-def create_combined_mask_mod(lengths: torch.Tensor, total_seq_len: int):
-    document_ids = torch.repeat_interleave(
-        torch.arange(lengths.shape[0], device=lengths.device, dtype=torch.long), lengths
-    )
-    # Pad ids with -1 to indicate padding
-    document_ids = torch.cat(
-        [
-            document_ids,
-            -1
-            * torch.ones(
-                total_seq_len - document_ids.shape[0],
-                device=lengths.device,
-                dtype=torch.long,
-            ),
-        ]
-    ).contiguous()
-
+def create_combined_mask_mod(document_ids: torch.Tensor, total_seq_len: int):
     def causal_mask_mod(_b, _h, q_idx, kv_idx):
         return q_idx >= kv_idx
 

src/speculators/models/eagle3/core.py

@@ -133,7 +133,7 @@ def forward(  # noqa: C901
         self,
         hidden_states: torch.Tensor,  # shape: [1, total_seq_len, 3 * hidden_size]
         input_ids: torch.Tensor,  # shape: [1, total_seq_len]
-        lengths: torch.Tensor | None = None,  # shape: [batch_size]
+        document_ids: torch.Tensor,  # shape: [1, total_seq_len]
         loss_mask: torch.Tensor | None = None,  # shape: [1, total_seq_len]
         position_ids: torch.Tensor | None = None,  # shape: [1, total_seq_len]
         verifier_last_hidden_states: torch.Tensor
@@ -148,8 +148,6 @@ def forward(  # noqa: C901
         device = hidden_states.device
         total_seq_len = hidden_states.shape[1]
 
-        if lengths is None:
-            lengths = torch.tensor([total_seq_len], dtype=torch.long, device=device)
         if position_ids is None:
             position_ids = 1 + torch.arange(
                 total_seq_len, dtype=torch.long, device=device
@@ -158,7 +156,9 @@ def forward(  # noqa: C901
 
         past_key_values = DynamicCache(config=self.config.transformer_layer_config)
 
-        combined_mask_mod = create_combined_mask_mod(lengths.to(device), total_seq_len)
+        combined_mask_mod = create_combined_mask_mod(
+            document_ids.squeeze(0).to(device), total_seq_len
+        )
         # Note: Attention mask is stored as a BlockMask object
         attention_mask = create_block_mask(
             combined_mask_mod,

src/speculators/models/peagle/attention.py

@@ -6,8 +6,7 @@
 def create_peagle_mask_mod(
     anchor_pos: torch.Tensor,  # shape: [total_sampled]
     depth: torch.Tensor,  # shape: [total_sampled]
-    lengths: torch.Tensor,  # shape: [batch_size]
-    total_seq_len: int,
+    document_ids: torch.Tensor,  # shape: [total_seq_len]
 ):
     """
     Create a flex attention mask modifier for P-EAGLE parallel groups.
@@ -23,9 +22,7 @@ def create_peagle_mask_mod(
         anchor_pos: The starting position in the original sequence the current
             sampling chain started from.
         depth: Which COD sampling round each element belongs to
-        lengths: The length of each document. Used to produce a document mask to prevent
-            cross contamination
-        total_seq_len: int, combined padded length of the original sequences
+        document_ids: Maps each position to its document index, -1 for padding
 
     Args example:
 
@@ -45,23 +42,6 @@ def create_peagle_mask_mod(
         A mask_mod function compatible with flex_attention create_block_mask
     """
 
-    # Generate sample_ids to prevent cross-sample attention
-    document_ids = torch.repeat_interleave(
-        torch.arange(lengths.shape[0], device=lengths.device, dtype=torch.long), lengths
-    )
-    # Pad ids with -1 to indicate padding
-    document_ids = torch.cat(
-        [
-            document_ids,
-            -1
-            * torch.ones(
-                total_seq_len - document_ids.shape[0],
-                device=lengths.device,
-                dtype=torch.long,
-            ),
-        ]
-    ).contiguous()
-
     def peagle_mask_mod(_b, _h, q_idx, kv_idx):
         q_anchor_pos = anchor_pos[q_idx]
         kv_anchor_pos = anchor_pos[kv_idx]

src/speculators/models/peagle/core.py

@@ -52,7 +52,7 @@ def forward(
         self,
         hidden_states: torch.Tensor,
         input_ids: torch.Tensor,
-        lengths: torch.Tensor | None = None,
+        document_ids: torch.Tensor,
         position_ids: torch.Tensor | None = None,
         loss_mask: torch.Tensor | None = None,
         verifier_last_hidden_states: torch.Tensor | None = None,
@@ -62,12 +62,10 @@ def forward(
         """
         Forward pass for P-EAGLE model training with parallel group prediction.
 
-        Matches p-eagle-train implementation but accepts standard EAGLE3 data format.
-
         Args:
             hidden_states: Verifier hidden states [batch, seq_len, 3*hidden_size]
             input_ids: Input token IDs [batch, seq_len]
-            lengths: Sequence lengths for each sample in batch [batch_size]
+            document_ids: Document IDs [1, seq_len], maps positions to document index, -1 for padding
             position_ids: Position IDs [batch, seq_len] (optional)
             loss_mask: Loss mask for which tokens to compute loss on
                 [batch, seq_len]
@@ -83,8 +81,6 @@ def forward(
         device = hidden_states.device
         seq_length = input_ids.shape[1]
 
-        if lengths is None:
-            lengths = torch.tensor([seq_length], dtype=torch.long, device=device)
         if loss_mask is None:
             loss_mask = torch.ones_like(input_ids, dtype=torch.float32)
 
@@ -133,8 +129,7 @@ def forward(
         mask_mod = create_peagle_mask_mod(
             anchor_pos=anchor_pos,
             depth=depth,
-            lengths=lengths,
-            total_seq_len=seq_length,
+            document_ids=document_ids.squeeze(0).to(device),
         )
 
         attention_mask = create_block_mask(  # type: ignore[assignment]

src/speculators/train/data.py

@@ -524,7 +524,7 @@ def collate_fn(batch: list[BatchType | None]) -> BatchType:
         # Include lengths until while they fit in max_len
         # The last included length is (if necessary) truncated
         # Any additional lengths are discarded
-        lengths = collated_data["lengths"]
+        lengths = collated_data.pop("lengths")
         new_lengths = []
         cum_length = 0
         for length in lengths:
@@ -534,9 +534,20 @@ def collate_fn(batch: list[BatchType | None]) -> BatchType:
             new_lengths.append(length)
             cum_length += length
         lengths = torch.tensor(new_lengths, dtype=torch.long)
-        # Pad to at least size 4 to avoid torch.compile recompiling on sizes 0/1
-        pad_to = max(4, lengths.shape[0])
-        collated_data["lengths"] = F.pad(lengths, (0, pad_to - lengths.shape[0]))
+
+        # Create document_ids: maps each position to its document index, -1 for padding
+        document_ids = torch.repeat_interleave(
+            torch.arange(lengths.shape[0], dtype=torch.long), lengths
+        )
+        document_ids = torch.cat(
+            [
+                document_ids,
+                -1 * torch.ones(max_len - document_ids.shape[0], dtype=torch.long),
+            ]
+        ).unsqueeze(0)
+        # shape: [1, max_len]
+        collated_data["document_ids"] = document_ids
+
         return collated_data
 
     return collate_fn

tests/unit/train/test_data.py

@@ -188,7 +188,9 @@ def test_collate_fn_basic():
             [[[2.0], [3.0], [10.0], [11.0], [12.0], [13.0], [14.0], [15.0], [-1], [-1]]]
         ),
         "loss_mask": torch.tensor([[0, 1, 0, 0, 1, 0, 1, 1, -1, -1]], dtype=torch.long),
-        "lengths": torch.tensor([2, 6], dtype=torch.long),
+        "document_ids": torch.tensor(
+            [[0, 0, 1, 1, 1, 1, 1, 1, -1, -1]], dtype=torch.long
+        ),
         "position_ids": torch.tensor(
             [[0, 1, 0, 1, 2, 3, 4, 5, -1, -1]], dtype=torch.long
         ),
@@ -235,11 +237,13 @@ def test_collate_fn_length_truncation():
 
     collated = collate_fn(batch)
 
-    # Last length is truncated to fit in max_len
-    expected_lengths = torch.tensor([5, 6], dtype=torch.long)
+    # document_ids: doc 0 has length 5, doc 1 truncated to length 6, rest is padding
+    expected_document_ids = torch.tensor(
+        [[0, 0, 0, 0, 0, 1, 1, 1, 1, 1, 1]], dtype=torch.long
+    )
+    assert torch.equal(collated["document_ids"], expected_document_ids)
+    assert "lengths" not in collated
 
-    # All tensors (other than lengths) are concatenated then truncated to max_len
-    assert torch.equal(collated["lengths"], expected_lengths)
     for key in [
         "input_ids",
         "hidden_states",

tests/unit/train/test_eagle3_attention.py

@@ -10,8 +10,11 @@
 
 def test_create_combined_mask_mod():
     lengths = torch.tensor([1, 2, 3])
+    document_ids = torch.repeat_interleave(
+        torch.arange(lengths.shape[0], dtype=torch.long), lengths
+    )
     mask_mod = create_combined_mask_mod(
-        lengths, total_seq_len=int(lengths.sum().item())
+        document_ids, total_seq_len=int(lengths.sum().item())
     )
 
     # Creates causal document mask mod that supports extended diagonals
@@ -44,7 +47,10 @@ def test_diagonal_draft_tokens_mask_mod(lengths):
     # If kv_idx > N (N = orig seq len = num query inds), only the diagonal tokens are
     # in the mask. Diagonal tokens are those where kv_idx % N == q_idx
 
-    mask_mod = create_combined_mask_mod(lengths, total_seq_len=lengths.sum().item())
+    document_ids = torch.repeat_interleave(
+        torch.arange(lengths.shape[0], dtype=torch.long), lengths
+    )
+    mask_mod = create_combined_mask_mod(document_ids, total_seq_len=lengths.sum().item())
 
     N = lengths.sum().item()