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

Fix SWA admission livelock on cached-prefix resumes (#32379)

Browse Source
This commit is contained in:
Ke Bao
2026-07-25 22:37:35 +08:00
committed by GitHub
parent e943e609dc
commit 69a3c54c70
2 changed files with 155 additions and 27 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/managers/schedule_policy.py
+62 -17
View File
@@ -620,7 +620,7 @@ class PrefillAdder:
return available_and_evictable - self.cur_rem_token_offset
def _swa_budget_for_req(
self, extend_input_len: int, swa_host_hit_length: int = 0
self, extend_input_len: int, max_new_tokens: int, swa_host_hit_length: int = 0
) -> int:
"""SWA pool budget per request. Only valid when is_hybrid_swa is True.
@@ -629,26 +629,52 @@ class PrefillAdder:
+ chunk N+1 (new allocation)
Since chunk N and locked tokens are already excluded from
swa_available + swa_evictable, the budget only needs to cover the
chunk N+1 allocation. We floor at sliding_window_size to reserve
room for the decode phase.
chunk N+1 allocation plus decode headroom:
budget = max(alloc - window, 0) + min(extend + max_new_tokens, window) + page
where alloc = min(extend, rem_chunk); the min() cap keeps the two terms
from double-counting extend, so budget <= extend + max_new_tokens + page.
"""
if self.rem_chunk_tokens is not None:
alloc = min(extend_input_len, self.rem_chunk_tokens)
else:
alloc = extend_input_len
window = self.tree_cache.sliding_window_size
return max(alloc - window, 0) + self._swa_reserved_tokens(swa_host_hit_length)
return max(alloc - window, 0) + self._swa_reserved_tokens(
extend_input_len, max_new_tokens, swa_host_hit_length
)
def _swa_reserved_tokens(self, swa_host_hit_length: int = 0) -> int:
"""SWA tokens a request needs regardless of extend length: the sliding
window (decode headroom) + allocator page slack + the load-back window
charge. Shared floor of _swa_budget_for_req and _swa_chunk_cap."""
reserved = self.tree_cache.sliding_window_size + self.page_size
def _swa_reserved_tokens(
self, extend_input_len: int, max_new_tokens: int, swa_host_hit_length: int = 0
) -> int:
"""SWA slots a request adds to its own sliding window + page slack + the
load-back charge. Shared floor of _swa_budget_for_req and _swa_chunk_cap.
The headroom is min(extend + decode, window), not a constant window: a
request contributes only extend + decode fresh tokens to its window and
a cached SWA prefix funds the rest. Charging a full window double-counted
a short cached-prefix resume and livelocked admission at a ~2-window
pool; keeping extend in the min() holds the reservation >= the prefill
allocation so an admitted request cannot OOM."""
window = self.tree_cache.sliding_window_size
headroom = min(extend_input_len + max_new_tokens, window)
reserved = headroom + self.page_size
if swa_host_hit_length > 0:
reserved += self.ceil_paged_tokens(swa_host_hit_length)
return reserved
def _swa_chunk_cap(self, swa_host_hit_length: int = 0) -> int:
def _swa_new_tokens(self, req: Req) -> int:
"""Tokens a request may still decode, for SWA headroom sizing. Mirrors
the remaining-then-clip order of add_one_req's max_new: clip-then-subtract
would zero out a request that has already generated >= CLIP tokens but
still has a long decode ahead, under-reserving its window."""
return min(
max(req.sampling_params.max_new_tokens - len(req.output_ids), 0),
CLIP_MAX_NEW_TOKENS,
)
def _swa_chunk_cap(self, max_new_tokens: int, swa_host_hit_length: int = 0) -> int:
"""Largest page-aligned extend chunk the SWA pool can admit right now,
keeping a sliding window of headroom below rem_swa_tokens; 0 if not
even one page fits. Only valid when is_hybrid_swa is True.
@@ -659,7 +685,11 @@ class PrefillAdder:
forever (head-of-line livelock). Shrinking is sound because past a
chunk boundary only the sliding window stays locked — the rest turns
evictable — so each pass's transient footprint fits the pool."""
cap = int(self.rem_swa_tokens) - self._swa_reserved_tokens(swa_host_hit_length)
# extend_input_len=0: this solves for the extend chunk itself, so the
# reserved headroom is the post-chunk decode window only.
cap = int(self.rem_swa_tokens) - self._swa_reserved_tokens(
0, max_new_tokens, swa_host_hit_length
)
if cap <= 0:
return 0
return cap // self.page_size * self.page_size
@@ -735,7 +765,9 @@ class PrefillAdder:
self.rem_input_tokens -= extend_input_len
if self.is_hybrid_swa:
self.rem_swa_token_offset += self._swa_budget_for_req(extend_input_len)
self.rem_swa_token_offset += self._swa_budget_for_req(
extend_input_len, max_new_tokens
)
if self.dllm_config is not None:
self.rem_dllm_tokens -= extend_input_len
@@ -906,7 +938,12 @@ class PrefillAdder:
if paged_input > min(self.cur_rem_tokens, self.rem_total_tokens):
return AddReqResult.NO_TOKEN
if self.is_hybrid_swa:
if self._swa_budget_for_req(cand_extend_input_len) > self.rem_swa_tokens:
if (
self._swa_budget_for_req(
cand_extend_input_len, self._swa_new_tokens(req)
)
> self.rem_swa_tokens
):
return AddReqResult.NO_TOKEN
def add_req_state(r, insert_sort=False):
@@ -1053,10 +1090,14 @@ class PrefillAdder:
# driven only by the freshly-prefilled tail (the loaded window is
# charged separately via swa_host_hit_length).
swa_needed = self._swa_budget_for_req(
real_input_tokens, swa_host_hit_length=req.swa_host_hit_length
real_input_tokens,
self._swa_new_tokens(req),
swa_host_hit_length=req.swa_host_hit_length,
)
if swa_needed >= self.rem_swa_tokens:
swa_cap = self._swa_chunk_cap(req.swa_host_hit_length)
swa_cap = self._swa_chunk_cap(
self._swa_new_tokens(req), req.swa_host_hit_length
)
if self.rem_chunk_tokens is None or swa_cap <= 0:
return AddReqResult.NO_TOKEN
chunk_tokens_limit = min(self.rem_chunk_tokens, swa_cap)
@@ -1079,10 +1120,14 @@ class PrefillAdder:
if self.is_hybrid_swa:
# self.rem_swa_tokens may decrease after the lock acquisition
swa_needed = self._swa_budget_for_req(
real_input_tokens, swa_host_hit_length=req.swa_host_hit_length
real_input_tokens,
self._swa_new_tokens(req),
swa_host_hit_length=req.swa_host_hit_length,
)
if swa_needed >= self.rem_swa_tokens:
swa_cap = self._swa_chunk_cap(req.swa_host_hit_length)
swa_cap = self._swa_chunk_cap(
self._swa_new_tokens(req), req.swa_host_hit_length
)
if self.rem_chunk_tokens is None or swa_cap <= 0:
return AddReqResult.NO_TOKEN
chunk_tokens_limit = min(self.rem_chunk_tokens, swa_cap)
test/registered/unit/managers/test_prefill_adder.py
+93 -10
View File
@@ -1,6 +1,6 @@
import unittest
from types import SimpleNamespace
from unittest.mock import MagicMock
from unittest.mock import MagicMock, patch
from sglang.srt.managers.schedule_batch import Req
from sglang.srt.managers.schedule_policy import AddReqResult, PrefillAdder
@@ -532,16 +532,26 @@ class TestPrefillAdder(CustomTestCase):
self.assertEqual(len(adder.can_run_list), 0)
def test_swa_budget_for_req(self):
# budget = max(alloc - window, 0) + min(extend + max_new, window) + page,
# where alloc = min(extend, rem_chunk). The decode headroom is the SWA the
# request adds to its own sliding window (extend + decode), capped at the
# window -- a cached prefix funds the rest -- rather than a constant
# window, which over-charged short requests into an admission livelock.
cases = [
# (extend, rem_chunk, window, page, expected, label)
(64, None, 128, 16, 128 + 16, "no_cap_floor_active"),
(200, None, 256, 32, 256 + 32, "no_cap_floor_active_other_dims"),
(300, None, 128, 16, 300 + 16, "no_cap_floor_inactive"),
(200, 50, 64, 8, 64 + 8, "cap_binds_then_floor"),
(300, 500, 64, 64, 300 + 64, "cap_does_not_bind"),
(0, None, 128, 16, 128 + 16, "extend_zero_floor_only"),
# (extend, max_new, rem_chunk, window, page, expected, label)
(64, 512, None, 128, 16, 128 + 16, "long_decode_hits_window_floor"),
(64, 32, None, 128, 16, 96 + 16, "short_req_reserves_below_window"),
(10, 20, None, 512, 8, 30 + 8, "short_resume_tiny_budget"),
(300, 512, None, 128, 16, 300 + 16, "big_extend_over_window"),
(200, 512, 50, 64, 8, 64 + 8, "chunk_capped_alloc_hits_floor"),
# Multi-chunk: alloc = rem_chunk (1024) caps a huge extend, and the
# chunk itself exceeds the window, so term1 = alloc - window is driven
# by the chunk, not the full extend -> budget = chunk + page.
(2000, 256, 1024, 512, 16, 1024 + 16, "multichunk_alloc_over_window"),
(0, 512, None, 128, 16, 128 + 16, "extend_zero_long_decode"),
(0, 40, None, 128, 16, 40 + 16, "extend_zero_short_decode"),
]
for extend, rem_chunk, window, page, expected, label in cases:
for extend, max_new, rem_chunk, window, page, expected, label in cases:
with self.subTest(label=label):
self.mock_tree_cache.sliding_window_size = window
adder = self.create_adder(
@@ -549,7 +559,80 @@ class TestPrefillAdder(CustomTestCase):
page_size=page,
rem_chunk_tokens=rem_chunk,
)
self.assertEqual(adder._swa_budget_for_req(extend), expected)
self.assertEqual(adder._swa_budget_for_req(extend, max_new), expected)
def test_swa_admission_admits_short_cached_resume_at_two_window_pool(self):
# Livelock regression (real incident). At an SWA pool ~= 2 sliding
# windows, a cached-prefix resume matches >= 1 window (locked, excluded
# from rem_swa) and has only a short uncached tail + a little decode
# left. The pre-fix constant-window reservation charged a second full
# window, so the admission gate (swa_needed >= rem_swa_tokens) rejected
# it every scheduler iteration while LPM kept it at the queue head --
# 100% scheduler CPU, idle GPU. Capping the reservation at
# min(extend + decode, window) admits it.
WINDOW, PAGE, REM_SWA = 128, 8, 100
PREFIX, EXTEND = 200, 16 # cached prefix > window; short uncached tail
self.mock_token_allocator.swa_available_size.return_value = REM_SWA
self.mock_token_allocator.full_available_size.return_value = 100_000
self.mock_token_allocator.available_size.return_value = 100_000
self.mock_tree_cache.sliding_window_size = WINDOW
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
req = self.create_mock_req(
"resume", priority=0, max_new_tokens=40, output_len=10
)
req.prefix_indices = list(range(PREFIX))
req.full_untruncated_fill_ids = list(range(PREFIX + EXTEND))
req.host_hit_length = 0
req.swa_host_hit_length = 0
req.last_node = MagicMock()
req.set_extend_range = MagicMock(
side_effect=lambda start, end: setattr(
req, "extend_range", Range(start, end)
)
)
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):
self.assertIs(
adder.add_one_req(
req, has_chunked_req=False, truncation_align_size=None
),
AddReqResult.NO_TOKEN,
)
self.assertEqual(len(adder.can_run_list), 0)
# Fix: min(extend + decode, window) reservation admits it.
adder.add_one_req(req, has_chunked_req=False, truncation_align_size=None)
self.assertIn(req, adder.can_run_list)
def test_swa_new_tokens_clamps_remaining_not_total(self):
# Remaining decode headroom must be min(max_new - generated, CLIP)
# (subtract-then-clip). The reversed order (clip-then-subtract) zeroes
# out a request that has already generated >= CLIP tokens but still has a
# long decode ahead, under-reserving its SWA window -> OOM risk on resume
# of a long-generation request.
from sglang.srt.managers.schedule_policy import CLIP_MAX_NEW_TOKENS as CLIP
adder = self.create_adder(self.create_running_batch())
cases = [
# (max_new, generated, expected, label)
(100, 10, 90, "below_clip_normal"),
(40, 100, 0, "already_finished"),
# clip-then-subtract would give 3996; subtract-then-clip caps at CLIP.
(CLIP + 6000, 100, CLIP, "long_gen_small_output_caps_at_clip"),
# clip-then-subtract would give 0; the long decode still needs CLIP.
(CLIP + 6000, CLIP + 100, CLIP, "long_gen_output_over_clip"),
]
for max_new, generated, expected, label in cases:
with self.subTest(label=label):
req = self.create_mock_req(
label, priority=0, max_new_tokens=max_new, output_len=generated
)
self.assertEqual(adder._swa_new_tokens(req), expected)
def test_delayer_not_consulted_when_kv_budget_rejects(self):
"""A rank whose first candidate fails the KV-budget gate must NOT