diff --git a/src/speculators/models/dflash/attention.py b/src/speculators/models/dflash/attention.py
index c65816dfd..722e7c7b6 100644
--- a/src/speculators/models/dflash/attention.py
+++ b/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 doc index, pad -1
         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,33 +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()}"
-        )
-
     # For each query position, which anchor does it belong to?
     # query q in [j*block_size, (j+1)*block_size) belongs to anchor_positions[j]
     query_anchor_positions = torch.repeat_interleave(anchor_positions, block_size)
diff --git a/src/speculators/models/dflash/core.py b/src/speculators/models/dflash/core.py
index 807b21902..504e6a7dd 100644
--- a/src/speculators/models/dflash/core.py
+++ b/src/speculators/models/dflash/core.py
@@ -203,9 +203,10 @@ def mask_token_id(self) -> int:
             )
         return self.config.mask_token_id
 
+    @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,
@@ -213,7 +214,7 @@ def _create_attention_mask(
         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,
@@ -230,7 +231,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(
@@ -240,7 +241,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,
@@ -250,7 +251,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,
@@ -267,7 +268,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,
@@ -276,15 +277,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
@@ -304,7 +303,7 @@ def forward(
         # shape: [1, total_seq_len, hidden_size]
 
         mask_position_ids = get_base_indices_for_anchored_blocks(
-            position_ids[:, anchor_positions], self.block_size, input_ids.numel()
+            position_ids[0, anchor_positions], self.block_size
         )
         position_ids = torch.cat([position_ids, mask_position_ids.unsqueeze(0)], dim=1)
         # shape: [1, total_seq_len + num_anchors*block_size]
@@ -314,7 +313,7 @@ def forward(
         position_embeddings = self.rotary_emb(hidden_states, position_ids)
 
         anchored_block_indices = get_base_indices_for_anchored_blocks(
-            anchor_positions, self.block_size, input_ids.numel()
+            anchor_positions, self.block_size
         )  # shape: [num_anchors*block_size]
 
         with torch.no_grad():
diff --git a/src/speculators/models/dflash/utils.py b/src/speculators/models/dflash/utils.py
index 192893815..0553750b9 100644
--- a/src/speculators/models/dflash/utils.py
+++ b/src/speculators/models/dflash/utils.py
@@ -6,7 +6,6 @@
 def get_base_indices_for_anchored_blocks(
     anchor_positions: torch.Tensor,  # shape: [1, num_anchors]
     block_size: int,
-    total_seq_len: int | None = None,
 ) -> torch.Tensor:  # shape: [num_anchors*block_size]
     anchor_positions = anchor_positions.to(dtype=torch.long).view(-1)
     # dtype: long, shape: [num_anchors]
@@ -16,12 +15,6 @@ def get_base_indices_for_anchored_blocks(
         anchor_positions[:, None] + offsets[None, :]
     )  # shape: [num_anchors, block_size]
 
-    if (idx < 0).any() or (total_seq_len and (idx >= total_seq_len).any()):
-        raise ValueError(
-            "Some anchor_positions + offsets are out of range for total_seq_len"
-            f"={total_seq_len}. Max={idx.max().item()}, min={idx.min().item()}"
-        )
-
     return idx.reshape(-1)
 
 
@@ -60,10 +53,14 @@ def select_anchors(
     anchor_valid = torch.zeros(num_anchors, dtype=torch.bool, device=device)
 
     k = min(num_anchors, valid_indices.numel())
-    if k > 0:
-        perm = torch.randperm(valid_indices.numel(), device=loss_mask.device)
-        anchors[:k] = valid_indices[perm[:k]]
-        anchor_valid[:k] = True
+
+    # Constrain value of k for torch dynamo
+    torch._check(k <= valid_indices.numel())  # noqa: SLF001
+    torch._check(k >= 0)  # noqa: SLF001
+
+    perm = torch.randperm(valid_indices.numel(), device=loss_mask.device)
+    anchors[:k] = torch.gather(valid_indices, 0, perm[:k])
+    anchor_valid[:k] = True
 
     return anchors, anchor_valid
     # shape: [num_anchors], [num_anchors]
diff --git a/src/speculators/models/eagle3/attention.py b/src/speculators/models/eagle3/attention.py
index 186a0d601..d365d16a5 100644
--- a/src/speculators/models/eagle3/attention.py
+++ b/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
 
diff --git a/src/speculators/models/eagle3/core.py b/src/speculators/models/eagle3/core.py
index 83ddbdefe..fc67fd57c 100644
--- a/src/speculators/models/eagle3/core.py
+++ b/src/speculators/models/eagle3/core.py
@@ -129,11 +129,11 @@ def load_verifier_weights(self):
             )
 
     @conditional_torch_compile
-    def forward(  # noqa: C901
+    def forward(
         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,
diff --git a/src/speculators/models/peagle/attention.py b/src/speculators/models/peagle/attention.py
index ee4e9f9ae..2377ad052 100644
--- a/src/speculators/models/peagle/attention.py
+++ b/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]
diff --git a/src/speculators/models/peagle/core.py b/src/speculators/models/peagle/core.py
index 1b9aaf5b1..cba6e480f 100644
--- a/src/speculators/models/peagle/core.py
+++ b/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 doc index, pad -1
             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]
diff --git a/src/speculators/train/data.py b/src/speculators/train/data.py
index 83120afab..a0f11052a 100644
--- a/src/speculators/train/data.py
+++ b/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:
@@ -533,7 +533,21 @@ def collate_fn(batch: list[BatchType | None]) -> BatchType:
                 break
             new_lengths.append(length)
             cum_length += length
-        collated_data["lengths"] = torch.tensor(new_lengths, dtype=torch.long)
+        lengths = torch.tensor(new_lengths, dtype=torch.long)
+
+        # 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
diff --git a/tests/unit/models/test_dflash_attention.py b/tests/unit/models/test_dflash_attention.py
index ea823012c..75d61f8cd 100644
--- a/tests/unit/models/test_dflash_attention.py
+++ b/tests/unit/models/test_dflash_attention.py
@@ -1,14 +1,21 @@
 """Unit tests for the DFlash anchor-block attention mask."""
 
-import pytest
 import torch
 from torch.nn.attention.flex_attention import create_mask
 
 from speculators.models.dflash.attention import create_anchor_block_mask_mod
 
 
+def _lengths_to_document_ids(lengths, total_seq_len):
+    document_ids = torch.full((total_seq_len,), -1, dtype=torch.long)
+    document_ids[: lengths.sum()] = torch.repeat_interleave(
+        torch.arange(lengths.shape[0], dtype=torch.long), lengths
+    )
+    return document_ids
+
+
 def _reference_dense_from_mask_mod(
-    lengths,
+    document_ids,
     total_seq_len,
     anchor_positions,
     block_size,
@@ -17,7 +24,7 @@ def _reference_dense_from_mask_mod(
 ):
     """Ground truth: evaluate the flex mask_mod element-wise over the q x kv grid."""
     mask_mod, q_len, kv_len = create_anchor_block_mask_mod(
-        lengths=lengths,
+        document_ids=document_ids,
         total_seq_len=total_seq_len,
         anchor_positions=anchor_positions,
         block_size=block_size,
@@ -33,7 +40,7 @@ def _reference_dense_from_mask_mod(
 
 
 def _dense_from_create_mask(
-    lengths,
+    document_ids,
     total_seq_len,
     anchor_positions,
     block_size,
@@ -41,7 +48,7 @@ def _dense_from_create_mask(
     sliding_window_non_causal=False,
 ):
     mask_mod, q_len, kv_len = create_anchor_block_mask_mod(
-        lengths=lengths,
+        document_ids=document_ids,
         total_seq_len=total_seq_len,
         anchor_positions=anchor_positions,
         block_size=block_size,
@@ -54,7 +61,7 @@ def _dense_from_create_mask(
         H=None,
         Q_LEN=q_len,
         KV_LEN=kv_len,
-        device=lengths.device,
+        device=document_ids.device,
     )
 
 
@@ -63,13 +70,14 @@ def test_create_mask_matches_mask_mod_full_attention():
     device = torch.device("cpu")
     total_seq_len, block_size = 16, 4
     lengths = torch.tensor([10, 6])  # two packed documents summing to total_seq_len
+    document_ids = _lengths_to_document_ids(lengths, total_seq_len)
     anchor_positions = torch.tensor([3, 8, 12])
 
     ref = _reference_dense_from_mask_mod(
-        lengths, total_seq_len, anchor_positions, block_size
+        document_ids, total_seq_len, anchor_positions, block_size
     )
     dense = _dense_from_create_mask(
-        lengths.to(device), total_seq_len, anchor_positions, block_size
+        document_ids.to(device), total_seq_len, anchor_positions, block_size
     )
 
     assert dense.shape == (1, 1, ref.shape[0], ref.shape[1])
@@ -81,18 +89,19 @@ def test_create_mask_matches_mask_mod_sliding_window():
     device = torch.device("cpu")
     total_seq_len, block_size = 16, 4
     lengths = torch.tensor([16])  # single document
+    document_ids = _lengths_to_document_ids(lengths, total_seq_len)
     anchor_positions = torch.tensor([5, 9, 14])
     sliding_window = 4
 
     ref = _reference_dense_from_mask_mod(
-        lengths,
+        document_ids,
         total_seq_len,
         anchor_positions,
         block_size,
         sliding_window=sliding_window,
     )
     dense = _dense_from_create_mask(
-        lengths.to(device),
+        document_ids.to(device),
         total_seq_len,
         anchor_positions,
         block_size,
@@ -107,32 +116,11 @@ def test_create_mask_each_query_sees_its_own_block():
     device = torch.device("cpu")
     total_seq_len, block_size = 12, 4
     lengths = torch.tensor([12])
+    document_ids = _lengths_to_document_ids(lengths, total_seq_len)
     anchor_positions = torch.tensor([2, 7, 10])
 
     dense = _dense_from_create_mask(
-        lengths.to(device), total_seq_len, anchor_positions, block_size
+        document_ids.to(device), total_seq_len, anchor_positions, block_size
     )
 
     assert bool(dense[0, 0].any(dim=-1).all())
-
-
-def test_mask_mod_rejects_lengths_overflow():
-    """sum(lengths) > total_seq_len raises."""
-    with pytest.raises(ValueError, match="exceeds total_seq_len"):
-        create_anchor_block_mask_mod(
-            lengths=torch.tensor([10, 10]),  # sum = 20 > total_seq_len
-            total_seq_len=16,
-            anchor_positions=torch.tensor([3, 8]),
-            block_size=4,
-        )
-
-
-def test_mask_mod_rejects_out_of_range_anchor():
-    """anchor_positions outside [0, total_seq_len) raises."""
-    with pytest.raises(ValueError, match="out of range"):
-        create_anchor_block_mask_mod(
-            lengths=torch.tensor([16]),
-            total_seq_len=16,
-            anchor_positions=torch.tensor([3, 20]),  # 20 >= total_seq_len
-            block_size=4,
-        )
diff --git a/tests/unit/models/test_dflash_utils.py b/tests/unit/models/test_dflash_utils.py
index 9910c5983..fdf286e3b 100644
--- a/tests/unit/models/test_dflash_utils.py
+++ b/tests/unit/models/test_dflash_utils.py
@@ -1,6 +1,5 @@
 """Unit tests for get_base_indices_for_anchored_blocks."""
 
-import pytest
 import torch
 
 from speculators.models.dflash.utils import get_base_indices_for_anchored_blocks
@@ -44,31 +43,3 @@ def test_output_dtype_is_long(self):
         anchor_positions = torch.tensor([[2.0, 5.0]])
         result = get_base_indices_for_anchored_blocks(anchor_positions, block_size=2)
         assert result.dtype == torch.long
-
-    def test_total_seq_len_valid(self):
-        anchor_positions = torch.tensor([[0, 3]])
-        result = get_base_indices_for_anchored_blocks(
-            anchor_positions, block_size=4, total_seq_len=7
-        )
-        expected = torch.tensor([0, 1, 2, 3, 3, 4, 5, 6])
-        assert torch.equal(result, expected)
-
-    def test_total_seq_len_out_of_range(self):
-        anchor_positions = torch.tensor([[0, 3]])
-        with pytest.raises(ValueError, match="out of range"):
-            get_base_indices_for_anchored_blocks(
-                anchor_positions, block_size=4, total_seq_len=6
-            )
-
-    def test_negative_anchor_raises(self):
-        anchor_positions = torch.tensor([[-1, 3]])
-        with pytest.raises(ValueError, match="out of range"):
-            get_base_indices_for_anchored_blocks(anchor_positions, block_size=2)
-
-    def test_no_total_seq_len_skips_upper_bound_check(self):
-        anchor_positions = torch.tensor([[100]])
-        result = get_base_indices_for_anchored_blocks(
-            anchor_positions, block_size=3, total_seq_len=None
-        )
-        expected = torch.tensor([100, 101, 102])
-        assert torch.equal(result, expected)
diff --git a/tests/unit/train/test_data.py b/tests/unit/train/test_data.py
index 5cdecc2d9..8cc23fa87 100644
--- a/tests/unit/train/test_data.py
+++ b/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",
diff --git a/tests/unit/train/test_eagle3_attention.py b/tests/unit/train/test_eagle3_attention.py
index de87dcbae..027f0566f 100644
--- a/tests/unit/train/test_eagle3_attention.py
+++ b/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,12 @@ 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()