Update torch requirement from <=2.11.0,>=2.9.0 to >=2.9.0,<=2.12.0

pyproject.toml

@@ -48,8 +48,7 @@ dependencies = [
     "pydantic-settings>=2.0.0",
     "rich",
     "safetensors",
-    "setuptools",
-    "torch>=2.9.0,<=2.11.0",
+    "torch>=2.9.0,<=2.12.0",
     "torchaudio",
     "torchvision",
     "tqdm>=4.66.3,<=4.67.3",

src/speculators/models/attention.py

@@ -4,8 +4,6 @@
 speculator architectures (EAGLE3, DFlash, etc.) to avoid code duplication.
 """
 
-from typing import cast
-
 import torch
 from torch.nn.attention.flex_attention import BlockMask, flex_attention
 from transformers.modeling_utils import AttentionInterface
@@ -55,7 +53,7 @@ def flex_attention_forward(
         enable_gqa=enable_gqa,
         scale=scaling,
     )
-    attention_output: torch.Tensor = cast("torch.Tensor", flex_attention_output)
+    attention_output: torch.Tensor = flex_attention_output
     attention_output = attention_output.transpose(1, 2).contiguous()
     return attention_output, None
 

src/speculators/models/dflash/core.py

@@ -185,6 +185,30 @@ def mask_token_id(self) -> int:
             )
         return self.config.mask_token_id
 
+    @torch.compiler.disable
+    def _build_attention_mask(self, loss_mask, lengths, device):
+        anchor_positions, anchor_valid = select_anchors(
+            loss_mask, self.config.max_anchors, self.block_size
+        )
+        # shape: [num_anchors], [num_anchors]
+
+        mask_mod, q_len, kv_len = create_anchor_block_mask_mod(
+            lengths=lengths.to(device),
+            total_seq_len=loss_mask.shape[1],
+            anchor_positions=anchor_positions,
+            block_size=self.block_size,
+        )
+
+        attention_mask = create_block_mask(
+            mask_mod,
+            B=None,
+            H=None,
+            Q_LEN=q_len,
+            KV_LEN=kv_len,
+            device=device,
+        )
+        return attention_mask, anchor_positions, anchor_valid
+
     @torch.compile
     def forward(
         self,
@@ -207,25 +231,8 @@ def forward(
                 total_seq_len, dtype=torch.long, device=device
             ).unsqueeze(0)
 
-        anchor_positions, anchor_valid = select_anchors(
-            loss_mask, num_anchors, self.block_size
-        )
-        # shape: [num_anchors], [num_anchors]
-
-        mask_mod, q_len, kv_len = create_anchor_block_mask_mod(
-            lengths=lengths.to(device),
-            total_seq_len=total_seq_len,
-            anchor_positions=anchor_positions,
-            block_size=self.block_size,
-        )
-
-        attention_mask = create_block_mask(
-            mask_mod,
-            B=None,
-            H=None,
-            Q_LEN=q_len,
-            KV_LEN=kv_len,
-            device=device,
+        attention_mask, anchor_positions, anchor_valid = self._build_attention_mask(
+            loss_mask, lengths, device
         )
 
         mask_tokens_size = num_anchors * self.block_size

src/speculators/models/eagle3/attention.py

@@ -43,6 +43,7 @@ def diagonal_draft_mask_mod(_b, _h, q_idx, kv_idx):
     )
 
 
+@torch.compiler.disable
 def extend_mask_for_draft_tokens(block_mask):
     """
     Extend the block mask to include new draft tokens. Concatenates a diagonal mask for