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

model: support Step-3.5-Flash (#18084)

Browse Source 
Co-authored-by: ltd0924 <ltd0924@sina.com>
This commit is contained in:
Yuhao Yang
2026-02-03 00:40:07 +08:00
committed by GitHub
parent c781db0f6c
commit 980d2936cd
15 changed files with 1557 additions and 12 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
python/sglang/srt/configs/__init__.py
View File

@@ -24,6 +24,7 @@ from sglang.srt.configs.step3_vl import (
Step3VisionEncoderConfig,
Step3VLConfig,
)
from sglang.srt.configs.step3p5 import Step3p5Config
__all__ = [
"AfmoeConfig",
@@ -50,4 +51,5 @@ __all__ = [
"NemotronH_Nano_VL_V2_Config",
"JetNemotronConfig",
"JetVLMConfig",
"Step3p5Config",
]

24
python/sglang/srt/configs/model_config.py
View File

@@ -302,6 +302,8 @@ class ModelConfig:
and self.hf_config.architectures[0] == "MiMoV2FlashForCausalLM"
):
self.hf_config.architectures[0] = "MiMoV2MTP"
if is_draft_model and self.hf_config.architectures[0] == "Step3p5ForCausalLM":
self.hf_config.architectures[0] = "Step3p5MTP"
if is_draft_model and self.hf_config.architectures[0] in [
"BailingMoeV2ForCausalLM",
"BailingMoeForCausalLM",
@@ -606,6 +608,11 @@ class ModelConfig:
if hasattr(self.hf_text_config, "swa_num_key_value_heads"):
total_num_kv_heads = self.hf_text_config.swa_num_key_value_heads
return max(1, total_num_kv_heads // tensor_parallel_size)
elif hasattr(self.hf_text_config, "attention_other_setting"): # For step3p5
total_num_kv_heads = self.hf_text_config.attention_other_setting.get(
"num_attention_groups"
)
return max(1, total_num_kv_heads // tensor_parallel_size)
else:
return self.get_num_kv_heads(tensor_parallel_size)
@@ -1268,6 +1275,8 @@ def is_hybrid_swa_model(model_architectures: List[str]):
"GptOssForCausalLM",
"MiMoV2FlashForCausalLM",
"MiMoV2MTP",
"Step3p5ForCausalLM",
"Step3p5MTP",
}
return any(arch in hybrid_swa_archs for arch in model_architectures)
@@ -1303,6 +1312,21 @@ def get_hybrid_layer_ids(
elif "MiMoV2MTP" in model_architectures:
swa_attention_layer_ids = [0]
full_attention_layer_ids = []
elif "Step3p5ForCausalLM" in model_architectures:
layer_types = hf_text_config.layer_types
swa_attention_layer_ids = [
i
for i, x in enumerate(layer_types)
if x == "sliding_attention" and i < num_hidden_layers
]
full_attention_layer_ids = [
i
for i, x in enumerate(layer_types)
if x == "full_attention" and i < num_hidden_layers
]
elif "Step3p5MTP" in model_architectures:
swa_attention_layer_ids = [0]
full_attention_layer_ids = []
else:
swa_attention_layer_ids = None
full_attention_layer_ids = None

97
python/sglang/srt/configs/step3p5.py Normal file
View File

@@ -0,0 +1,97 @@
from typing import Any, Optional
from transformers.configuration_utils import PretrainedConfig
class Step3p5Config(PretrainedConfig):
model_type = "step3p5"
architectures = ["Step3p5ForCausalLM"]
def __init__(
self,
hidden_size: int = 4096,
intermediate_size: int = 11264,
num_attention_heads: int = 64,
num_attention_groups: int = 8,
num_hidden_layers: int = 45,
max_seq_len: int = 128000,
vocab_size: int = 128815,
rms_norm_eps: float = 1e-5,
moe_intermediate_size: int = 1280,
moe_num_experts: int = 288,
moe_top_k: int = 8,
rope_theta: float = 10000,
rope_scaling: Optional[dict[str, Any]] = None,
max_position_embeddings: int = 128000,
share_expert_dims: int = 1280,
head_dim: int = 128,
norm_expert_weight: bool = True,
layer_types: list[str] = None,
sliding_window: Optional[int] = None,
moe_layers_enum: tuple[int] = (
3,
4,
5,
6,
7,
8,
9,
10,
11,
12,
13,
14,
15,
16,
17,
18,
19,
20,
21,
22,
23,
24,
25,
26,
27,
28,
29,
30,
31,
32,
33,
34,
35,
36,
37,
38,
39,
40,
41,
42,
43,
44,
),
**kwargs,
) -> None:
self.hidden_size = hidden_size
self.intermediate_size = intermediate_size
self.num_attention_heads = num_attention_heads
self.num_attention_groups = num_attention_groups
self.num_hidden_layers = num_hidden_layers
self.max_seq_len = max_seq_len
self.vocab_size = vocab_size
self.rms_norm_eps = rms_norm_eps
self.moe_intermediate_size = moe_intermediate_size
self.moe_num_experts = moe_num_experts
self.moe_top_k = moe_top_k
self.rope_theta = rope_theta
self.rope_scaling = rope_scaling
self.max_position_embeddings = max_position_embeddings
self.share_expert_dim = share_expert_dims
self.head_dim = head_dim
self.norm_expert_weight = norm_expert_weight
self.moe_layers_enum = moe_layers_enum
self.layer_types = layer_types
self.sliding_window = sliding_window
super().__init__(**kwargs)

1
python/sglang/srt/function_call/function_call_parser.py
View File

@@ -62,6 +62,7 @@ class FunctionCallParser:
"qwen25": Qwen25Detector,
"qwen3_coder": Qwen3CoderDetector,
"step3": Step3Detector,
"step3p5": Qwen3CoderDetector,
"minimax-m2": MinimaxM2Detector,
"trinity": TrinityDetector,
"interns1": InternlmDetector,

25
python/sglang/srt/layers/moe/fused_moe_triton/fused_moe.py
View File

@@ -278,7 +278,18 @@ def moe_sum_reduce_torch_compile(x, out, routed_scaling_factor):
@torch.compile
def swiglu_with_alpha_and_limit(x, gemm1_alpha, gemm1_limit):
def _swiglu_silu_clamp_mul(x, gemm1_limit):
gate, up = x.chunk(2, dim=-1)
gate = F.silu(gate)
gate = gate.clamp(min=None, max=gemm1_limit)
up = up.clamp(min=-gemm1_limit, max=gemm1_limit)
return gate * up
@torch.compile
def _swiglu_gpt_oss_sigmoid_alpha(x, gemm1_alpha, gemm1_limit):
# NOTE: This variant uses gemm1_alpha, unlike _swiglu_silu_clamp_mul.
# At present, only GPT-OSS uses this variant.
gate, up = x[..., ::2], x[..., 1::2]
gate = gate.clamp(min=None, max=gemm1_limit)
up = up.clamp(min=-gemm1_limit, max=gemm1_limit)
@@ -471,12 +482,16 @@ def fused_experts_impl(
# Activation function with multiplication
if activation == "silu" and is_gated:
# - gemm1_alpha != None: GPT-OSS-style swiglu(alpha, limit)
# - gemm1_alpha == None and gemm1_limit != None: silu+clamp+mul(limit-only)
if gemm1_alpha is not None:
assert gemm1_limit is not None
intermediate_cache2 = swiglu_with_alpha_and_limit(
intermediate_cache1.view(-1, N),
gemm1_alpha,
gemm1_limit,
intermediate_cache2 = _swiglu_gpt_oss_sigmoid_alpha(
intermediate_cache1.view(-1, N), gemm1_alpha, gemm1_limit
)
elif gemm1_limit is not None:
intermediate_cache2 = _swiglu_silu_clamp_mul(
intermediate_cache1.view(-1, N), gemm1_limit
)
elif _is_cuda or _is_hip:
if not filter_expert:

13
python/sglang/srt/layers/moe/moe_runner/triton.py
View File

@@ -117,10 +117,11 @@ class TritonRunnerCore(MoeRunnerCore):
# TODO: move these functions to the triton runner
from sglang.srt.layers.moe.fused_moe_triton.fused_moe import (
_swiglu_gpt_oss_sigmoid_alpha,
_swiglu_silu_clamp_mul,
invoke_fused_moe_kernel,
moe_sum_reduce_torch_compile,
moe_sum_reduce_triton,
swiglu_with_alpha_and_limit,
)
hidden_states = runner_input.hidden_states
@@ -203,10 +204,12 @@ class TritonRunnerCore(MoeRunnerCore):
if activation == "silu":
if gemm1_alpha is not None:
assert gemm1_limit is not None
intermediate_cache2 = swiglu_with_alpha_and_limit(
intermediate_cache1.view(-1, N),
gemm1_alpha,
gemm1_limit,
intermediate_cache2 = _swiglu_gpt_oss_sigmoid_alpha(
intermediate_cache1.view(-1, N), gemm1_alpha, gemm1_limit
)
elif gemm1_limit is not None:
intermediate_cache2 = _swiglu_silu_clamp_mul(
intermediate_cache1.view(-1, N), gemm1_limit
)
elif _is_cuda or _is_hip:
silu_and_mul(intermediate_cache1.view(-1, N), intermediate_cache2)

2
python/sglang/srt/model_executor/model_runner.py
View File

@@ -493,6 +493,8 @@ class ModelRunner(ModelRunnerKVCacheMixin):
)
if self.model_config.hf_config.architectures[0] == "MiMoV2MTP":
model_num_layers = 1
elif self.model_config.hf_config.architectures[0] == "Step3p5MTP":
model_num_layers = 1
self.start_layer = getattr(self.model, "start_layer", 0)
self.end_layer = getattr(self.model, "end_layer", model_num_layers)
self.num_effective_layers = self.end_layer - self.start_layer

3
python/sglang/srt/model_loader/loader.py
View File

@@ -275,6 +275,9 @@ def _initialize_model(
kwargs["sparse_head"] = envs.SGLANG_EMBEDDINGS_SPARSE_HEAD.get()
kwargs["model_path"] = model_config.model_path
if load_config.draft_model_idx is not None:
kwargs["draft_model_idx"] = load_config.draft_model_idx
return model_class(**kwargs)

1
python/sglang/srt/models/mimo_v2_flash_nextn.py
View File

@@ -229,6 +229,7 @@ class MiMoV2MTP(MiMoV2FlashForCausalLM):
self,
config: PretrainedConfig,
quant_config: Optional[QuantizationConfig] = None,
draft_model_idx: Optional[int] = None,
prefix: str = "",
) -> None:
nn.Module.__init__(self)

1037
python/sglang/srt/models/step3p5.py Normal file
View File

File diff suppressed because it is too large Load Diff

336
python/sglang/srt/models/step3p5_mtp.py Normal file
View File

@@ -0,0 +1,336 @@
import logging
from collections.abc import Iterable
from typing import Optional
import torch
import torch.nn as nn
from transformers import PretrainedConfig
from sglang.srt.distributed import get_tensor_model_parallel_world_size
from sglang.srt.layers.layernorm import GemmaRMSNorm
from sglang.srt.layers.logits_processor import LogitsProcessor
from sglang.srt.layers.quantization.base_config import QuantizationConfig
from sglang.srt.layers.vocab_parallel_embedding import (
ParallelLMHead,
VocabParallelEmbedding,
)
from sglang.srt.model_executor.forward_batch_info import ForwardBatch
from sglang.srt.model_loader.weight_utils import default_weight_loader
from sglang.srt.models.step3p5 import Step3p5DecoderLayer, Step3p5ForCausalLM
from sglang.srt.utils import add_prefix
logger = logging.getLogger(__name__)
def get_spec_layer_idx_from_weight_name(
config: PretrainedConfig, weight_name: str
) -> Optional[int]:
"""Return MTP/nextn layer index if this weight belongs to spec layers.
Step3p5 MTP/nextn checkpoints append extra layers after the main decoder:
model.layers.[num_hidden_layers ... num_hidden_layers + num_nextn_predict_layers)
"""
if hasattr(config, "num_nextn_predict_layers") and (
getattr(config, "num_nextn_predict_layers", 0) > 0
):
base = config.num_hidden_layers
for i in range(config.num_nextn_predict_layers):
if weight_name.startswith(f"model.layers.{base + i}."):
return base + i
return None
class SharedHead(nn.Module):
def __init__(
self,
config,
quant_config=None,
) -> None:
super().__init__()
self.norm = GemmaRMSNorm(config.hidden_size, config.rms_norm_eps)
self.head = ParallelLMHead(
config.vocab_size, config.hidden_size, quant_config=quant_config
)
self.lm_head = self.head
def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
return self.norm(hidden_states)
class Step3p5AMultiTokenPredictor(nn.Module):
def __init__(
self,
config: PretrainedConfig,
quant_config: Optional[QuantizationConfig] = None,
prefix: str = "",
) -> None:
super().__init__()
self.config = config
self.embed_tokens = VocabParallelEmbedding(
config.vocab_size,
config.hidden_size,
)
self.mtp_start_layer_idx = config.num_hidden_layers
self.num_mtp_layers = config.num_nextn_predict_layers
layer_id = 45 # FIXME
self.enorm = GemmaRMSNorm(config.hidden_size, config.rms_norm_eps)
self.hnorm = GemmaRMSNorm(config.hidden_size, config.rms_norm_eps)
self.eh_proj = nn.Linear(config.hidden_size * 2, config.hidden_size, bias=False)
self.shared_head = SharedHead(config=config, quant_config=quant_config)
self.mtp_block = Step3p5DecoderLayer(
config=config, layer_id=layer_id, prefix=f"{prefix}.mtp_block"
)
self.lm_head = self.shared_head.head
def forward(
self,
input_ids: torch.Tensor,
positions: torch.Tensor,
forward_batch: ForwardBatch,
input_embeds: torch.Tensor = None,
) -> torch.Tensor:
if input_embeds is None:
hidden_states = self.embed_tokens(input_ids)
else:
hidden_states = input_embeds
if hidden_states.shape[0] > 0:
hidden_states = self.eh_proj(
torch.cat(
(
self.enorm(hidden_states),
self.hnorm(forward_batch.spec_info.hidden_states),
),
dim=-1,
)
)
hidden_states, residual = self.mtp_block(
positions=positions,
hidden_states=hidden_states,
forward_batch=forward_batch,
residual=None,
)
hidden_states_before_norm = None
if not forward_batch.forward_mode.is_idle():
# if forward_batch.return_hidden_states_before_norm:
hidden_states_before_norm = (
hidden_states if residual is None else hidden_states + residual
)
if residual is not None:
hidden_states, _ = self.shared_head.norm(hidden_states, residual)
else:
hidden_states = self.shared_head.norm(hidden_states)
return hidden_states, hidden_states_before_norm
def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:
return self.embed_tokens(input_ids)
# The current implementation differs slightly from the standard MTP implementation in Step3.5 Flash.
# In the standard multi-layer MTP design of Step3.5 Flash,
# the hidden states of each MTP layer are passed from the preceding MTP layer
# (the hidden states of the initial (layer-0) MTP still being provided by the target model).
# In contrast, the current SGL implementation obtains hidden states directly from the target model for all MTP layers.
# Empirical evaluations indicate that the overall performance remains strong;
# however, this design choice may lead to a slight reduction in acceptance rate in certain scenarios.
# This behavior will be corrected shortly, and we expect to implement the standard multi-layer MTP design of Step3.5 Flash in the near future.
# FIXME(yhyang201)
class Step3p5MTP(Step3p5ForCausalLM):
def __init__(
self,
config: PretrainedConfig,
quant_config: Optional[QuantizationConfig] = None,
draft_model_idx: Optional[int] = None,
prefix: str = "",
) -> None:
nn.Module.__init__(self)
self.config = config
self.tp_size = get_tensor_model_parallel_world_size()
self.quant_config = quant_config
self.draft_model_idx = draft_model_idx
self.model = Step3p5AMultiTokenPredictor(
config=config, quant_config=quant_config, prefix=add_prefix("model", prefix)
)
self.logits_processor = LogitsProcessor(config)
self.lm_head = self.model.lm_head
def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:
return self.model.embed_input_ids(input_ids)
def forward(
self,
input_ids: torch.Tensor,
positions: torch.Tensor,
forward_batch: ForwardBatch,
) -> torch.Tensor:
hidden_states, hidden_states_before_norm = self.model(
input_ids, positions, forward_batch
)
return self.logits_processor(
input_ids,
hidden_states,
self.model.shared_head.head,
forward_batch,
hidden_states_before_norm=hidden_states_before_norm,
)
def get_embed_and_head(self):
return self.model.embed_tokens.weight, self.model.shared_head.head.weight
def set_embed_and_head(self, embed, head):
return
def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
stacked_params_mapping = [
# (param_name, shard_name, shard_id)
("qkv_proj", "q_proj", "q"),
("qkv_proj", "k_proj", "k"),
("qkv_proj", "v_proj", "v"),
("gate_up_proj", "gate_proj", 0),
("gate_up_proj", "up_proj", 1),
]
expert_params_mapping = [
(".moe.experts.w13_weight", ".moe.gate_proj.weight", "w1"),
(".moe.experts.w13_weight", ".moe.up_proj.weight", "w3"),
(".moe.experts.w2_weight", ".moe.down_proj.weight", "w2"),
]
params_dict = dict(self.named_parameters())
loaded_params: set[str] = set()
for name, loaded_weight in weights:
if "rotary_emb.inv_freq" in name:
continue
spec_layer = get_spec_layer_idx_from_weight_name(self.config, name)
if spec_layer is not None and spec_layer != (
self.config.num_hidden_layers + self.draft_model_idx
):
continue
if "embed_tokens" not in name and spec_layer is None:
continue
name = self._rewrite_spec_layer_name(spec_layer, name)
for param_name, weight_name, shard_id in stacked_params_mapping:
# Skip non-stacked layers and experts (experts handled below).
if weight_name not in name:
continue
# We have mlp.experts[0].gate_proj in the checkpoint.
# Since we handle the experts below in expert_params_mapping,
# we need to skip here BEFORE we update the name, otherwise
# name will be updated to mlp.experts[0].gate_up_proj, which
# will then be updated below in expert_params_mapping
# for mlp.experts[0].gate_gate_up_proj, which breaks load.
if ("mlp.experts." in name) and name not in params_dict:
continue
if "experts" in name or "moe" in name:
continue
name = name.replace(weight_name, param_name)
# Skip loading extra bias for GPTQ models.
if name.endswith(".bias") and name not in params_dict:
continue
param = params_dict[name]
weight_loader = param.weight_loader
weight_loader(param, loaded_weight, shard_id)
break
else:
for mapping in expert_params_mapping:
param_name, weight_name, shard_id = mapping
if weight_name not in name:
continue
name = name.replace(weight_name, param_name)
# Skip loading extra bias for GPTQ models.
if (
name.endswith(".bias") or name.endswith("_bias")
) and name not in params_dict:
continue
param = params_dict[name]
weight_loader = param.weight_loader
for expert_id in range(loaded_weight.shape[0]):
loaded_weight_expert = loaded_weight[expert_id]
weight_loader(
param,
loaded_weight_expert,
name,
shard_id=shard_id,
expert_id=expert_id,
)
loaded_params.add(name)
break
else:
# Skip loading extra bias for GPTQ models.
if (
name.endswith(".bias")
and name not in params_dict
or "tok_embeddings" in name
):
continue
if "shared_head" in name:
name = name.replace("shared_head.output", "shared_head.head")
if "embed_tokens" in name:
assert (
hasattr(self.config, "num_nextn_predict_layers")
and self.config.num_nextn_predict_layers > 0
)
name = "model.embed_tokens.weight"
param = params_dict[name]
weight_loader = getattr(
param, "weight_loader", default_weight_loader
)
weight_loader(param, loaded_weight)
loaded_params.add(name)
params_need_to_load = set(params_dict.keys())
if params_need_to_load != loaded_params:
missing_params = list(params_need_to_load - loaded_params)
param_name_example = missing_params[0]
raise RuntimeError(
f"Some parameters like {param_name_example} are not in the checkpoint and will falsely use random initialization"
)
return loaded_params
def _rewrite_spec_layer_name(self, spec_layer: Optional[int], name: str) -> str:
"""
Rewrite the weight name to match the format of the original model.
Add .mtp_block for modules in transformer layer block for spec layer
"""
if spec_layer is None:
return name
# Some checkpoints place MTP weights under "model.layers.<id>.transformer.*".
# Our modules use "model.layers.<id>.*", so drop the ".transformer." segment.
transformer_prefix = f"model.layers.{spec_layer}.transformer."
if name.startswith(transformer_prefix):
name = name.replace(".transformer.", ".", 1)
spec_layer_weight_names = [
"embed_tokens",
"enorm",
"hnorm",
"eh_proj",
"shared_head",
]
spec_layer_weight = False
for weight_name in spec_layer_weight_names:
if weight_name in name:
spec_layer_weight = True
break
if not spec_layer_weight:
# treat rest weights as weights for transformer layer block
name = name.replace(
f"model.layers.{spec_layer}.", f"model.layers.{spec_layer}.mtp_block."
)
# NEW: drop "layers.<idx>." from the rewritten name (minimal change).
layers_prefix = f"model.layers.{spec_layer}."
if name.startswith(layers_prefix):
name = name.replace(layers_prefix, "model.", 1)
return name
EntryClass = [Step3p5MTP]

1
python/sglang/srt/parser/reasoning_parser.py
View File

@@ -384,6 +384,7 @@ class ReasoningParser:
"minimax": Qwen3Detector,
"minimax-append-think": MiniMaxAppendThinkDetector,
"step3": DeepSeekR1Detector,
"step3p5": DeepSeekR1Detector,
"nano_v3": NanoV3Detector,
"interns1": Qwen3Detector,
}

20
python/sglang/srt/server_args.py
View File

@@ -1406,6 +1406,26 @@ class ServerArgs:
logger.warning(
"Disable hybrid SWA memory for MiMoV2FlashForCausalLM model with hierarchical cache"
)
elif "Step3p5ForCausalLM" in model_arch:
if self.speculative_algorithm == "EAGLE":
self.enable_multi_layer_eagle = True
logger.info(
"Enable multi-layer EAGLE speculative decoding for Step3p5ForCausalLM model."
)
if not envs.SGLANG_ENABLE_SPEC_V2.get():
envs.SGLANG_ENABLE_SPEC_V2.set(True)
logger.warning(
"Spec v2 is enabled for multi-layer EAGLE speculative decoding."
)
if self.enable_hierarchical_cache:
self.swa_full_tokens_ratio = 1.0
logger.warning(
"Reset swa_full_tokens_ratio to 1.0 for Step3p5ForCausalLM model with hierarchical cache"
)
self.disable_hybrid_swa_memory = True
logger.warning(
"Disable hybrid SWA memory for Step3p5ForCausalLM model with hierarchical cache"
)
elif "Llama4" in model_arch and self.device != "cpu":
# Auto-select attention backend for Llama4 if not specified
if self.attention_backend is None:

5
python/sglang/srt/speculative/multi_layer_eagle_worker_v2.py
View File

@@ -466,11 +466,12 @@ class MultiLayerEagleDraftWorker(BaseDraftWorker):
draft_logits_output.topk_index,
)
else:
draft_logits_output, _ = self.draft_runner_list[step].forward(
draft_logits_output = self.draft_runner_list[step].forward(
forward_batch, skip_attn_backend_init=True
)
probs = torch.softmax(
draft_logits_output.next_token_logits[select_index], dim=-1
draft_logits_output.logits_output.next_token_logits[select_index],
dim=-1,
)
ret_topk_p, ret_topk_index = fast_topk(probs, self.topk, dim=-1)
if forward_batch.extend_seq_lens is not None:

2
python/sglang/srt/utils/hf_transformers_utils.py
View File

@@ -63,6 +63,7 @@ from sglang.srt.configs import (
NemotronHConfig,
Olmo3Config,
Qwen3NextConfig,
Step3p5Config,
Step3VLConfig,
)
from sglang.srt.configs.deepseek_ocr import DeepseekVLV2Config
@@ -95,6 +96,7 @@ _CONFIG_REGISTRY: List[Type[PretrainedConfig]] = [
JetNemotronConfig,
JetVLMConfig,
KimiK25Config,
Step3p5Config,
]
_CONFIG_REGISTRY = {