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

[Core] Replace server_args mutation hack with explicit MemoryPoolConfig for draft worker init (#20183)

Browse Source
This commit is contained in:
Liangsheng Yin
2026-03-09 11:45:54 -07:00
committed by GitHub
parent ecca8c553d
commit ffb4b6f4c1
9 changed files with 93 additions and 49 deletions
Download Patch File Download Diff File Expand all files Collapse all files

5
python/sglang/srt/managers/tp_worker.py
View File

@@ -54,6 +54,7 @@ from sglang.srt.weight_sync.tensor_bucket import FlattenedTensorBucket
if TYPE_CHECKING:
from sglang.srt.managers.cache_controller import LayerDoneCounter
from sglang.srt.model_executor.model_runner import ModelRunner
from sglang.srt.model_executor.model_runner_kv_cache_mixin import MemoryPoolConfig
logger = logging.getLogger(__name__)
@@ -231,6 +232,7 @@ class TpModelWorker(BaseTpWorker):
is_draft_worker: bool = False,
req_to_token_pool: Optional[ReqToTokenPool] = None,
token_to_kv_pool_allocator: Optional[BaseTokenToKVPoolAllocator] = None,
memory_pool_config: Optional[MemoryPoolConfig] = None,
is_multi_layer_eagle: bool = False,
):
# Parse args
@@ -248,6 +250,7 @@ class TpModelWorker(BaseTpWorker):
self.is_multi_layer_eagle = is_multi_layer_eagle
self.req_to_token_pool = req_to_token_pool
self.token_to_kv_pool_allocator = token_to_kv_pool_allocator
self.memory_pool_config = memory_pool_config
self.attn_cp_rank = attn_cp_rank
self.moe_dp_rank = moe_dp_rank
@@ -354,6 +357,7 @@ class TpModelWorker(BaseTpWorker):
is_draft_worker=self.is_draft_worker,
req_to_token_pool=self.req_to_token_pool,
token_to_kv_pool_allocator=self.token_to_kv_pool_allocator,
memory_pool_config=self.memory_pool_config,
draft_model_idx=0 if self.is_multi_layer_eagle else None,
)
@@ -379,6 +383,7 @@ class TpModelWorker(BaseTpWorker):
is_draft_worker=self.is_draft_worker,
req_to_token_pool=self.req_to_token_pool,
token_to_kv_pool_allocator=self.token_to_kv_pool_allocator,
memory_pool_config=self.memory_pool_config,
draft_model_idx=i,
)
)

5
python/sglang/srt/model_executor/model_runner.py
View File

@@ -13,6 +13,8 @@
# ==============================================================================
"""ModelRunner runs the forward passes of the models."""
from __future__ import annotations
import datetime
import gc
import inspect
@@ -132,6 +134,7 @@ from sglang.srt.model_executor.forward_batch_info import (
)
from sglang.srt.model_executor.hook_manager import register_forward_hooks
from sglang.srt.model_executor.model_runner_kv_cache_mixin import (
MemoryPoolConfig,
ModelRunnerKVCacheMixin,
)
from sglang.srt.model_executor.piecewise_cuda_graph_runner import (
@@ -301,6 +304,7 @@ class ModelRunner(ModelRunnerKVCacheMixin):
is_draft_worker: bool = False,
req_to_token_pool: Optional[ReqToTokenPool] = None,
token_to_kv_pool_allocator: Optional[BaseTokenToKVPoolAllocator] = None,
memory_pool_config: Optional[MemoryPoolConfig] = None,
draft_model_idx: Optional[int] = None,
):
# Parse args
@@ -322,6 +326,7 @@ class ModelRunner(ModelRunnerKVCacheMixin):
self.dist_port = nccl_port
self.server_args = server_args
self.is_draft_worker = is_draft_worker
self.memory_pool_config = memory_pool_config
self.is_generation = model_config.is_generation
self.is_multimodal = model_config.is_multimodal
self.is_multimodal_chunked_prefill_supported = (

126
python/sglang/srt/model_executor/model_runner_kv_cache_mixin.py
View File

@@ -1,7 +1,8 @@
from __future__ import annotations
import logging
from typing import TYPE_CHECKING
from dataclasses import dataclass
from typing import TYPE_CHECKING, Optional, Tuple
import torch
@@ -33,6 +34,26 @@ from sglang.srt.utils.common import (
if TYPE_CHECKING:
from sglang.srt.model_executor.model_runner import ModelRunner
@dataclass
class MemoryPoolConfig:
"""Resolved memory pool config, shared between target and draft workers."""
max_total_num_tokens: int
max_running_requests: int
full_max_total_num_tokens: Optional[int] = None
swa_max_total_num_tokens: Optional[int] = None
mem_fraction_static: Optional[float] = None
def __post_init__(self):
if self.max_total_num_tokens <= 0:
msg = "Not enough memory. Please try to increase --mem-fraction-static."
if self.mem_fraction_static is not None:
msg += f" Current value: mem_fraction_static={self.mem_fraction_static}"
raise RuntimeError(msg)
# the ratio of mamba cache pool size to max_running_requests
MAMBA_CACHE_SIZE_MAX_RUNNING_REQUESTS_RATIO = 3
MAMBA_CACHE_V2_ADDITIONAL_RATIO_OVERLAP = 2
@@ -248,9 +269,13 @@ class ModelRunnerKVCacheMixin:
return kv_cache_dim
def set_num_tokens_hybrid_swa(self: ModelRunner, token_capacity: int) -> int:
"""Split token_capacity into full/swa pools. Returns the effective
max_total_num_tokens (= full pool size)."""
def _resolve_hybrid_swa_tokens(
self: ModelRunner, token_capacity: int
) -> Tuple[int, int, int]:
"""Split token_capacity into full/swa pools.
Returns (effective_capacity, full_max_total_num_tokens, swa_max_total_num_tokens).
"""
page_size = self.server_args.page_size
assert self.sliding_window_size is not None and self.sliding_window_size > 0
@@ -264,12 +289,11 @@ class ModelRunnerKVCacheMixin:
if full_layers_num == 0:
# all layers are SWA
self.swa_max_total_num_tokens = align_page_size(token_capacity)
self.full_max_total_num_tokens = 0
swa_tokens = align_page_size(token_capacity)
logger.info(
f"Use sliding window memory pool (all SWA). swa_layer_tokens={self.swa_max_total_num_tokens}"
f"Use sliding window memory pool (all SWA). swa_layer_tokens={swa_tokens}"
)
return self.swa_max_total_num_tokens
return swa_tokens, 0, swa_tokens
swa_full_tokens_ratio = self.server_args.swa_full_tokens_ratio
@@ -323,17 +347,13 @@ class ModelRunnerKVCacheMixin:
denominator > 0
), f"Invalid denominator={denominator} for memory-based allocation. full_per_token={full_per_token}, full_layers_num={full_layers_num}, swa_per_token={swa_per_token}, swa_layers_num={swa_layers_num}, swa_full_tokens_ratio={swa_full_tokens_ratio}"
self.full_max_total_num_tokens = align_page_size(
int(total_memory / denominator)
)
self.swa_max_total_num_tokens = align_page_size(
int(self.full_max_total_num_tokens * swa_full_tokens_ratio)
)
full_tokens = align_page_size(int(total_memory / denominator))
swa_tokens = align_page_size(int(full_tokens * swa_full_tokens_ratio))
logger.info(
f"Use sliding window memory pool. full_layer_tokens={self.full_max_total_num_tokens}, swa_layer_tokens={self.swa_max_total_num_tokens}"
f"Use sliding window memory pool. full_layer_tokens={full_tokens}, swa_layer_tokens={swa_tokens}"
)
return self.full_max_total_num_tokens
return full_tokens, full_tokens, swa_tokens
def _calculate_mamba_ratio(self: ModelRunner) -> int:
if self.server_args.disable_radix_cache:
@@ -349,7 +369,10 @@ class ModelRunnerKVCacheMixin:
return MAMBA_CACHE_SIZE_MAX_RUNNING_REQUESTS_RATIO + additional_ratio
def _init_pools(self: ModelRunner, max_num_reqs: int):
def _init_pools(self: ModelRunner):
"""Initialize the memory pools."""
max_num_reqs = self.max_running_requests
# Initialize req_to_token_pool
if self.req_to_token_pool is None:
# FIXME(lsyin): this is the temporary fix for the context length issue when using speculative decoding
@@ -728,44 +751,49 @@ class ModelRunnerKVCacheMixin:
return max_num_reqs
def init_memory_pool(self: ModelRunner, pre_model_load_memory: int):
# Profile the maximum number of tokens
profiled_tokens = self.profile_max_num_token(pre_model_load_memory)
def _apply_memory_pool_config(self: ModelRunner, config: MemoryPoolConfig):
"""Apply a resolved MemoryPoolConfig and initialize pools."""
self.max_total_num_tokens = config.max_total_num_tokens
self.max_running_requests = config.max_running_requests
if self.is_hybrid_swa:
self.full_max_total_num_tokens = config.full_max_total_num_tokens
self.swa_max_total_num_tokens = config.swa_max_total_num_tokens
# Resolve the token capacity
self._init_pools()
def _resolve_memory_pool_config(
self: ModelRunner, pre_model_load_memory: int
) -> MemoryPoolConfig:
"""Profile GPU memory and resolve all pool parameters into a config."""
profiled_tokens = self.profile_max_num_token(pre_model_load_memory)
token_capacity = self._resolve_token_capacity(profiled_tokens)
# HACK: spec decode uses server_args as a mutable channel to pass
# resolved values between target and draft workers. Target writes first,
# draft reads later. Should be replaced with an explicit handoff.
# NOTE: draft worker override must happen BEFORE SWA splitting so that
# swa_max_total_num_tokens is computed from the correct base value.
if not self.spec_algorithm.is_none() and self.is_draft_worker:
token_capacity = self.server_args.draft_runner_cache_size
# Hybrid SWA: split capacity into full/swa pools, adjust effective capacity
full_tokens = None
swa_tokens = None
if self.is_hybrid_swa:
token_capacity = self.set_num_tokens_hybrid_swa(token_capacity)
# Commit the resolved token capacity & max number of requests
self.max_total_num_tokens = token_capacity
if not self.spec_algorithm.is_none() and self.is_draft_worker:
self.max_running_requests = self.server_args.max_num_reqs
else:
self.max_running_requests = self._resolve_max_num_reqs(token_capacity)
# Target worker stores resolved values for draft worker to read later
if not self.spec_algorithm.is_none() and not self.is_draft_worker:
self.server_args.draft_runner_cache_size = self.max_total_num_tokens
self.server_args.max_num_reqs = self.max_running_requests
if self.max_total_num_tokens <= 0:
raise RuntimeError(
f"Not enough memory. Please try to increase --mem-fraction-static. "
f"Current value: {self.server_args.mem_fraction_static=}"
token_capacity, full_tokens, swa_tokens = self._resolve_hybrid_swa_tokens(
token_capacity
)
self._init_pools(self.max_running_requests)
return MemoryPoolConfig(
max_total_num_tokens=token_capacity,
max_running_requests=self._resolve_max_num_reqs(token_capacity),
full_max_total_num_tokens=full_tokens,
swa_max_total_num_tokens=swa_tokens,
mem_fraction_static=self.server_args.mem_fraction_static,
)
def init_memory_pool(self: ModelRunner, pre_model_load_memory: int):
if not self.spec_algorithm.is_none() and self.is_draft_worker:
assert (
self.memory_pool_config is not None
), "Draft worker requires memory_pool_config"
else:
self.memory_pool_config = self._resolve_memory_pool_config(
pre_model_load_memory
)
self._apply_memory_pool_config(self.memory_pool_config)
logger.info(
f"Memory pool end. "

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

@@ -152,6 +152,7 @@ class EAGLEWorker(TpModelWorker):
is_draft_worker=True,
req_to_token_pool=self.req_to_token_pool,
token_to_kv_pool_allocator=self.token_to_kv_pool_allocator,
memory_pool_config=target_worker.model_runner.memory_pool_config,
)
embed, head = self.target_worker.model_runner.model.get_embed_and_head()

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

@@ -145,6 +145,7 @@ class EagleDraftWorker(BaseDraftWorker):
is_draft_worker=True,
req_to_token_pool=self.req_to_token_pool,
token_to_kv_pool_allocator=self.token_to_kv_pool_allocator,
memory_pool_config=target_worker.model_runner.memory_pool_config,
)
# Alias for better readability

1
python/sglang/srt/speculative/multi_layer_eagle_worker.py
View File

@@ -142,6 +142,7 @@ class MultiLayerEagleWorker(TpModelWorker):
is_draft_worker=True,
req_to_token_pool=self.req_to_token_pool,
token_to_kv_pool_allocator=self.token_to_kv_pool_allocator,
memory_pool_config=target_worker.model_runner.memory_pool_config,
is_multi_layer_eagle=True,
)

1
python/sglang/srt/speculative/multi_layer_eagle_worker_v2.py
View File

@@ -126,6 +126,7 @@ class MultiLayerEagleDraftWorker(BaseDraftWorker):
is_draft_worker=True,
req_to_token_pool=self.req_to_token_pool,
token_to_kv_pool_allocator=self.token_to_kv_pool_allocator,
memory_pool_config=target_worker.model_runner.memory_pool_config,
is_multi_layer_eagle=True,
)

1
python/sglang/srt/speculative/standalone_worker.py
View File

@@ -86,6 +86,7 @@ class StandaloneWorker(EAGLEWorker):
is_draft_worker=True,
req_to_token_pool=self.req_to_token_pool,
token_to_kv_pool_allocator=self.token_to_kv_pool_allocator,
memory_pool_config=target_worker.model_runner.memory_pool_config,
)
# Init attention backend and cuda graphs

1
python/sglang/srt/speculative/standalone_worker_v2.py
View File

@@ -99,6 +99,7 @@ class StandaloneDraftWorker(EagleDraftWorker):
is_draft_worker=True,
req_to_token_pool=self.req_to_token_pool,
token_to_kv_pool_allocator=self.token_to_kv_pool_allocator,
memory_pool_config=target_worker.model_runner.memory_pool_config,
)
# Alias for better readability