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

Fix unified SWA: size a non-owner's v2p by the id space it must address (#37560)

Browse Source 
Co-authored-by: Claude Opus 5 <noreply@anthropic.com>
This commit is contained in:
Cheng Wan
2026-09-02 16:56:18 -07:00
committed by GitHub
co-authored by Claude Opus 5
parent 5a1275a519
commit 5ddca6819e
3 changed files with 232 additions and 6 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/mem_cache/multi_ended_allocator.py
+13 -6
View File
@@ -267,6 +267,7 @@ class MultiEndedAllocator(BaseTokenToKVPoolAllocator):
sub_pool_name: str,
device: str,
is_id_owner: bool,
virtual_num_pages: Optional[int] = None,
page_size: int = 1,
shards_under_dcp: bool = False,
need_sort: bool = False,
@@ -332,10 +333,15 @@ class MultiEndedAllocator(BaseTokenToKVPoolAllocator):
) // self.pool_page_size
self.entry_bytes_per_page = self.entry_bytes * self.pool_page_size
# v2p / p2v sized by PAGES. Page 0 is the padding anchor; trailing row is
# the -1 sentinel.
# v2p is indexed by VIRTUAL page id, p2v by PHYSICAL page id. A
# non-owner consumes the owner's ids, so its v2p spans the owner's
# count; the two are unrelated and either can be the larger.
self.num_virtual_ids = (
self.num_pages if virtual_num_pages is None else virtual_num_pages
)
# Page 0 is the padding anchor; the trailing row is the -1 sentinel.
self.virtual_to_physical = torch.full(
(self.num_pages + 1,),
(self.num_virtual_ids + 1,),
-1,
dtype=torch.int64,
device=device,
@@ -346,8 +352,6 @@ class MultiEndedAllocator(BaseTokenToKVPoolAllocator):
dtype=torch.int64,
device=device,
)
# Back-compat alias (count of virtual PAGES) consulted by is_slot_allocated.
self.num_virtual_ids = self.num_pages
# Chain neighbours: `low_peer` toward byte 0, `high_peer` toward
# `total_bytes`. Ends have one (`bind_peer`), float middles have both.
@@ -540,7 +544,7 @@ class MultiEndedAllocator(BaseTokenToKVPoolAllocator):
def is_slot_allocated(self, slot: int) -> bool:
"""Whether the PAGE containing this virtual id is in use."""
virt_page = slot // self.page_size
if virt_page < 0 or virt_page >= self.num_pages:
if virt_page < 0 or virt_page >= self.num_virtual_ids:
return False
return int(self.virtual_to_physical[virt_page].item()) != -1
@@ -3212,6 +3216,9 @@ class UnifiedSWATokenToKVPoolAllocator(SWATokenToKVPoolAllocator):
need_sort=need_sort,
forward_stream=forward_stream,
lazy_compaction=lazy_compaction,
# swa binds the virtual pages full mints, so it must address
# full's whole id space.
virtual_num_pages=self.full_attn_allocator.num_virtual_ids,
)
self._wire_peers()
test/registered/attention/test_gemma4_unified_swa_virtual_ids.py
+100
View File
@@ -0,0 +1,100 @@
# 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.
# ==============================================================================
"""Unified memory on a hybrid sliding-window model, which nothing else covered.
The unified SWA composite mints one virtual page id per allocation and binds it
on both sides, so the swa sub-pool's `virtual_to_physical` is indexed by the
FULL side's ids while it used to be sized by its own page count. A model with
few full-attention layers and many sliding ones gives the id owner more pages,
and the top of that id space then lands off the end of the swa table. On GPU
the write is unchecked (`alloc_bind_inplace`'s `tl.store`), so the symptom
surfaces later and elsewhere: a device-side index assert on the read in
`_swa_write_loc_unified`, and a dead scheduler.
The reachability of it is a KV-budget property, not just a model property: ids
come off the front of the owner's free list and freed ids return to the back, so
the cursor sweeps the owner's range over time and the failure needs cumulative
churn past `swa.num_pages`. `--max-total-tokens 60000` puts that within one
200-example GSM8K run; at this model's default budget the swa table is 3.8M
entries and it would take dozens.
Every argument below is load-bearing: a larger pool, or synthetic prompts in
place of this eval, and the narrow table passes.
python -m pytest test/registered/attention/test_gemma4_unified_swa_virtual_ids.py -v
"""
import types
import unittest
from sglang.srt.utils import kill_process_tree
from sglang.test.ci.ci_register import register_cuda_ci
from sglang.test.run_eval import run_eval
from sglang.test.test_utils import (
DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
DEFAULT_URL_FOR_TEST,
CustomTestCase,
popen_launch_server,
)
register_cuda_ci(est_time=65, stage="base-b", runner_config="1-gpu-large")
# Healthy 0.865-0.880 (static pool 0.870); a died-mid-run server scores
# 0.05-0.21. 1 sigma over 200 examples is ~0.024.
SCORE_THRESHOLD = 0.82
class TestGemma4UnifiedSwaVirtualIds(CustomTestCase):
@classmethod
def setUpClass(cls):
cls.model = "google/gemma-4-E2B-it"
cls.base_url = DEFAULT_URL_FOR_TEST
cls.process = popen_launch_server(
cls.model,
cls.base_url,
timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
other_args=[
"--attention-backend",
"triton",
"--enable-unified-memory",
"--disable-radix-cache",
"--mem-fraction-static",
"0.8",
# The trigger: this is what brings `swa.num_pages` within one
# eval's worth of churn.
"--max-total-tokens",
"60000",
],
)
@classmethod
def tearDownClass(cls):
if hasattr(cls, "process") and cls.process:
kill_process_tree(cls.process.pid)
def test_gsm8k_survives_churn_past_the_swa_page_count(self):
metrics = run_eval(
types.SimpleNamespace(
base_url=self.base_url,
model=self.model,
eval_name="gsm8k",
num_examples=200,
num_threads=128,
)
)
self.assertGreaterEqual(metrics["score"], SCORE_THRESHOLD)
if __name__ == "__main__":
unittest.main()
test/registered/unit/mem_cache/test_unified_swa_shared_virtual_ids.py
+119
View File
@@ -0,0 +1,119 @@
# 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.
# ==============================================================================
"""The unified SWA composite mints one virtual page id and binds it on both
sides, so the swa side's `virtual_to_physical` is indexed by the FULL side's
ids. Its own page count says nothing about how wide that has to be: which side
gets more pages out of a shared byte budget depends on the layer split, and a
model with few full-attention layers and many sliding ones (gemma-4: 10 and 50)
gives the full side an order of magnitude more.
Reaching an id that high takes cumulative allocation, so a narrow table fails
after churn rather than at once -- and on GPU it fails as a `tl.store` past the
end of the table: an unchecked write, not a raised index.
python -m pytest test/registered/unit/mem_cache/test_unified_swa_shared_virtual_ids.py -v
"""
import unittest
import torch
from test_swa_locked_full_recover_unified import _DEV, _FakeUnifiedSWAKVPool
from sglang.srt.mem_cache.multi_ended_allocator import (
UnifiedSWATokenToKVPoolAllocator,
)
from sglang.srt.mem_cache.unified_memory_pool import MHASubPoolSpec, UnifiedKVPool
from sglang.test.ci.ci_register import register_cpu_ci
register_cpu_ci(est_time=10, suite="base-a-test-cpu")
def _build(n_full: int, n_swa: int, full_layers: int, swa_layers: int):
"""A composite whose sides have different per-page byte costs. `full_layers
< swa_layers` is the gemma-4 shape: the cheap side is the id owner and ends
up with far more pages than the side that has to address them."""
full_spec = MHASubPoolSpec(
name="full",
layer_num=full_layers,
head_num=2,
head_dim=4,
store_dtype=torch.float16,
grow_direction="up",
)
swa_spec = MHASubPoolSpec(
name="swa",
layer_num=swa_layers,
head_num=2,
head_dim=4,
store_dtype=torch.float16,
grow_direction="down",
)
pool = UnifiedKVPool(
total_bytes=n_full * full_spec.entry_bytes() + n_swa * swa_spec.entry_bytes(),
sub_pool_specs=[full_spec, swa_spec],
device=_DEV,
enable_memory_saver=False,
)
return UnifiedSWATokenToKVPoolAllocator(
unified_buffer=pool,
kvcache=_FakeUnifiedSWAKVPool(pool),
device=_DEV,
full_max_total_num_tokens=n_full,
swa_max_total_num_tokens=n_swa,
need_sort=False,
forward_stream=None,
)
class TestSharedVirtualIdSpace(unittest.TestCase):
def test_swa_table_spans_the_owners_id_space(self):
"""Static form: the table has to be wide enough before any alloc runs."""
for full_layers, swa_layers in ((1, 5), (5, 1), (2, 2)):
with self.subTest(full_layers=full_layers, swa_layers=swa_layers):
alloc = _build(200, 20, full_layers, swa_layers)
owner = alloc.full_attn_allocator
swa = alloc.swa_attn_allocator
self.assertGreaterEqual(
int(swa.virtual_to_physical.shape[0]),
owner.num_virtual_ids + 1,
"swa v2p cannot address every id the owner can mint",
)
# p2v stays this pool's own business: it is indexed by physical id.
self.assertEqual(
int(swa.physical_to_virtual.shape[0]), swa.num_pages + 1
)
def test_churn_past_the_swa_page_count_binds_cleanly(self):
"""Dynamic form: alloc/free until the owner's cursor passes the swa
side's page count, which is where the narrow table used to be written
off the end."""
alloc = _build(200, 20, full_layers=1, swa_layers=5)
swa_pages = alloc.swa_attn_allocator.num_pages
highest = 0
for _ in range(40):
v = alloc.alloc(4)
if v is None:
break
highest = max(highest, int(v.max()) // alloc.page_size)
alloc.free(v)
self.assertGreater(
highest,
swa_pages,
f"churn never reached past the swa side's {swa_pages} pages, so this "
"test would pass on a narrow table too",
)
if __name__ == "__main__":
unittest.main()