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

Extract spec aux-hidden-state resolution into a module (#31157)

Browse Source
This commit is contained in:
fzyzcjy
2026-07-14 15:59:35 +08:00
committed by GitHub
parent 39e508b7fc
commit 17c04602c6
3 changed files with 188 additions and 113 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/model_executor/model_runner.py
+20 -110
View File
@@ -140,6 +140,10 @@ from sglang.srt.model_executor.model_runner_components.ngram_embedding_manager i
from sglang.srt.model_executor.model_runner_components.remote_instance_weight_transporter import (
RemoteInstanceWeightTransporter,
)
from sglang.srt.model_executor.model_runner_components.spec_aux_hidden_state import (
SpecAuxHiddenStateConfig,
resolve_spec_aux_hidden_state_config,
)
from sglang.srt.model_executor.model_runner_components.weight_exporter import (
WeightExporter,
)
@@ -317,111 +321,7 @@ class ModelRunner(ModelRunnerKVCacheMixin):
self.init_msprobe()
# 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_family_use_aux_hidden_state = False
self.dflash_family_target_layer_ids = None
self.dflash_family_draft_num_layers = None
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 = ModelConfig.from_server_args(
server_args,
model_path=server_args.speculative_draft_model_path,
model_revision=server_args.speculative_draft_model_revision,
is_draft_model=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:
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_family() and not self.is_draft_worker:
from sglang.srt.speculative.dflash_utils import parse_dflash_draft_config
# Select target layers to capture for building draft context features.
draft_model_config = ModelConfig.from_server_args(
server_args,
model_path=(server_args.speculative_draft_model_path),
model_revision=server_args.speculative_draft_model_revision,
is_draft_model=True,
)
dflash_draft_config = parse_dflash_draft_config(
draft_hf_config=draft_model_config.hf_config
)
draft_num_layers = dflash_draft_config.require_num_layers()
trained_target_layers = dflash_draft_config.num_target_layers
target_num_layers = getattr(
self.model_config.hf_text_config, "num_hidden_layers", None
)
if target_num_layers is None:
raise ValueError(
"Block-draft-with-target-kv spec requires target num_hidden_layers "
f"in config. Got target={target_num_layers}."
)
target_num_layers = int(target_num_layers)
if (
trained_target_layers is not None
and trained_target_layers != target_num_layers
):
logger.warning(
"Draft config num_target_layers=%s differs from runtime target num_hidden_layers=%s; "
"selecting capture layers based on the runtime target model.",
trained_target_layers,
target_num_layers,
)
target_layer_ids = dflash_draft_config.resolve_target_layer_ids(
target_num_layers=int(target_num_layers),
draft_num_layers=int(draft_num_layers),
)
if self.spec_algorithm.is_dspark():
from sglang.srt.speculative.dspark_components.dspark_config import (
parse_dspark_draft_config,
)
dspark_draft_config = parse_dspark_draft_config(
draft_hf_config=draft_model_config.hf_config
)
if not dspark_draft_config.require_markov():
raise ValueError(
"DSPARK requires markov_rank > 0 in the draft config, "
f"got markov_rank={dspark_draft_config.markov_rank}."
)
if dspark_draft_config.target_layer_ids is not None:
target_layer_ids = list(dspark_draft_config.target_layer_ids)
self.dflash_family_use_aux_hidden_state = True
self.dflash_family_draft_num_layers = int(draft_num_layers)
self.dflash_family_target_layer_ids = target_layer_ids
self.init_spec_aux_hidden_state()
# Apply the rank zero filter to logger
if server_args.show_time_cost:
@@ -538,6 +438,16 @@ class ModelRunner(ModelRunnerKVCacheMixin):
get_model_runner=lambda: self,
)
def init_spec_aux_hidden_state(self):
self.spec_aux_config: SpecAuxHiddenStateConfig = (
resolve_spec_aux_hidden_state_config(
server_args=self.server_args,
model_config=self.model_config,
spec_algorithm=self.spec_algorithm,
is_draft_worker=self.is_draft_worker,
)
)
def init_weight_exporter(self):
self.weight_exporter = WeightExporter(
tp_rank=self.tp_rank,
@@ -921,20 +831,20 @@ class ModelRunner(ModelRunnerKVCacheMixin):
Must be called before CUDA graph capture so the captured graphs
include aux hidden state output paths.
"""
if self.eagle_use_aux_hidden_state:
if self.spec_aux_config.eagle_use_aux_hidden_state:
self.model.set_eagle3_layers_to_capture(
self.eagle_aux_hidden_state_layer_ids
self.spec_aux_config.eagle_aux_hidden_state_layer_ids
)
if self.dflash_family_use_aux_hidden_state:
if self.spec_aux_config.dflash_use_aux_hidden_state:
if self.spec_algorithm.is_dspark() and hasattr(
self.model, "set_dspark_layers_to_capture"
):
self.model.set_dspark_layers_to_capture(
self.dflash_family_target_layer_ids
self.spec_aux_config.dflash_target_layer_ids
)
elif hasattr(self.model, "set_dflash_layers_to_capture"):
self.model.set_dflash_layers_to_capture(
self.dflash_family_target_layer_ids
self.spec_aux_config.dflash_target_layer_ids
)
else:
raise ValueError(
python/sglang/srt/model_executor/model_runner_components/spec_aux_hidden_state.py
+165
View File
@@ -0,0 +1,165 @@
from __future__ import annotations
import logging
from typing import TYPE_CHECKING, Any, Optional
import msgspec
from sglang.srt.configs.model_config import ModelConfig
if TYPE_CHECKING:
from sglang.srt.server_args import ServerArgs
from sglang.srt.speculative.spec_info import SpeculativeAlgorithm
logger = logging.getLogger(__name__)
class SpecAuxHiddenStateConfig(msgspec.Struct, kw_only=True):
eagle_use_aux_hidden_state: bool = False
eagle_draft_num_layers: Optional[int] = None
eagle_aux_hidden_state_layer_ids: Any = None
dflash_use_aux_hidden_state: bool = False
dflash_draft_num_layers: Optional[int] = None
dflash_target_layer_ids: Any = None
def resolve_spec_aux_hidden_state_config(
*,
server_args: ServerArgs,
model_config: ModelConfig,
spec_algorithm: SpeculativeAlgorithm,
is_draft_worker: bool,
) -> SpecAuxHiddenStateConfig:
config = SpecAuxHiddenStateConfig()
_resolve_eagle_aux_hidden_state(
config=config,
server_args=server_args,
spec_algorithm=spec_algorithm,
is_draft_worker=is_draft_worker,
)
_resolve_dflash_aux_hidden_state(
config=config,
server_args=server_args,
model_config=model_config,
spec_algorithm=spec_algorithm,
is_draft_worker=is_draft_worker,
)
return config
def _resolve_eagle_aux_hidden_state(
*,
config: SpecAuxHiddenStateConfig,
server_args: ServerArgs,
spec_algorithm: SpeculativeAlgorithm,
is_draft_worker: bool,
) -> None:
if (
(spec_algorithm.is_eagle() or spec_algorithm.is_standalone())
and not is_draft_worker
and server_args.speculative_draft_model_path
):
# Load draft config to get layer count for KV cache sizing
draft_model_config = ModelConfig.from_server_args(
server_args,
model_path=server_args.speculative_draft_model_path,
model_revision=server_args.speculative_draft_model_revision,
is_draft_model=True,
)
num_nextn_predict_layers = draft_model_config.num_nextn_predict_layers
if num_nextn_predict_layers is not None:
config.eagle_draft_num_layers = int(num_nextn_predict_layers)
else:
config.eagle_draft_num_layers = int(
max(
draft_model_config.num_hidden_layers,
draft_model_config.num_attention_layers,
)
)
if spec_algorithm.is_eagle3():
config.eagle_use_aux_hidden_state = True
try:
eagle_config = getattr(
draft_model_config.hf_config, "eagle_config", None
)
config.eagle_use_aux_hidden_state = eagle_config.get(
"use_aux_hidden_state", True
)
config.eagle_aux_hidden_state_layer_ids = eagle_config[
"eagle_aux_hidden_state_layer_ids"
]
except:
# if there is no aux layer, set to None
config.eagle_aux_hidden_state_layer_ids = None
def _resolve_dflash_aux_hidden_state(
*,
config: SpecAuxHiddenStateConfig,
server_args: ServerArgs,
model_config: ModelConfig,
spec_algorithm: SpeculativeAlgorithm,
is_draft_worker: bool,
) -> None:
if spec_algorithm.is_dflash_family() and not is_draft_worker:
from sglang.srt.speculative.dflash_utils import parse_dflash_draft_config
# Select target layers to capture for building draft context features.
draft_model_config = ModelConfig.from_server_args(
server_args,
model_path=(server_args.speculative_draft_model_path),
model_revision=server_args.speculative_draft_model_revision,
is_draft_model=True,
)
dflash_draft_config = parse_dflash_draft_config(
draft_hf_config=draft_model_config.hf_config
)
draft_num_layers = dflash_draft_config.require_num_layers()
trained_target_layers = dflash_draft_config.num_target_layers
target_num_layers = getattr(
model_config.hf_text_config, "num_hidden_layers", None
)
if target_num_layers is None:
raise ValueError(
"Block-draft-with-target-kv spec requires target num_hidden_layers "
f"in config. Got target={target_num_layers}."
)
target_num_layers = int(target_num_layers)
if (
trained_target_layers is not None
and trained_target_layers != target_num_layers
):
logger.warning(
"Draft config num_target_layers=%s differs from runtime target num_hidden_layers=%s; "
"selecting capture layers based on the runtime target model.",
trained_target_layers,
target_num_layers,
)
target_layer_ids = dflash_draft_config.resolve_target_layer_ids(
target_num_layers=int(target_num_layers),
draft_num_layers=int(draft_num_layers),
)
if spec_algorithm.is_dspark():
from sglang.srt.speculative.dspark_components.dspark_config import (
parse_dspark_draft_config,
)
dspark_draft_config = parse_dspark_draft_config(
draft_hf_config=draft_model_config.hf_config
)
if not dspark_draft_config.require_markov():
raise ValueError(
"DSPARK requires markov_rank > 0 in the draft config, "
f"got markov_rank={dspark_draft_config.markov_rank}."
)
if dspark_draft_config.target_layer_ids is not None:
target_layer_ids = list(dspark_draft_config.target_layer_ids)
config.dflash_use_aux_hidden_state = True
config.dflash_draft_num_layers = int(draft_num_layers)
config.dflash_target_layer_ids = target_layer_ids
python/sglang/srt/model_executor/pool_configurator.py
+3 -3
View File
@@ -144,7 +144,7 @@ class DefaultPoolConfigurator(MemoryPoolConfigurator):
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)
eagle_draft_num_layers = mr.spec_aux_config.eagle_draft_num_layers
if (
eagle_draft_num_layers is not None
and int(eagle_draft_num_layers) > 0
@@ -161,7 +161,7 @@ class DefaultPoolConfigurator(MemoryPoolConfigurator):
scale_kv_cell_size_per_token_for_dflash,
)
draft_num_layers = mr.dflash_family_draft_num_layers
draft_num_layers = mr.spec_aux_config.dflash_draft_num_layers
if (
draft_num_layers is not None
and int(draft_num_layers) > 0
@@ -332,7 +332,7 @@ class HybridSWAPoolConfigurator(MemoryPoolConfigurator):
if (
mr.spec_algorithm.is_eagle() or mr.spec_algorithm.is_standalone()
) and not mr.is_draft_worker:
draft_layers = getattr(mr, "eagle_draft_num_layers", None)
draft_layers = mr.spec_aux_config.eagle_draft_num_layers
if draft_layers is not None and int(draft_layers) > 0:
self._draft_full_layers_num = int(draft_layers)