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

[Spec] Frozen-KV MTP: delay target KV binding to pool init + reset stale draft out_cache_loc (#29616)

Browse Source 
Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Co-authored-by: bef0rewind <490285+ronhuafeng@users.noreply.github.com>
Co-authored-by: AlejandroParedesLT <99832715+AlejandroParedesLT@users.noreply.github.com>
Co-authored-by: Cursor <cursoragent@cursor.com>
This commit is contained in:
Khoa Pham
2026-06-29 11:57:10 -07:00
committed by GitHub
co-authored by Claude Opus 4.8 bef0rewind AlejandroParedesLT Cursor
parent b0be644133
commit f480c5f1f9
2 changed files with 21 additions and 28 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/managers/scheduler.py
-3
View File
@@ -850,9 +850,6 @@ class Scheduler(
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.
python/sglang/srt/speculative/frozen_kv_mtp_worker_v2.py
+21 -25
View File
@@ -114,16 +114,11 @@ class FrozenKVMTPDraftWorker(EagleDraftWorkerBase, TpModelWorker):
f"{self.speculative_algorithm.name}."
)
# 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
self.draft_pool_config = MemoryPoolConfig(
max_total_num_tokens=64, # Dummy value
max_running_requests=target_cfg.max_running_requests,
)
# Target pools (read-only) are bound in alloc_memory_pool(), not here, so
# the worker can be built before the target pool exists (see #29021).
self.req_to_token_pool = None
self.token_to_kv_pool_allocator = None
self.draft_pool_config: Optional[MemoryPoolConfig] = None
self.hot_token_id = None
@@ -144,9 +139,6 @@ class FrozenKVMTPDraftWorker(EagleDraftWorkerBase, TpModelWorker):
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=self.draft_pool_config,
)
embed, head = self.target_worker.model_runner.model.get_embed_and_head()
@@ -160,8 +152,6 @@ class FrozenKVMTPDraftWorker(EagleDraftWorkerBase, TpModelWorker):
)
self.kv_context: Optional[FrozenKVMTPContext] = None
if hasattr(self.draft_model_runner.model, "bind_frozen_kv_context"):
self._bind_kv_context()
self.draft_tp_context = (
draft_tp_context if server_args.enable_dp_attention else empty_context
@@ -180,23 +170,26 @@ class FrozenKVMTPDraftWorker(EagleDraftWorkerBase, TpModelWorker):
req_to_token_pool=None,
token_to_kv_pool_allocator=None,
):
self.req_to_token_pool = req_to_token_pool
self.token_to_kv_pool_allocator = token_to_kv_pool_allocator
self.draft_pool_config = MemoryPoolConfig(
max_total_num_tokens=64, # Dummy value
max_running_requests=memory_pool_config.max_running_requests,
)
# NOTE: call TpModelWorker explicitly -- EagleDraftWorkerBase 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
),
req_to_token_pool=req_to_token_pool,
token_to_kv_pool_allocator=token_to_kv_pool_allocator,
)
if hasattr(self.draft_model_runner.model, "bind_frozen_kv_context"):
self._bind_kv_context()
def init_attention_backends(self):
with (
self.draft_tp_context(self.draft_model_runner.tp_group),
@@ -442,6 +435,9 @@ class FrozenKVMTPDraftWorker(EagleDraftWorkerBase, TpModelWorker):
self._set_positions(forward_batch)
self._expand_for_topk_draft(forward_batch)
# Frozen draft never writes KV; None signals fill_from to skip the slot.
forward_batch.out_cache_loc = None
can_run_cuda_graph = (
self.cuda_graph_runner
and self.cuda_graph_runner.can_run_graph(forward_batch)