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

Use a shared byte budget for unified hybrid-SWA memory (#36729)

Browse Source 
Co-authored-by: yhzhuang <yhzhuang@fb.com>
Co-authored-by: Cheng Wan <cheng.wan@radixark.ai>
This commit is contained in:
Yonghao Zhuang
2026-09-15 15:27:00 -07:00
committed by GitHub
co-authored by yhzhuang Cheng Wan
parent 4da5599e93
commit 2929a39927
32 changed files with 2220 additions and 588 deletions
Download Patch File Download Diff File Expand all files Collapse all files
test/registered/unit/mem_cache/test_hisparse_max_token_pool_size.py
+6 -2
View File
@@ -14,7 +14,9 @@ from types import SimpleNamespace
from unittest.mock import MagicMock
from sglang.srt.disaggregation.decode import DecodePreallocQueue
from sglang.srt.mem_cache.kv_cache_configurator import KVCacheConfigurator
from sglang.srt.model_executor.model_runner import ModelRunner
from sglang.srt.runtime_context import get_context
from sglang.test.ci.ci_register import register_cpu_ci
from sglang.test.test_utils import CustomTestCase
@@ -28,6 +30,7 @@ def _make_model_runner(**attrs):
raise AttributeError on the internal `self.effective_max_total_num_tokens`
read inside `max_token_pool_size`."""
instance = object.__new__(ModelRunner)
instance.kv_cache_configurator = object.__new__(KVCacheConfigurator)
for name, value in attrs.items():
setattr(instance, name, value)
return instance
@@ -85,8 +88,9 @@ class TestMaxTokenPoolSize(CustomTestCase):
full_max_total_num_tokens=3000,
swa_max_total_num_tokens=500,
)
self.assertEqual(instance.max_token_pool_size, 3000)
self.assertEqual(instance.effective_max_total_num_tokens, 3000)
with get_context().override_server_args(enable_unified_memory=False):
self.assertEqual(instance.max_token_pool_size, 3000)
self.assertEqual(instance.effective_max_total_num_tokens, 3000)
def _make_prealloc_queue(
test/registered/unit/mem_cache/test_kv_index_translator.py
+119 -1
View File
@@ -35,6 +35,8 @@ from sglang.srt.mem_cache.allocator.unified_hybrid_swa import (
from sglang.srt.mem_cache.kv_index_translator import KVIndexTranslator, KVReadTables
from sglang.srt.mem_cache.swa_memory_pool import SWAKVPool
from sglang.srt.mem_cache.unified_memory_pool import MHASubPoolSpec, UnifiedKVPool
from sglang.srt.state_capturer.base import BaseTopkCapturer
from sglang.test.test_utils import CustomTestCase
_DEV = "cpu"
_FULL_L = 2
@@ -135,6 +137,7 @@ class TestPassthrough(unittest.TestCase):
device=_DEV,
)
self.assertFalse(src.is_translating)
self.assertEqual(src.capture_token_capacity(17), 18)
rows = torch.tensor([2, 0])
view = src.build_index_table(
req_pool_indices=rows, seq_lens=torch.tensor([5, 3])
@@ -594,7 +597,7 @@ class TestViewMemo(unittest.TestCase):
self.assertEqual(v2.ids.shape[0], 1)
class TestWriteLoc(unittest.TestCase):
class TestWriteLoc(CustomTestCase):
"""The two-phase write contract: `rebind_write_loc` rebinds the full side
once at ForwardBatch construction, and the sliding-window write loc derives
POINTWISE from the full-side values -- pads, slices, and fresh copies
@@ -630,6 +633,121 @@ class TestWriteLoc(unittest.TestCase):
self.assertTrue(torch.equal(fb.out_cache_loc, want_full))
self.assertTrue(torch.equal(virt, keep))
def test_capture_capacity_covers_dcp_widened_mamba_ids(self):
"""A capturer must store the highest virtual IDs issued under DCP."""
from sglang.srt.mem_cache.allocator.unified_mamba import (
UnifiedMambaTokenToKVPoolAllocator,
)
from sglang.srt.mem_cache.unified_memory_pool import MambaSubPoolSpec
from sglang.srt.runtime_context import get_parallel
for dcp_size in (1, 2, 4):
with (
self.subTest(dcp_size=dcp_size),
get_parallel().override(attn_dcp_size=dcp_size),
):
pool = UnifiedKVPool(
total_bytes=2048,
sub_pool_specs=[
MHASubPoolSpec(
name="full",
layer_num=1,
head_num=1,
head_dim=4,
store_dtype=torch.float16,
grow_direction="up",
),
MambaSubPoolSpec(
name="mamba",
layer_num=1,
conv_state_shapes=((2, 2),),
conv_dtype=torch.float16,
temporal_state_shape=(2, 2),
temporal_dtype=torch.float16,
grow_direction="down",
),
],
device="cpu",
enable_memory_saver=False,
page_size=4,
)
allocator = UnifiedMambaTokenToKVPoolAllocator(
unified_buffer=pool,
kvcache=SimpleNamespace(full_kv_pool=None, mamba_pool=None),
device="cpu",
page_size=4,
)
virt = allocator.alloc(allocator.available_size())
self.assertIsNotNone(virt)
src = _make_source(allocator, virt[None, :], 4)
cap = object.__new__(BaseTopkCapturer)
cap.topk_size = 1
expected = torch.arange(len(virt), dtype=torch.int32).reshape(-1, 1, 1)
cap.device_cache = SimpleNamespace(buffer=expected)
cap.host_cache = SimpleNamespace(
buffer=torch.zeros(
src.capture_token_capacity(1), 1, 1, dtype=torch.int32
)
)
fb = _FakeForwardBatch(out_cache_loc=virt)
fb.out_cache_loc_virtual = virt
cap.on_forward_end(fb, False, None, no_copy_to_cpu=False)
req_pool = SimpleNamespace(req_to_token=virt[None, :])
self.assertTrue(
torch.equal(cap.get_topk(0, len(virt) + 1, req_pool), expected)
)
def test_topk_capture_round_trips_request_token_ids(self):
for ps in (1, 4, 64):
for translating in (False, True):
for overlap in (False, True):
with self.subTest(
page_size=ps, translating=translating, overlap=overlap
):
src, _, _, _, virt, _, _ = self._built(ps=ps, n=3 * ps)
req_pool = SimpleNamespace(req_to_token=virt.clone()[None, :])
fb = _FakeForwardBatch(out_cache_loc=virt.clone())
if translating:
src.rebind_write_loc(fb)
fb.out_cache_loc = torch.cat(
[fb.out_cache_loc, virt.new_zeros(2)]
)
else:
fb.out_cache_loc_virtual = None
# Admission may be capped below IDs issued after reuse.
capacity = src.capture_token_capacity(ps)
self.assertGreater(capacity, int(virt.max()))
expected = (
torch.arange(len(virt) * 4, dtype=torch.int32).reshape(
-1, 2, 2
)
+ 1
)
cap = object.__new__(BaseTopkCapturer)
cap.topk_size = 2
cap.device_cache = SimpleNamespace(
buffer=torch.cat([expected, expected.new_zeros(2, 2, 2)])
)
cap.host_cache = SimpleNamespace(
buffer=torch.zeros(capacity, 2, 2, dtype=torch.int32)
)
result = cap.on_forward_end(
fb, False, None, no_copy_to_cpu=overlap
)
if overlap:
fb.out_cache_loc.zero_()
if translating:
fb.out_cache_loc_virtual.zero_()
cap.device_cache.buffer.zero_()
result.map_device_tensors(lambda value: value.cpu())
result.finalize()
self.assertTrue(
torch.equal(
cap.get_topk(0, len(virt) + 1, req_pool), expected
)
)
def test_swa_write_loc_round_trips_from_full_side(self):
"""Derived property: `field(full(t)) == swa(t)` for any virtual run t,
across page sizes and multipliers."""
test/registered/unit/mem_cache/test_multi_ended_allocator.py
+167
View File
@@ -40,6 +40,7 @@ from sglang.srt.mem_cache.allocator.unified_sub_pool import (
MultiEndedAllocator,
)
from sglang.srt.mem_cache.base_prefix_cache import EvictParams
from sglang.srt.mem_cache.prefill_budget import estimate_swa_kv_tokens
from sglang.srt.mem_cache.unified_cache.components import ComponentType
from sglang.srt.mem_cache.unified_memory_pool import (
MambaSubPoolSpec,
@@ -558,6 +559,7 @@ class TestUnifiedSWATokenToKVPoolAllocator(unittest.TestCase):
swa_layer_num=2,
head_num=2,
head_dim=4,
page_size=1,
):
full_spec = MHASubPoolSpec(
name="full",
@@ -584,6 +586,7 @@ class TestUnifiedSWATokenToKVPoolAllocator(unittest.TestCase):
sub_pool_specs=[full_spec, swa_spec],
device=_DEV,
enable_memory_saver=False,
page_size=page_size,
)
kvcache = _FakeUnifiedSWAKVPool(pool)
allocator = UnifiedSWATokenToKVPoolAllocator(
@@ -592,11 +595,175 @@ class TestUnifiedSWATokenToKVPoolAllocator(unittest.TestCase):
device=_DEV,
full_max_total_num_tokens=n_full_slots,
swa_max_total_num_tokens=n_swa_slots,
page_size=page_size,
need_sort=False,
forward_stream=None,
)
return pool, allocator, kvcache
def test_reclaim_plan_matches_exhaustive_page_targets(self):
page_size = 4
_, allocator, _ = self._build(
n_full_slots=40, n_swa_slots=24, page_size=page_size
)
allocator.lazy_compaction = True
for sub_pool in (allocator.full_attn_allocator, allocator.swa_attn_allocator):
sub_pool.lazy_compaction = True
sub_pool.disagg_move_gate = lambda: False
live = allocator.alloc(16)
self.assertIsNotNone(live)
allocator.free(live[4:8])
allocator.free_swa(live[8:12])
for compacted in (False, True):
for sub_pool in (
allocator.full_attn_allocator,
allocator.swa_attn_allocator,
):
sub_pool.disagg_move_gate = lambda: compacted
for full_evictable, swa_evictable in ((0, 0), (7, 5), (12, 8), (100, 100)):
max_full = min(12, full_evictable) // page_size
max_swa = min(8, swa_evictable) // page_size
for full_pages in range(9):
for swa_pages in range(9):
feasible = [
(full * page_size, swa * page_size)
for swa in range(max_swa + 1)
for full in range(max_full + 1)
if allocator._fits_page_demand(
full_pages,
swa_pages,
full_reclaim_pages=full,
swa_reclaim_pages=swa,
compacted=compacted,
)
]
with self.subTest(
compacted=compacted,
evictable=(full_evictable, swa_evictable),
pages=(full_pages, swa_pages),
):
self.assertEqual(
allocator.reclaim_plan(
full_pages * page_size,
swa_pages * page_size,
full_evictable_tokens=full_evictable,
swa_evictable_tokens=swa_evictable,
),
feasible[0] if feasible else None,
)
def test_restore_swa_without_allocating_more_full(self):
_, allocator, _ = self._build(page_size=4)
indices = allocator.alloc(8)
full_before = allocator.translate_kv_indices_for_transfer(indices).clone()
allocator.free_swa(indices)
self.assertEqual(allocator.reclaim_plan(0, 8), (0, 0))
self.assertTrue(allocator.can_reserve(0, 8))
self.assertTrue(allocator.ensure_capacity(0, 8))
allocator.swa_attn_allocator.alloc_with_virtual((indices // 4).unique())
self.assertTrue(
torch.equal(
allocator.translate_kv_indices_for_transfer(indices), full_before
)
)
self.assertTrue(
bool((allocator.swa_attn_allocator.translate_kv_loc(indices) > 0).all())
)
def test_empty_pool_reservation_matches_packed_byte_boundary(self):
page_size = 4
_, allocator, _ = self._build(
n_full_slots=40,
n_swa_slots=24,
full_layer_num=4,
swa_layer_num=2,
page_size=page_size,
)
full_allocator = allocator.full_attn_allocator
swa_allocator = allocator.swa_attn_allocator
swa_pages = 2
full_pages = (
allocator._empty_shared_gap_bytes
- swa_pages * swa_allocator.entry_bytes_per_page
) // full_allocator.entry_bytes_per_page
packed_bytes = (
full_pages * full_allocator.entry_bytes_per_page
+ swa_pages * swa_allocator.entry_bytes_per_page
)
self.assertLessEqual(
full_pages + 1, full_allocator.num_pages - full_allocator.min_page_index
)
self.assertLessEqual(
swa_pages, swa_allocator.num_pages - swa_allocator.min_page_index
)
self.assertLessEqual(packed_bytes, allocator._empty_shared_gap_bytes)
self.assertGreater(
packed_bytes + full_allocator.entry_bytes_per_page,
allocator._empty_shared_gap_bytes,
)
self.assertTrue(
allocator.can_reserve(
full_pages * page_size,
swa_pages * page_size,
empty_pool=True,
)
)
self.assertFalse(
allocator.can_reserve(
full_pages * page_size + 1,
swa_pages * page_size,
empty_pool=True,
)
)
extend_tokens = 32
max_new_tokens = 0
reservation_full_tokens = extend_tokens + max_new_tokens + page_size
reservation_swa_tokens = estimate_swa_kv_tokens(
extend_tokens,
max_new_tokens,
sliding_window_size=16,
page_size=page_size,
allocation_limit=16,
)
reservation_swa_with_tail = estimate_swa_kv_tokens(
extend_tokens,
max_new_tokens,
sliding_window_size=16,
page_size=page_size,
)
reservation_bytes = (
reservation_full_tokens // page_size
) * full_allocator.entry_bytes_per_page + (
reservation_swa_tokens // page_size
) * swa_allocator.entry_bytes_per_page
reservation_bytes_with_tail = (
reservation_full_tokens // page_size
) * full_allocator.entry_bytes_per_page + (
reservation_swa_with_tail // page_size
) * swa_allocator.entry_bytes_per_page
self.assertLessEqual(reservation_bytes, allocator._empty_shared_gap_bytes)
self.assertGreater(
reservation_bytes_with_tail, allocator._empty_shared_gap_bytes
)
self.assertTrue(
allocator.can_reserve(
reservation_full_tokens,
reservation_swa_tokens,
empty_pool=True,
)
)
self.assertFalse(
allocator.can_reserve(
reservation_full_tokens,
reservation_swa_with_tail,
empty_pool=True,
)
)
def _alloc(self, allocator, kvcache, n):
"""Allocate N virtual ids; stamp the data marker on both sub-pools."""
v = allocator.alloc(n)
test/registered/unit/mem_cache/test_prefill_memory_budget.py
+300
View File
@@ -0,0 +1,300 @@
"""CPU regressions for allocator-owned prefill admission and pending demand."""
import unittest
from array import array
from types import SimpleNamespace
from unittest.mock import MagicMock
import torch
from sglang.srt.managers.schedule_batch import Req
from sglang.srt.managers.schedule_policy import PrefillAdder
from sglang.srt.managers.scheduler import Scheduler
from sglang.srt.mem_cache.allocator.hisparse import (
DeepSeekV4HiSparseTokenToKVPoolAllocator,
)
from sglang.srt.mem_cache.allocator.swa import (
PureSWATokenToKVPoolAllocator,
SWATokenToKVPoolAllocator,
)
from sglang.srt.mem_cache.allocator.unified_hybrid_swa import (
UnifiedMambaSWATokenToKVPoolAllocator,
)
from sglang.srt.mem_cache.common import evict_from_tree_cache
from sglang.srt.mem_cache.prefill_budget import SWAPrefillBudget
from sglang.srt.mem_cache.unified_memory_pool import init_unified_swa_pools
from sglang.srt.runtime_context import get_parallel
from sglang.srt.sampling.sampling_params import SamplingParams
from sglang.test.ci.ci_register import register_cpu_ci
register_cpu_ci(est_time=5, suite="base-a-test-cpu")
def _shared_allocator(*, page_size=4, total_bytes=1024):
return init_unified_swa_pools(
device="cpu",
kv_cache_dtype=torch.float16,
head_num=1,
head_dim=4,
v_head_dim=4,
swa_head_num=1,
swa_head_dim=4,
swa_v_head_dim=4,
page_size=page_size,
start_layer=0,
end_layer=2,
swa_attention_layer_ids=[1],
full_attention_layer_ids=[0],
total_bytes=total_bytes,
enable_memory_saver=False,
need_sort=False,
lazy_compaction=True,
).token_to_kv_pool_allocator
def _cache():
return SimpleNamespace(
sliding_window_size=8,
full_evictable_size=lambda: 0,
swa_evictable_size=lambda: 0,
is_chunk_cache=lambda: False,
)
class TestSharedPrefillMemoryBudget(unittest.TestCase):
def setUp(self):
self.allocator = _shared_allocator()
self.cache = _cache()
self.budget = self.allocator.create_prefill_budget(self.cache)
self.request = dict(
extend_input_len=12,
total_tokens=20,
max_new_tokens=4,
input_tokens=12,
swa_host_hit_length=0,
chunk_limit=16,
)
def test_pending_batch_cannot_spend_shared_bytes_twice(self):
self.assertEqual(self.budget.check_prefill(**self.request), (True, 16))
self.budget.reserve(12, 4, chunk_limit=16)
# Each side separately has room, but their combined reservation does not.
self.assertGreater(self.budget.remaining_total, 20)
self.assertGreater(self.budget.remaining_swa, 16)
self.assertEqual(self.budget.check_prefill(**self.request), (False, None))
self.assertEqual(self.allocator.full_attn_allocator.allocated_count(), 0)
self.assertEqual(self.allocator.swa_attn_allocator.allocated_count(), 0)
def test_mixed_decode_reserves_both_sides(self):
budget = self.allocator.create_prefill_budget(
self.cache, num_mixed_decode_tokens=4
)
budget.reserve(12, 4, chunk_limit=16)
self.assertEqual(
(budget.total_offset, budget.current_offset, budget.swa_offset),
(24, 20, 20),
)
self.assertEqual(budget.check_prefill(**self.request), (False, None))
def test_prefix_lock_changes_admission_without_rebuilding_budget(self):
self.assertIsNotNone(self.allocator.alloc(24))
self.cache.full_evictable_size = lambda: 24
self.cache.swa_evictable_size = lambda: 24
self.assertEqual(self.budget.check_prefill(**self.request), (True, 16))
# Locking the cached prefix removes its eviction credit.
self.cache.full_evictable_size = lambda: 0
self.cache.swa_evictable_size = lambda: 0
self.assertEqual(self.budget.check_prefill(**self.request), (False, None))
def test_final_chunk_reserves_decode_headroom(self):
limit = self.budget.fit_chunk(
extend_input_len=12, max_new_tokens=80, chunk_limit=16
)
self.assertEqual(limit, 8)
self.budget.reserve(limit, 0, chunk_limit=16, is_chunked_continuation=True)
self.assertIsNotNone(self.allocator.alloc(limit))
def test_host_swa_load_is_part_of_joint_demand(self):
self.assertEqual(self.budget.check_prefill(**self.request), (True, 16))
request = {**self.request, "swa_host_hit_length": 32}
self.assertEqual(self.budget.check_prefill(**request), (False, None))
def test_prompt_clipping_uses_the_empty_pool(self):
kwargs = dict(token_capacity=1, sliding_window_size=8, chunk_size=16)
limit = self.allocator.max_new_tokens_for_memory(12, 80, **kwargs)
self.assertIsNotNone(limit)
self.assertGreater(limit, 0)
self.assertIsNotNone(self.allocator.alloc(24))
self.assertEqual(
self.allocator.max_new_tokens_for_memory(12, 80, **kwargs), limit
)
self.assertIsNone(self.allocator.max_new_tokens_for_memory(100, 0, **kwargs))
def test_shared_stats_pair_available_tokens_with_current_capacity(self):
self.assertIsNotNone(self.allocator.alloc(12))
(full_capacity, full_free), (swa_capacity, swa_free) = (
self.allocator.swa_capacity_and_available(full_capacity=1, swa_capacity=1)
)
self.assertEqual(full_capacity - full_free, 12)
self.assertEqual(swa_capacity - swa_free, 12)
def test_common_eviction_dispatches_joint_reclaim(self):
self.cache.token_to_kv_pool_allocator = self.allocator
self.allocator.evict_to_free_tokens = MagicMock()
evict_from_tree_cache(self.cache, 8)
self.allocator.evict_to_free_tokens.assert_called_once_with(self.cache, 8)
class TestSharedPrefillAdmission(unittest.TestCase):
def _new_admission(self, page_size, pool_pages, *, ignore_eos=False):
allocator = _shared_allocator(
page_size=page_size, total_bytes=pool_pages * page_size * 16
)
req = Req(
rid="unaligned-prompt",
origin_input_text=None,
origin_input_ids=array("q", [1] * (page_size + 1)),
sampling_params=SamplingParams(max_new_tokens=1, ignore_eos=ignore_eos),
)
scheduler = Scheduler.__new__(Scheduler)
scheduler.max_req_len = 16 * page_size
scheduler.max_total_num_tokens = allocator.size_full
scheduler.page_size = page_size
scheduler.max_new_tokens_limit = None
scheduler.sliding_window_size = page_size
scheduler.chunked_prefill_size = page_size
scheduler.token_to_kv_pool_allocator = allocator
with get_parallel().override(attn_dcp_size=1):
scheduler.init_req_max_new_tokens(req)
self.assertEqual(req.sampling_params.max_new_tokens, 1)
req._refresh_fill_ids()
cache = SimpleNamespace(
sliding_window_size=page_size,
disable=True,
full_evictable_size=lambda: 0,
swa_evictable_size=lambda: 0,
is_chunk_cache=lambda: True,
supports_mamba=lambda: False,
)
adder = PrefillAdder(
page_size=page_size,
tree_cache=cache,
token_to_kv_pool_allocator=allocator,
running_batch=None,
new_token_ratio=1.0,
rem_input_tokens=16 * page_size,
rem_chunk_tokens=page_size,
)
return allocator, req, adder
def test_unaligned_final_chunk_makes_progress(self):
for page_size in (4, 64):
with self.subTest(page_size=page_size):
allocator, req, adder = self._new_admission(page_size, pool_pages=7)
req.prefix_indices = allocator.alloc(page_size)
self.assertIsNotNone(req.prefix_indices)
self.assertTrue(allocator.can_reserve(page_size + 2, page_size + 2))
self.assertIsNone(adder.add_chunked_req(req))
self.assertEqual(adder.can_run_list, [req])
self.assertEqual(req.extend_range.length, 1)
def test_unaligned_ignore_eos_enters_empty_pool(self):
for page_size in (4, 64):
with self.subTest(page_size=page_size):
allocator, req, adder = self._new_admission(
page_size, pool_pages=6, ignore_eos=True
)
self.assertEqual(len(req.prefix_indices), 0)
self.assertTrue(allocator.can_reserve(2 * page_size + 2, 2 * page_size))
adder.add_one_req(
req, has_chunked_req=False, truncation_align_size=None
)
self.assertEqual(adder.can_run_list, [req])
class TestFixedPrefillMemoryBudget(unittest.TestCase):
def _allocator(self, cls=SWATokenToKVPoolAllocator):
allocator = object.__new__(cls)
allocator.page_size = 4
allocator._size_full = 128
allocator._size_swa = 64
allocator.full_available_size = lambda: 128
allocator.swa_available_size = lambda: 64
return allocator
def test_ring_slot_reserved_once_and_evictable_tokens_give_no_credit(self):
allocator = self._allocator()
allocator._swa_req_ring = True
allocator._swa_ring_cost = 32
cache = _cache()
cache.swa_evictable_size = lambda: 1000
budget = allocator.create_prefill_budget(cache)
budget.reserve(12, 4, chunk_limit=16)
self.assertEqual(budget.remaining_swa, 32)
budget.reserve(12, 4, chunk_limit=16, is_chunked_continuation=True)
self.assertEqual(budget.remaining_swa, 32)
# The last exact ring slot remains admissible.
self.assertEqual(
budget.check_prefill(
extend_input_len=4,
total_tokens=12,
max_new_tokens=4,
input_tokens=4,
swa_host_hit_length=0,
chunk_limit=16,
),
(True, 16),
)
def test_pure_swa_budget_reads_swa_capacity(self):
allocator = self._allocator(PureSWATokenToKVPoolAllocator)
allocator.full_available_size = lambda: 0
budget = allocator.create_prefill_budget(_cache())
self.assertEqual(budget.remaining_total, 64)
self.assertTrue(budget.has_capacity())
def test_hisparse_budget_reads_wrapper_capacity(self):
allocator = self._allocator(DeepSeekV4HiSparseTokenToKVPoolAllocator)
allocator.full_available_size = lambda: 12
budget = allocator.create_prefill_budget(_cache())
self.assertEqual(budget.remaining_total, 12)
self.assertEqual(budget.remaining_swa, 64)
def test_tri_pool_keeps_fixed_admission_and_clipping(self):
allocator = self._allocator(UnifiedMambaSWATokenToKVPoolAllocator)
allocator.can_reserve = MagicMock(
side_effect=AssertionError("two-pool reservation")
)
budget = allocator.create_prefill_budget(_cache())
self.assertIs(type(budget), SWAPrefillBudget)
self.assertTrue(budget.has_capacity())
self.assertEqual(
allocator.max_new_tokens_for_memory(
5,
100,
token_capacity=32,
sliding_window_size=8,
chunk_size=16,
),
19,
)
def test_tri_pool_eviction_does_not_reenter_common(self):
allocator = self._allocator(UnifiedMambaSWATokenToKVPoolAllocator)
allocator.available_size = lambda: 0
allocator.full_available_size = lambda: 0
allocator.swa_available_size = lambda: 0
cache = _cache()
cache.token_to_kv_pool_allocator = allocator
cache.evict_for_alloc = MagicMock()
evict_from_tree_cache(cache, 8)
cache.evict_for_alloc.assert_called_once()
params = cache.evict_for_alloc.call_args.args[0]
self.assertEqual((params.num_tokens, params.swa_num_tokens), (8, 8))
if __name__ == "__main__":
unittest.main()
test/registered/unit/mem_cache/test_unified_byte_budget_sizing.py
+6 -6
View File
@@ -13,9 +13,9 @@
# ==============================================================================
"""Byte-budget buffer sizing for the unified 2-pool factories.
With ``unified_total_bytes`` set the buffer is that many bytes exactly (the
mamba pair adds the state pool's bytes on top -- the budget is captured AFTER
the state carve-out); without it, sizing falls back to the token-count re-sum.
The SWA factory's ``total_bytes`` sets the exact buffer size. The Mamba pair's
``unified_total_bytes`` adds state bytes because its budget excludes that carve-out.
Without an explicit budget, sizing falls back to the token-count re-sum.
Sizing from the ratio-derived token counts instead would re-introduce the
configurator's rounding, which floors by the cell size and then page-aligns
EACH side, losing up to about a page of tokens per side.
@@ -91,7 +91,7 @@ class TestBudgetSizing(unittest.TestCase):
(64 + 32) * e + (e - 2), # almost one more entry
):
with self.subTest(budget=budget):
bundle = _swa_factory(unified_total_bytes=budget)
bundle = _swa_factory(total_bytes=budget)
self.assertEqual(bundle.unified_memory_pool.total_bytes, budget)
def test_fallback_is_the_token_count_resum(self):
@@ -166,7 +166,7 @@ class TestBs1FeasibilityFloor(unittest.TestCase):
must fail loud instead of livelocking later."""
with self.assertRaises(RuntimeError) as ctx:
_swa_factory(
unified_total_bytes=8 * _entry_bytes(),
total_bytes=8 * _entry_bytes(),
model_context_len=4096,
sliding_window_size=4096,
)
@@ -184,7 +184,7 @@ class TestBs1FeasibilityFloor(unittest.TestCase):
):
with self.subTest(case=case):
bundle = _swa_factory(
unified_total_bytes=200 * e,
total_bytes=200 * e,
model_context_len=model_context_len,
sliding_window_size=sliding_window_size,
)
test/registered/unit/mem_cache/test_unified_free_no_host_sync.py
+6 -1
View File
@@ -327,6 +327,7 @@ class TestEveryUnifiedAllocatorOverridesFreeSegment(unittest.TestCase):
mea.MultiEndedAllocator,
unified_mamba.UnifiedMambaTokenToKVPoolAllocator,
unified_hybrid_swa.UnifiedSWATokenToKVPoolAllocator,
unified_hybrid_swa.UnifiedMambaSWATokenToKVPoolAllocator,
):
with self.subTest(cls=cls.__name__):
self.assertIsNot(
@@ -346,9 +347,13 @@ class TestEveryUnifiedAllocatorOverridesFreeSegment(unittest.TestCase):
mea.MultiEndedAllocator,
unified_mamba.UnifiedMambaTokenToKVPoolAllocator,
unified_hybrid_swa.UnifiedSWATokenToKVPoolAllocator,
unified_hybrid_swa.UnifiedMambaSWATokenToKVPoolAllocator,
):
with self.subTest(cls=cls.__name__):
self.assertIn("free_page_reps_group", inspect.getsource(cls))
alloc = object.__new__(cls)
alloc.free_group = None
alloc.free_group_begin()
self.assertEqual(alloc.free_page_reps_group, [])
class TestUnifiedSwaFullSideGroup(unittest.TestCase):
test/registered/unit/mem_cache/test_unified_radix_allocation_eviction.py
+5 -1
View File
@@ -554,7 +554,11 @@ class TestUnifiedRadixAllocationEviction(CustomTestCase):
def test_common_helper_uses_allocation_aware_entry_point(self):
tree_cache = MagicMock()
tree_cache.is_chunk_cache.return_value = False
tree_cache.token_to_kv_pool_allocator.available_size.return_value = 30
from sglang.srt.mem_cache.allocator.token import TokenToKVPoolAllocator
allocator = object.__new__(TokenToKVPoolAllocator)
allocator.available_size = lambda: 30
tree_cache.token_to_kv_pool_allocator = allocator
evict_from_tree_cache(tree_cache, num_tokens=100)
test/registered/unit/mem_cache/test_unified_radix_cache_unittest.py
+1 -1
View File
@@ -4319,7 +4319,7 @@ class UnifiedRadixCacheSuite:
def test_buffer_load_back_swa_window_charged_at_admission(self):
"""Admission contract: a request the SWA budget gate accepts must be
allocatable at batch time (_swa_reserved_tokens: "an admitted request
allocatable at batch time (estimate_swa_kv_tokens: "an admitted request
cannot OOM"). Regression: buffer mode surfaced a staged prefetch as
host_hit_length only, so the gate never charged the SWA window that
consumption (init_load_back -> cc.load) allocates and the request
test/registered/unit/mem_cache/test_unified_tri_pool.py
+36 -1
View File
@@ -20,7 +20,7 @@ Pure CPU; fakes stand in for the KV pools (data markers verify moves).
import inspect
import unittest
from unittest.mock import MagicMock
from unittest.mock import MagicMock, patch
import torch
@@ -175,6 +175,33 @@ class TestUnifiedTriPool(unittest.TestCase):
self.assertIs(kvcache._full_allocator, fa)
self.assertIs(kvcache._swa_allocator, sa)
def test_pd_preallocation_binds_only_swa_tail_pages(self):
_, allocator, _, _ = self._build(page_size=4)
before = allocator.available_size()
prefix = torch.tensor([0], dtype=torch.int64)
seq = torch.tensor([12], dtype=torch.int64)
# With an empty prefix, ordinary allocation supplies the same virtual
# pages without launching the GPU extend kernel. All page binding and
# capacity accounting still run through the real sub-allocators.
full = allocator.full_attn_allocator
with patch.object(
full, "alloc_extend", side_effect=lambda *a, **kw: full.alloc(12)
):
virtual = allocator.alloc_extend_swa_tail(
prefix, prefix, seq, seq, torch.tensor([-1]), 12, swa_tail_len=5
)
self.assertIsNotNone(virtual)
self.assertEqual(len(virtual), 12)
self.assertEqual(allocator.full_attn_allocator.allocated_count(), 12)
self.assertEqual(allocator.swa_attn_allocator.allocated_count(), 8)
swa_pages = allocator.swa_v2p_page_table[virtual[::4] // 4]
self.assertLessEqual(swa_pages[0].item(), 0)
self.assertTrue(torch.all(swa_pages[1:] > 0).item())
allocator.free(virtual)
self.assertEqual(allocator.full_attn_allocator.allocated_count(), 0)
self.assertEqual(allocator.swa_attn_allocator.allocated_count(), 0)
self.assertEqual(allocator.available_size(), before)
def test_empty_float_is_transparent_to_the_ends(self):
_, allocator, _, _ = self._build()
fa = allocator.full_attn_allocator
@@ -961,11 +988,19 @@ class TestTriFactorySizing(unittest.TestCase):
return kw
def test_budget_sizing_and_boot_signature(self):
from sglang.srt.mem_cache.allocator.unified_hybrid_swa import (
UnifiedSWAAllocatorBase,
UnifiedSWATokenToKVPoolAllocator,
)
budget = 1 << 20
bundle = init_unified_mamba_swa_pools(
**self._factory_kwargs(unified_total_bytes=budget)
)
pool = bundle.unified_memory_pool
allocator = bundle.token_to_kv_pool_allocator
self.assertIsInstance(allocator, UnifiedSWAAllocatorBase)
self.assertNotIsInstance(allocator, UnifiedSWATokenToKVPoolAllocator)
# Buffer = budget + the state pool's bytes (budget captured AFTER the
# state carve-out), never the token-count re-sum.
state_bytes = 4 * pool.spec("mamba").entry_bytes()