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

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,

python/sglang/srt/configs/model_config.py

@@ -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.head_dim = getattr(self.hf_text_config, "head_dim", None)
-            self.hf_text_config,
+        if self.head_dim is None:
-            "head_dim",
+            self.head_dim = (
-            self.hf_text_config.hidden_size // self.hf_text_config.num_attention_heads,
+                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,
             )
+            setattr(self.hf_text_config, "head_dim", self.head_dim)
 
-        self.swa_head_dim = getattr(
+        self.v_head_dim = getattr(self.hf_text_config, "v_head_dim", None)
-            self.hf_text_config,
+        if self.v_head_dim is None:
-            "swa_head_dim",
+            self.v_head_dim = self.head_dim
-            self.head_dim,
+            setattr(self.hf_text_config, "v_head_dim", self.v_head_dim)
-        )
+
-        self.swa_v_head_dim = getattr(
+        self.swa_head_dim = getattr(self.hf_text_config, "swa_head_dim", None)
-            self.hf_text_config,
+        if self.swa_head_dim is None:
-            "swa_v_head_dim",
+            self.swa_head_dim = self.head_dim
-            self.swa_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

@@ -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

@@ -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.model_runner.max_total_num_tokens,
-            self.max_prefill_tokens,
+            self.server_args.max_prefill_tokens,
-            self.max_running_requests,
+            self.model_runner.max_running_requests,
-            self.max_queued_requests,
+            self.server_args.max_queued_requests,
-            self.max_req_len,
+            max_req_len,
-            self.max_req_input_len,
+            max_req_len - 5,
             self.random_seed,
             self.device,
             self.model_runner.forward_stream,

python/sglang/srt/model_executor/model_runner.py

@@ -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:
+        if (
-            # load draft config
+            (self.spec_algorithm.is_eagle() or self.spec_algorithm.is_standalone())
-            draft_model_config = ModelConfig.from_server_args(
+            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 (
+            from sglang.srt.speculative.dflash_utils import parse_dflash_draft_config
-                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
+        # Get available memory before model loading.
-        pre_model_load_memory = self.init_torch_distributed()
+        # 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
+        # Load model weights and configure
-        self.initialize(pre_model_load_memory)
+        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
+        # Snapshot free memory at the end of the weight-load phase. KV-pool
-        self.init_memory_pool(pre_model_load_memory)
+        # 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 (
+        from sglang.srt.layers.flashinfer_comm_fusion import pre_initialize_workspaces
-            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

@@ -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

@@ -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

@@ -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

@@ -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
+        mtp_config = copy.deepcopy(config)
-        config.full_attention_interval = 1
+        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

@@ -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
+        mtp_config = copy.deepcopy(config)
-        config.full_attention_interval = 1
+        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

@@ -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

@@ -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

@@ -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

@@ -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

@@ -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,
+        # Load draft model weights only.
-        # 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
         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.tree_mask_mode = TreeMaskMode.FULL_MASK
-            self.draft_tp_context(self.draft_runner.tp_group),
+
-            speculative_moe_backend_context(),
+        self.plan_stream, self.plan_stream_ctx = _get_plan_stream(self.device)
-            speculative_moe_a2a_backend_context(),
+
+    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

@@ -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 (
+from sglang.srt.model_executor.cuda_graph_config import cuda_graph_fully_disabled
-    Backend,
-    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_attn_backend = None
-        self.draft_model_runner.draft_attn_backend = self.draft_attn_backend
         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

@@ -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,
+        # Load draft model weights only.
-        # 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()
-        )
         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
+    def init_backends(self):
-        for i in range(self.speculative_num_steps):
+        with self.draft_tp_context(
-            self.draft_runner_list[i].server_args.cuda_graph_config.decode.backend = (
+            self.draft_runner_list[0].tp_group
-                backup_decode_mode
+        ), speculative_moe_backend_context():
-            )
+            super().init_backends()
-        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 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

@@ -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 = (
+        # req_to_token_pool / token_to_kv_pool_allocator are set in
-            target_worker.get_memory_pool()
+        # alloc_memory_pool(), after the target pools are allocated.
-        )
-
-        self.max_batch_size = target_worker.max_running_requests
         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

@@ -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,
+        # Load draft model weights only.
-        # 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()
-        )
         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

@@ -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

@@ -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

@@ -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()