From ca07917c58bdf87ae5aa299ab93fb7695c70432b Mon Sep 17 00:00:00 2001 From: Liangsheng Yin Date: Fri, 31 Jul 2026 17:50:40 -0700 Subject: [PATCH] [Fix] Bound FULL_MASK verify-mask reuse by the captured max_bs (#33127) --- .../srt/layers/attention/verify_mask.py | 19 ++++++------ .../srt/speculative/eagle_worker_common.py | 4 +-- .../sglang/srt/speculative/eagle_worker_v2.py | 2 +- .../unit/layers/attention/test_verify_mask.py | 31 +++++++++++++------ 4 files changed, 33 insertions(+), 23 deletions(-) diff --git a/python/sglang/srt/layers/attention/verify_mask.py b/python/sglang/srt/layers/attention/verify_mask.py index c0864a29f..bb156c448 100644 --- a/python/sglang/srt/layers/attention/verify_mask.py +++ b/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, ) diff --git a/python/sglang/srt/speculative/eagle_worker_common.py b/python/sglang/srt/speculative/eagle_worker_common.py index 89a371f3e..8e0c01826 100644 --- a/python/sglang/srt/speculative/eagle_worker_common.py +++ b/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. diff --git a/python/sglang/srt/speculative/eagle_worker_v2.py b/python/sglang/srt/speculative/eagle_worker_v2.py index 60f4dad89..f6e86eeb3 100644 --- a/python/sglang/srt/speculative/eagle_worker_v2.py +++ b/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: diff --git a/test/registered/unit/layers/attention/test_verify_mask.py b/test/registered/unit/layers/attention/test_verify_mask.py index 6a62e7a20..3386c697c 100644 --- a/test/registered/unit/layers/attention/test_verify_mask.py +++ b/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 - layout -- it has no context-dimension slack to absorb the overflow.""" + """A batch past the captured max_bs must not silently reuse the buffer.""" 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(1, _DRAFT)) + self.assertTrue(mask.fits(_MAX_BS)) + 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): - """FULL_MASK is exempt from the check -- see fits().""" + def test_full_mask_does_not_fit_beyond_max_bs(self): + """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):