diff --git a/python/sglang/srt/managers/scheduler.py b/python/sglang/srt/managers/scheduler.py index 9426c6671..358f30a5d 100644 --- a/python/sglang/srt/managers/scheduler.py +++ b/python/sglang/srt/managers/scheduler.py @@ -1124,16 +1124,11 @@ class Scheduler( ) # todo: should we fix this when enabling mtp or it doesn't matter since we only enable mtp in decode node thus we don't transfer draft kvs between P and D? - draft_token_to_kv_pool, model_config = kv_cache_builder.get_draft_kv_pool( + draft_token_to_kv_pool = kv_cache_builder.get_draft_kv_pool( draft_worker=self.draft_worker, spec_algorithm=self.spec_algorithm, server_args=self.server_args, ) - # Default to the target model_config so the MetadataBuffers branches - # below can always access it; overridden by the draft model_config - # when this node runs a spec module. - if model_config is None: - model_config = self.model_config if self.spec_algorithm.carries_draft_hidden_states(): # `draft_runner` aliases `draft_runner_list[0]` in the multi-layer diff --git a/python/sglang/srt/mem_cache/kv_cache_builder.py b/python/sglang/srt/mem_cache/kv_cache_builder.py index 70785d35e..5788a67bc 100644 --- a/python/sglang/srt/mem_cache/kv_cache_builder.py +++ b/python/sglang/srt/mem_cache/kv_cache_builder.py @@ -52,17 +52,17 @@ def get_draft_kv_pool( spec_algorithm: SpeculativeAlgorithm, server_args: ServerArgs, ): - """Return (draft_token_to_kv_pool, draft_model_config) for the current - draft worker, or (None, None) when no draft KV pool is available.""" + """Return the draft token-to-KV pool for the current draft worker, + or None when no draft KV pool is available.""" if draft_worker is None or spec_algorithm.is_ngram(): - return None, None + return None # V2 workers nest the draft runner under `.draft_worker`. if server_args.enable_multi_layer_eagle: draft_runner = draft_worker.draft_worker.draft_runner_list[0] else: draft_runner = draft_worker.draft_worker.draft_runner - return draft_runner.token_to_kv_pool, draft_runner.model_config + return draft_runner.token_to_kv_pool def maybe_register_hicache_draft( @@ -78,7 +78,7 @@ def maybe_register_hicache_draft( if not enable_hierarchical_cache: return - draft_kv_pool, _ = get_draft_kv_pool( + draft_kv_pool = get_draft_kv_pool( draft_worker=draft_worker, spec_algorithm=spec_algorithm, server_args=server_args, diff --git a/python/sglang/srt/model_executor/runner/prefill_cuda_graph_runner.py b/python/sglang/srt/model_executor/runner/prefill_cuda_graph_runner.py index 0a9347609..2574cd297 100644 --- a/python/sglang/srt/model_executor/runner/prefill_cuda_graph_runner.py +++ b/python/sglang/srt/model_executor/runner/prefill_cuda_graph_runner.py @@ -27,6 +27,7 @@ Backend selection comes from cuda_graph_config.prefill: from __future__ import annotations +import inspect import logging import warnings from typing import TYPE_CHECKING, Dict, Optional, Union @@ -191,13 +192,6 @@ class PrefillCudaGraphRunner(BaseCudaGraphRunner): hidden_size=self.model_runner.model_config.hidden_size, embed_dtype=self.model_runner.dtype, enable_mamba_track=self.mamba_track_enabled, - # Register the multimodal input_embeds slot for every prefill - # backend (default True). The slot is only added when is_multimodal, - # so text-only models are unaffected. Both tc_piecewise (outer MM - # wrapper passes composed input_embeds as an argument) and breakable - # (captures the input_embeds path; general_mm_embed_routine fills the - # slot) need it, otherwise the captured graph re-embeds input_ids and - # drops the scattered vision embeddings. source=self.buffers, ) @@ -294,6 +288,10 @@ class PrefillCudaGraphRunner(BaseCudaGraphRunner): language_model.model, "layers" ): self.layer_model = language_model.model + params = list(inspect.signature(self.layer_model.forward).parameters) + self._input_embeds_arg_idx = ( + params.index("input_embeds") if "input_embeds" in params else None + ) else: raise RuntimeError( f"BCG could not resolve inner layer_model on " @@ -880,6 +878,8 @@ class PrefillCudaGraphRunner(BaseCudaGraphRunner): shape_key = ShapeKey(size=self._static_num_tokens) static_n = self._static_num_tokens + ie_idx = self._input_embeds_arg_idx + def replay_layer_forward(*args, **layer_kwargs): # The captured BCG graph reads activations from the static # input_embeds slot. The outer model.forward (run eagerly) @@ -891,8 +891,8 @@ class PrefillCudaGraphRunner(BaseCudaGraphRunner): # current request's embeddings (mirrors main's BCG closure). if self.buffer_registry.has_slot("input_embeds"): ie = layer_kwargs.get("input_embeds") - if ie is None and len(args) > 3: - ie = args[3] + if ie is None and ie_idx is not None and len(args) > ie_idx: + ie = args[ie_idx] if ie is not None: self.buffer_registry.get_slot("input_embeds").slice_for( 1, static_n