diff --git a/python/sglang/bench_one_batch.py b/python/sglang/bench_one_batch.py index 1cfb78b56..1aa7cf7fd 100644 --- a/python/sglang/bench_one_batch.py +++ b/python/sglang/bench_one_batch.py @@ -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, diff --git a/python/sglang/srt/configs/model_config.py b/python/sglang/srt/configs/model_config.py index 2356270c0..abc43d354 100644 --- a/python/sglang/srt/configs/model_config.py +++ b/python/sglang/srt/configs/model_config.py @@ -12,6 +12,7 @@ # limitations under the License. # ============================================================================== +import copy import json import logging import math @@ -198,13 +199,18 @@ class ModelConfig: kwargs = {} if override_config_file and override_config_file.strip(): kwargs["_configuration_file"] = override_config_file.strip() - self.hf_config = get_config( - self.model_path, - trust_remote_code=trust_remote_code, - revision=revision, - model_override_args=self.model_override_args, - model_config_parser=model_config_parser, - **kwargs, + # 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, + model_override_args=self.model_override_args, + 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( @@ -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 diff --git a/python/sglang/srt/managers/scheduler.py b/python/sglang/srt/managers/scheduler.py index 99fdf2c28..fa95f870c 100644 --- a/python/sglang/srt/managers/scheduler.py +++ b/python/sglang/srt/managers/scheduler.py @@ -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 diff --git a/python/sglang/srt/managers/tp_worker.py b/python/sglang/srt/managers/tp_worker.py index 93b669389..bfc13609d 100644 --- a/python/sglang/srt/managers/tp_worker.py +++ b/python/sglang/srt/managers/tp_worker.py @@ -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, diff --git a/python/sglang/srt/model_executor/model_runner.py b/python/sglang/srt/model_executor/model_runner.py index 4117bc87b..9f43032d8 100644 --- a/python/sglang/srt/model_executor/model_runner.py +++ b/python/sglang/srt/model_executor/model_runner.py @@ -449,41 +449,54 @@ 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, + ) + ) - try: - # get the aux layer from draft model config - eagle_config = getattr( - draft_model_config.hf_config, "eagle_config", None - ) - self.eagle_use_aux_hidden_state = eagle_config.get( - "use_aux_hidden_state", True - ) - self.eagle_aux_hidden_state_layer_ids = eagle_config[ - "eagle_aux_hidden_state_layer_ids" - ] - except: - # if there is no aux layer, set to None - self.eagle_aux_hidden_state_layer_ids = None + if self.spec_algorithm.is_eagle3(): + self.eagle_use_aux_hidden_state = True + try: + eagle_config = getattr( + draft_model_config.hf_config, "eagle_config", None + ) + self.eagle_use_aux_hidden_state = eagle_config.get( + "use_aux_hidden_state", True + ) + self.eagle_aux_hidden_state_layer_ids = eagle_config[ + "eagle_aux_hidden_state_layer_ids" + ] + except: + # if there is no aux layer, set to None + 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,11 +838,42 @@ class ModelRunner(ModelRunnerKVCacheMixin): # Init ngram embedding token table self.maybe_init_ngram_embedding() - # Init routed experts capturer - self.init_routed_experts_capturer() + if self.enable_hisparse: + from sglang.srt.managers.hisparse_coordinator import HiSparseCoordinator + from sglang.srt.mem_cache.sparsity import parse_hisparse_config + hisparse_cfg = parse_hisparse_config(self.server_args) + hisparse_top_k = getattr( + self.model_config.hf_text_config, "index_topk", hisparse_cfg.top_k + ) + self.hisparse_coordinator = HiSparseCoordinator( + req_to_token_pool=self.req_to_token_pool, + token_to_kv_pool_allocator=self.token_to_kv_pool_allocator, + top_k=hisparse_top_k, + device_buffer_size=hisparse_cfg.device_buffer_size, + device=self.device, + tp_group=( + self.attention_tp_group.cpu_group + if self.server_args.enable_dp_attention + else self.tp_group.cpu_group + ), + 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. @@ -810,34 +881,15 @@ class ModelRunner(ModelRunnerKVCacheMixin): 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 - - hisparse_cfg = parse_hisparse_config(self.server_args) - hisparse_top_k = getattr( - self.model_config.hf_text_config, "index_topk", hisparse_cfg.top_k - ) - self.hisparse_coordinator = HiSparseCoordinator( - req_to_token_pool=self.req_to_token_pool, - token_to_kv_pool_allocator=self.token_to_kv_pool_allocator, - top_k=hisparse_top_k, - device_buffer_size=hisparse_cfg.device_buffer_size, - device=self.device, - tp_group=( - self.attention_tp_group.cpu_group - if self.server_args.enable_dp_attention - else self.tp_group.cpu_group - ), - host_to_device_ratio=hisparse_cfg.host_to_device_ratio, - ) self.init_attention_backend() self.kernel_warmup() self._pre_initialize_flashinfer_allreduce_workspace() - self.init_decode_cuda_graph() + if not disable_cuda_graph: + self.init_decode_cuda_graph() elif self.device == "cpu": self.init_attention_backend() - self.init_decode_cuda_graph() + if not disable_cuda_graph: + self.init_decode_cuda_graph() elif self.device == "npu": self.init_attention_backend() # lazy init for zbal with mix mode (before graph capture when enable_cuda_graph) @@ -851,10 +903,11 @@ class ModelRunner(ModelRunnerKVCacheMixin): get_world_group().world_size, get_world_group().cpu_group, ) - self.init_decode_cuda_graph() + 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 diff --git a/python/sglang/srt/model_executor/model_runner_kv_cache_mixin.py b/python/sglang/srt/model_executor/model_runner_kv_cache_mixin.py index e81b649f9..fcba8ad6c 100644 --- a/python/sglang/srt/model_executor/model_runner_kv_cache_mixin.py +++ b/python/sglang/srt/model_executor/model_runner_kv_cache_mixin.py @@ -64,12 +64,17 @@ _is_hip = is_hip() class ModelRunnerKVCacheMixin: def _profile_available_bytes(self: ModelRunner, pre_model_load_memory: int) -> int: - 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, - ) + # 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, + distributed=get_world_group().world_size > 1, + cpu_group=get_world_group().cpu_group, + ) rest_memory = post_model_load_memory - pre_model_load_memory * ( 1 - self.mem_fraction_static diff --git a/python/sglang/srt/model_executor/pool_configurator.py b/python/sglang/srt/model_executor/pool_configurator.py index 5803e0bcf..166c973b4 100644 --- a/python/sglang/srt/model_executor/pool_configurator.py +++ b/python/sglang/srt/model_executor/pool_configurator.py @@ -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 ( diff --git a/python/sglang/srt/model_executor/runner/decode_cuda_graph_runner.py b/python/sglang/srt/model_executor/runner/decode_cuda_graph_runner.py index 87cf0d986..4278efacf 100644 --- a/python/sglang/srt/model_executor/runner/decode_cuda_graph_runner.py +++ b/python/sglang/srt/model_executor/runner/decode_cuda_graph_runner.py @@ -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 ( diff --git a/python/sglang/srt/models/qwen3_5_mtp.py b/python/sglang/srt/models/qwen3_5_mtp.py index c699a79aa..8e89b689f 100644 --- a/python/sglang/srt/models/qwen3_5_mtp.py +++ b/python/sglang/srt/models/qwen3_5_mtp.py @@ -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, diff --git a/python/sglang/srt/models/qwen3_next_mtp.py b/python/sglang/srt/models/qwen3_next_mtp.py index c9e8064f5..f239339c6 100644 --- a/python/sglang/srt/models/qwen3_next_mtp.py +++ b/python/sglang/srt/models/qwen3_next_mtp.py @@ -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, diff --git a/python/sglang/srt/server_args.py b/python/sglang/srt/server_args.py index 98d68f101..913601cfa 100644 --- a/python/sglang/srt/server_args.py +++ b/python/sglang/srt/server_args.py @@ -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 diff --git a/python/sglang/srt/speculative/base_spec_worker.py b/python/sglang/srt/speculative/base_spec_worker.py index 3d08dcec4..b1ba3d3bc 100644 --- a/python/sglang/srt/speculative/base_spec_worker.py +++ b/python/sglang/srt/speculative/base_spec_worker.py @@ -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: diff --git a/python/sglang/srt/speculative/dflash_worker_v2.py b/python/sglang/srt/speculative/dflash_worker_v2.py index d476a52fd..605516575 100644 --- a/python/sglang/srt/speculative/dflash_worker_v2.py +++ b/python/sglang/srt/speculative/dflash_worker_v2.py @@ -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: diff --git a/python/sglang/srt/speculative/eagle_draft_extend_cuda_graph_runner.py b/python/sglang/srt/speculative/eagle_draft_extend_cuda_graph_runner.py index 80e60e4da..752916e31 100644 --- a/python/sglang/srt/speculative/eagle_draft_extend_cuda_graph_runner.py +++ b/python/sglang/srt/speculative/eagle_draft_extend_cuda_graph_runner.py @@ -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 diff --git a/python/sglang/srt/speculative/eagle_worker_v2.py b/python/sglang/srt/speculative/eagle_worker_v2.py index 8532df6e0..bd22c68fa 100644 --- a/python/sglang/srt/speculative/eagle_worker_v2.py +++ b/python/sglang/srt/speculative/eagle_worker_v2.py @@ -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 diff --git a/python/sglang/srt/speculative/frozen_kv_mtp_worker_v2.py b/python/sglang/srt/speculative/frozen_kv_mtp_worker_v2.py index a15c9e4ad..07afe21c1 100644 --- a/python/sglang/srt/speculative/frozen_kv_mtp_worker_v2.py +++ b/python/sglang/srt/speculative/frozen_kv_mtp_worker_v2.py @@ -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 diff --git a/python/sglang/srt/speculative/multi_layer_eagle_worker_v2.py b/python/sglang/srt/speculative/multi_layer_eagle_worker_v2.py index 1f77b4269..38e91f3a4 100644 --- a/python/sglang/srt/speculative/multi_layer_eagle_worker_v2.py +++ b/python/sglang/srt/speculative/multi_layer_eagle_worker_v2.py @@ -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,9 +223,10 @@ class MultiLayerEagleDraftWorker(BaseDraftWorker): self.draft_extend_attn_backend_list.append( draft_backend_factory.create_draft_extend_backend() ) - self.draft_runner_list[step].attn_backend = ( - self.draft_extend_attn_backend_list[-1] - ) + 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] + ) def init_cuda_graphs(self): """Capture cuda graphs.""" @@ -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): diff --git a/python/sglang/srt/speculative/ngram_worker.py b/python/sglang/srt/speculative/ngram_worker.py index 0fef94b2c..f2d7c405a 100644 --- a/python/sglang/srt/speculative/ngram_worker.py +++ b/python/sglang/srt/speculative/ngram_worker.py @@ -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). diff --git a/python/sglang/srt/speculative/standalone_worker_v2.py b/python/sglang/srt/speculative/standalone_worker_v2.py index 60a826a12..e07e5b83c 100644 --- a/python/sglang/srt/speculative/standalone_worker_v2.py +++ b/python/sglang/srt/speculative/standalone_worker_v2.py @@ -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 diff --git a/test/registered/unit/configs/test_model_config_shapes.py b/test/registered/unit/configs/test_model_config_shapes.py new file mode 100644 index 000000000..ec505f378 --- /dev/null +++ b/test/registered/unit/configs/test_model_config_shapes.py @@ -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() diff --git a/test/registered/unit/model_executor/test_pool_configurator.py b/test/registered/unit/model_executor/test_pool_configurator.py index 761efeb88..9d33f38ee 100644 --- a/test/registered/unit/model_executor/test_pool_configurator.py +++ b/test/registered/unit/model_executor/test_pool_configurator.py @@ -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) diff --git a/test/registered/unit/spec/test_eagle_worker_v2_topk1_fastpath.py b/test/registered/unit/spec/test_eagle_worker_v2_topk1_fastpath.py index a42cd162c..e94732d46 100644 --- a/test/registered/unit/spec/test_eagle_worker_v2_topk1_fastpath.py +++ b/test/registered/unit/spec/test_eagle_worker_v2_topk1_fastpath.py @@ -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()