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

Fix --mem-fraction-static not accounting for EAGLE draft model KV cache (#23862)

Browse Source
This commit is contained in:
cctry
2026-06-12 10:35:48 -07:00
committed by GitHub
parent 533b59d00c
commit 75998d0421
22 changed files with 688 additions and 295 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/bench_one_batch.py
+2
View File
@@ -321,6 +321,8 @@ def load_model(server_args, port_args, gpu_id, tp_rank):
model_runner = MlxModelRunnerStub(**runner_kwargs)
else:
model_runner = ModelRunner(**runner_kwargs)
model_runner.alloc_memory_pool()
model_runner.init_backends()
rank_print(f"max_total_num_tokens={model_runner.max_total_num_tokens}")
tokenizer = get_tokenizer(
server_args.tokenizer_path,
python/sglang/srt/configs/model_config.py
+27 -20
View File
@@ -12,6 +12,7 @@
# limitations under the License.
# ==============================================================================
import copy
import json
import logging
import math
@@ -198,7 +199,11 @@ class ModelConfig:
kwargs = {}
if override_config_file and override_config_file.strip():
kwargs["_configuration_file"] = override_config_file.strip()
self.hf_config = get_config(
# get_config() is cached. ModelConfig mutates hf_config for draft-model
# remapping and architecture-specific normalization, so each instance
# must own an isolated copy.
self.hf_config = copy.deepcopy(
get_config(
self.model_path,
trust_remote_code=trust_remote_code,
revision=revision,
@@ -206,6 +211,7 @@ class ModelConfig:
model_config_parser=model_config_parser,
**kwargs,
)
)
self.hf_text_config = get_hf_text_config(self.hf_config)
self.hf_generation_config = get_generation_config(
self.model_path,
@@ -589,27 +595,28 @@ class ModelConfig:
def _derive_model_shapes(self):
# Unify the config keys for hf_text_config
self.head_dim = getattr(
self.hf_text_config,
"head_dim",
self.hf_text_config.hidden_size // self.hf_text_config.num_attention_heads,
)
self.v_head_dim = getattr(
self.hf_text_config,
"v_head_dim",
self.head_dim,
self.head_dim = getattr(self.hf_text_config, "head_dim", None)
if self.head_dim is None:
self.head_dim = (
self.hf_text_config.hidden_size
// self.hf_text_config.num_attention_heads
)
setattr(self.hf_text_config, "head_dim", self.head_dim)
self.swa_head_dim = getattr(
self.hf_text_config,
"swa_head_dim",
self.head_dim,
)
self.swa_v_head_dim = getattr(
self.hf_text_config,
"swa_v_head_dim",
self.swa_head_dim,
)
self.v_head_dim = getattr(self.hf_text_config, "v_head_dim", None)
if self.v_head_dim is None:
self.v_head_dim = self.head_dim
setattr(self.hf_text_config, "v_head_dim", self.v_head_dim)
self.swa_head_dim = getattr(self.hf_text_config, "swa_head_dim", None)
if self.swa_head_dim is None:
self.swa_head_dim = self.head_dim
setattr(self.hf_text_config, "swa_head_dim", self.swa_head_dim)
self.swa_v_head_dim = getattr(self.hf_text_config, "swa_v_head_dim", None)
if self.swa_v_head_dim is None:
self.swa_v_head_dim = self.swa_head_dim
setattr(self.hf_text_config, "swa_v_head_dim", self.swa_v_head_dim)
# FIXME: temporary special judge for MLA architecture
if (
"DeepseekV2ForCausalLM" in self.hf_config.architectures
python/sglang/srt/managers/scheduler.py
+39
View File
@@ -817,10 +817,49 @@ class Scheduler(
else:
self.external_corpus_manager = None
def init_target_memory_pool(self):
"""Allocate target KV cache pools if they have not been allocated yet."""
if (
self.tp_worker.model_runner.memory_pool_config is not None
and self.tp_worker.model_runner.req_to_token_pool is not None
and self.tp_worker.model_runner.token_to_kv_pool_allocator is not None
):
return
self.tp_worker.alloc_memory_pool()
def init_memory_pools(self):
"""Allocate KV cache pools for target and draft workers."""
self.init_target_memory_pool()
if self.draft_worker is not None:
pool, allocator = self.tp_worker.get_memory_pool()
self.draft_worker.alloc_memory_pool(
memory_pool_config=self.tp_worker.model_runner.memory_pool_config,
req_to_token_pool=pool,
token_to_kv_pool_allocator=allocator,
)
def init_all_backends(self):
"""Initialize attention backends and capture cuda graphs for all workers."""
self.tp_worker.init_backends()
if self.draft_worker is not None:
self.draft_worker.init_backends()
def init_model_worker(self):
# Load model weights.
self.init_tp_model_worker()
if self.spec_algorithm.is_frozen_kv_mtp():
# Frozen-KV MTP draft construction needs the target KV pool.
self.init_target_memory_pool()
self.maybe_init_draft_worker()
# Allocate KV cache pools for all workers.
# Memory profiling now sees all loaded weights.
self.init_memory_pools()
# Initialize attention backends and capture cuda graphs.
# TODO: make memory profile consider cuda graph memory as well
self.init_all_backends()
# Dispatch the model worker
if self.spec_algorithm.is_none():
self.model_worker = self.tp_worker
python/sglang/srt/managers/tp_worker.py
+43 -24
View File
@@ -292,24 +292,6 @@ class TpModelWorker(BaseTpWorker):
self.pp_group = get_pp_group()
self.world_group = get_world_group()
# Profile number of tokens
self.max_total_num_tokens = self.model_runner.max_total_num_tokens
self.max_prefill_tokens = server_args.max_prefill_tokens
self.max_running_requests = self.model_runner.max_running_requests
assert self.max_running_requests > 0, "max_running_request is zero"
self.max_queued_requests = server_args.max_queued_requests
assert (
self.max_queued_requests is None or self.max_queued_requests >= 1
), "If configured, max_queued_requests must be at least 1 for any work to be scheduled."
self.max_req_len = min(
self.model_config.context_len - 1,
self.model_runner.max_token_pool_size - 1,
)
self.max_req_input_len = self.max_req_len - 5
assert (
self.max_req_len > 0 and self.max_req_input_len > 0
), "Memory pool size is too small"
# Sync random seed across TP workers
self.random_seed = broadcast_pyobj(
[server_args.random_seed],
@@ -323,6 +305,39 @@ class TpModelWorker(BaseTpWorker):
self.enable_spec = server_args.speculative_algorithm is not None
self.hicache_layer_transfer_counter = None
def alloc_memory_pool(
self,
memory_pool_config: Optional[MemoryPoolConfig] = None,
req_to_token_pool: Optional[ReqToTokenPool] = None,
token_to_kv_pool_allocator: Optional[BaseTokenToKVPoolAllocator] = None,
):
"""Allocate KV cache pools only (no backends or cuda graphs)."""
if req_to_token_pool is not None:
self.req_to_token_pool = req_to_token_pool
self.model_runner.req_to_token_pool = req_to_token_pool
if token_to_kv_pool_allocator is not None:
self.token_to_kv_pool_allocator = token_to_kv_pool_allocator
self.model_runner.token_to_kv_pool_allocator = token_to_kv_pool_allocator
self.model_runner.alloc_memory_pool(memory_pool_config)
for mr in self.model_runner_list[1:]:
mr.req_to_token_pool = self.req_to_token_pool
mr.token_to_kv_pool_allocator = self.token_to_kv_pool_allocator
mr.alloc_memory_pool(memory_pool_config)
# Validation
assert self.model_runner.max_running_requests > 0, "max_running_request is zero"
max_req_len = min(
self.model_config.context_len - 1,
self.model_runner.max_token_pool_size - 1,
)
assert max_req_len > 0, "Memory pool size is too small"
def init_backends(self, disable_cuda_graph: bool = False):
"""Initialize attention backends and capture cuda graphs."""
self.model_runner.init_backends(disable_cuda_graph=disable_cuda_graph)
for mr in self.model_runner_list[1:]:
mr.init_backends(disable_cuda_graph=disable_cuda_graph)
def _init_model_config(self):
from sglang.srt.configs.model_config import ModelConfig
@@ -414,13 +429,17 @@ class TpModelWorker(BaseTpWorker):
self.model_runner.hisparse_coordinator = coordinator
def get_worker_info(self):
max_req_len = min(
self.model_config.context_len - 1,
self.model_runner.max_token_pool_size - 1,
)
return (
self.max_total_num_tokens,
self.max_prefill_tokens,
self.max_running_requests,
self.max_queued_requests,
self.max_req_len,
self.max_req_input_len,
self.model_runner.max_total_num_tokens,
self.server_args.max_prefill_tokens,
self.model_runner.max_running_requests,
self.server_args.max_queued_requests,
max_req_len,
max_req_len - 5,
self.random_seed,
self.device,
self.model_runner.forward_stream,
python/sglang/srt/model_executor/model_runner.py
+97 -36
View File
@@ -449,21 +449,36 @@ class ModelRunner(ModelRunnerKVCacheMixin):
# auxiliary hidden capture mode. TODO: expose this to server args?
self.eagle_use_aux_hidden_state = False
self.eagle_draft_num_layers = None
self.dflash_use_aux_hidden_state = False
self.dflash_target_layer_ids = None
self.dflash_draft_num_layers = None
if self.spec_algorithm.is_eagle3() and not self.is_draft_worker:
# load draft config
draft_model_config = ModelConfig.from_server_args(
if (
(self.spec_algorithm.is_eagle() or self.spec_algorithm.is_standalone())
and not self.is_draft_worker
and server_args.speculative_draft_model_path
):
# Load draft config to get layer count for KV cache sizing
draft_model_config = self._build_model_config(
server_args,
model_path=(server_args.speculative_draft_model_path),
model_path=server_args.speculative_draft_model_path,
model_revision=server_args.speculative_draft_model_revision,
is_draft_model=True,
)
self.eagle_use_aux_hidden_state = True
num_nextn_predict_layers = draft_model_config.num_nextn_predict_layers
if num_nextn_predict_layers is not None:
self.eagle_draft_num_layers = int(num_nextn_predict_layers)
else:
self.eagle_draft_num_layers = int(
max(
draft_model_config.num_hidden_layers,
draft_model_config.num_attention_layers,
)
)
if self.spec_algorithm.is_eagle3():
self.eagle_use_aux_hidden_state = True
try:
# get the aux layer from draft model config
eagle_config = getattr(
draft_model_config.hf_config, "eagle_config", None
)
@@ -478,12 +493,10 @@ class ModelRunner(ModelRunnerKVCacheMixin):
self.eagle_aux_hidden_state_layer_ids = None
if self.spec_algorithm.is_dflash() and not self.is_draft_worker:
from sglang.srt.speculative.dflash_utils import (
parse_dflash_draft_config,
)
from sglang.srt.speculative.dflash_utils import parse_dflash_draft_config
# Select target layers to capture for building DFlash context features.
draft_model_config = ModelConfig.from_server_args(
draft_model_config = self._build_model_config(
server_args,
model_path=(server_args.speculative_draft_model_path),
model_revision=server_args.speculative_draft_model_revision,
@@ -541,8 +554,9 @@ class ModelRunner(ModelRunnerKVCacheMixin):
if self.device == "cpu":
self.init_threads_binding()
# Get available memory before model loading
pre_model_load_memory = self.init_torch_distributed()
# Get available memory before model loading.
# Stored for later use by alloc_memory_pool().
self.pre_model_load_memory = self.init_torch_distributed()
# Initialize MooncakeTransferEngine
self.init_shared_mooncake_transfer_engine()
@@ -570,8 +584,8 @@ class ModelRunner(ModelRunnerKVCacheMixin):
self._linear_attn_registry_cache: Any = _UNSET
# Initialize the model runner
self.initialize(pre_model_load_memory)
# Load model weights and configure
self.initialize()
self.check_quantized_moe_compatibility()
if (
@@ -603,6 +617,16 @@ class ModelRunner(ModelRunnerKVCacheMixin):
self._model_update_group = {}
self._weights_send_group = {}
def _build_model_config(
self, server_args, model_path=None, model_revision=None, is_draft_model=False
):
return ModelConfig.from_server_args(
server_args,
model_path=model_path,
model_revision=model_revision,
is_draft_model=is_draft_model,
)
def init_msprobe(self):
# Init the msprobe
try:
@@ -632,7 +656,7 @@ class ModelRunner(ModelRunnerKVCacheMixin):
port=self.dist_port,
)
def initialize(self, pre_model_load_memory: float):
def initialize(self):
server_args = self.server_args
self.memory_saver_adapter = TorchMemorySaverAdapter.create(
@@ -783,14 +807,30 @@ class ModelRunner(ModelRunnerKVCacheMixin):
# Deduce KV cache dtype
self.configure_kv_cache_dtype()
# Init memory pool and attention backends
self.init_memory_pool(pre_model_load_memory)
# Snapshot free memory at the end of the weight-load phase. KV-pool
# profiling uses this instead of measuring at alloc_memory_pool()
# time: draft-model weights load between the two phases and must stay
# outside the --mem-fraction-static budget (deployments tune the
# fraction assuming draft weights live in the non-static slack).
self.post_model_load_memory = get_available_gpu_memory(
self.device,
self.gpu_id,
distributed=get_world_group().world_size > 1,
cpu_group=get_world_group().cpu_group,
)
def alloc_memory_pool(self, memory_pool_config: Optional[MemoryPoolConfig] = None):
"""Allocate KV cache memory pools only (no backends or cuda graphs)."""
if memory_pool_config is not None:
self.memory_pool_config = memory_pool_config
self.init_memory_pool(self.pre_model_load_memory)
# Must be called AFTER init_memory_pool so the pool object exists for
# canary to monkey-patch, and BEFORE init_device_graphs so warmup
# canary to monkey-patch, and BEFORE init_decode_cuda_graph so warmup
# forwards captured into the graph see the patched pool methods.
self.canary_manager = install_canary(
server_args=server_args,
server_args=self.server_args,
model_runner=self,
token_oracle_manager=self._token_oracle_manager,
)
@@ -798,18 +838,6 @@ class ModelRunner(ModelRunnerKVCacheMixin):
# Init ngram embedding token table
self.maybe_init_ngram_embedding()
# Init routed experts capturer
self.init_routed_experts_capturer()
self.init_indexer_capturer()
# TODO: Refactor device-specific init branches into platform interface (separate PR).
# Must be called BEFORE init_decode_cuda_graph() so CUDA graph capture
# runs with aux hidden state capture enabled.
self.init_aux_hidden_state_capture()
if self.device == "cuda" or self.device == "musa":
self.init_cublas()
if self.enable_hisparse:
from sglang.srt.managers.hisparse_coordinator import HiSparseCoordinator
from sglang.srt.mem_cache.sparsity import parse_hisparse_config
@@ -831,12 +859,36 @@ class ModelRunner(ModelRunnerKVCacheMixin):
),
host_to_device_ratio=hisparse_cfg.host_to_device_ratio,
)
self.init_routed_experts_capturer()
self.init_indexer_capturer()
self.attn_backend = None
self.decode_attn_backend = None
self.decode_attn_backend_group = []
self.decode_cuda_graph_runner = None
self.graph_mem_usage = 0
self.prefill_cuda_graph_runner = None
def init_backends(self, disable_cuda_graph: bool = False):
"""Initialize attention backends and capture cuda graphs."""
server_args = self.server_args
# TODO: Refactor device-specific init branches into platform interface (separate PR).
# Must be called BEFORE init_decode_cuda_graph() so CUDA graph capture
# runs with aux hidden state capture enabled.
self.init_aux_hidden_state_capture()
if self.device == "cuda" or self.device == "musa":
self.init_cublas()
self.init_attention_backend()
self.kernel_warmup()
self._pre_initialize_flashinfer_allreduce_workspace()
if not disable_cuda_graph:
self.init_decode_cuda_graph()
elif self.device == "cpu":
self.init_attention_backend()
if not disable_cuda_graph:
self.init_decode_cuda_graph()
elif self.device == "npu":
self.init_attention_backend()
@@ -851,10 +903,11 @@ class ModelRunner(ModelRunnerKVCacheMixin):
get_world_group().world_size,
get_world_group().cpu_group,
)
if not disable_cuda_graph:
self.init_decode_cuda_graph()
elif current_platform.is_out_of_tree():
self.init_attention_backend()
if current_platform.support_cuda_graph():
if current_platform.support_cuda_graph() and not disable_cuda_graph:
self.init_decode_cuda_graph()
else:
self.decode_cuda_graph_runner = None
@@ -864,10 +917,13 @@ class ModelRunner(ModelRunnerKVCacheMixin):
self.graph_mem_usage = 0
self.init_attention_backend()
if disable_cuda_graph:
self.decode_cuda_graph_runner = None
self.graph_mem_usage = 0
if server_args.forward_hooks:
register_forward_hooks(self.model, server_args.forward_hooks)
# Initialize piecewise CUDA graph
self.init_prefill_cuda_graph()
self.prealloc_symmetric_memory_pool()
@@ -2407,9 +2463,7 @@ class ModelRunner(ModelRunnerKVCacheMixin):
return
from sglang.srt.layers.communicator import FUSE_ALLREDUCE_MAX_BATCH_SIZE
from sglang.srt.layers.flashinfer_comm_fusion import (
pre_initialize_workspaces,
)
from sglang.srt.layers.flashinfer_comm_fusion import pre_initialize_workspaces
pre_initialize_workspaces(
max_token_num=FUSE_ALLREDUCE_MAX_BATCH_SIZE,
@@ -2706,6 +2760,13 @@ class ModelRunner(ModelRunnerKVCacheMixin):
seq_lens_sum=None,
seq_lens_cpu=None,
)
# MTP models (e.g. deepseek_nextn) read spec_info.hidden_states
# during forward; provide a dummy so warmup doesn't crash.
spec_info.hidden_states = torch.zeros(
(num_tokens, self.model_config.hidden_size),
dtype=self.dtype,
device=self.device,
)
elif self.spec_algorithm.is_dflash():
from sglang.srt.speculative.dflash_info import DFlashVerifyInput
python/sglang/srt/model_executor/model_runner_kv_cache_mixin.py
+5
View File
@@ -64,6 +64,11 @@ _is_hip = is_hip()
class ModelRunnerKVCacheMixin:
def _profile_available_bytes(self: ModelRunner, pre_model_load_memory: int) -> int:
# Use the snapshot taken at the end of this runner's weight-load phase,
# not the current free memory: draft-model weights loaded after that
# point are charged to the non-static slack, not the static budget.
post_model_load_memory = getattr(self, "post_model_load_memory", None)
if post_model_load_memory is None:
post_model_load_memory = get_available_gpu_memory(
self.device,
self.gpu_id,
python/sglang/srt/model_executor/pool_configurator.py
+18
View File
@@ -104,6 +104,24 @@ class DefaultPoolConfigurator(MemoryPoolConfigurator):
self._cell_size = self._compute_cell_size(mr, num_layers)
# EAGLE/STANDALONE: scale cell_size to account for draft model KV cache.
# Assumes draft and target share the same per-layer KV size (head_dim,
# num_kv_heads, dtype), which holds for EAGLE/MTP draft models that
# reuse the target architecture's attention config.
if (
mr.spec_algorithm.is_eagle() or mr.spec_algorithm.is_standalone()
) and not mr.is_draft_worker:
eagle_draft_num_layers = getattr(mr, "eagle_draft_num_layers", None)
if (
eagle_draft_num_layers is not None
and int(eagle_draft_num_layers) > 0
and int(num_layers) > 0
):
self._cell_size = int(
self._cell_size
* (1 + int(eagle_draft_num_layers) / int(num_layers))
)
# DFLASH: scale cell_size to account for draft model KV cache
if mr.spec_algorithm.is_dflash() and not mr.is_draft_worker:
from sglang.srt.speculative.dflash_utils import (
python/sglang/srt/model_executor/runner/decode_cuda_graph_runner.py
+6
View File
@@ -1118,6 +1118,12 @@ class DecodeCudaGraphRunner(BaseCudaGraphRunner):
seq_lens_sum=None,
seq_lens_cpu=None,
)
# MTP models (e.g. deepseek_nextn) read spec_info.hidden_states
spec_info.hidden_states = torch.zeros(
(num_tokens, self.model_runner.model_config.hidden_size),
dtype=self.model_runner.dtype,
device=self.model_runner.device,
)
elif self.model_runner.spec_algorithm.is_dflash():
from sglang.srt.speculative.dflash_info import DFlashVerifyInput
from sglang.srt.speculative.dflash_utils import (
python/sglang/srt/models/qwen3_5_mtp.py
+4 -3
View File
@@ -89,10 +89,11 @@ class Qwen3_5ForCausalLMMTP(nn.Module):
config.hidden_size, config.rms_norm_eps
)
self.pre_fc_norm_hidden = RMSNorm_cls(config.hidden_size, config.rms_norm_eps)
config.num_hidden_layers = 1
config.full_attention_interval = 1
mtp_config = copy.deepcopy(config)
mtp_config.num_hidden_layers = 1
mtp_config.full_attention_interval = 1
self.model = Qwen3_5ForCausalLM(
config,
mtp_config,
quant_config,
prefix=add_prefix("mtp", prefix),
is_nextn=True,
python/sglang/srt/models/qwen3_next_mtp.py
+4 -3
View File
@@ -69,10 +69,11 @@ class Qwen3NextForCausalLMMTP(Qwen3NextForCausalLM):
config.hidden_size, config.rms_norm_eps
)
self.pre_fc_norm_hidden = RMSNorm_cls(config.hidden_size, config.rms_norm_eps)
config.num_hidden_layers = 1
config.full_attention_interval = 1
mtp_config = copy.deepcopy(config)
mtp_config.num_hidden_layers = 1
mtp_config.full_attention_interval = 1
self.model = Qwen3NextModel(
config,
mtp_config,
quant_config,
prefix=add_prefix("model", prefix),
is_nextn=True,
python/sglang/srt/server_args.py
-8
View File
@@ -1777,14 +1777,6 @@ class ServerArgs:
if gpu_mem is not None and gpu_mem > 60 * 1024:
reserved_mem = max(reserved_mem, 10 * 1024)
if self.speculative_algorithm is not None:
if self.speculative_algorithm == "STANDALONE":
# standalonedraft model and cuda graphs
reserved_mem += 6 * 1024
elif self.speculative_algorithm not in {"NGRAM", "DFLASH"}:
# eagle draft models and cuda graphs
reserved_mem += 4 * 1024
self.mem_fraction_static = (
round((gpu_mem - reserved_mem) / gpu_mem, 3)
if gpu_mem is not None
python/sglang/srt/speculative/base_spec_worker.py
+19
View File
@@ -16,6 +16,19 @@ class BaseDraftWorker(ABC):
def draft_extend():
pass
def alloc_memory_pool(self, **kwargs):
pass
def init_backends(self):
"""Initialize standard backends (no cuda graphs) then draft-specific backends.
Subclasses should wrap this with their context managers (draft_tp_context,
speculative_moe_backend_context, etc.) rather than reimplementing the logic.
"""
self.draft_worker.init_backends(disable_cuda_graph=True)
self.init_attention_backend()
self.init_cuda_graphs()
class BaseSpecWorker(ABC):
@property
@@ -39,6 +52,12 @@ class BaseSpecWorker(ABC):
# TODO: move this abstract method to BaseTpWorker and call through self.model_runner
pass
def alloc_memory_pool(self, **kwargs):
pass
def init_backends(self):
pass
def on_verify_complete_cpu(
self, num_correct_drafts_per_req: list[int], batch_size: int = 0
) -> None:
python/sglang/srt/speculative/dflash_worker_v2.py
+33 -15
View File
@@ -40,7 +40,7 @@ from sglang.srt.speculative.triton_ops.dflash_accept_bonus import (
from sglang.srt.speculative.triton_ops.dflash_prepare_block import (
_prepare_dflash_draft_block_unchecked,
)
from sglang.srt.utils import is_cuda, is_hip, is_npu
from sglang.srt.utils import get_available_gpu_memory, is_cuda, is_hip, is_npu
_is_npu = is_npu()
@@ -102,17 +102,6 @@ class DFlashWorkerV2(BaseSpecWorker):
self._logged_first_verify = False
# Draft runner (separate KV cache + attention backend).
# Without draft windowing, the draft worker aliases the target request->token
# mapping and allocation state. With draft windowing enabled, the draft worker
# keeps a private compact req->token table over the same global KV index space,
# so radix-cache/prefix-hit KV remains reusable while draft attention sees only
# the recent window.
target_req_to_token_pool, target_token_to_kv_pool_allocator = (
target_worker.get_memory_pool()
)
shared_req_to_token_pool = (
None if self.use_compact_draft_cache else target_req_to_token_pool
)
draft_server_args = deepcopy(server_args)
draft_server_args.skip_tokenizer_init = True
draft_backend = draft_server_args.speculative_draft_attention_backend
@@ -168,9 +157,6 @@ class DFlashWorkerV2(BaseSpecWorker):
dp_rank=dp_rank,
nccl_port=nccl_port,
is_draft_worker=True,
req_to_token_pool=shared_req_to_token_pool,
token_to_kv_pool_allocator=target_token_to_kv_pool_allocator,
memory_pool_config=target_worker.model_runner.memory_pool_config,
)
set_global_server_args_for_scheduler(saved_server_args)
self.draft_model_runner = self._draft_worker.model_runner
@@ -288,6 +274,38 @@ class DFlashWorkerV2(BaseSpecWorker):
self.draft_model_runner.attn_backend,
)
def alloc_memory_pool(
self,
memory_pool_config=None,
req_to_token_pool=None,
token_to_kv_pool_allocator=None,
):
# Without draft windowing, the draft worker aliases the target
# request->token mapping and allocation state. With draft windowing
# enabled, the draft worker keeps a private compact req->token table
# over the same global KV index space, so radix-cache/prefix-hit KV
# remains reusable while draft attention sees only the recent window.
self._draft_worker.alloc_memory_pool(
memory_pool_config=memory_pool_config,
req_to_token_pool=(
None if self.use_compact_draft_cache else req_to_token_pool
),
token_to_kv_pool_allocator=token_to_kv_pool_allocator,
)
def init_backends(self):
disable_cuda_graph = False
if is_cuda() and not self.server_args.disable_cuda_graph:
available_mem = get_available_gpu_memory(self.device, self.gpu_id)
disable_cuda_graph = available_mem < 1.0
if disable_cuda_graph:
logger.warning(
"Disable DFLASH draft cuda graph because only %.2f GB GPU "
"memory is available after target backend initialization.",
available_mem,
)
self._draft_worker.init_backends(disable_cuda_graph=disable_cuda_graph)
def _init_fused_kv_helper(self) -> None:
"""Initialize the fused KV materialization helper with pre-stacked weights."""
try:
python/sglang/srt/speculative/eagle_draft_extend_cuda_graph_runner.py
+1 -1
View File
@@ -165,7 +165,7 @@ class EAGLEDraftExtendCudaGraphRunner(DecodeCudaGraphRunner):
else None
)
self.seq_len_fill_value = (
self.model_runner.attn_backend.get_cuda_graph_seq_len_fill_value()
self.draft_extend_attn_backend.get_cuda_graph_seq_len_fill_value()
)
seq_lens = torch.full(
(self.max_bs,), self.seq_len_fill_value, dtype=torch.int64
python/sglang/srt/speculative/eagle_worker_v2.py
+52 -46
View File
@@ -155,18 +155,7 @@ class EagleDraftWorker(BaseDraftWorker):
self._topk1_score_indices_prealloc = None
self._rebuild_topk1_chain_buffers()
# Do not capture cuda graph in `TpModelWorker` init,
# will capture later with init_cuda_graphs()
backup_decode_mode = server_args.cuda_graph_config.decode.backend
server_args.cuda_graph_config.decode.backend = Backend.DISABLED
# Share the allocator with a target worker.
# Draft and target worker own their own KV cache pools.
self.req_to_token_pool, self.token_to_kv_pool_allocator = (
target_worker.get_memory_pool()
)
# Init draft worker
# Load draft model weights only.
if server_args.enable_dp_attention and self.speculative_algorithm.is_eagle3():
ctx = draft_tp_context(get_attention_tp_group())
else:
@@ -174,7 +163,6 @@ class EagleDraftWorker(BaseDraftWorker):
with (
ctx
), speculative_moe_backend_context(), speculative_moe_a2a_backend_context():
# Init draft worker
self.draft_worker = TpModelWorker(
server_args=server_args,
gpu_id=gpu_id,
@@ -186,9 +174,6 @@ class EagleDraftWorker(BaseDraftWorker):
moe_dp_rank=moe_dp_rank,
nccl_port=nccl_port,
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
@@ -201,21 +186,35 @@ class EagleDraftWorker(BaseDraftWorker):
self.eagle_use_aux_hidden_state = eagle_config.get(
"use_aux_hidden_state", True
)
self.init_token_map()
self.init_lm_head()
# Init attention backend and cuda graphs
self.draft_runner.server_args.cuda_graph_config.decode.backend = (
backup_decode_mode
)
self.draft_tp_context = (
draft_tp_context if server_args.enable_dp_attention else empty_context
)
with (
self.draft_tp_context(self.draft_runner.tp_group),
speculative_moe_backend_context(),
speculative_moe_a2a_backend_context(),
self.tree_mask_mode = TreeMaskMode.FULL_MASK
self.plan_stream, self.plan_stream_ctx = _get_plan_stream(self.device)
def alloc_memory_pool(
self,
memory_pool_config=None,
req_to_token_pool=None,
token_to_kv_pool_allocator=None,
):
"""Allocate draft KV cache pools (called by scheduler)."""
self.req_to_token_pool = req_to_token_pool
self.token_to_kv_pool_allocator = token_to_kv_pool_allocator
self.draft_worker.alloc_memory_pool(
memory_pool_config=memory_pool_config,
req_to_token_pool=req_to_token_pool,
token_to_kv_pool_allocator=token_to_kv_pool_allocator,
)
self.init_token_map()
self.init_lm_head()
def init_backends(self):
with self.draft_tp_context(
self.draft_runner.tp_group
), speculative_moe_backend_context(), speculative_moe_a2a_backend_context():
self.draft_worker.init_backends(disable_cuda_graph=True)
self.init_attention_backend()
if check_cuda_graph_backend(Phase.PREFILL, Backend.BREAKABLE):
self.draft_runner.init_prefill_cuda_graph(force_for_draft_worker=True)
@@ -224,10 +223,6 @@ class EagleDraftWorker(BaseDraftWorker):
if (c := self.draft_runner.canary_manager) is not None:
c.mark_init_finished()
self.tree_mask_mode = TreeMaskMode.FULL_MASK
self.plan_stream, self.plan_stream_ctx = _get_plan_stream(self.device)
def _rebuild_topk1_chain_buffers(self) -> None:
# For topk=1 the draft tree degenerates to a chain, so parent_list and
# top_scores_index are runtime-invariant. Must be rebuilt after any
@@ -844,10 +839,6 @@ class EAGLEWorkerV2(BaseSpecWorker):
server_args.speculative_algorithm
)
self.req_to_token_pool, self.token_to_kv_pool_allocator = (
target_worker.get_memory_pool()
)
# Override the context length of the draft model to be the same as the target model.
server_args.context_length = target_worker.model_runner.model_config.context_len
@@ -879,7 +870,32 @@ class EAGLEWorkerV2(BaseSpecWorker):
self.plan_stream, self.plan_stream_ctx = _get_plan_stream(self.device)
# Build adaptive runtime states (must be after draft worker is fully initialized)
@property
def spec_v2_attn_backends(self) -> tuple:
# Every attn backend a spec_v2 forward touches; consumed by
# decide_needs_cpu_seq_lens to gate the seq_lens_cpu D2H.
return (
self._target_worker.model_runner.attn_backend,
self._draft_worker.draft_attn_backend,
self._draft_worker.draft_extend_attn_backend
or self._draft_worker.draft_runner.attn_backend,
)
def alloc_memory_pool(
self,
memory_pool_config=None,
req_to_token_pool=None,
token_to_kv_pool_allocator=None,
):
self._draft_worker.alloc_memory_pool(
memory_pool_config, req_to_token_pool, token_to_kv_pool_allocator
)
self.req_to_token_pool = req_to_token_pool
self.token_to_kv_pool_allocator = token_to_kv_pool_allocator
def init_backends(self):
self._draft_worker.init_backends()
# Build adaptive runtime states after target and draft backends exist.
if self.adaptive_controller is not None:
with (
self._draft_worker.draft_tp_context(
@@ -908,16 +924,6 @@ class EAGLEWorkerV2(BaseSpecWorker):
),
)
@property
def spec_v2_attn_backends(self) -> tuple:
# Every attn backend a spec_v2 forward touches; consumed by
# decide_needs_cpu_seq_lens to gate the seq_lens_cpu D2H.
return (
self._target_worker.model_runner.attn_backend,
self._draft_worker.draft_attn_backend,
self._draft_worker.draft_extend_attn_backend,
)
@property
def target_worker(self):
return self._target_worker
python/sglang/srt/speculative/frozen_kv_mtp_worker_v2.py
+31 -16
View File
@@ -32,10 +32,7 @@ from sglang.srt.layers.moe.utils import (
)
from sglang.srt.managers.schedule_batch import ScheduleBatch
from sglang.srt.managers.tp_worker import TpModelWorker
from sglang.srt.model_executor.cuda_graph_config import (
Backend,
cuda_graph_fully_disabled,
)
from sglang.srt.model_executor.cuda_graph_config import cuda_graph_fully_disabled
from sglang.srt.model_executor.forward_batch_info import (
CaptureHiddenMode,
ForwardBatch,
@@ -117,17 +114,13 @@ class FrozenKVMTPDraftWorker(BaseDraftWorker, TpModelWorker):
f"{self.speculative_algorithm.name}."
)
# Defer cuda graph capture; we do it ourselves below.
backup_decode_mode = server_args.cuda_graph_config.decode.backend
server_args.cuda_graph_config.decode.backend = Backend.DISABLED
# Draft attention uses target req_to_token + KV allocator (read-only).
self.req_to_token_pool, self.token_to_kv_pool_allocator = (
target_worker.get_memory_pool()
)
target_cfg = target_worker.model_runner.memory_pool_config
draft_pool_config = MemoryPoolConfig(
self.draft_pool_config = MemoryPoolConfig(
max_total_num_tokens=64, # Dummy value
max_running_requests=target_cfg.max_running_requests,
)
@@ -153,7 +146,7 @@ class FrozenKVMTPDraftWorker(BaseDraftWorker, 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=draft_pool_config,
memory_pool_config=self.draft_pool_config,
)
embed, head = self.target_worker.model_runner.model.get_embed_and_head()
@@ -170,27 +163,49 @@ class FrozenKVMTPDraftWorker(BaseDraftWorker, TpModelWorker):
if hasattr(self.draft_model_runner.model, "bind_frozen_kv_context"):
self._bind_kv_context()
self.draft_model_runner.server_args.cuda_graph_config.decode.backend = (
backup_decode_mode
)
self.draft_tp_context = (
draft_tp_context if server_args.enable_dp_attention else empty_context
)
self.draft_attn_backend = self._init_draft_attn_backend()
self.draft_model_runner.draft_attn_backend = self.draft_attn_backend
self.draft_attn_backend = None
self.cuda_graph_runner = None
# Frozen draft has no draft-extend forward (seed-select only); keep these
# None so inherited probes (spec_v2_attn_backends, adaptive) stay typed.
self.draft_extend_attn_backend = None
self.cuda_graph_runner_for_draft_extend = None
def alloc_memory_pool(
self,
memory_pool_config=None,
req_to_token_pool=None,
token_to_kv_pool_allocator=None,
):
# NOTE: call TpModelWorker explicitly -- BaseDraftWorker precedes it in
# the MRO and its alloc_memory_pool is a no-op stub.
TpModelWorker.alloc_memory_pool(
self,
memory_pool_config=self.draft_pool_config,
req_to_token_pool=(
req_to_token_pool
if req_to_token_pool is not None
else self.req_to_token_pool
),
token_to_kv_pool_allocator=(
token_to_kv_pool_allocator
if token_to_kv_pool_allocator is not None
else self.token_to_kv_pool_allocator
),
)
def init_backends(self):
with (
self.draft_tp_context(self.draft_model_runner.tp_group),
speculative_moe_backend_context(),
speculative_moe_a2a_backend_context(),
):
TpModelWorker.init_backends(self, disable_cuda_graph=True)
self.draft_attn_backend = self._init_draft_attn_backend()
self.draft_model_runner.draft_attn_backend = self.draft_attn_backend
self.init_cuda_graphs()
@property
python/sglang/srt/speculative/multi_layer_eagle_worker_v2.py
+48 -37
View File
@@ -135,18 +135,8 @@ class MultiLayerEagleDraftWorker(BaseDraftWorker):
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_decode_mode = server_args.cuda_graph_config.decode.backend
server_args.cuda_graph_config.decode.backend = Backend.DISABLED
# Share the allocator with a target worker.
# Draft and target worker own their own KV cache pools.
self.req_to_token_pool, self.token_to_kv_pool_allocator = (
target_worker.get_memory_pool()
)
# Load draft model weights only.
with empty_context(), speculative_moe_backend_context():
# Init draft worker
self.draft_worker = TpModelWorker(
server_args=server_args,
gpu_id=gpu_id,
@@ -158,9 +148,6 @@ class MultiLayerEagleDraftWorker(BaseDraftWorker):
moe_dp_rank=moe_dp_rank,
nccl_port=nccl_port,
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,
)
@@ -174,7 +161,26 @@ class MultiLayerEagleDraftWorker(BaseDraftWorker):
# next step. Non-chain: each step uses the target model's hidden states.
draft_arch = self.draft_worker.model_config.hf_config.architectures[0]
self.chain_mtp_hidden_states = draft_arch in ["Step3p5MTP"]
self.draft_tp_context = (
draft_tp_context if server_args.enable_dp_attention else empty_context
)
self.tree_mask_mode = TreeMaskMode.FULL_MASK
self.plan_stream, self.plan_stream_ctx = _get_plan_stream(self.device)
def alloc_memory_pool(
self,
memory_pool_config=None,
req_to_token_pool=None,
token_to_kv_pool_allocator=None,
):
"""Allocate draft KV cache pools (called by scheduler)."""
self.req_to_token_pool = req_to_token_pool
self.token_to_kv_pool_allocator = token_to_kv_pool_allocator
self.draft_worker.alloc_memory_pool(
memory_pool_config=memory_pool_config,
req_to_token_pool=req_to_token_pool,
token_to_kv_pool_allocator=token_to_kv_pool_allocator,
)
self.init_lm_head()
# KV cache reversion buffer; sized to mirror req_to_token (indexed by
@@ -189,24 +195,11 @@ class MultiLayerEagleDraftWorker(BaseDraftWorker):
device=self.device,
)
# Init attention backend and cuda graphs
for i in range(self.speculative_num_steps):
self.draft_runner_list[i].server_args.cuda_graph_config.decode.backend = (
backup_decode_mode
)
self.draft_tp_context = (
draft_tp_context if server_args.enable_dp_attention else empty_context
)
with (
self.draft_tp_context(self.draft_runner_list[0].tp_group),
speculative_moe_backend_context(),
):
self.init_attention_backend()
self.init_cuda_graphs()
self.tree_mask_mode = TreeMaskMode.FULL_MASK
self.plan_stream, self.plan_stream_ctx = _get_plan_stream(self.device)
def init_backends(self):
with self.draft_tp_context(
self.draft_runner_list[0].tp_group
), speculative_moe_backend_context():
super().init_backends()
def mtp_model_runner(self, step: int):
return self.draft_runner_list[step]
@@ -230,6 +223,7 @@ class MultiLayerEagleDraftWorker(BaseDraftWorker):
self.draft_extend_attn_backend_list.append(
draft_backend_factory.create_draft_extend_backend()
)
if self.draft_extend_attn_backend_list[-1] is not None:
self.draft_runner_list[step].attn_backend = (
self.draft_extend_attn_backend_list[-1]
)
@@ -685,10 +679,6 @@ class MultiLayerEagleWorkerV2(BaseSpecWorker):
server_args.speculative_algorithm
)
self.req_to_token_pool, self.token_to_kv_pool_allocator = (
target_worker.get_memory_pool()
)
# Override the context length of the draft model to be the same as the target model.
server_args.context_length = target_worker.model_runner.model_config.context_len
@@ -712,6 +702,21 @@ class MultiLayerEagleWorkerV2(BaseSpecWorker):
self.plan_stream, self.plan_stream_ctx = _get_plan_stream(self.device)
def alloc_memory_pool(
self,
memory_pool_config=None,
req_to_token_pool=None,
token_to_kv_pool_allocator=None,
):
self._draft_worker.alloc_memory_pool(
memory_pool_config, req_to_token_pool, token_to_kv_pool_allocator
)
self.req_to_token_pool = req_to_token_pool
self.token_to_kv_pool_allocator = token_to_kv_pool_allocator
def init_backends(self):
self._draft_worker.init_backends()
@property
def target_worker(self):
return self._target_worker
@@ -724,7 +729,13 @@ class MultiLayerEagleWorkerV2(BaseSpecWorker):
def spec_v2_attn_backends(self) -> tuple:
return (
self._target_worker.model_runner.attn_backend,
*self._draft_worker.draft_extend_attn_backend_list,
*(
backend or runner.attn_backend
for backend, runner in zip(
self._draft_worker.draft_extend_attn_backend_list,
self._draft_worker.draft_runner_list,
)
),
)
def clear_cache_pool(self):
python/sglang/srt/speculative/ngram_worker.py
+10 -7
View File
@@ -34,6 +34,14 @@ USE_FULL_MASK = True
class NGRAMWorker(BaseSpecWorker):
def alloc_memory_pool(self, **kwargs):
# The target memory pool does not exist yet when __init__ runs.
self.req_to_token_pool, self.token_to_kv_pool_allocator = (
self._target_worker.get_memory_pool()
)
self.max_batch_size = self.model_runner.max_running_requests
self._init_preallocated_tensors()
def __init__(
self,
server_args: ServerArgs,
@@ -57,15 +65,10 @@ class NGRAMWorker(BaseSpecWorker):
self.speculative_num_draft_tokens = server_args.speculative_num_draft_tokens
self.topk = server_args.speculative_eagle_topk
self.speculative_num_steps = server_args.speculative_num_steps
self.req_to_token_pool, self.token_to_kv_pool_allocator = (
target_worker.get_memory_pool()
)
self.max_batch_size = target_worker.max_running_requests
# req_to_token_pool / token_to_kv_pool_allocator are set in
# alloc_memory_pool(), after the target pools are allocated.
self.device = f"cuda:{gpu_id}" if gpu_id >= 0 else "cuda"
self._init_preallocated_tensors()
self.adaptive_controller = None
# rids of the last decode batch; used to erase corpus match state for
# requests that left the batch (see forward_batch_generation).
python/sglang/srt/speculative/standalone_worker_v2.py
+24 -34
View File
@@ -7,7 +7,6 @@ import torch
from sglang.srt.environ import envs
from sglang.srt.layers.moe.utils import speculative_moe_backend_context
from sglang.srt.managers.tp_worker import TpModelWorker
from sglang.srt.model_executor.cuda_graph_config import Backend
from sglang.srt.server_args import ServerArgs
from sglang.srt.speculative.adaptive_runtime_state import (
AdaptiveController,
@@ -83,18 +82,8 @@ class StandaloneDraftWorker(EagleDraftWorker):
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_decode_mode = server_args.cuda_graph_config.decode.backend
server_args.cuda_graph_config.decode.backend = Backend.DISABLED
# Share the allocator with a target worker.
# Draft and target worker own their own KV cache pools.
self.req_to_token_pool, self.token_to_kv_pool_allocator = (
target_worker.get_memory_pool()
)
# Load draft model weights only.
with empty_context():
# Init draft worker
self.draft_worker = TpModelWorker(
server_args=server_args,
gpu_id=gpu_id,
@@ -106,34 +95,39 @@ class StandaloneDraftWorker(EagleDraftWorker):
moe_dp_rank=moe_dp_rank,
nccl_port=nccl_port,
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
self.draft_runner = self.draft_worker.model_runner
self.init_token_map()
self.init_lm_head()
# Init attention backend and cuda graphs
self.draft_runner.server_args.cuda_graph_config.decode.backend = (
backup_decode_mode
)
self.draft_tp_context = (
draft_tp_context if server_args.enable_dp_attention else empty_context
)
with (
self.draft_tp_context(self.draft_runner.tp_group),
speculative_moe_backend_context(),
):
self.init_attention_backend()
self.init_cuda_graphs()
self.tree_mask_mode = TreeMaskMode.FULL_MASK
self.plan_stream, self.plan_stream_ctx = _get_plan_stream(self.device)
def alloc_memory_pool(
self,
memory_pool_config=None,
req_to_token_pool=None,
token_to_kv_pool_allocator=None,
):
"""Standalone: allocate pools without sharing embeddings."""
self.req_to_token_pool = req_to_token_pool
self.token_to_kv_pool_allocator = token_to_kv_pool_allocator
self.draft_worker.alloc_memory_pool(
memory_pool_config=memory_pool_config,
req_to_token_pool=req_to_token_pool,
token_to_kv_pool_allocator=token_to_kv_pool_allocator,
)
self.init_token_map()
self.init_lm_head()
def init_backends(self):
with self.draft_tp_context(
self.draft_runner.tp_group
), speculative_moe_backend_context():
super().init_backends()
def init_lm_head(self):
"""Override to prevent sharing embeddings and lm_head with target model."""
# For standalone worker, we don't share embeddings and lm_head
@@ -168,10 +162,6 @@ class StandaloneWorkerV2(EAGLEWorkerV2):
server_args.speculative_algorithm
)
self.req_to_token_pool, self.token_to_kv_pool_allocator = (
target_worker.get_memory_pool()
)
# Override the context length of the draft model to be the same as the target model.
server_args.context_length = target_worker.model_runner.model_config.context_len
test/registered/unit/configs/test_model_config_shapes.py
+71
View File
@@ -0,0 +1,71 @@
"""Unit tests for ModelConfig shape normalization."""
import unittest
from types import SimpleNamespace
from sglang.srt.configs.model_config import ModelConfig
from sglang.test.ci.ci_register import register_cpu_ci
from sglang.test.test_utils import CustomTestCase
register_cpu_ci(est_time=1, suite="base-a-test-cpu")
def _make_text_config(**overrides):
defaults = dict(
architectures=["MixtralForCausalLM"],
model_type="mixtral",
hidden_size=4096,
num_attention_heads=32,
num_hidden_layers=2,
vocab_size=32000,
num_key_value_heads=8,
)
defaults.update(overrides)
return SimpleNamespace(**defaults)
class TestModelConfigShapes(CustomTestCase):
def _derive_shapes(self, text_config):
model_config = ModelConfig.__new__(ModelConfig)
model_config.hf_config = text_config
model_config.hf_text_config = text_config
model_config._derive_model_shapes()
return model_config
def test_optional_head_dims_default_when_none(self):
text_config = _make_text_config(
head_dim=None,
v_head_dim=None,
swa_head_dim=None,
swa_v_head_dim=None,
)
model_config = self._derive_shapes(text_config)
self.assertEqual(model_config.head_dim, 128)
self.assertEqual(model_config.v_head_dim, 128)
self.assertEqual(model_config.swa_head_dim, 128)
self.assertEqual(model_config.swa_v_head_dim, 128)
self.assertEqual(text_config.head_dim, 128)
self.assertEqual(text_config.v_head_dim, 128)
self.assertEqual(text_config.swa_head_dim, 128)
self.assertEqual(text_config.swa_v_head_dim, 128)
def test_explicit_head_dims_are_preserved(self):
text_config = _make_text_config(
head_dim=128,
v_head_dim=96,
swa_head_dim=64,
swa_v_head_dim=48,
)
model_config = self._derive_shapes(text_config)
self.assertEqual(model_config.head_dim, 128)
self.assertEqual(model_config.v_head_dim, 96)
self.assertEqual(model_config.swa_head_dim, 64)
self.assertEqual(model_config.swa_v_head_dim, 48)
if __name__ == "__main__":
unittest.main()
test/registered/unit/model_executor/test_pool_configurator.py
+31
View File
@@ -83,6 +83,8 @@ def _make_model_runner(
mr.server_args = sa
spec = MagicMock()
spec.is_eagle.return_value = False
spec.is_standalone.return_value = False
spec.is_dflash.return_value = False
spec.is_none.return_value = True
mr.spec_algorithm = spec
@@ -303,6 +305,35 @@ class TestAllSWAConfigurator(unittest.TestCase):
self.assertEqual(config.swa_max_total_num_tokens, 500)
class TestEagleConfigurator(unittest.TestCase):
"""EAGLE: draft KV cache must be accounted for so total allocation fits in budget."""
def test_eagle_does_not_exceed_budget(self):
"""Total memory (target + draft KV cache) must not exceed available."""
available = 10_000_000
num_layers = 32
eagle_draft_num_layers = 4
mr = _make_model_runner(num_layers=num_layers)
mr.spec_algorithm.is_eagle.return_value = True
mr.spec_algorithm.is_standalone.return_value = False
mr.spec_algorithm.is_none.return_value = False
mr.eagle_draft_num_layers = eagle_draft_num_layers
with mock_cpu_env():
from sglang.srt.model_executor.pool_configurator import (
create_memory_pool_configurator,
)
cfg = create_memory_pool_configurator(mr)
config = cfg.calculate_pool_sizes(available, 1)
full_pt = _full_per_token(mr)
total_layers = num_layers + eagle_draft_num_layers
used = config.max_total_num_tokens * full_pt * total_layers
self.assertLessEqual(used, available)
class TestFactory(unittest.TestCase):
def test_default_for_non_swa(self):
mr = _make_model_runner(is_hybrid_swa=False)
test/registered/unit/spec/test_eagle_worker_v2_topk1_fastpath.py
+79 -1
View File
@@ -8,11 +8,12 @@ slow path (`organize_draft_results`) for num_steps in {1, 2, 3, 4}.
import unittest
from types import SimpleNamespace
from unittest.mock import patch
import torch
from sglang.srt.speculative.eagle_utils import organize_draft_results
from sglang.srt.speculative.eagle_worker_v2 import EagleDraftWorker
from sglang.srt.speculative.eagle_worker_v2 import EagleDraftWorker, EAGLEWorkerV2
from sglang.srt.utils import get_device
from sglang.test.ci.ci_register import register_cuda_ci
from sglang.test.test_utils import CustomTestCase
@@ -57,6 +58,20 @@ def _make_worker(num_steps: int, num_draft_tokens: int):
return worker
def _make_backend_factory(decode_backend, draft_extend_backend):
class FakeDraftBackendFactory:
def __init__(self, *args, **kwargs):
pass
def create_decode_backend(self):
return decode_backend
def create_draft_extend_backend(self):
return draft_extend_backend
return FakeDraftBackendFactory
class TestEagleWorkerV2Topk1FastPath(CustomTestCase):
def test_fast_path_matches_slow_path(self):
bs = 3
@@ -93,5 +108,68 @@ class TestEagleWorkerV2Topk1FastPath(CustomTestCase):
worker._rebuild_topk1_chain_buffers()
class TestEagleWorkerV2BackendFallback(CustomTestCase):
def test_preserves_initialized_backend_when_draft_extend_backend_is_unset(self):
worker = object.__new__(EagleDraftWorker)
existing_backend = object()
decode_backend = object()
worker.server_args = SimpleNamespace()
worker.draft_runner = SimpleNamespace(attn_backend=existing_backend)
worker.topk = 1
worker.speculative_num_steps = 2
with patch(
"sglang.srt.speculative.eagle_worker_v2.DraftBackendFactory",
_make_backend_factory(decode_backend, None),
):
worker.init_attention_backend()
self.assertIs(worker.draft_attn_backend, decode_backend)
self.assertIsNone(worker.draft_extend_attn_backend)
self.assertIs(worker.draft_runner.draft_attn_backend, decode_backend)
self.assertIs(worker.draft_runner.attn_backend, existing_backend)
def test_uses_draft_extend_backend_when_available(self):
worker = object.__new__(EagleDraftWorker)
existing_backend = object()
decode_backend = object()
draft_extend_backend = object()
worker.server_args = SimpleNamespace()
worker.draft_runner = SimpleNamespace(attn_backend=existing_backend)
worker.topk = 1
worker.speculative_num_steps = 2
with patch(
"sglang.srt.speculative.eagle_worker_v2.DraftBackendFactory",
_make_backend_factory(decode_backend, draft_extend_backend),
):
worker.init_attention_backend()
self.assertIs(worker.draft_attn_backend, decode_backend)
self.assertIs(worker.draft_extend_attn_backend, draft_extend_backend)
self.assertIs(worker.draft_runner.draft_attn_backend, decode_backend)
self.assertIs(worker.draft_runner.attn_backend, draft_extend_backend)
def test_spec_v2_attn_backends_include_draft_extend_fallback(self):
target_backend = object()
decode_backend = object()
fallback_backend = object()
worker = object.__new__(EAGLEWorkerV2)
worker._target_worker = SimpleNamespace(
model_runner=SimpleNamespace(attn_backend=target_backend)
)
worker._draft_worker = SimpleNamespace(
draft_attn_backend=decode_backend,
draft_extend_attn_backend=None,
draft_runner=SimpleNamespace(attn_backend=fallback_backend),
)
self.assertEqual(
worker.spec_v2_attn_backends,
(target_backend, decode_backend, fallback_backend),
)
if __name__ == "__main__":
unittest.main()