minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

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

Browse Source
This commit is contained in:
Liangsheng Yin
2026-07-31 17:50:40 -07:00
committed by GitHub
parent 1496bfee93
commit ca07917c58
4 changed files with 33 additions and 23 deletions
Download Patch File Download Diff File Expand all files Collapse all files
test/registered/unit/layers/attention/test_verify_mask.py
+21 -10
View File
@@ -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):