[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
+        ``tree_mask_numel`` is ``bs * per_req`` and per_req is fixed at
-        because its runtime bound depends on sequence lengths not passed to fits().
+        allocation (FULL_MASK's spans max_context_len), so max_bs bounds it.
         """
-        if self.mode != TreeMaskMode.QLEN_ONLY:
+        return bs <= self.max_bs
-            return True
-        return self.buffer.numel() >= bs * num_draft_tokens * num_draft_tokens
 
 
 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 = (
+        tree_mask_buf = verify_mask.buffer if verify_mask.fits(bs) else None
-            verify_mask.buffer if verify_mask.fits(bs, num_draft_tokens) 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:

test/registered/unit/layers/attention/test_verify_mask.py

@@ -61,25 +61,35 @@ class TestVerifyMaskSizing(CustomTestCase):
 
 
 class TestVerifyMaskCapacity(CustomTestCase):
-    """A batch past the captured max_bs must not silently reuse the compact
+    """A batch past the captured max_bs must not silently reuse the buffer."""
-    layout -- it has no context-dimension slack to absorb the overflow."""
 
     def test_compact_layout_fits_up_to_max_bs(self):
         # is_read=False pins QLEN_ONLY; the read layout is build-dependent.
         mask = _create(is_read=False)
-        self.assertTrue(mask.fits(_MAX_BS, _DRAFT))
+        self.assertTrue(mask.fits(_MAX_BS))
-        self.assertTrue(mask.fits(1, _DRAFT))
+        self.assertTrue(mask.fits(1))
 
     def test_compact_layout_does_not_fit_beyond_max_bs(self):
         mask = _create(is_read=False)
-        self.assertFalse(mask.fits(_MAX_BS + 1, _DRAFT))
+        self.assertFalse(mask.fits(_MAX_BS + 1))
 
-    def test_full_mask_always_fits(self):
+    def test_full_mask_does_not_fit_beyond_max_bs(self):
-        """FULL_MASK is exempt from the check -- see fits()."""
+        """FULL_MASK's context dimension is per-request slack, not spare room
+        for extra requests -- it must not be exempt from the check. Built
+        explicitly because default_tree_mask_mode() is host-dependent."""
         mask = VerifyMask(
-            buffer=torch.zeros(8, dtype=torch.bool), mode=TreeMaskMode.FULL_MASK
+            buffer=torch.zeros(
+                tree_mask_numel(
+                    TreeMaskMode.FULL_MASK, _MAX_BS, _DRAFT, _MAX_CONTEXT_LEN
+                ),
+                dtype=torch.bool,
+            ),
+            mode=TreeMaskMode.FULL_MASK,
+            max_bs=_MAX_BS,
         )
-        self.assertTrue(mask.fits(_MAX_BS * 1000, _DRAFT))
+
+        self.assertTrue(mask.fits(_MAX_BS))
+        self.assertFalse(mask.fits(_MAX_BS + 1))
 
 
 class TestVerifyMaskGate(CustomTestCase):
@@ -107,6 +117,7 @@ def _mask(numel, **kwargs):
     return VerifyMask(
         buffer=torch.zeros(numel, dtype=torch.bool),
         mode=TreeMaskMode.QLEN_ONLY,
+        max_bs=_MAX_BS,
         **kwargs,
     )
 
@@ -148,7 +159,7 @@ class TestHybridAttnBackendHandsOutSelectedChildMask(CustomTestCase):
     def test_capacity_check_needs_nothing_from_the_backend(self):
         backend = _make_hybrid_backend("prefill", _mask(64, is_read=False), None)
 
-        self.assertTrue(backend.verify_mask.fits(_MAX_BS, _DRAFT))
+        self.assertTrue(backend.verify_mask.fits(_MAX_BS))
 
 
 class TestTreeMaskNumel(CustomTestCase):