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

Introduce global alloc_len_per_decode & clean check decode memory (#15115)

Browse Source
This commit is contained in:
Liangsheng Yin
2026-01-26 10:26:20 -08:00
committed by GitHub
parent 8643fb2f52
commit 85d077f44d
10 changed files with 76 additions and 101 deletions
Download Patch File Download Diff File Expand all files Collapse all files

10
python/sglang/srt/disaggregation/decode.py
View File

@@ -616,15 +616,7 @@ class DecodePreallocQueue:
and len(self.scheduler.running_batch.reqs) > 0
else 0
)
if self.scheduler.model_config.is_hybrid_swa:
available_size = min(
self.token_to_kv_pool_allocator.full_available_size(),
self.token_to_kv_pool_allocator.swa_available_size(),
)
else:
available_size = self.token_to_kv_pool_allocator.available_size()
available_size = self.token_to_kv_pool_allocator.available_size()
allocatable_tokens = available_size - max(
# preserve some space for future decode
self.num_reserved_decode_tokens

75
python/sglang/srt/managers/schedule_batch.py
View File

@@ -4,6 +4,7 @@ import enum
from sglang.srt.dllm.config import DllmConfig
from sglang.srt.model_executor.forward_batch_info import ForwardBatch
from sglang.srt.utils.common import ceil_align
# Copyright 2023-2024 SGLang Team
# Licensed under the Apache License, Version 2.0 (the "License");
@@ -1794,46 +1795,43 @@ class ScheduleBatch(ScheduleBatchDisaggregationDecodeMixin):
self.extend_logprob_start_lens.extend([0] * running_bs)
self.is_prefill_only = False
def new_page_count_next_decode(self, selected_indices: Optional[List[int]] = None):
def new_tokens_required_next_decode(
self, selected_indices: Optional[List[int]] = None
):
page_size = self.token_to_kv_pool_allocator.page_size
requests = (
self.reqs
if selected_indices is None
else [self.reqs[i] for i in selected_indices]
)
if page_size == 1:
return len(requests)
if not self.spec_algorithm.is_none():
# A loose bound that err towards safety
server_args = get_global_server_args()
thresh = server_args.speculative_num_draft_tokens + (
(server_args.speculative_eagle_topk or 1)
* (server_args.speculative_num_steps or 1)
)
return sum(
1 for req in requests if ((req.seqlen + thresh) % page_size) <= thresh
)
if self.spec_algorithm.is_none():
new_pages = sum(1 for r in requests if r.kv_committed_len % page_size == 0)
return new_pages * page_size
# In the decoding phase, the length of a request's KV cache should be
# the total length of the request minus 1
return (
sum(1 for req in requests if req.seqlen % page_size == 0)
if self.enable_overlap
else sum(1 for req in requests if (req.seqlen - 1) % page_size == 0)
)
server_args = get_global_server_args()
len_per_topk = server_args.speculative_num_steps or 1
spec_topk = server_args.speculative_eagle_topk or 1
spec_tokens = server_args.speculative_num_draft_tokens
def check_decode_mem(
self, buf_multiplier=1, selected_indices: Optional[List[int]] = None
):
num_tokens = (
self.new_page_count_next_decode(selected_indices)
* buf_multiplier
* self.token_to_kv_pool_allocator.page_size
)
if page_size > 1 and spec_topk > 1:
# last partial page and ceil alignment
len_per_topk = ceil_align(len_per_topk + page_size, page_size)
spec_tokens = ceil_align(spec_tokens, page_size)
elif page_size > 1:
# only page alignment
len_per_topk = ceil_align(len_per_topk, page_size)
spec_tokens = ceil_align(spec_tokens, page_size)
num_tokens = max(len_per_topk * spec_topk, spec_tokens) * len(requests)
# v2 eagle has over-allocation
return num_tokens * (1 + self.is_spec_v2)
def check_decode_mem(self, selected_indices: Optional[List[int]] = None):
num_tokens = self.new_tokens_required_next_decode(selected_indices)
evict_from_tree_cache(self.tree_cache, num_tokens)
return self._is_available_size_sufficient(num_tokens)
return self.token_to_kv_pool_allocator.available_size() >= num_tokens
def retract_all(self, server_args: ServerArgs):
retracted_reqs = self.reqs
@@ -1844,9 +1842,7 @@ class ScheduleBatch(ScheduleBatchDisaggregationDecodeMixin):
return retracted_reqs
def retract_decode(
self,
server_args: ServerArgs,
buf_multiplier: int = 1,
self, server_args: ServerArgs
) -> Tuple[List[Req], float, List[Req]]:
"""Retract the decoding requests when there is not enough memory."""
sorted_indices = list(range(len(self.reqs)))
@@ -1868,9 +1864,7 @@ class ScheduleBatch(ScheduleBatchDisaggregationDecodeMixin):
retracted_reqs = []
first_iter = True
while first_iter or (
not self.check_decode_mem(
selected_indices=sorted_indices, buf_multiplier=buf_multiplier
)
not self.check_decode_mem(selected_indices=sorted_indices)
):
if len(sorted_indices) == 1:
# Always keep at least one request
@@ -1884,7 +1878,7 @@ class ScheduleBatch(ScheduleBatchDisaggregationDecodeMixin):
self.release_req(idx, len(sorted_indices), server_args)
if len(sorted_indices) <= 1 and not self.check_decode_mem(
selected_indices=sorted_indices, buf_multiplier=buf_multiplier
selected_indices=sorted_indices
):
# Retracting loops ends and still not enough memory
raise ValueError(
@@ -2320,15 +2314,6 @@ class ScheduleBatch(ScheduleBatchDisaggregationDecodeMixin):
self.token_to_kv_pool_allocator.free_swa(free_slots)
req.swa_evicted_seqlen = new_swa_evicted_seqlen
def _is_available_size_sufficient(self, num_tokens: int) -> bool:
if self.is_hybrid_swa:
return (
self.token_to_kv_pool_allocator.full_available_size() >= num_tokens
and self.token_to_kv_pool_allocator.swa_available_size() >= num_tokens
)
else:
return self.token_to_kv_pool_allocator.available_size() >= num_tokens
def __str__(self):
return (
f"ScheduleBatch(forward_mode={self.forward_mode.name if self.forward_mode else 'None'}, "

38
python/sglang/srt/managers/scheduler.py
View File

@@ -704,18 +704,6 @@ class Scheduler(
else:
self.decode_offload_manager = None
self.decode_mem_cache_buf_multiplier = (
1
if self.spec_algorithm.is_none()
else (
server_args.speculative_num_draft_tokens
+ (
(server_args.speculative_eagle_topk or 1)
* (server_args.speculative_num_steps or 1)
)
)
)
embedding_cache_size = envs.SGLANG_VLM_CACHE_SIZE_MB.get()
init_mm_embedding_cache(embedding_cache_size * 1024 * 1024)
@@ -2162,29 +2150,15 @@ class Scheduler(
return batch
# Check if decode out of memory
if (
kv_full_retract_flag := not batch.check_decode_mem(
self.decode_mem_cache_buf_multiplier
)
) or (TEST_RETRACT and self.forward_ct % TEST_RETRACT_INTERVAL == 0):
if self.is_hybrid_swa:
old_available_tokens = min(
self.token_to_kv_pool_allocator.full_available_size(),
self.token_to_kv_pool_allocator.swa_available_size(),
)
else:
old_available_tokens = self.token_to_kv_pool_allocator.available_size()
if (kv_full_retract_flag := not batch.check_decode_mem()) or (
TEST_RETRACT and self.forward_ct % TEST_RETRACT_INTERVAL == 0
):
old_available_tokens = self.token_to_kv_pool_allocator.available_size()
old_ratio = self.new_token_ratio
retracted_reqs, new_token_ratio, reqs_to_abort = batch.retract_decode(
self.server_args, self.decode_mem_cache_buf_multiplier
self.server_args
)
if self.is_hybrid_swa:
new_available_tokens = min(
self.token_to_kv_pool_allocator.full_available_size(),
self.token_to_kv_pool_allocator.swa_available_size(),
)
else:
new_available_tokens = self.token_to_kv_pool_allocator.available_size()
new_available_tokens = self.token_to_kv_pool_allocator.available_size()
new_token_gained = new_available_tokens - old_available_tokens
self.num_retracted_reqs = len(retracted_reqs)

28
python/sglang/srt/managers/utils.py
View File

@@ -11,6 +11,7 @@ from sglang.srt.layers.logits_processor import LogitsProcessorOutput
from sglang.srt.managers.overlap_utils import FutureIndices
from sglang.srt.managers.schedule_batch import Req
from sglang.srt.model_executor.forward_batch_info import PPProxyTensors
from sglang.srt.server_args import ServerArgs
if TYPE_CHECKING:
from sglang.srt.managers.scheduler import GenerationBatchResult
@@ -176,3 +177,30 @@ def get_logprob_from_pp_outputs(
]
return logits_output, extend_input_len_per_req, extend_logprob_start_len_per_req
def get_alloc_len_per_decode(server_args: Optional[ServerArgs] = None) -> int:
if server_args is None:
from sglang.srt.server_args import get_global_server_args
server_args = get_global_server_args()
if server_args.speculative_algorithm is None:
return 1
# Spec v1:
# 1) alloc topk * num_steps when draft decoding and then restore the allocation
# 2) alloc num_draft_tokens when verifying the drafts
# Sepc v2: allocate max(topk * num_steps, num_draft_tokens)
spec_steps = server_args.speculative_num_steps or 1
spec_topk = server_args.speculative_eagle_topk or 1
spec_tokens = server_args.speculative_num_draft_tokens
page_size = server_args.page_size
if page_size == 1 or spec_topk == 1:
return max(spec_steps * spec_topk, spec_tokens)
else:
raise NotImplementedError(
"get_alloc_len_per_decode not implemented for page_size > 1 and spec_topk > 1"
)

3
python/sglang/srt/mem_cache/common.py
View File

@@ -235,8 +235,7 @@ def evict_from_tree_cache(tree_cache: BasePrefixCache | None, num_tokens: int):
allocator = tree_cache.token_to_kv_pool_allocator
# Check if this is a hybrid allocator
if hasattr(allocator, "full_available_size"):
if isinstance(allocator, SWATokenToKVPoolAllocator):
# Hybrid allocator
full_available_size = allocator.full_available_size()
swa_available_size = allocator.swa_available_size()

6
python/sglang/srt/mem_cache/swa_memory_pool.py
View File

@@ -297,8 +297,10 @@ class SWATokenToKVPoolAllocator(BaseTokenToKVPoolAllocator):
self._kvcache.register_mapping(weakref.proxy(self.full_to_swa_index_mapping))
def available_size(self):
# Note: use full_available_size() and swa_available_size() instead.
raise NotImplementedError()
return min(
self.full_attn_allocator.available_size(),
self.swa_attn_allocator.available_size(),
)
def full_available_size(self):
return self.full_attn_allocator.available_size()

5
python/sglang/srt/speculative/eagle_info.py
View File

@@ -1,7 +1,7 @@
import logging
from copy import copy
from dataclasses import dataclass
from typing import ClassVar, List, Optional, Tuple
from typing import List, Optional, Tuple
import torch
import torch.nn.functional as F
@@ -614,9 +614,6 @@ class EagleVerifyInput(SpecInput, EagleVerifyInputV2Mixin):
@dataclass
class EagleDraftInput(SpecInput, EagleDraftInputV2Mixin):
# Constant: alloc length per decode step
ALLOC_LEN_PER_DECODE: ClassVar[int] = None
# The inputs for decode
# shape: (b, topk)
topk_p: torch.Tensor = None

4
python/sglang/srt/speculative/eagle_info_v2.py
View File

@@ -10,6 +10,7 @@ import triton.language as tl
from sglang.srt.layers.logits_processor import LogitsProcessorOutput
from sglang.srt.managers.schedule_batch import ModelWorkerBatch, ScheduleBatch
from sglang.srt.managers.utils import get_alloc_len_per_decode
from sglang.srt.mem_cache.common import (
alloc_paged_token_slots_extend,
alloc_token_slots,
@@ -92,9 +93,10 @@ class EagleDraftInputV2Mixin:
cur_kv_lens_cpu = []
nxt_kv_lens_cpu = []
num_needed_tokens = 0
alloc_len_per_decode = get_alloc_len_per_decode()
for r in batch.reqs:
# Over-allocation happens here
x = r.kv_committed_len + 2 * self.ALLOC_LEN_PER_DECODE - r.kv_allocated_len
x = r.kv_committed_len + 2 * alloc_len_per_decode - r.kv_allocated_len
cur_kv_lens_cpu.append(r.kv_allocated_len)
nxt_kv_lens_cpu.append(r.kv_allocated_len + x)
num_needed_tokens += x

3
python/sglang/srt/speculative/eagle_worker.py
View File

@@ -394,10 +394,11 @@ class EAGLEWorker(TpModelWorker):
# [ topk 0 ] [ topk 1 ]
# [iter=0, iter=1, iter=2] [iter=0, iter=1, iter=2]
if self.page_size == 1:
alloc_len_per_decode = self.speculative_num_steps * self.topk
# TODO: We only need self.speculative_num_steps - 1 * topk cache loc
out_cache_loc, token_to_kv_pool_state_backup = alloc_token_slots(
batch.tree_cache,
num_seqs * self.speculative_num_steps * self.topk,
num_seqs * alloc_len_per_decode,
backup_state=True,
)
else:

5
python/sglang/srt/speculative/eagle_worker_v2.py
View File

@@ -107,11 +107,6 @@ class EagleDraftWorker(BaseDraftWorker):
server_args.speculative_algorithm
)
# Set constant
EagleDraftInput.ALLOC_LEN_PER_DECODE = max(
self.speculative_num_steps * self.topk, self.speculative_num_draft_tokens
)
# Do not capture cuda graph in `TpModelWorker` init,
# will capture later with init_cuda_graphs()
backup_disable_cuda_graph = server_args.disable_cuda_graph