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

python/sglang/srt/disaggregation/decode.py

@@ -616,15 +616,7 @@ class DecodePreallocQueue:
             and len(self.scheduler.running_batch.reqs) > 0
             else 0
         )
-
+        available_size = self.token_to_kv_pool_allocator.available_size()
-        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()
-
         allocatable_tokens = available_size - max(
             # preserve some space for future decode
             self.num_reserved_decode_tokens

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():
+        if self.spec_algorithm.is_none():
-            # A loose bound that err towards safety
+            new_pages = sum(1 for r in requests if r.kv_committed_len % page_size == 0)
-            server_args = get_global_server_args()
+            return new_pages * page_size
-            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
-            )
 
-        # In the decoding phase, the length of a request's KV cache should be
+        server_args = get_global_server_args()
-        # the total length of the request minus 1
+        len_per_topk = server_args.speculative_num_steps or 1
-        return (
+        spec_topk = server_args.speculative_eagle_topk or 1
-            sum(1 for req in requests if req.seqlen % page_size == 0)
+        spec_tokens = server_args.speculative_num_draft_tokens
-            if self.enable_overlap
-            else sum(1 for req in requests if (req.seqlen - 1) % page_size == 0)
-        )
 
-    def check_decode_mem(
+        if page_size > 1 and spec_topk > 1:
-        self, buf_multiplier=1, selected_indices: Optional[List[int]] = None
+            # last partial page and ceil alignment
-    ):
+            len_per_topk = ceil_align(len_per_topk + page_size, page_size)
-        num_tokens = (
+            spec_tokens = ceil_align(spec_tokens, page_size)
-            self.new_page_count_next_decode(selected_indices)
+        elif page_size > 1:
-            * buf_multiplier
+            # only page alignment
-            * self.token_to_kv_pool_allocator.page_size
+            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,
+        self, server_args: ServerArgs
-        server_args: ServerArgs,
-        buf_multiplier: int = 1,
     ) -> 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(
+            not self.check_decode_mem(selected_indices=sorted_indices)
-                selected_indices=sorted_indices, buf_multiplier=buf_multiplier
-            )
         ):
             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'}, "

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 (
+        if (kv_full_retract_flag := not batch.check_decode_mem()) or (
-            kv_full_retract_flag := not batch.check_decode_mem(
+            TEST_RETRACT and self.forward_ct % TEST_RETRACT_INTERVAL == 0
-                self.decode_mem_cache_buf_multiplier
+        ):
-            )
+            old_available_tokens = self.token_to_kv_pool_allocator.available_size()
-        ) 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()
             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 = self.token_to_kv_pool_allocator.available_size()
-                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_token_gained = new_available_tokens - old_available_tokens
 
             self.num_retracted_reqs = len(retracted_reqs)

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"
+        )

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 isinstance(allocator, SWATokenToKVPoolAllocator):
-    if hasattr(allocator, "full_available_size"):
         # Hybrid allocator
         full_available_size = allocator.full_available_size()
         swa_available_size = allocator.swa_available_size()

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.
+        return min(
-        raise NotImplementedError()
+            self.full_attn_allocator.available_size(),
+            self.swa_attn_allocator.available_size(),
+        )
 
     def full_available_size(self):
         return self.full_attn_allocator.available_size()

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

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

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:

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