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

Optimize SWA memory preallocation for disaggregated decode (#24857)

Browse Source 
Co-authored-by: weireweire <weiliangl@nvidia.com>
Co-authored-by: Cheng Wan <chwan@rice.edu>
This commit is contained in:
Yuhao Yang
2026-05-13 09:09:34 +08:00
committed by GitHub
co-authored by weireweire Cheng Wan
parent 4fb40bffac
commit d6d3d0f599
3 changed files with 391 additions and 58 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/disaggregation/decode.py
+270 -52
View File
@@ -62,11 +62,13 @@ from sglang.srt.mem_cache.common import (
page_align_floor,
release_kv_cache,
)
from sglang.srt.mem_cache.deepseek_v4_memory_pool import DeepSeekV4TokenToKVPool
from sglang.srt.mem_cache.memory_pool import (
HybridReqToTokenPool,
KVCache,
ReqToTokenPool,
)
from sglang.srt.mem_cache.swa_memory_pool import SWAKVPool
from sglang.srt.observability.req_time_stats import (
set_schedule_time_batch,
set_time_batch,
@@ -175,7 +177,6 @@ class DecodeReqToTokenPool:
class HybridMambaDecodeReqToTokenPool(HybridReqToTokenPool):
def __init__(
self,
size: int,
@@ -316,13 +317,55 @@ class DecodePreallocQueue:
if self.enable_staging:
self.transfer_queue._init_staging_handler(self.kv_manager)
if self.scheduler.tp_worker.is_hybrid_swa:
# FIXME: current SWA allocation allocate full kv cache size in prefill
if (
self.scheduler.tp_worker.is_hybrid_swa
and not self._uses_swa_tail_prealloc()
):
# Fallback for SWA allocators that still allocate the SWA pool at
# full prompt length.
self.max_total_num_tokens = min(
self.max_total_num_tokens,
self.scheduler.tp_worker.model_runner.swa_max_total_num_tokens,
)
def _uses_swa_tail_prealloc(self) -> bool:
return (
isinstance(self.token_to_kv_pool, (SWAKVPool, DeepSeekV4TokenToKVPool))
and self.token_to_kv_pool_allocator.page_size > 1
and hasattr(self.token_to_kv_pool_allocator, "alloc_extend_swa_tail")
)
def _swa_tail_len(self, seq_len: int) -> int:
if not self._uses_swa_tail_prealloc() or seq_len <= 0:
return max(seq_len, 0)
window_size = self.scheduler.sliding_window_size
if window_size is None or window_size <= 0:
return seq_len
page_size = self.token_to_kv_pool_allocator.page_size
window_start = max(0, seq_len - window_size)
window_start = (window_start // page_size) * page_size
return seq_len - window_start
def _swa_retractable_len(self, req: Req) -> int:
if not self._uses_swa_tail_prealloc():
return len(req.origin_input_ids) + len(req.output_ids)
return self._swa_tail_len(len(req.origin_input_ids)) + len(req.output_ids)
def _prealloc_kv_lens(self, req: Req) -> Tuple[int, int]:
allocated_kv_len = len(req.origin_input_ids) + max(len(req.output_ids) - 1, 0)
if self._uses_swa_tail_prealloc():
return allocated_kv_len, self._swa_tail_len(allocated_kv_len)
return allocated_kv_len, allocated_kv_len
def _prealloc_required_tokens(self, req: Req) -> Tuple[int, int]:
full_len, swa_len = self._prealloc_kv_lens(req)
return (
full_len + self.num_reserved_decode_tokens,
swa_len + self.num_reserved_decode_tokens,
)
def _init_kv_manager(self) -> CommonKVManager:
kv_args_class = get_kv_class(self.transfer_backend, KVClassType.KVARGS)
kv_args = kv_args_class()
@@ -486,6 +529,18 @@ class DecodePreallocQueue:
prepare_abort(req, message, status_code=HTTPStatus.BAD_REQUEST)
self.scheduler.stream_output([req], req.return_logprob)
return True
if self._uses_swa_tail_prealloc():
_, swa_required = self._prealloc_required_tokens(req)
swa_capacity = self.token_to_kv_pool_allocator.size_swa
if swa_required > swa_capacity:
message = (
f"Request {req.rid} requires too many SWA KV tokens for "
f"decode preallocation: {swa_required} > {swa_capacity}"
)
logger.error(message)
prepare_abort(req, message, status_code=HTTPStatus.BAD_REQUEST)
self.scheduler.stream_output([req], req.return_logprob)
return True
return False
def extend(self, reqs: List[Req], is_retracted: bool = False) -> None:
@@ -501,7 +556,15 @@ class DecodePreallocQueue:
# allocate memory
resumed_reqs = []
indices_to_remove = set()
allocatable_tokens = self._allocatable_tokens(count_retracted=False)
uses_swa_tail_prealloc = self._uses_swa_tail_prealloc()
if uses_swa_tail_prealloc:
full_allocatable_tokens, swa_allocatable_tokens = (
self._swa_aware_allocatable_token_budgets(count_retracted=False)
)
else:
full_allocatable_tokens = self._allocatable_token_budgets(
count_retracted=False
)
for i, req in enumerate(self.retracted_queue):
if rids_to_check is not None and req.rid not in rids_to_check:
@@ -510,19 +573,19 @@ class DecodePreallocQueue:
if self.req_to_token_pool.available_size() <= 0:
break
required_tokens_for_request = (
len(req.origin_input_ids)
+ len(req.output_ids)
+ self.num_reserved_decode_tokens
)
if required_tokens_for_request > allocatable_tokens:
full_required, swa_required = self._prealloc_required_tokens(req)
if full_required > full_allocatable_tokens:
break
if uses_swa_tail_prealloc and swa_required > swa_allocatable_tokens:
break
resumed_reqs.append(req)
indices_to_remove.add(i)
req.is_retracted = False
self._pre_alloc(req)
allocatable_tokens -= required_tokens_for_request
full_allocatable_tokens -= full_required
if uses_swa_tail_prealloc:
swa_allocatable_tokens -= swa_required
# load from cpu, release the cpu copy
req.load_kv_cache(self.req_to_token_pool, self.token_to_kv_pool_allocator)
@@ -677,9 +740,24 @@ class DecodePreallocQueue:
len(r.origin_input_ids) + len(r.output_ids)
for r in self.scheduler.running_batch.reqs
)
allocatable_tokens = self._allocatable_tokens(
retractable_tokens=retractable_tokens, count_retracted=True
)
uses_swa_tail_prealloc = self._uses_swa_tail_prealloc()
swa_allocatable_tokens = 0
if uses_swa_tail_prealloc:
retractable_swa_tokens = sum(
self._swa_retractable_len(r) for r in self.scheduler.running_batch.reqs
)
full_allocatable_tokens, swa_allocatable_tokens = (
self._swa_aware_allocatable_token_budgets(
retractable_tokens=retractable_tokens,
retractable_swa_tokens=retractable_swa_tokens,
count_retracted=True,
)
)
else:
retractable_swa_tokens = 0
full_allocatable_tokens = self._allocatable_token_budgets(
retractable_tokens=retractable_tokens, count_retracted=True
)
# First, remove all failed requests from the queue
for i, decode_req in enumerate(self.queue):
if rids_to_check is not None and decode_req.req.rid not in rids_to_check:
@@ -746,7 +824,7 @@ class DecodePreallocQueue:
# Matching may lock previously-evictable radix pages, so refresh
# the admission budget against the post-lock pool state before we
# decide whether this request still fits.
allocatable_tokens = self._allocatable_tokens(
full_allocatable_tokens = self._allocatable_token_budgets(
retractable_tokens=retractable_tokens,
count_retracted=True,
extra_reserved_reqs=len(preallocated_reqs),
@@ -771,25 +849,47 @@ class DecodePreallocQueue:
)
- retractable_tokens,
)
> allocatable_tokens
> full_allocatable_tokens
):
if prefix_len > 0:
self.tree_cache.dec_lock_ref(decode_req.req.last_node)
break
if required_tokens_for_request > allocatable_tokens:
if required_tokens_for_request > full_allocatable_tokens:
if prefix_len > 0:
self.tree_cache.dec_lock_ref(decode_req.req.last_node)
break
if uses_swa_tail_prealloc:
_, swa_required = self._prealloc_required_tokens(decode_req.req)
_, swa_len = self._prealloc_kv_lens(decode_req.req)
max_new_tokens = min(
decode_req.req.sampling_params.max_new_tokens,
CLIP_MAX_NEW_TOKEN,
)
if (
max(
swa_required,
swa_len + max_new_tokens - retractable_swa_tokens,
)
> swa_allocatable_tokens
):
if prefix_len > 0:
self.tree_cache.dec_lock_ref(decode_req.req.last_node)
break
dst_kv_indices = self._pre_alloc(decode_req.req, prefix_indices, prefix_len)
hisparse_req_budget -= 1
# Recompute from actual pool state for the next queue entry.
# This accounts for page rounding and newly locked evictable cache.
allocatable_tokens = self._allocatable_tokens(
full_allocatable_tokens = self._allocatable_token_budgets(
retractable_tokens=retractable_tokens,
count_retracted=True,
extra_reserved_reqs=len(preallocated_reqs) + 1,
)
if uses_swa_tail_prealloc:
# SWA budget uses simple decrement (no radix cache eviction in
# the SWA pool, so page-rounding drift is negligible).
swa_allocatable_tokens -= swa_required
decode_req.req.cache_protected_len = prefix_len
if self.scheduler.enable_hisparse:
@@ -896,11 +996,8 @@ class DecodePreallocQueue:
len(decode_req.req.fill_ids) for decode_req in self.transfer_queue.queue
)
def _allocatable_tokens(
self,
retractable_tokens: Optional[int] = None,
count_retracted: bool = True,
extra_reserved_reqs: int = 0,
def _need_space_for_single_req(
self, retractable_tokens: Optional[int] = None
) -> int:
need_space_for_single_req = (
max(
@@ -915,12 +1012,69 @@ class DecodePreallocQueue:
and len(self.scheduler.running_batch.reqs) > 0
else 0
)
return need_space_for_single_req
def _active_req_count(self, extra_reserved_reqs: int = 0) -> int:
return (
len(self.scheduler.running_batch.reqs)
+ len(self.transfer_queue.queue)
+ len(self.scheduler.waiting_queue)
+ extra_reserved_reqs
)
def _active_reserved_tokens(
self, n_active: Optional[int] = None, extra_reserved_reqs: int = 0
) -> int:
if n_active is None:
n_active = self._active_req_count(extra_reserved_reqs)
return self.num_reserved_decode_tokens * n_active
def _swa_aware_allocatable_token_budgets(
self,
retractable_tokens: Optional[int] = None,
retractable_swa_tokens: Optional[int] = None,
count_retracted: bool = True,
) -> Tuple[int, int]:
n_active = self._active_req_count()
reserved_tokens = self._active_reserved_tokens(n_active)
full_allocatable_tokens = self._allocatable_token_budgets(
retractable_tokens=retractable_tokens,
count_retracted=count_retracted,
reserved_tokens=reserved_tokens,
)
return full_allocatable_tokens, self._swa_tail_allocatable_token_budget(
retractable_tokens=retractable_tokens,
retractable_swa_tokens=retractable_swa_tokens,
count_retracted=count_retracted,
n_active=n_active,
reserved_tokens=reserved_tokens,
)
def _allocatable_token_budgets(
self,
retractable_tokens: Optional[int] = None,
count_retracted: bool = True,
extra_reserved_reqs: int = 0,
reserved_tokens: Optional[int] = None,
) -> int:
need_space_for_single_req = self._need_space_for_single_req(retractable_tokens)
if reserved_tokens is None:
reserved_tokens = self._active_reserved_tokens(
extra_reserved_reqs=extra_reserved_reqs
)
if self.scheduler.enable_hisparse:
# HiSparse pre-alloc only allocates logical indices (alloc_logical_only),
# so the logical pool is the binding constraint for admission control.
available_size = (
self.token_to_kv_pool_allocator.logical_attn_allocator.available_size()
)
elif self._uses_swa_tail_prealloc():
available_size = self.token_to_kv_pool_allocator.full_available_size()
if self.scheduler.server_args.disaggregation_decode_enable_radix_cache:
available_size += self.tree_cache.evictable_size()
else:
available_size = self.token_to_kv_pool_allocator.available_size()
# Include evictable decode-radix cache entries in the budget -- they
@@ -928,16 +1082,7 @@ class DecodePreallocQueue:
if self.scheduler.server_args.disaggregation_decode_enable_radix_cache:
available_size += self.tree_cache.evictable_size()
allocatable_tokens = available_size - max(
# preserve some space for future decode
self.num_reserved_decode_tokens
* (
len(self.scheduler.running_batch.reqs)
+ len(self.transfer_queue.queue)
+ len(self.scheduler.waiting_queue)
+ extra_reserved_reqs
),
# make sure each request can finish if reach max_tokens with all other requests retracted
need_space_for_single_req,
reserved_tokens, need_space_for_single_req
)
# Note: if the last prebuilt extend just finishes, and we enter `pop_preallocated` immediately in the next iteration
@@ -951,16 +1096,75 @@ class DecodePreallocQueue:
)
if count_retracted:
allocatable_tokens -= sum(
[
len(req.origin_input_ids)
+ len(req.output_ids)
+ self.num_reserved_decode_tokens
for req in self.retracted_queue
]
)
for req in self.retracted_queue:
full_required, _ = self._prealloc_required_tokens(req)
allocatable_tokens -= full_required
return allocatable_tokens
def _swa_tail_allocatable_token_budget(
self,
retractable_tokens: Optional[int] = None,
retractable_swa_tokens: Optional[int] = None,
count_retracted: bool = True,
n_active: Optional[int] = None,
reserved_tokens: Optional[int] = None,
) -> int:
need_swa_space_for_single_req = self._need_space_for_single_req(
retractable_tokens
)
if (
retractable_swa_tokens is not None
and len(self.scheduler.running_batch.reqs) > 0
):
need_swa_space_for_single_req = max(
self._swa_tail_len(len(x.origin_input_ids))
+ min(x.sampling_params.max_new_tokens, CLIP_MAX_NEW_TOKEN)
- retractable_swa_tokens
for x in self.scheduler.running_batch.reqs
)
if n_active is None:
n_active = self._active_req_count()
if reserved_tokens is None:
reserved_tokens = self._active_reserved_tokens(n_active)
# SWA growth is bounded by the sliding window: once a req's SWA
# footprint reaches `sliding_window_size`, further decode tokens
# evict old ones and net growth is zero. The linear reservation
# `num_reserved_decode_tokens * n_active` (correct for the full
# pool) over-reserves SWA in steady state. Cap by the actual
# remaining headroom up to per-req window cap.
window_size = self.scheduler.sliding_window_size or 0
swa_total = self.token_to_kv_pool_allocator.size_swa
swa_used = swa_total - self.token_to_kv_pool_allocator.swa_available_size()
swa_growth_potential = max(0, n_active * window_size - swa_used)
swa_reserved_tokens = min(reserved_tokens, swa_growth_potential)
swa_allocatable_tokens = (
self.token_to_kv_pool_allocator.swa_available_size()
- max(swa_reserved_tokens, need_swa_space_for_single_req)
)
# Note: if the last prebuilt extend just finishes, and we enter `pop_preallocated` immediately in the next iteration
# the extend batch is not in any queue, so we need to explicitly add the tokens slots here
if (
self.scheduler.last_batch
and self.scheduler.last_batch.forward_mode.is_prebuilt()
):
prebuilt_reserved_tokens = self.num_reserved_decode_tokens * len(
self.scheduler.last_batch.reqs
)
prebuilt_n = len(self.scheduler.last_batch.reqs)
prebuilt_swa_growth = max(0, prebuilt_n * window_size - swa_used)
swa_allocatable_tokens -= min(prebuilt_reserved_tokens, prebuilt_swa_growth)
if count_retracted:
for req in self.retracted_queue:
_, swa_required = self._prealloc_required_tokens(req)
swa_allocatable_tokens -= swa_required
return swa_allocatable_tokens
def _required_alloc_tokens(self, *, fill_len: int, prefix_len: int) -> int:
page_size = self.token_to_kv_pool_allocator.page_size
if page_size == 1:
@@ -1062,16 +1266,31 @@ class DecodePreallocQueue:
if prefix_len > 0
else torch.tensor([-1], dtype=torch.int64, device=device)
)
kv_loc = self.token_to_kv_pool_allocator.alloc_extend(
prefix_lens=torch.tensor(
[prefix_len], dtype=torch.int64, device=device
),
prefix_lens_cpu=torch.tensor([prefix_len], dtype=torch.int64),
seq_lens=torch.tensor([fill_len], dtype=torch.int64, device=device),
seq_lens_cpu=torch.tensor([fill_len], dtype=torch.int64),
last_loc=last_loc,
extend_num_tokens=delta_len,
)
if self._uses_swa_tail_prealloc() and prefix_len == 0:
# Tail-only SWA allocation: only valid when prefix_len == 0.
# When prefix_len > 0 (radix cache hit), we fall back to
# alloc_extend which allocates SWA at full page count; the
# SWA budget in that case may slightly under-estimate.
kv_loc = self.token_to_kv_pool_allocator.alloc_extend_swa_tail(
prefix_lens=torch.tensor([0], dtype=torch.int64, device=device),
prefix_lens_cpu=torch.tensor([0], dtype=torch.int64),
seq_lens=torch.tensor([fill_len], dtype=torch.int64, device=device),
seq_lens_cpu=torch.tensor([fill_len], dtype=torch.int64),
last_loc=last_loc,
extend_num_tokens=fill_len,
swa_tail_len=self._swa_tail_len(fill_len),
)
else:
kv_loc = self.token_to_kv_pool_allocator.alloc_extend(
prefix_lens=torch.tensor(
[prefix_len], dtype=torch.int64, device=device
),
prefix_lens_cpu=torch.tensor([prefix_len], dtype=torch.int64),
seq_lens=torch.tensor([fill_len], dtype=torch.int64, device=device),
seq_lens_cpu=torch.tensor([fill_len], dtype=torch.int64),
last_loc=last_loc,
extend_num_tokens=delta_len,
)
assert kv_loc is not None, (
f"KV cache is full! Bug in memory estimation. "
@@ -1333,7 +1552,6 @@ class DecodeTransferQueue:
class SchedulerDisaggregationDecodeMixin:
@torch.no_grad()
def event_loop_normal_disagg_decode(self: Scheduler):
"""A normal scheduler loop for decode worker in disaggregation mode."""
python/sglang/srt/mem_cache/swa_memory_pool.py
+118 -4
View File
@@ -214,20 +214,55 @@ class SWAKVPool(BaseSWAKVPool):
src_loc_swa = self.translate_loc_from_full_to_swa(src_loc)
self.swa_kv_pool.move_kv_cache(tgt_loc_swa, src_loc_swa)
def _filter_swa_cpu_copy(self, swa_kv_cpu, row_mask: torch.Tensor):
if swa_kv_cpu is None:
return None
if row_mask is None or bool(torch.all(row_mask).item()):
return swa_kv_cpu
chunk_size = getattr(
self.swa_kv_pool, "cpu_offloading_chunk_size", len(row_mask)
)
filtered = []
for layer_chunks in swa_kv_cpu:
if len(layer_chunks) == 0:
filtered.append([])
continue
k_cpu = torch.cat([chunk[0] for chunk in layer_chunks], dim=0)
v_cpu = torch.cat([chunk[1] for chunk in layer_chunks], dim=0)
k_cpu = k_cpu[row_mask]
v_cpu = v_cpu[row_mask]
filtered_layer = []
for i in range(0, len(k_cpu), chunk_size):
filtered_layer.append(
[k_cpu[i : i + chunk_size], v_cpu[i : i + chunk_size]]
)
filtered.append(filtered_layer)
return filtered
def get_cpu_copy(self, indices, mamba_indices=None):
# For SWA, we need to copy KV cache from both full and SWA pools
# The indices are for the full pool, and we use mapping to get SWA indices
full_kv_cpu = self.full_kv_pool.get_cpu_copy(indices)
# Get SWA indices through the mapping
# Note: SWA allocation always creates 1:1 mapping, so no need to filter
swa_mask = None
if self.full_to_swa_index_mapping is not None:
swa_indices = self.full_to_swa_index_mapping[indices]
swa_kv_cpu = self.swa_kv_pool.get_cpu_copy(swa_indices)
# Slot 0 is reserved as a dummy slot. Tail-only SWA allocations leave
# the out-of-window full KV indices unmapped, so only copy mapped SWA
# tokens and keep their positions for load_cpu_copy().
swa_mask = swa_indices > 0
if torch.any(swa_mask):
swa_kv_cpu = self.swa_kv_pool.get_cpu_copy(swa_indices[swa_mask])
swa_mask = swa_mask.cpu()
else:
swa_kv_cpu = None
else:
swa_kv_cpu = None
return {"full": full_kv_cpu, "swa": swa_kv_cpu}
return {"full": full_kv_cpu, "swa": swa_kv_cpu, "swa_mask": swa_mask}
def load_cpu_copy(self, kv_cache_cpu, indices, mamba_indices=None):
# Load KV cache back from CPU to both full and SWA pools
@@ -241,6 +276,20 @@ class SWAKVPool(BaseSWAKVPool):
# Load SWA KV cache if it exists
if swa_kv_cpu is not None and self.full_to_swa_index_mapping is not None:
swa_indices = self.full_to_swa_index_mapping[indices]
new_swa_mask = swa_indices > 0
old_swa_mask = kv_cache_cpu.get("swa_mask")
if old_swa_mask is not None:
old_swa_mask = old_swa_mask.to(indices.device)
row_mask = new_swa_mask[old_swa_mask].cpu()
swa_indices = swa_indices[old_swa_mask][row_mask.to(indices.device)]
else:
row_mask = new_swa_mask.cpu()
swa_indices = swa_indices[new_swa_mask]
if swa_indices.numel() == 0:
return
swa_kv_cpu = self._filter_swa_cpu_copy(swa_kv_cpu, row_mask)
self.swa_kv_pool.load_cpu_copy(swa_kv_cpu, swa_indices)
@@ -437,6 +486,71 @@ class SWATokenToKVPoolAllocator(BaseTokenToKVPoolAllocator):
return alloc_full_indices
def alloc_extend_swa_tail(
self,
prefix_lens: torch.Tensor,
prefix_lens_cpu: torch.Tensor,
seq_lens: torch.Tensor,
seq_lens_cpu: torch.Tensor,
last_loc: torch.Tensor, # last_loc for full layers
extend_num_tokens: int,
swa_tail_len: int,
):
"""Allocate full KV for the whole extend and SWA KV only for the tail.
This is used by disaggregated decode preallocation: decode receives full
prompt KV for full-attention layers, but only the sliding-window state is
transferred for SWA layers.
"""
assert self.page_size > 1
assert len(seq_lens_cpu) == 1, "SWA tail allocation currently supports bs=1"
assert len(prefix_lens_cpu) == 1
assert 0 <= swa_tail_len <= extend_num_tokens
num_full_pages = get_num_new_pages(
seq_lens=seq_lens_cpu, page_size=self.page_size, prefix_lens=prefix_lens_cpu
)
num_swa_pages = (swa_tail_len + self.page_size - 1) // self.page_size
if num_full_pages > self.full_attn_allocator.available_size() // self.page_size:
return None
if num_swa_pages > self.swa_attn_allocator.available_size() // self.page_size:
return None
alloc_full_indices = self.full_attn_allocator.alloc_extend(
prefix_lens,
prefix_lens_cpu,
seq_lens,
seq_lens_cpu,
last_loc,
extend_num_tokens,
)
assert alloc_full_indices is not None
if swa_tail_len == 0:
return alloc_full_indices
device = self.device
swa_prefix_lens = torch.zeros((1,), dtype=torch.int64, device=device)
swa_prefix_lens_cpu = torch.zeros((1,), dtype=torch.int64)
swa_seq_lens = torch.tensor([swa_tail_len], dtype=torch.int64, device=device)
swa_seq_lens_cpu = torch.tensor([swa_tail_len], dtype=torch.int64)
swa_last_loc = torch.tensor([-1], dtype=torch.int64, device=device)
alloc_swa_indices = self.swa_attn_allocator.alloc_extend(
swa_prefix_lens,
swa_prefix_lens_cpu,
swa_seq_lens,
swa_seq_lens_cpu,
swa_last_loc,
swa_tail_len,
)
assert alloc_swa_indices is not None
self.full_to_swa_index_mapping[alloc_full_indices[-swa_tail_len:]] = (
alloc_swa_indices
)
return alloc_full_indices
def alloc_decode(
self,
seq_lens: torch.Tensor,