Remove dflash graph breaks

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

Author: fynnsu

src/speculators/models/dflash/attention.py

@@ -54,11 +54,6 @@ def create_anchor_block_mask_mod(
     q_len = n_anchors * block_size
     kv_len = total_seq_len + q_len
 
-    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

@@ -304,7 +304,7 @@ def forward(
         # shape: [1, total_seq_len, hidden_size]
 
         mask_position_ids = get_base_indices_for_anchored_blocks(
-            position_ids[0, 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 +314,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())
+    torch._check(k >= 0)
+
+    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]

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)