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/disaggregation/test_unified_memory_move_gate.py
+49 -59
View File
@@ -11,13 +11,17 @@ corruption with no crash.
"""
import unittest
from types import SimpleNamespace
from typing import List, Optional, Set
import torch
from sglang.srt.disaggregation.utils import (
DisaggregationMode,
unified_memory_disagg_move_gate,
)
from sglang.srt.mem_cache.allocator.unified_sub_pool import MultiEndedAllocator
from sglang.srt.runtime_context import get_parallel
from sglang.test.ci.ci_register import register_cpu_ci
from sglang.test.test_utils import CustomTestCase
@@ -187,58 +191,41 @@ class TestMoveGateRejectsNonPdNode(CustomTestCase):
class TestUnifiedAllocatorsPublishTheTransferContract(CustomTestCase):
"""Every unified composite allocator must OVERRIDE the two PD hooks.
"""Unified composites must translate virtual IDs before PD transfer.
`BaseTokenToKVPoolAllocator.translate_kv_indices_for_transfer` is the
IDENTITY, and `set_disagg_move_gate` exists only where a composite defines
it. Inheriting either is silent, not loud: identity puts VIRTUAL ids on the
wire (they address real bytes, so the peer gets plausible garbage), and a
missing gate lets lazy compaction relocate pages under in-flight RDMA.
An AST-level check because instantiating these composites needs a GPU.
The implementation may be inherited from a shared unified allocator base,
but inheriting the static allocator's identity would put virtual IDs on the
wire and silently corrupt KV. Gate installation must reach every member.
"""
# Composites that own the full-side virtual ids and so must define the
# transfer translate themselves.
_COMPOSITES = (
"UnifiedMambaTokenToKVPoolAllocator",
"UnifiedSWATokenToKVPoolAllocator",
)
# Every composite must define the gate setter, including the tri-pool,
# which inherits the SWA translates (same full side) but has a THIRD
# member the 2-pool setter does not reach.
_GATE_COMPOSITES = _COMPOSITES + ("UnifiedMambaSWATokenToKVPoolAllocator",)
@staticmethod
def _own_methods(cls_name: str) -> Set[str]:
"""Names this class defines ITSELF, inheritance excluded.
Resolved off the class object rather than by parsing a named module:
these composites have already been moved once (out of
`multi_ended_allocator` into `allocator/unified_*`), and a hardcoded
module path turns that kind of move into a test failure that says
nothing about the contract. `__dict__` needs no GPU -- it is the class
body, not an instance.
"""
from sglang.srt.mem_cache.allocator import (
unified_hybrid_swa,
unified_mamba,
def _allocator_class(name):
from sglang.srt.mem_cache.allocator.unified_hybrid_swa import (
UnifiedMambaSWATokenToKVPoolAllocator,
UnifiedSWATokenToKVPoolAllocator,
)
from sglang.srt.mem_cache.allocator.unified_mamba import (
UnifiedMambaTokenToKVPoolAllocator,
)
for mod in (unified_mamba, unified_hybrid_swa):
cls = getattr(mod, cls_name, None)
if cls is not None:
return set(vars(cls))
raise AssertionError(f"class {cls_name} not found in the unified allocators")
classes = (
UnifiedMambaTokenToKVPoolAllocator,
UnifiedSWATokenToKVPoolAllocator,
UnifiedMambaSWATokenToKVPoolAllocator,
)
return {cls.__name__: cls for cls in classes}[name]
def test_transfer_translate_is_not_inherited_identity(self):
for name in self._COMPOSITES:
with self.subTest(composite=name):
self.assertIn(
"translate_kv_indices_for_transfer",
self._own_methods(name),
f"{name} inherits the identity transfer translate; PD would "
"ship VIRTUAL ids and corrupt KV without any error",
virtual = torch.tensor([1, 3], dtype=torch.int32)
for name in self._EXPECTED_COVERAGE:
with self.subTest(composite=name), get_parallel().override(attn_dcp_size=1):
alloc = object.__new__(self._allocator_class(name))
alloc.full_attn_allocator = SimpleNamespace(
translate_kv_loc=lambda ids: ids + 16
)
physical = alloc.translate_kv_indices_for_transfer(virtual)
self.assertEqual(physical.dtype, torch.int64)
self.assertEqual(physical.tolist(), [17, 19])
# Every sub-allocator attribute a composite can hold. The stub carries all
# of them regardless of composite, so the assertion is on what installation
@@ -270,11 +257,7 @@ class TestUnifiedAllocatorsPublishTheTransferContract(CustomTestCase):
`object.__new__` skips `__init__` (which needs a GPU); the setter reads
only `lazy_compaction` and the member attributes.
"""
from sglang.srt.mem_cache.allocator import unified_hybrid_swa, unified_mamba
cls = getattr(unified_mamba, cls_name, None) or getattr(
unified_hybrid_swa, cls_name
)
cls = self._allocator_class(cls_name)
alloc = object.__new__(cls)
alloc.lazy_compaction = True
for attr in self._MEMBER_ATTRS:
@@ -313,14 +296,8 @@ class TestUnifiedAllocatorsPublishTheTransferContract(CustomTestCase):
"""
import inspect
from sglang.srt.mem_cache.allocator import unified_hybrid_swa, unified_mamba
for name in self._EXPECTED_COVERAGE:
cls = getattr(unified_mamba, name, None) or getattr(
unified_hybrid_swa, name
)
if "set_disagg_move_gate" not in vars(cls):
continue # inherited, and the inherited one is checked above
cls = self._allocator_class(name)
with self.subTest(composite=name):
body = inspect.getsource(cls.set_disagg_move_gate)
self.assertIn("install_move_gate", body)
@@ -331,10 +308,23 @@ class TestUnifiedAllocatorsPublishTheTransferContract(CustomTestCase):
does not name the SWA page holding the same virtual token. The read-path
`translate_loc_from_full_to_swa` cannot stand in either: it returns
kernel-facing ids, and the transfer addresses raw page envelopes."""
self.assertIn(
"translate_swa_indices_for_transfer",
self._own_methods("UnifiedSWATokenToKVPoolAllocator"),
)
virtual = torch.tensor([1, 3], dtype=torch.int32)
for name in (
"UnifiedSWATokenToKVPoolAllocator",
"UnifiedMambaSWATokenToKVPoolAllocator",
):
with self.subTest(composite=name), get_parallel().override(attn_dcp_size=1):
alloc = object.__new__(self._allocator_class(name))
alloc.full_attn_allocator = SimpleNamespace(
translate_kv_loc=lambda ids: ids + 16
)
alloc.swa_attn_allocator = SimpleNamespace(
translate_kv_loc=lambda ids: ids + 32,
translate_kv_loc_for_kernel=lambda ids: ids + 64,
)
physical = alloc.translate_swa_indices_for_transfer(virtual)
self.assertEqual(physical.dtype, torch.int64)
self.assertEqual(physical.tolist(), [33, 35])
class TestEverySwaAllocatorAnswersTheTransferTranslate(CustomTestCase):
test/registered/unit/managers/test_prefill_adder.py
+172 -25
View File
@@ -16,6 +16,12 @@ from sglang.srt.mem_cache.base_prefix_cache import (
DecLockRefResult,
IncLockRefResult,
)
from sglang.srt.mem_cache.prefill_budget import (
PrefillBudget,
SWAPrefillBudget,
estimate_swa_kv_tokens,
)
from sglang.srt.mem_cache.unified_memory_pool import init_unified_swa_pools
from sglang.srt.runtime_context import get_context
from sglang.srt.server_args import ServerArgs, set_global_server_args_for_scheduler
from sglang.srt.utils.common import Range
@@ -75,6 +81,10 @@ class TestPrefillAdder(CustomTestCase):
allocator.swa_available_size.return_value = swa_available_size
allocator.available_size.return_value = available_size
allocator.size_swa = size_swa
allocator.swa_req_ring = False
allocator.create_prefill_budget.side_effect = lambda tree_cache, **kwargs: (
PrefillBudget(allocator, tree_cache, **kwargs)
)
return allocator
def create_running_batch(self, reqs=None) -> MagicMock:
@@ -129,8 +139,115 @@ class TestPrefillAdder(CustomTestCase):
priority_scheduling_preemption_threshold=0,
)
defaults.update(kwargs)
defaults["token_to_kv_pool_allocator"].page_size = defaults["page_size"]
return PrefillAdder(**defaults)
def create_shared_adder(self, *, num_mixed_decode_tokens=0):
self.mock_tree_cache.supports_mamba.return_value = False
self.mock_tree_cache.sliding_window_size = 8
self.mock_tree_cache.is_tree_cache.return_value = False
allocator = 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=4,
start_layer=0,
end_layer=2,
swa_attention_layer_ids=[1],
full_attention_layer_ids=[0],
total_bytes=1024,
enable_memory_saver=False,
need_sort=False,
lazy_compaction=True,
).token_to_kv_pool_allocator
return self.create_adder(
self.create_running_batch(),
page_size=4,
rem_chunk_tokens=16,
num_mixed_decode_tokens=num_mixed_decode_tokens,
token_to_kv_pool_allocator=allocator,
)
def create_shared_req(self, rid, max_new_tokens=4):
req = self.create_mock_req(rid, priority=0, max_new_tokens=max_new_tokens)
req.sampling_params.ignore_eos = False
req.swa_host_hit_length = 0
req.last_node = MagicMock()
req.full_untruncated_fill_ids = list(range(12))
req.set_extend_range = MagicMock(
side_effect=lambda start, end: setattr(
req, "extend_range", Range(start, end)
)
)
return req
def test_shared_admission_reserves_all_pending_requests(self):
adder = self.create_shared_adder()
first, second = (
self.create_shared_req("first"),
self.create_shared_req("second"),
)
adder.add_one_req(first, has_chunked_req=False, truncation_align_size=None)
self.assertEqual(adder.can_run_list, [first])
self.assertEqual(
adder.add_one_req(
second, has_chunked_req=False, truncation_align_size=None
),
AddReqResult.NO_TOKEN,
)
self.assertEqual(adder.can_run_list, [first])
def test_shared_admission_rechecks_after_prefix_lock(self):
adder = self.create_shared_adder()
self.assertIsNotNone(adder.token_to_kv_pool_allocator.alloc(24))
self.mock_tree_cache.full_evictable_size.return_value = 24
self.mock_tree_cache.swa_evictable_size.return_value = 24
def lock_prefix(_):
self.mock_tree_cache.full_evictable_size.return_value = 0
self.mock_tree_cache.swa_evictable_size.return_value = 0
return IncLockRefResult()
self.mock_tree_cache.inc_lock_ref.side_effect = lock_prefix
req = self.create_shared_req("locked-prefix")
self.assertEqual(
adder.add_one_req(req, has_chunked_req=False, truncation_align_size=None),
AddReqResult.NO_TOKEN,
)
self.mock_tree_cache.inc_lock_ref.assert_called_once()
self.assertEqual(adder.can_run_list, [])
def test_shared_continuation_defers_when_decode_consumes_chunk_budget(self):
"""Mixed decode must not commit an empty or negatively sliced prompt."""
for decode_tokens in (16, 28):
with self.subTest(decode_tokens=decode_tokens):
adder = self.create_shared_adder(num_mixed_decode_tokens=decode_tokens)
req = self.create_shared_req("continuation")
before = (
adder.memory_budget.total_offset,
adder.memory_budget.swa_offset,
)
self.assertIs(adder.add_chunked_req(req), req)
self.assertEqual(adder.can_run_list, [])
req.set_extend_range.assert_not_called()
self.assertEqual(
(adder.memory_budget.total_offset, adder.memory_budget.swa_offset),
before,
)
def test_shared_continuation_uses_memory_chunk_limit(self):
adder = self.create_shared_adder()
req = self.create_shared_req("continuation", max_new_tokens=80)
self.assertIs(adder.add_chunked_req(req), req)
self.assertEqual(req.extend_range.length, 8)
self.assertEqual(adder.memory_budget.total_offset, 12)
self.assertEqual(adder.memory_budget.swa_offset, 12)
def test_storage_prefetch_fulfillment_resolves_at_admission(self):
adder = self.create_adder(self.create_running_batch())
req = self.create_mock_req("storage-hit", priority=0, max_new_tokens=1)
@@ -208,7 +325,7 @@ class TestPrefillAdder(CustomTestCase):
running_batch = self.create_running_batch(running_reqs)
adder = self.create_adder(running_batch)
self.assertEqual(adder.rem_total_token_offset, 225)
self.assertEqual(adder.memory_budget.total_offset, 225)
self.mock_token_allocator.full_available_size.return_value = (
225 # full occupation of GRam
@@ -221,7 +338,9 @@ class TestPrefillAdder(CustomTestCase):
self.assertTrue(success)
self.assertIn(running_reqs[0], adder.preempt_list)
self.assertEqual(adder.rem_total_token_offset, 175) # 50 + 75 + 100 - 50 = 175
self.assertEqual(
adder.memory_budget.total_offset, 175
) # 50 + 75 + 100 - 50 = 175
running_batch.release_req.assert_called_once()
def test_preempt_success_low_priority_values_first(self):
@@ -238,7 +357,7 @@ class TestPrefillAdder(CustomTestCase):
running_batch = self.create_running_batch(running_reqs)
adder = self.create_adder(running_batch)
self.assertEqual(adder.rem_total_token_offset, 225)
self.assertEqual(adder.memory_budget.total_offset, 225)
self.mock_token_allocator.full_available_size.return_value = (
225 # full occupation of GRam
@@ -251,7 +370,9 @@ class TestPrefillAdder(CustomTestCase):
self.assertTrue(success)
self.assertIn(running_reqs[2], adder.preempt_list)
self.assertEqual(adder.rem_total_token_offset, 125) # 50 + 75 + 100 - 100 = 125
self.assertEqual(
adder.memory_budget.total_offset, 125
) # 50 + 75 + 100 - 100 = 125
running_batch.release_req.assert_called_once()
def test_preempt_fail_low_priority_values_first(self):
@@ -268,7 +389,7 @@ class TestPrefillAdder(CustomTestCase):
running_batch = self.create_running_batch(running_reqs)
adder = self.create_adder(running_batch)
self.assertEqual(adder.rem_total_token_offset, 225)
self.assertEqual(adder.memory_budget.total_offset, 225)
self.mock_token_allocator.full_available_size.return_value = (
225 # full occupation of GRam
@@ -306,7 +427,7 @@ class TestPrefillAdder(CustomTestCase):
running_batch = self.create_running_batch(running_reqs)
adder = self.create_adder(running_batch)
self.assertEqual(adder.rem_total_token_offset, 225)
self.assertEqual(adder.memory_budget.total_offset, 225)
self.mock_token_allocator.full_available_size.return_value = (
225 # full occupation of GRam
@@ -344,7 +465,7 @@ class TestPrefillAdder(CustomTestCase):
running_batch = self.create_running_batch(running_reqs)
adder = self.create_adder(running_batch)
self.assertEqual(adder.rem_total_token_offset, 225)
self.assertEqual(adder.memory_budget.total_offset, 225)
self.mock_token_allocator.full_available_size.return_value = 225
self.mock_token_allocator.available_size.return_value = 225
@@ -356,7 +477,7 @@ class TestPrefillAdder(CustomTestCase):
first_success = adder.preempt_to_schedule(first_req)
self.assertTrue(first_success)
self.assertIn(running_reqs[0], adder.preempt_list)
self.assertEqual(adder.rem_total_token_offset, 175)
self.assertEqual(adder.memory_budget.total_offset, 175)
running_batch.release_req.assert_called_once()
# Second call needs more tokens than currently free, so it would need to
@@ -367,7 +488,7 @@ class TestPrefillAdder(CustomTestCase):
second_success = adder.preempt_to_schedule(second_req)
self.assertFalse(second_success)
self.assertEqual(adder.rem_total_token_offset, 175)
self.assertEqual(adder.memory_budget.total_offset, 175)
self.assertEqual(adder.preempt_list.count(running_reqs[0]), 1)
running_batch.release_req.assert_called_once()
@@ -387,7 +508,7 @@ class TestPrefillAdder(CustomTestCase):
running_batch = self.create_running_batch(running_reqs)
adder = self.create_adder(running_batch)
self.assertEqual(adder.rem_total_token_offset, 475)
self.assertEqual(adder.memory_budget.total_offset, 475)
self.mock_token_allocator.full_available_size.return_value = (
475 # full occupation of GRam
@@ -400,7 +521,7 @@ class TestPrefillAdder(CustomTestCase):
self.assertTrue(success)
self.assertIn(running_reqs[2], adder.preempt_list)
self.assertEqual(
adder.rem_total_token_offset, 375
adder.memory_budget.total_offset, 375
) # 50 + 75 + 100 + 125 + 125 - 100 = 375
running_batch.release_req.assert_called_once()
@@ -420,7 +541,7 @@ class TestPrefillAdder(CustomTestCase):
running_batch = self.create_running_batch(running_reqs)
adder = self.create_adder(running_batch)
self.assertEqual(adder.rem_total_token_offset, 475)
self.assertEqual(adder.memory_budget.total_offset, 475)
self.mock_token_allocator.full_available_size.return_value = (
475 # full occupation of GRam
@@ -434,7 +555,7 @@ class TestPrefillAdder(CustomTestCase):
self.assertIn(running_reqs[2], adder.preempt_list)
self.assertIn(running_reqs[3], adder.preempt_list)
self.assertEqual(
adder.rem_total_token_offset, 250
adder.memory_budget.total_offset, 250
) # 50 + 75 + 100 + 125 + 125 - 100 - 125 = 250
self.assertEqual(running_batch.release_req.call_count, 2)
@@ -456,8 +577,8 @@ class TestPrefillAdder(CustomTestCase):
self.assertEqual(adder.rem_input_tokens, 192) # 200 - 8
self.assertEqual(adder.rem_chunk_tokens, 56) # 64 - 8
self.assertEqual(adder.rem_total_token_offset, 408) # 8 + 8 * 50
self.assertEqual(adder.cur_rem_token_offset, 8)
self.assertEqual(adder.memory_budget.total_offset, 408) # 8 + 8 * 50
self.assertEqual(adder.memory_budget.current_offset, 8)
self.assertEqual(adder.budget_state(), AddReqResult.CONTINUE)
# Add a prefill that exactly consumes the chunk budget
@@ -497,7 +618,7 @@ class TestPrefillAdder(CustomTestCase):
self.assertEqual(adder2.rem_input_tokens, 195) # 200 - 5
self.assertEqual(adder2.rem_chunk_tokens, 59) # 64 - 5
self.assertEqual(adder2.rem_total_token_offset, 255) # 5 + 5 * 50
self.assertEqual(adder2.memory_budget.total_offset, 255) # 5 + 5 * 50
self.assertEqual(adder2.budget_state(), AddReqResult.CONTINUE)
# Same prefill no longer exhausts the chunk budget
@@ -562,6 +683,10 @@ class TestPrefillAdder(CustomTestCase):
rem_chunk_tokens=rem_chunk,
)
adder.is_hybrid_swa = is_hybrid_swa
if is_hybrid_swa:
adder.memory_budget = SWAPrefillBudget(
self.mock_token_allocator, self.mock_tree_cache
)
req = self.create_mock_req("chunked", priority=0, max_new_tokens=128)
req.prefix_indices = []
@@ -640,7 +765,16 @@ class TestPrefillAdder(CustomTestCase):
page_size=page,
rem_chunk_tokens=rem_chunk,
)
self.assertEqual(adder._swa_budget_for_req(extend, max_new), expected)
self.assertEqual(
estimate_swa_kv_tokens(
extend,
max_new,
sliding_window_size=window,
page_size=page,
allocation_limit=rem_chunk,
),
expected,
)
def test_swa_admission_admits_short_cached_resume_at_two_window_pool(self):
# Livelock regression (real incident). At an SWA pool ~= 2 sliding
@@ -660,6 +794,9 @@ class TestPrefillAdder(CustomTestCase):
self.mock_tree_cache.is_tree_cache.return_value = False
adder = self.create_adder(self.create_running_batch(), page_size=PAGE)
adder.is_hybrid_swa = True
adder.memory_budget = SWAPrefillBudget(
self.mock_token_allocator, self.mock_tree_cache
)
req = self.create_mock_req(
"resume", priority=0, max_new_tokens=40, output_len=10
@@ -677,7 +814,9 @@ class TestPrefillAdder(CustomTestCase):
req.sampling_params = SimpleNamespace(max_new_tokens=40, ignore_eos=False)
# Pre-fix: a constant sliding-window reservation rejects the resume.
with patch.object(adder, "_swa_reserved_tokens", return_value=WINDOW + PAGE):
with patch.object(
adder.memory_budget, "swa_tokens", return_value=WINDOW + PAGE
):
self.assertIs(
adder.add_one_req(
req, has_chunked_req=False, truncation_align_size=None
@@ -708,6 +847,9 @@ class TestPrefillAdder(CustomTestCase):
self.mock_token_allocator.swa_available_size.return_value = 400
adder = self.create_adder(self.create_running_batch(), page_size=PAGE)
adder.is_hybrid_swa = True
adder.memory_budget = SWAPrefillBudget(
self.mock_token_allocator, self.mock_tree_cache
)
req = self.create_mock_req("dropped-fetch", priority=0, max_new_tokens=8)
req.prefix_indices = torch.empty(0, dtype=torch.int64)
req.full_untruncated_fill_ids = list(range(SPAN))
@@ -844,7 +986,7 @@ class TestPrefillAdder(CustomTestCase):
extend if req.retracted_stain else 0,
)
self.assertEqual(
adder.rem_total_token_offset,
adder.memory_budget.total_offset,
adder.ceil_paged_tokens(extend) + decode + 2,
)
self.assertEqual(
@@ -1092,6 +1234,11 @@ class TestPrefillAdder(CustomTestCase):
) -> PrefillAdder:
self.mock_tree_cache.sliding_window_size = sliding_window
self.mock_token_allocator = self.create_token_allocator(size_swa=size_swa)
self.mock_token_allocator.create_prefill_budget.side_effect = (
lambda tree_cache, **kwargs: SWAPrefillBudget(
self.mock_token_allocator, tree_cache, **kwargs
)
)
return self.create_adder(
self.create_running_batch(),
page_size=page_size,
@@ -1103,7 +1250,7 @@ class TestPrefillAdder(CustomTestCase):
# once running decodes drain -> must wait, not take the hatch.
adder = self.create_swa_adder(size_swa=1024, sliding_window=128)
self.assertFalse(
adder._swa_req_never_fits(extend_input_len=256, max_new_tokens=64)
adder.memory_budget.swa_never_fits(extend_input_len=256, max_new_tokens=64)
)
def test_swa_never_fits_true_when_budget_exceeds_whole_pool(self):
@@ -1111,7 +1258,7 @@ class TestPrefillAdder(CustomTestCase):
# pool: it can never fit however far the pool drains -> hatch.
adder = self.create_swa_adder(size_swa=1024, sliding_window=128)
self.assertTrue(
adder._swa_req_never_fits(
adder.memory_budget.swa_never_fits(
extend_input_len=256, max_new_tokens=64, swa_host_hit_length=4096
)
)
@@ -1121,14 +1268,14 @@ class TestPrefillAdder(CustomTestCase):
# the budget against size_swa (guards against a wrong-accessor bug).
req = dict(extend_input_len=256, max_new_tokens=64, swa_host_hit_length=600)
self.assertTrue(
self.create_swa_adder(size_swa=512, sliding_window=128)._swa_req_never_fits(
**req
)
self.create_swa_adder(
size_swa=512, sliding_window=128
).memory_budget.swa_never_fits(**req)
)
self.assertFalse(
self.create_swa_adder(
size_swa=4096, sliding_window=128
)._swa_req_never_fits(**req)
).memory_budget.swa_never_fits(**req)
)
test/registered/unit/managers/test_scheduler_init_req_max_new_tokens.py
+45
View File
@@ -2,8 +2,12 @@ import logging
import unittest
from types import SimpleNamespace
import torch
from sglang.srt.environ import envs
from sglang.srt.managers.scheduler import Scheduler
from sglang.srt.mem_cache.allocator.token import TokenToKVPoolAllocator
from sglang.srt.mem_cache.unified_memory_pool import init_unified_swa_pools
from sglang.srt.runtime_context import get_parallel
from sglang.test.ci.ci_register import register_cpu_ci
@@ -56,6 +60,16 @@ class TestSchedulerInitReqMaxNewTokens(unittest.TestCase):
scheduler.max_total_num_tokens = max_total_num_tokens
scheduler.page_size = page_size
scheduler.max_new_tokens_limit = envs.SGLANG_MAX_NEW_TOKENS_LIMIT.get()
scheduler.sliding_window_size = None
scheduler.chunked_prefill_size = None
scheduler.token_to_kv_pool_allocator = TokenToKVPoolAllocator(
size=max_total_num_tokens,
dtype=torch.int64,
device="cpu",
kvcache=None,
need_sort=False,
)
scheduler.token_to_kv_pool_allocator.page_size = page_size
return scheduler
def _new_req(self, max_new_tokens, input_len: int = 8, min_new_tokens: int = 0):
@@ -180,6 +194,37 @@ class TestSchedulerInitReqMaxNewTokens(unittest.TestCase):
)
self._init_and_check(scheduler, req)
def test_unified_budget_rounds_prompt_and_decode_together(self):
bundle = 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=4,
start_layer=0,
end_layer=2,
swa_attention_layer_ids=[1],
full_attention_layer_ids=[0],
total_bytes=384,
enable_memory_saver=False,
need_sort=False,
lazy_compaction=True,
)
scheduler = self._new_scheduler(page_size=4)
scheduler.token_to_kv_pool_allocator = bundle.token_to_kv_pool_allocator
scheduler.sliding_window_size = 4
scheduler.chunked_prefill_size = 4
scheduler.max_new_tokens_limit = None
for prompt_len in (4, 5, 6, 7):
with self.subTest(prompt_len=prompt_len):
req = self._new_req(max_new_tokens=1, input_len=prompt_len)
scheduler.init_req_max_new_tokens(req)
self.assertEqual(req.sampling_params.max_new_tokens, 1)
if __name__ == "__main__":
unittest.main()
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()
test/registered/unit/model_executor/test_pool_configurator.py
+97 -1
View File
@@ -83,6 +83,7 @@ def _make_model_runner(
disaggregation_mode="null",
max_running_requests=None,
disaggregation_decode_extra_slots=0,
enable_unified_memory=False,
kv_lora_rank=512,
qk_rope_head_dim=64,
swa_kv_lora_rank=128,
@@ -150,6 +151,7 @@ def _make_model_runner(
disaggregation_mode=disaggregation_mode,
max_running_requests=max_running_requests,
disaggregation_decode_extra_slots=disaggregation_decode_extra_slots,
enable_unified_memory=enable_unified_memory,
enable_hisparse=False,
enable_hierarchical_cache=False,
enable_dsa_cache_layer_split=False,
@@ -299,7 +301,14 @@ class TestDefaultConfigurator(CustomTestCase):
class TestHybridSWAConfigurator(CustomTestCase):
"""Hybrid SWA: full/swa split, ratio, memory invariant."""
def _make_swa_runner(self, full_layers=16, swa_layers=16, ratio=0.5, page_size=1):
def _make_swa_runner(
self,
full_layers=16,
swa_layers=16,
ratio=0.5,
page_size=1,
enable_unified_memory=False,
):
return _make_model_runner(
self,
is_hybrid_swa=True,
@@ -308,6 +317,7 @@ class TestHybridSWAConfigurator(CustomTestCase):
swa_num_kv_heads=4,
page_size=page_size,
swa_full_tokens_ratio=ratio,
enable_unified_memory=enable_unified_memory,
)
def _run(self, available_bytes, **kwargs):
@@ -329,6 +339,92 @@ class TestHybridSWAConfigurator(CustomTestCase):
self.assertLessEqual(used, available)
self.assertGreater(used, available * 0.99)
def test_draft_does_not_inherit_target_shared_byte_budget(self):
"""A separate draft pool must not allocate the target's byte envelope again."""
from sglang.srt.mem_cache.kv_cache_configurator import KVCacheConfigurator
mr, _, config = self._run(1 << 20, enable_unified_memory=True)
self.assertIsNotNone(config.unified_memory_pool_bytes)
configurator = object.__new__(KVCacheConfigurator)
configurator.model_config = mr.model_config
configurator.is_hybrid_swa = True
configurator.is_draft_worker = False
target = configurator._derive_pool_sizes(config=config)
configurator.is_draft_worker = True
draft = configurator._derive_pool_sizes(config=config)
self.assertEqual(
target.unified_memory_pool_bytes, config.unified_memory_pool_bytes
)
self.assertIsNone(draft.unified_memory_pool_bytes)
self.assertEqual(
draft.full_max_total_num_tokens, config.full_max_total_num_tokens
)
self.assertEqual(
draft.swa_max_total_num_tokens, config.swa_max_total_num_tokens
)
def test_unified_capacity_is_maximal_with_draft_pool(self):
page_size = 8
full_layers = 2
swa_layers = 1
draft_layers = 2
draft_swa_layers = 1
ratio = 0.5
mr = _make_model_runner(
self,
is_hybrid_swa=True,
full_attention_layer_ids=list(range(full_layers)),
swa_attention_layer_ids=list(range(full_layers, full_layers + swa_layers)),
swa_num_kv_heads=4,
swa_full_tokens_ratio=ratio,
page_size=page_size,
enable_unified_memory=True,
speculative_algorithm="EAGLE",
)
mr.spec_algorithm.is_eagle.return_value = True
mr.spec_algorithm.is_none.return_value = False
mr.spec_aux_config.eagle_draft_num_layers = draft_layers
mr.spec_aux_config.eagle_draft_swa_num_layers = draft_swa_layers
full_bytes_per_token = _full_per_token(mr)
swa_bytes_per_token = _swa_per_token(mr)
target_full_bytes_per_token = full_bytes_per_token * full_layers
draft_bytes_per_token = (
full_bytes_per_token * (draft_layers - draft_swa_layers)
+ swa_bytes_per_token * draft_swa_layers
)
def allocation_bytes(full_tokens, *, include_reserved_draft_page=True):
swa_tokens = int(full_tokens * ratio) // page_size * page_size
target_bytes = (
full_tokens * target_full_bytes_per_token
+ swa_tokens * swa_bytes_per_token * swa_layers
)
virtual_span = max(target_bytes // target_full_bytes_per_token - 1, 0)
draft_tokens = (virtual_span + page_size - 1) // page_size * page_size
if include_reserved_draft_page:
draft_tokens += page_size
return target_bytes + draft_tokens * draft_bytes_per_token
expected_full_tokens = 10 * page_size
available = allocation_bytes(
expected_full_tokens + page_size,
include_reserved_draft_page=False,
)
with mock_cpu_env():
from sglang.srt.model_executor.pool_configurator import (
create_memory_pool_configurator,
)
cfg = create_memory_pool_configurator(mr)
config = cfg.calculate_pool_sizes(available, page_size)
full_tokens = config.full_max_total_num_tokens
self.assertEqual(full_tokens % page_size, 0)
self.assertEqual(full_tokens, expected_full_tokens)
self.assertLessEqual(allocation_bytes(full_tokens), available)
self.assertGreater(allocation_bytes(full_tokens + page_size), available)
@patch(
"sglang.srt.mem_cache.kv_cache_configurator.calculate_mla_kv_cache_dim",
return_value=576,