[Fix] Bound FULL_MASK verify-mask reuse by the captured max_bs (#33127)

python/sglang/srt/layers/attention/verify_mask.py

@@ -26,14 +26,15 @@ def tree_mask_numel(
     return bs * per_req
 
 
-class VerifyMask(msgspec.Struct):
+class VerifyMask(msgspec.Struct, frozen=True):
     """The target-verify mask.
 
     ``build_tree_kernel_efficient`` writes the buffer in place after draft, which
-    is what lets the worker skip the ``seq_lens_sum`` D2H sync. Keep the three
+    is what lets the worker skip the ``seq_lens_sum`` D2H sync. Keep them
     together: taking the buffer without its layout has the kernel write a shape
     the reader does not expect. The kernel writes every cell even when unread, so
-    the buffer is always allocated.
+    the buffer is always allocated. Frozen -- resize by swapping the whole struct,
+    never a field, so ``max_bs`` cannot go stale against ``buffer``.
 
     Temporary home -- a phase-level buffer with no owner today (``spec_info`` is a
     per-phase union the graph registry cannot slot).
@@ -41,17 +42,16 @@ class VerifyMask(msgspec.Struct):
 
     buffer: torch.Tensor
     mode: TreeMaskMode
+    max_bs: int
     is_read: bool = True
 
-    def fits(self, bs: int, num_draft_tokens: int) -> bool:
+    def fits(self, bs: int) -> bool:
         """Whether this batch's writes stay inside the buffer.
 
-        Only the compact layout is checked. FULL_MASK keeps unconditional reuse
-        because its runtime bound depends on sequence lengths not passed to fits().
+        ``tree_mask_numel`` is ``bs * per_req`` and per_req is fixed at
+        allocation (FULL_MASK's spans max_context_len), so max_bs bounds it.
         """
-        if self.mode != TreeMaskMode.QLEN_ONLY:
-            return True
-        return self.buffer.numel() >= bs * num_draft_tokens * num_draft_tokens
+        return bs <= self.max_bs
 
 
 def maybe_create_verify_mask(
@@ -76,5 +76,6 @@ def maybe_create_verify_mask(
             device=device,
         ),
         mode=mode,
+        max_bs=max_bs,
         is_read=is_read,
     )

python/sglang/srt/speculative/eagle_worker_common.py

@@ -351,9 +351,7 @@ def build_eagle_verify_input(
         tree_mask_buf, mask_mode, fill_mask = None, tree_mask_mode, True
     else:
         mask_mode, fill_mask = verify_mask.mode, verify_mask.is_read
-        tree_mask_buf = (
-            verify_mask.buffer if verify_mask.fits(bs, num_draft_tokens) else None
-        )
+        tree_mask_buf = verify_mask.buffer if verify_mask.fits(bs) else None
 
     # build_tree_kernel uses seq_lens_sum only to size the (non-preallocated)
     # FULL_MASK tree mask; over-size is safe. Skip per-iter .sum().item() D2H via UB.

python/sglang/srt/speculative/eagle_worker_v2.py

@@ -1210,7 +1210,7 @@ class EAGLEWorkerV2(BaseSpecWorker):
         verify_mask = attn_backend.verify_mask
         # Every position in a 1-node tree is visible, so an all-True fill is
         # correct under either layout.
-        if verify_mask is not None and verify_mask.fits(bs, 1):
+        if verify_mask is not None and verify_mask.fits(bs):
             custom_mask = verify_mask.buffer
             custom_mask.fill_(True)
         else: