Quality

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

Author: fynnsu

src/speculators/models/dflash/attention.py

@@ -27,7 +27,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 Any, 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

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

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

@@ -43,7 +43,7 @@ def conditional_torch_compile(func):
         def wrapper(*args, **kwargs):
             lengths = kwargs.get("lengths")
             if lengths is not None:
-                torch._dynamo.mark_dynamic(lengths, 0)
+                torch._dynamo.mark_dynamic(lengths, 0)  # noqa: SLF001
             return compiled(*args, **kwargs)
 
         return wrapper

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()