diff --git a/python/sglang/srt/disaggregation/decode.py b/python/sglang/srt/disaggregation/decode.py index 9e9633d4d..26df1bcb7 100644 --- a/python/sglang/srt/disaggregation/decode.py +++ b/python/sglang/srt/disaggregation/decode.py @@ -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 diff --git a/python/sglang/srt/managers/schedule_batch.py b/python/sglang/srt/managers/schedule_batch.py index f3ca299bd..9681a1d70 100644 --- a/python/sglang/srt/managers/schedule_batch.py +++ b/python/sglang/srt/managers/schedule_batch.py @@ -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'}, " diff --git a/python/sglang/srt/managers/scheduler.py b/python/sglang/srt/managers/scheduler.py index 2c835b92d..18fd50130 100644 --- a/python/sglang/srt/managers/scheduler.py +++ b/python/sglang/srt/managers/scheduler.py @@ -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) diff --git a/python/sglang/srt/managers/utils.py b/python/sglang/srt/managers/utils.py index 056671d51..9eea302b7 100644 --- a/python/sglang/srt/managers/utils.py +++ b/python/sglang/srt/managers/utils.py @@ -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" + ) diff --git a/python/sglang/srt/mem_cache/common.py b/python/sglang/srt/mem_cache/common.py index 5f7589998..f77a654ef 100644 --- a/python/sglang/srt/mem_cache/common.py +++ b/python/sglang/srt/mem_cache/common.py @@ -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() diff --git a/python/sglang/srt/mem_cache/swa_memory_pool.py b/python/sglang/srt/mem_cache/swa_memory_pool.py index e323b0e31..0faf201cb 100644 --- a/python/sglang/srt/mem_cache/swa_memory_pool.py +++ b/python/sglang/srt/mem_cache/swa_memory_pool.py @@ -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() diff --git a/python/sglang/srt/speculative/eagle_info.py b/python/sglang/srt/speculative/eagle_info.py index e72c9d725..e22eeaee4 100644 --- a/python/sglang/srt/speculative/eagle_info.py +++ b/python/sglang/srt/speculative/eagle_info.py @@ -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 diff --git a/python/sglang/srt/speculative/eagle_info_v2.py b/python/sglang/srt/speculative/eagle_info_v2.py index d9d5147e0..b542e9615 100644 --- a/python/sglang/srt/speculative/eagle_info_v2.py +++ b/python/sglang/srt/speculative/eagle_info_v2.py @@ -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 diff --git a/python/sglang/srt/speculative/eagle_worker.py b/python/sglang/srt/speculative/eagle_worker.py index 5a6cc4b03..0086e2aa7 100644 --- a/python/sglang/srt/speculative/eagle_worker.py +++ b/python/sglang/srt/speculative/eagle_worker.py @@ -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: diff --git a/python/sglang/srt/speculative/eagle_worker_v2.py b/python/sglang/srt/speculative/eagle_worker_v2.py index 92f7886da..1c90a3041 100644 --- a/python/sglang/srt/speculative/eagle_worker_v2.py +++ b/python/sglang/srt/speculative/eagle_worker_v2.py @@ -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