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

feat(unified-memory): byte-budget sizing, feasibility floor, and a conservation verifier (#35158)

Browse Source 
Co-authored-by: Caihua Li <caihua.li@bytedance.com>
Co-authored-by: Claude Fable 5 <noreply@anthropic.com>
Co-authored-by: Cheng Wan <cheng.wan@radixark.ai>
This commit is contained in:
caihuali95
2026-08-31 15:09:28 -07:00
committed by GitHub
co-authored by Caihua Li Claude Fable 5 Cheng Wan
parent 961beee9e5
commit 98cb3535b7
9 changed files with 658 additions and 21 deletions
Download Patch File Download Diff File Expand all files Collapse all files
test/registered/unit/mem_cache/test_unified_byte_accounting.py
+157
View File
@@ -0,0 +1,157 @@
# Copyright 2023-2026 SGLang Team
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
"""Byte-conservation verifier for the unified 2-pool composites.
`verify_byte_accounting` is the idle-time tripwire the token-identity leak
check cannot provide: the unified pool's correctness rests on BYTE bookkeeping
(watermark spans, holes, pending compaction, frontier ordering inside one
shared buffer), and a drifted counter admits requests into memory that is not
actually free — silent corruption territory, not a crash.
Derived properties pinned here:
* Conservation: on a lazy end pool the watermark span must equal
live + holes + pending pages at EVERY point of a healthy lifecycle
(alloc, partial free, group free, flush) — not just at rest.
* The check is not vacuous: drifting any single term (live count, watermark,
a leaked hole) reports loudly, naming the sub-pool.
* Chain order: one member's low frontier clearing the other's high frontier
is what "two pools share one buffer without overlap" MEANS; the pair check
must hold regardless of which member grows up.
* The strict escalation env defaults OFF: promoting the diagnostic to a
RuntimeError is a validation posture, not the production one.
python -m pytest test/registered/unit/mem_cache/test_unified_byte_accounting.py -v
"""
import unittest
from test_multi_ended_allocator import TestPagedMultiEndedAllocator as _PagedFixture
from test_multi_ended_allocator import (
TestUnifiedSWATokenToKVPoolAllocator as _SwaFixture,
)
from sglang.srt.mem_cache import multi_ended_allocator as mea
from sglang.test.ci.ci_register import register_cpu_ci
register_cpu_ci(est_time=15, suite="base-a-test-cpu")
def _swa_composite():
inst = _SwaFixture([m for m in dir(_SwaFixture) if m.startswith("test_")][0])
pool, allocator, kvcache = inst._build()
return inst, allocator, kvcache
def _paged_pair(lazy: bool):
inst = _PagedFixture([m for m in dir(_PagedFixture) if m.startswith("test_")][0])
_pool, full, swa, _fkv, _skv = inst._build()
full.lazy_compaction = lazy
return full, swa
class TestHealthyLifecycleReportsClean(unittest.TestCase):
def test_swa_composite_clean_at_every_step(self):
inst, allocator, kvcache = _swa_composite()
self.assertEqual(allocator.verify_byte_accounting(), [])
v = inst._alloc(allocator, kvcache, 8)
self.assertEqual(allocator.verify_byte_accounting(), [])
allocator.free_swa(v[:4]) # tombstone half the swa side
self.assertEqual(allocator.verify_byte_accounting(), [])
inst._free(allocator, kvcache, v)
self.assertEqual(allocator.verify_byte_accounting(), [])
allocator.clear()
self.assertEqual(allocator.verify_byte_accounting(), [])
def test_lazy_end_pool_clean_through_free_and_flush(self):
full, _swa = _paged_pair(lazy=True)
self.assertEqual(full._byte_accounting_violations(), [])
v = full.alloc(full.page_size * 4)
self.assertEqual(full._byte_accounting_violations(), [])
full.free(v[: full.page_size * 2]) # lazy: holes, no compaction yet
self.assertEqual(full._byte_accounting_violations(), [])
full._flush(urgent=True)
self.assertEqual(full._byte_accounting_violations(), [])
class TestDriftReportsLoudly(unittest.TestCase):
"""Each mutation below models a distinct bookkeeping bug; the verifier
must name the drifted sub-pool. Without these, a regression in any single
counter passes every other test (the pool still 'works' — it just lies
about capacity)."""
def _lazy_full(self):
full, _swa = _paged_pair(lazy=True)
v = full.alloc(full.page_size * 4)
full.free(v[: full.page_size]) # one hole so all three terms are live
self.assertEqual(full._byte_accounting_violations(), [])
return full
def test_drifted_live_count(self):
full = self._lazy_full()
full.live_page_count += 1
self.assertTrue(any("span" in s for s in full._byte_accounting_violations()))
def test_leaked_hole(self):
full = self._lazy_full()
full._free_phys_pages = full._free_phys_pages[:-1] # hole vanished
self.assertTrue(any("span" in s for s in full._byte_accounting_violations()))
def test_drifted_watermark(self):
full = self._lazy_full()
full.watermark_physical += 1
self.assertTrue(any("span" in s for s in full._byte_accounting_violations()))
def test_composite_report_names_the_sub_pool(self):
"""Frontier-bounds drift (checked in BOTH lazy and eager modes): push
the swa band's watermark outside the buffer."""
inst, allocator, kvcache = _swa_composite()
inst._alloc(allocator, kvcache, 8)
swa = allocator.swa_attn_allocator
# grow-down member: low frontier = (wm+1)*bytes; wm == num_pages puts
# it past the buffer top.
self.assertEqual(swa.grow_direction, "down")
swa.watermark_physical = swa.num_pages
out = allocator.verify_byte_accounting()
self.assertTrue(out and any("[swa]" in s for s in out), out)
class TestChainFrontierOrder(unittest.TestCase):
def test_overlapping_frontiers_report(self):
"""Both bands hold pages, then the up member's watermark is pushed past
the down member's LIVE low frontier: the two bands now claim the same
bytes of one buffer. (An empty down band cannot overlap — its low
frontier IS the buffer top — so both sides must be populated for the
scenario to be a real corruption.)"""
full, swa = _paged_pair(lazy=False)
chain = mea._end_pair_chain(full, swa)
up, down = chain
self.assertEqual(up.grow_direction, "up")
self.assertIsNotNone(down.alloc(down.page_size * 2)) # down side live
self.assertLess(down._byte_low_frontier(), up.unified_buffer.total_bytes)
up.watermark_physical = up.num_pages # up band swallows the buffer
out = mea._chain_byte_accounting_violations(chain)
self.assertTrue(any("overlap" in s for s in out), out)
def test_pair_order_is_direction_agnostic(self):
"""The factories and the unit fixtures orient the pair differently;
the check must order by grow direction, not by argument position."""
full, swa = _paged_pair(lazy=False)
a = mea._end_pair_chain(full, swa)
b = mea._end_pair_chain(swa, full)
self.assertEqual([x.sub_pool_name for x in a], [x.sub_pool_name for x in b])
if __name__ == "__main__":
unittest.main()
test/registered/unit/mem_cache/test_unified_byte_budget_sizing.py
+244
View File
@@ -0,0 +1,244 @@
# Copyright 2023-2026 SGLang Team
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
"""Byte-budget buffer sizing for the unified 2-pool factories.
Derived properties pinned here:
* Budget honored EXACTLY: with ``unified_total_bytes`` set, the swa pair's
buffer is that many bytes (the mamba pair adds the state pool's bytes on
top — the budget is captured AFTER the state carve-out). Sizing from the
ratio-derived token counts instead re-introduces the configurator's
rounding: the swa split floors the budget by the cell size and then
page-aligns EACH side's token count, so the re-sum reconstructs less
than the profiled budget by up to about one page of tokens per side.
* Fallback: without the budget, sizing is the historical token-count re-sum,
bit-for-bit.
* bs=1 feasibility floor: a budget that cannot fit ONE worst-case request
(full KV at max context, plus one SWA window / the state slots a single
running request locks) raises at BOOT, before any pool construction —
under-sizing is a retract LIVELOCK at runtime, not a perf bug.
* The 4096-byte alignment exists because the factories ``.view()`` the whole
uint8 buffer as the KV dtype; an unaligned budget must be floored, never
rounded up (rounding up overcommits profiled memory).
python -m pytest test/registered/unit/mem_cache/test_unified_byte_budget_sizing.py -v
"""
import unittest
import torch
from sglang.srt.mem_cache.unified_memory_pool import (
MambaSubPoolSpec,
MHASubPoolSpec,
UnifiedKVPool,
_check_bs1_feasibility_floor,
_reserved_floor_bytes,
init_unified_swa_pools,
)
from sglang.test.ci.ci_register import register_cpu_ci
register_cpu_ci(est_time=15, suite="base-a-test-cpu")
_DEV = "cpu"
def _swa_factory(**over):
kw = dict(
device=_DEV,
kv_cache_dtype=torch.float16,
head_num=2,
head_dim=8,
v_head_dim=8,
swa_head_num=2,
swa_head_dim=8,
swa_v_head_dim=8,
page_size=1,
start_layer=0,
end_layer=4,
swa_attention_layer_ids=[1, 3],
full_attention_layer_ids=[0, 2],
full_max_total_num_tokens=64,
swa_max_total_num_tokens=32,
enable_memory_saver=False,
need_sort=False,
)
kw.update(over)
return init_unified_swa_pools(**kw)
def _entry_bytes():
full = MHASubPoolSpec(
name="full",
layer_num=2,
head_num=2,
head_dim=8,
store_dtype=torch.float16,
grow_direction="up",
)
return full.entry_bytes()
class TestBudgetSizing(unittest.TestCase):
def test_swa_factory_honors_the_budget_exactly(self):
e = _entry_bytes()
budget = 96 * e + 512 # deliberately NOT a token-count multiple
bundle = _swa_factory(unified_total_bytes=budget)
self.assertEqual(bundle.unified_memory_pool.total_bytes, budget)
def test_fallback_is_the_token_count_resum(self):
e = _entry_bytes()
bundle = _swa_factory()
self.assertEqual(bundle.unified_memory_pool.total_bytes, (64 + 32) * e)
def test_budget_beats_resum_on_rounding(self):
"""The property that motivates the whole phase: the re-sum cannot
represent a budget that is not a whole-token multiple per side, so it
strands bytes the buffer could have held."""
e = _entry_bytes()
budget = (64 + 32) * e + (e - 2) # almost one more entry
bundle = _swa_factory(unified_total_bytes=budget)
self.assertEqual(bundle.unified_memory_pool.total_bytes, budget)
self.assertGreater(budget, (64 + 32) * e)
class TestReservedFloorIsOneSourceOfTruth(unittest.TestCase):
"""The bs=1 floor charges the slot-0 sink, and MUST charge exactly what
`UnifiedKVPool` actually reserves.
Regression (GPU eval_434/436, Falcon-H1 boot): the floor hand-copied the
formula as `page_size * max(entry_bytes)`, applying the page multiplier to
the MAMBA spec. The pool deliberately excludes mamba (it is page_size=1),
so with page_size=256 and a ~139 MB state entry the floor over-charged the
sink by 256x — ~33 GiB of phantom requirement — and a healthy config
failed to boot with 25 GiB of real headroom.
"""
def _specs(self, page_size):
full = MHASubPoolSpec(
name="full",
layer_num=2,
head_num=2,
head_dim=8,
store_dtype=torch.float16,
grow_direction="down",
)
# A state entry vastly larger than a KV token entry — the real ratio
# (~139 MB vs ~45 KB) is what made the over-charge fatal.
mamba = MambaSubPoolSpec(
name="mamba",
layer_num=2,
conv_state_shapes=((4, 256),),
conv_dtype=torch.float16,
temporal_state_shape=(4, 256, 64),
temporal_dtype=torch.float16,
grow_direction="up",
)
return full, mamba
def test_mamba_entry_is_not_multiplied_by_page_size(self):
full, mamba = self._specs(page_size=256)
got = _reserved_floor_bytes([full, mamba], 256)
self.assertEqual(got, max(mamba.entry_bytes(), 256 * full.entry_bytes()))
self.assertLess(got, 256 * mamba.entry_bytes()) # the bug's value
def test_floor_sink_equals_what_the_pool_reserves(self):
"""Pin the two against each other so the formula cannot drift again."""
for page_size in (1, 4, 256):
with self.subTest(page_size=page_size):
full, mamba = self._specs(page_size)
floor = _reserved_floor_bytes([full, mamba], page_size)
pool = UnifiedKVPool(
total_bytes=floor + 64 * mamba.entry_bytes(),
sub_pool_specs=[full, mamba],
device=_DEV,
enable_memory_saver=False,
page_size=page_size,
)
# min_slot_index is ceil(reserved_floor / entry_bytes) per side.
for spec in (full, mamba):
self.assertEqual(
pool.min_slot_index(spec.name),
-(-floor // spec.entry_bytes()),
)
class TestBs1FeasibilityFloor(unittest.TestCase):
def test_infeasible_budget_raises_before_construction(self):
"""The buffer cannot hold one sliding window plus the sink, so boot
must fail loud instead of livelocking later."""
with self.assertRaises(RuntimeError) as ctx:
_swa_factory(
unified_total_bytes=8 * _entry_bytes(),
model_context_len=4096,
sliding_window_size=4096,
)
self.assertIn("bs=1 floor", str(ctx.exception))
self.assertIn("swa_window_kv", str(ctx.exception))
def test_context_longer_than_the_pool_is_not_rejected(self):
"""REGRESSION: the floor must NOT charge the full-attention token side.
`TpModelWorker.get_worker_info` clamps max_req_len to the pool, so a
context far larger than the buffer is refused at admission, not a
livelock -- and it is an ordinary way to serve a long-context model on
one GPU. Charging it here made such configs fail at boot."""
e = _entry_bytes()
bundle = _swa_factory(
unified_total_bytes=200 * e,
model_context_len=1_000_000, # far beyond what the buffer holds
sliding_window_size=16,
)
self.assertEqual(bundle.unified_memory_pool.total_bytes, 200 * e)
def test_feasible_config_boots_with_floor_inputs_present(self):
e = _entry_bytes()
bundle = _swa_factory(
unified_total_bytes=200 * e,
model_context_len=64,
sliding_window_size=16,
)
self.assertEqual(bundle.unified_memory_pool.total_bytes, 200 * e)
def test_window_term_is_clamped_to_context(self):
"""A window larger than the context must charge at most the context —
otherwise short-context models over-raise."""
e = _entry_bytes()
bundle = _swa_factory(
unified_total_bytes=200 * e,
model_context_len=64,
sliding_window_size=10_000, # window >> context
)
self.assertIsNotNone(bundle)
def test_floor_message_itemizes_terms(self):
with self.assertRaises(RuntimeError) as ctx:
_check_bs1_feasibility_floor(
total_bytes=10,
floor_terms=[("a", 8), ("b", 8)],
factory="test",
)
msg = str(ctx.exception)
self.assertIn("a=8", msg)
self.assertIn("b=8", msg)
self.assertIn("16", msg)
def test_exact_floor_passes(self):
"""Boundary: total == floor must NOT raise (>= is the contract)."""
_check_bs1_feasibility_floor(
total_bytes=16, floor_terms=[("a", 8), ("b", 8)], factory="test"
)
if __name__ == "__main__":
unittest.main()