Reduce graph breaks and recompiles

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)

src/speculators/models/dflash/core.py

@@ -203,17 +203,18 @@ 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,
         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,
@@ -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():

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]

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

@@ -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,

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 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]

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