Quality

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

Author: fynnsu

src/speculators/models/dflash/attention.py

@@ -29,7 +29,7 @@ def create_anchor_block_mask_mod(
         - may not attend to other synthetic blocks or later base tokens
 
     Args:
-        document_ids: [total_seq_len] maps each position to its document index, -1 for padding
+        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

src/speculators/models/dflash/core.py

@@ -1,4 +1,3 @@
-import functools
 from typing import ClassVar
 
 import torch

src/speculators/models/dflash/utils.py

@@ -55,8 +55,8 @@ def select_anchors(
     k = min(num_anchors, valid_indices.numel())
 
     # Constrain value of k for torch dynamo
-    torch._check(k <= valid_indices.numel())
-    torch._check(k >= 0)
+    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])

src/speculators/models/eagle3/core.py

@@ -129,7 +129,7 @@ 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]

src/speculators/models/peagle/core.py

@@ -65,7 +65,7 @@ def forward(
         Args:
             hidden_states: Verifier hidden states [batch, seq_len, 3*hidden_size]
             input_ids: Input token IDs [batch, seq_len]
-            document_ids: Document IDs [1, seq_len], maps positions to document index, -1 for padding
+            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]

src/speculators/models/utils.py

@@ -1,4 +1,3 @@
-import functools
 import warnings
 from functools import partial
 

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,15 +40,15 @@ def _reference_dense_from_mask_mod(
 
 
 def _dense_from_create_mask(
-    lengths,
+    document_ids,
     total_seq_len,
     anchor_positions,
     block_size,
     sliding_window=None,
     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,
-        )

tests/unit/train/test_eagle3_attention.py

@@ -50,7 +50,9 @@ def test_diagonal_draft_tokens_mask_mod(lengths):
     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())
+    mask_mod = create_combined_mask_mod(
+        document_ids, total_seq_len=lengths.sum().item()
+    )
 
     N = lengths.sum().item()