Skip to content

fix(unified-memory): propagate unified-memory lazy checkpoint policy and handle unallocated session slots #42072

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
ch-wan merged 2 commits into from
Oct 4, 2026
Merged

fix(unified-memory): propagate unified-memory lazy checkpoint policy and handle unallocated session slots #42072

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
1fdeb52
fix(mem-cache): propagate unified lazy policy and filter session sent…
seokwoosong Oct 1, 2026
26cce36
test(mem-cache): pass the lazy checkpoint flag in the shared allocato…
ch-wan Oct 4, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 2 additions & 0 deletions python/sglang/srt/mem_cache/kv_cache_configurator.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -723,6 +723,7 @@ def _init_unified_mamba_pools(
max_num_reqs=max_num_reqs,
enable_memory_saver=get_exec().features.enable_memory_saver,
enable_mamba_extra_buffer=get_exec().mamba.enable_mamba_extra_buffer,
enable_mamba_extra_buffer_lazy=get_exec().mamba.enable_mamba_extra_buffer_lazy,
speculative_num_draft_tokens=get_spec().speculative_num_draft_tokens,
disable_overlap_schedule=get_schedule().disable_overlap_schedule,
need_sort=get_disagg().disaggregation_mode in ("decode", "prefill"),
Expand Down Expand Up @@ -835,6 +836,7 @@ def _init_unified_mamba_swa_pools(
max_num_reqs=max_num_reqs,
enable_memory_saver=get_exec().features.enable_memory_saver,
enable_mamba_extra_buffer=get_exec().mamba.enable_mamba_extra_buffer,
enable_mamba_extra_buffer_lazy=get_exec().mamba.enable_mamba_extra_buffer_lazy,
disable_overlap_schedule=get_schedule().disable_overlap_schedule,
need_sort=get_disagg().disaggregation_mode in ("decode", "prefill"),
speculative_num_draft_tokens=get_spec().speculative_num_draft_tokens,
Expand Down
6 changes: 6 additions & 0 deletions python/sglang/srt/mem_cache/unified_memory_pool.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1148,6 +1148,7 @@ def __init__(
cache_params,
mamba_layer_ids: List[int],
enable_mamba_extra_buffer: bool,
enable_mamba_extra_buffer_lazy: bool,
speculative_num_draft_tokens: Optional[int] = None,
enable_overlap_schedule: bool = True,
start_layer: Optional[int] = None,
Expand All @@ -1171,6 +1172,7 @@ def __init__(
cache_params=cache_params,
mamba_layer_ids=mamba_layer_ids,
enable_mamba_extra_buffer=enable_mamba_extra_buffer,
enable_mamba_extra_buffer_lazy=enable_mamba_extra_buffer_lazy,
speculative_num_draft_tokens=speculative_num_draft_tokens,
enable_overlap_schedule=enable_overlap_schedule,
start_layer=start_layer,
Expand Down Expand Up @@ -1367,6 +1369,7 @@ def init_unified_mamba_pools(
max_num_reqs: int,
enable_memory_saver: bool,
enable_mamba_extra_buffer: bool,
enable_mamba_extra_buffer_lazy: bool,
speculative_num_draft_tokens: Optional[int],
disable_overlap_schedule: bool,
need_sort: bool,
Expand Down Expand Up @@ -1466,6 +1469,7 @@ def init_unified_mamba_pools(
cache_params=mamba2_cache_params,
mamba_layer_ids=mamba_layer_ids,
enable_mamba_extra_buffer=enable_mamba_extra_buffer,
enable_mamba_extra_buffer_lazy=enable_mamba_extra_buffer_lazy,
speculative_num_draft_tokens=speculative_num_draft_tokens,
enable_overlap_schedule=not disable_overlap_schedule,
start_layer=start_layer,
Expand Down Expand Up @@ -2062,6 +2066,7 @@ def init_unified_mamba_swa_pools(
max_num_reqs: int,
enable_memory_saver: bool,
enable_mamba_extra_buffer: bool,
enable_mamba_extra_buffer_lazy: bool,
disable_overlap_schedule: bool,
need_sort: bool,
speculative_num_draft_tokens: Optional[int] = None,
Expand Down Expand Up @@ -2194,6 +2199,7 @@ def init_unified_mamba_swa_pools(
cache_params=mamba2_cache_params,
mamba_layer_ids=mamba_layer_ids,
enable_mamba_extra_buffer=enable_mamba_extra_buffer,
enable_mamba_extra_buffer_lazy=enable_mamba_extra_buffer_lazy,
speculative_num_draft_tokens=speculative_num_draft_tokens,
enable_overlap_schedule=not disable_overlap_schedule,
start_layer=start_layer,
Expand Down
5 changes: 3 additions & 2 deletions python/sglang/srt/session/streaming_session.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -323,7 +323,7 @@ def session_held_mamba_slots(self, active_pool_idxs: Optional[set] = None) -> in
if slot.kv.holds_mamba:
total += slot.kv.mamba_pool_idx.numel()
if slot.kv.mamba_ping_pong_track_buffer is not None:
total += slot.kv.mamba_ping_pong_track_buffer.numel()
total += int((slot.kv.mamba_ping_pong_track_buffer != -1).sum().item())
return total

def _free_slot_mamba(self, slot: SessionSlot) -> None:
Expand All @@ -335,7 +335,8 @@ def _free_slot_mamba(self, slot: SessionSlot) -> None:
mamba_allocator.free(slot.kv.mamba_pool_idx.unsqueeze(0))
slot.kv.mamba_pool_idx = None
if slot.kv.mamba_ping_pong_track_buffer is not None:
mamba_allocator.free(slot.kv.mamba_ping_pong_track_buffer)
indices = slot.kv.mamba_ping_pong_track_buffer
mamba_allocator.free(indices[indices != -1])
slot.kv.mamba_ping_pong_track_buffer = None

# -- Internal helpers (streaming body bits) --
Expand Down
1 change: 1 addition & 0 deletions python/sglang/test/unified_allocator_fixtures.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -79,6 +79,7 @@ def build_tri_pool(*, page_size=1):
extra_max_context_len=1,
max_num_reqs=4,
enable_mamba_extra_buffer=False,
enable_mamba_extra_buffer_lazy=False,
disable_overlap_schedule=True,
sliding_window_size=32,
)
Expand Down
188 changes: 188 additions & 0 deletions test/registered/unit/mem_cache/test_unified_tri_pool.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -975,12 +975,200 @@ def _factory_kwargs(self, **over):
max_num_reqs=4,
enable_memory_saver=False,
enable_mamba_extra_buffer=False,
enable_mamba_extra_buffer_lazy=False,
disable_overlap_schedule=True,
need_sort=False,
)
kw.update(over)
return kw

def test_streaming_session_lazy_checkpoint_cleanup(self):
from types import SimpleNamespace

from sglang.srt.managers.schedule_batch import ReqKvInfo
from sglang.srt.mem_cache.allocator.mamba import MambaSlotAllocator
from sglang.srt.runtime_context import get_parallel
from sglang.srt.session.streaming_session import SessionSlot, StreamingSession

for lazy_compaction in (None, False, True):
for ids in ((0, -1), (-1, 0), (-1, -1), (0, 1)):
with self.subTest(compaction=lazy_compaction, ids=ids):
kw = self._factory_kwargs(
enable_mamba_extra_buffer=True,
enable_mamba_extra_buffer_lazy=True,
lazy_compaction=lazy_compaction,
)
with get_parallel().override(attn_dcp_size=1):
pool = init_unified_mamba_swa_pools(**kw).req_to_token_pool
if lazy_compaction is None:
pool = SimpleNamespace(
mamba_allocator=MambaSlotAllocator(16, "cpu")
)
allocator = pool.mamba_allocator
available = allocator.available_size()
count = sum(i != -1 for i in ids)
live = allocator.alloc(count + 1)
track = live.new_full((2,), -1)
for position, index in enumerate(ids):
if index != -1:
track[position] = live[index + 1]
slot = SessionSlot(
kv=ReqKvInfo(
req_pool_idx=1,
mamba_pool_idx=live[0],
mamba_ping_pong_track_buffer=track,
)
)
session = StreamingSession(SimpleNamespace(req_to_token_pool=pool))
session.slots = {"test": slot}
self.assertEqual(session.session_held_mamba_slots(), count + 1)
self.assertEqual(session.session_held_mamba_slots({1}), 0)
session._free_slot_mamba(slot)
self.assertEqual(session.session_held_mamba_slots(), 0)
self.assertEqual(allocator.available_size(), available)
session._free_slot_mamba(slot)
self.assertEqual(allocator.available_size(), available)

def test_mamba_lazy_checkpoint_cleanup(self):
from types import SimpleNamespace

from sglang.srt.runtime_context import get_parallel

for tri_pool in (False, True):
for lazy in (False, True):
for keep_checkpoint in (False, True):
with self.subTest(
tri_pool=tri_pool, lazy=lazy, keep=keep_checkpoint
):
from sglang.srt.mem_cache import kv_cache_configurator as cfg

kw = self._factory_kwargs()
model = SimpleNamespace(
get_num_kv_heads=lambda tp, dcp: 2,
head_dim=4,
context_len=16,
full_attention_layer_ids=[0],
swa_attention_layer_ids=[1],
sliding_window_size=8,
)
configurator = SimpleNamespace(
mambaish_config=SimpleNamespace(
mamba2_cache_params=kw["mamba2_cache_params"],
full_attention_layer_ids=[0],
),
model_config=model,
layer_info=SimpleNamespace(
swa_attention_layer_ids=[1],
full_attention_layer_ids=[0],
start_layer=0,
end_layer=2,
),
device=_DEV,
kv_cache_dtype=torch.float16,
page_size=1,
is_draft_worker=False,
use_mla_backend=False,
is_hybrid_swa=tri_pool,
is_hybrid_swa_compress=False,
forward_stream=None,
)
# Run the production configurator AND factory. Reverting
# either top-level flag forwarding must break cleanup.
with (
get_parallel().override(attn_dcp_size=1, attn_tp_size=1),
patch.object(
cfg,
"get_exec",
return_value=SimpleNamespace(
features=SimpleNamespace(enable_memory_saver=False),
mamba=SimpleNamespace(
enable_mamba_extra_buffer=True,
enable_mamba_extra_buffer_lazy=lazy,
),
),
),
patch.object(
cfg,
"get_spec",
return_value=SimpleNamespace(
speculative_num_draft_tokens=None,
),
),
patch.object(
cfg,
"get_schedule",
return_value=SimpleNamespace(
max_mamba_cache_size=4,
disable_overlap_schedule=False,
mamba_full_memory_ratio=None,
),
),
patch.object(
cfg,
"get_disagg",
return_value=SimpleNamespace(
disaggregation_mode="null",
),
),
patch.object(
cfg, "_should_enable_lazy_compaction", return_value=True
),
):
if tri_pool:
bundle = cfg.KVCacheConfigurator._init_unified_mamba_swa_pools(
configurator,
max_num_reqs=4,
full_max_total_num_tokens=64,
swa_max_total_num_tokens=32,
)
else:
bundle = (
cfg.KVCacheConfigurator._init_unified_mamba_pools(
configurator,
max_num_reqs=4,
max_total_num_tokens=64,
)
)
pool = bundle.req_to_token_pool
self.assertEqual(pool.enable_mamba_extra_buffer_lazy, lazy)
allocator = pool.mamba_allocator
initial_available = allocator.available_size()
req = SimpleNamespace(
kv=SimpleNamespace(
req_pool_idx=0, mamba_pool_idx=allocator.alloc(1)[0]
)
)
pool._alloc_ping_pong_buffer(req)
buf = req.kv.mamba_ping_pong_track_buffer
self.assertEqual(int((buf != -1).sum()), 1 if lazy else 2)
if lazy:
# A decode boundary replaces the old checkpoint and
# leaves its ping-pong entry unallocated (-1).
replacement = allocator.alloc(1)
allocator.free(buf[:1].clone())
pool.set_mamba_ping_pong_slot(req, 0, -1)
pool.set_mamba_ping_pong_slot(req, 1, replacement[0])
else:
pool.set_mamba_ping_pong_slot(req, 1, buf[1])
retained = buf[1:2].clone()
with patch.object(
allocator, "free", wraps=allocator.free
) as free:
pool.free_mamba_cache(req, 1 if keep_checkpoint else None)
for call in free.call_args_list:
self.assertTrue(bool((call.args[0] >= 0).all()))
self.assertEqual(
allocator.available_size(),
initial_available - int(keep_checkpoint),
)
self.assertIsNone(req.kv.mamba_ping_pong_track_buffer)
if keep_checkpoint:
self.assertTrue(
bool((allocator.translate(retained) >= 0).all())
)
allocator.free(retained)
self.assertEqual(allocator.available_size(), initial_available)

def test_budget_sizing_and_boot_signature(self):
from sglang.srt.mem_cache.allocator.unified_hybrid_swa import (
UnifiedSWAAllocatorBase,
Expand Down
Loading