@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / c3b40dc3e7

Fork 0

A high-throughput and memory-efficient inference and serving engine for LLMs

Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git

Synced Nov 9, 03:45 UTC

[Models] Step-3.5-Flash (#33523)

Signed-off-by: Jee Jee Li <pandaleefree@gmail.com> Co-authored-by: i-zhangmingming <i-zhangmingming@stepfun.com> Co-authored-by: xiewuxun <xiewuxun@stepfun.com> Co-authored-by: zetaohong <i-hongzetao@stepfun.com> Co-authored-by: Jee Jee Li <pandaleefree@gmail.com>

csy0225 committed Feb 2, 2026 at 10:21 UTC c3b40dc3e74dc0f552f76a01ae38b4f1385ad0af

18 files changed +3107 -4

docs/models/supported_models.md

@@ -456,6 +456,7 @@ th {

458 | `Step1ForCausalLM` | Step-Audio | `stepfun-ai/Step-Audio-EditX`, etc. | ✅︎ | ✅︎ |

460 | `TeleChat2ForCausalLM` | TeleChat2 | `Tele-AI/TeleChat2-3B`, `Tele-AI/TeleChat2-7B`, `Tele-AI/TeleChat2-35B`, etc. | ✅︎ | ✅︎ |

461 | `TeleFLMForCausalLM` | TeleFLM | `CofeAI/FLM-2-52B-Instruct-2407`, `CofeAI/Tele-FLM`, etc. | ✅︎ | ✅︎ |

462 | `XverseForCausalLM` | XVERSE | `xverse/XVERSE-7B-Chat`, `xverse/XVERSE-13B-Chat`, `xverse/XVERSE-65B-Chat`, etc. | ✅︎ | ✅︎ |

tests/kernels/core/test_activation.py

+8 -2

@@ -17,6 +17,8 @@ from vllm.model_executor.layers.activation import (

17 QuickGELU,

18 SiluAndMul,

19 SwigluOAIAndMul,

20 + SwigluStepAndMul,

21 + swiglustep_and_mul_triton,

22 )

23 from vllm.utils.torch_utils import set_random_seed

@@ -36,6 +38,7 @@ CUDA_DEVICES = [f"cuda:{i}" for i in range(1 if torch.cuda.device_count() == 1 e

38 "gelu_tanh",

39 "fatrelu",

40 "swigluoai_and_mul",

41 + "swiglustep_and_mul",

42 ],

43 )

44 @pytest.mark.parametrize("num_tokens", NUM_TOKENS)

@@ -75,9 +78,12 @@ def test_act_and_mul(

78 elif activation == "swigluoai_and_mul":

79 layer = SwigluOAIAndMul()

80 fn = torch.ops._C.swigluoai_and_mul

81 + elif activation == "swiglustep_and_mul":

82 + layer = SwigluStepAndMul()

83 + fn = swiglustep_and_mul_triton

84 out = layer(x)

85 ref_out = layer.forward_native(x)

80 - if activation == "swigluoai_and_mul":

86 + if activation in ["swigluoai_and_mul", "swiglustep_and_mul"]:

87 rtol = {

88 # For fp16, change the relative tolerance from 1e-3 to 2e-3

89 torch.float16: 2e-3,

@@ -104,7 +110,7 @@ def test_act_and_mul(

110 opcheck(fn, (out, x, threshold))

111 elif activation == "swigluoai_and_mul":

112 opcheck(fn, (out, x, layer.alpha, layer.limit))

107 - else:

113 + elif activation != "swiglustep_and_mul":

114 opcheck(fn, (out, x))

115

116

tests/models/registry.py

@@ -484,6 +484,9 @@ _TEXT_GENERATION_EXAMPLE_MODELS = {

484 "Step1ForCausalLM": _HfExamplesInfo(

485 "stepfun-ai/Step-Audio-EditX", trust_remote_code=True

486 ),

487 + "Step3p5ForCausalLM": _HfExamplesInfo(

488 + "stepfun-ai/step-3.5-flash", is_available_online=False

489 + ),

490 "SmolLM3ForCausalLM": _HfExamplesInfo("HuggingFaceTB/SmolLM3-3B"),

491 "StableLMEpochForCausalLM": _HfExamplesInfo("stabilityai/stablelm-zephyr-3b"),

492 "StableLmForCausalLM": _HfExamplesInfo("stabilityai/stablelm-3b-4e1t"),

@@ -1113,6 +1116,12 @@ _SPECULATIVE_DECODING_EXAMPLE_MODELS = {

1116 "Qwen3NextMTP": _HfExamplesInfo(

1117 "Qwen/Qwen3-Next-80B-A3B-Instruct", min_transformers_version="4.56.3"

1118 ),

1119 + "Step3p5MTP": _HfExamplesInfo(

1120 + "stepfun-ai/Step-3.5-Flash",

1121 + trust_remote_code=True,

1122 + speculative_model="stepfun-ai/Step-3.5-Flash",

1123 + is_available_online=False,

1124 + ),

1125 }

1126

1127 _TRANSFORMERS_BACKEND_MODELS = {

vllm/config/speculative.py

@@ -41,6 +41,7 @@ MTPModelTypes = Literal[

41 "longcat_flash_mtp",

42 "mtp",

43 "pangu_ultra_moe_mtp",

44 + "step3p5_mtp",

45 ]

46 EagleModelTypes = Literal["eagle", "eagle3", MTPModelTypes]

47 SpeculativeMethod = Literal[

@@ -264,6 +265,11 @@ class SpeculativeConfig:

265 {"n_predict": n_predict, "architectures": ["LongCatFlashMTPModel"]}

266 )

267

268 + if hf_config.model_type == "step3p5":

269 + hf_config.model_type = "step3p5_mtp"

270 + n_predict = getattr(hf_config, "num_nextn_predict_layers", 1)

271 + hf_config.update({"n_predict": n_predict, "architectures": ["Step3p5MTP"]})

272 +

273 if initial_architecture == "MistralLarge3ForCausalLM":

274 hf_config.update({"architectures": ["EagleMistralLarge3ForCausalLM"]})

275

vllm/model_executor/layers/activation.py

+90

@@ -17,11 +17,63 @@ from vllm.logger import init_logger

17 from vllm.model_executor.custom_op import CustomOp

18 from vllm.model_executor.utils import set_weight_attrs

19 from vllm.platforms import current_platform

20 +from vllm.triton_utils import tl, triton

21 from vllm.utils.collection_utils import LazyDict

23 logger = init_logger(__name__)

26 +@triton.jit

27 +def _swiglustep_and_mul_kernel(

28 + o_ptr,

29 + o_stride,

30 + x_ptr,

31 + x_stride,

32 + limit: tl.constexpr,

33 + d: tl.constexpr,

34 + BLOCK_SIZE: tl.constexpr,

35 +) -> None:

36 + i = tl.program_id(axis=0).to(tl.int64)

37 + j = tl.program_id(axis=1)

38 + o_row_ptr = o_ptr + o_stride * i

39 + x_row_ptr = x_ptr + x_stride * i

40 + offsets = j * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)

41 + mask = offsets < d

42 +

43 + gate = tl.load(x_row_ptr + offsets, mask=mask).to(tl.float32)

44 + up = tl.load(x_row_ptr + offsets + d, mask=mask).to(tl.float32)

45 +

46 + gate_silu = tl.sigmoid(gate) * gate

47 + gate_clamped = tl.minimum(gate_silu, limit)

48 + up_clamped = tl.minimum(tl.maximum(up, -limit), limit)

49 +

50 + result = gate_clamped * up_clamped

51 + result = result.to(x_ptr.dtype.element_ty)

52 + tl.store(o_row_ptr + offsets, result, mask=mask)

53 +

54 +

55 +def swiglustep_and_mul_triton(

56 + output: torch.Tensor, input: torch.Tensor, limit: float = 7.0

57 +):

58 + b, n = input.shape

59 + assert input.ndim == 2

60 + assert n % 2 == 0

61 + d = n // 2

62 +

63 + def grid(meta):

64 + return (b, triton.cdiv(d, meta["BLOCK_SIZE"]))

65 +

66 + _swiglustep_and_mul_kernel[grid](

67 + output,

68 + output.stride(0),

69 + input,

70 + input.stride(0),

71 + limit=limit,

72 + d=d,

73 + BLOCK_SIZE=1024,

74 + )

75 +

76 +

77 # --8<-- [start:fatrelu_and_mul]

78 @CustomOp.register("fatrelu_and_mul")

79 class FatreluAndMul(CustomOp):

@@ -304,6 +356,44 @@ class SwigluOAIAndMul(CustomOp):

356 return f"alpha={repr(self.alpha)}, limit={repr(self.limit)}"

357

358

359 +# --8<-- [start:swiglustep_and_mul]

360 +@CustomOp.register("swiglustep_and_mul")

361 +class SwigluStepAndMul(CustomOp):

362 + """An activation function for SwiGLU with clamping.

363 +

364 + Computes x -> silu(x[:d]).clamp(max=limit) * x[d:].clamp(-limit, limit)

365 + where d = x.shape[-1] // 2.

366 +

367 + Shapes:

368 + x: (num_tokens, 2 * d) or (batch_size, seq_len, 2 * d)

369 + return: (num_tokens, d) or (batch_size, seq_len, d)

370 + """

371 +

372 + def __init__(self, limit: float = 7.0):

373 + super().__init__()

374 + if limit is None:

375 + raise ValueError("SwigluStepAndMul requires limit to be set.")

376 + self.limit = limit

377 +

378 + def forward_native(self, x: torch.Tensor) -> torch.Tensor:

379 + """PyTorch-native implementation equivalent to forward()."""

380 + gate, up = x.chunk(2, dim=-1)

381 + gate = F.silu(gate)

382 + gate = gate.clamp(max=self.limit)

383 + up = up.clamp(min=-self.limit, max=self.limit)

384 + return gate * up

385 +

386 + def forward_cuda(self, x: torch.Tensor) -> torch.Tensor:

387 + d = x.shape[-1] // 2

388 + output_shape = x.shape[:-1] + (d,)

389 + out = torch.empty(output_shape, dtype=x.dtype, device=x.device)

390 + swiglustep_and_mul_triton(out, x, self.limit)

391 + return out

392 +

393 + def extra_repr(self) -> str:

394 + return f"limit={repr(self.limit)}"

395 +

396 +

397 # --8<-- [start:gelu_new]

398 @CustomOp.register("gelu_new")

399 class NewGELU(CustomOp):

vllm/model_executor/layers/fused_moe/deep_gemm_moe.py

+1 -1

@@ -144,7 +144,7 @@ class DeepGemmExperts(mk.FusedMoEPermuteExpertsUnpermute):

144

145 @staticmethod

146 def _supports_activation(activation: str) -> bool:

147 - return activation in ["silu"]

147 + return activation in ["silu", "swiglustep"]

148

149 @staticmethod

150 def _supports_parallel_config(moe_parallel_config: FusedMoEParallelConfig) -> bool:

vllm/model_executor/layers/fused_moe/fused_moe.py

+1 -1

@@ -1939,7 +1939,7 @@ class TritonExperts(mk.FusedMoEPermuteExpertsUnpermute):

1939

1940 @staticmethod

1941 def _supports_activation(activation: str) -> bool:

1942 - return activation in ["silu", "gelu", "swigluoai"]

1942 + return activation in ["silu", "gelu", "swigluoai", "swiglustep"]

1943

1944 @staticmethod

1945 def _supports_parallel_config(moe_parallel_config: FusedMoEParallelConfig) -> bool:

vllm/model_executor/layers/fused_moe/utils.py

@@ -358,6 +358,11 @@ def apply_moe_activation(

358 torch.ops._C.gelu_and_mul(output, input)

359 elif activation == "swigluoai":

360 torch.ops._C.swigluoai_and_mul(output, input)

361 + elif activation == "swiglustep":

362 + from vllm.model_executor.layers.activation import swiglustep_and_mul_triton

363 +

364 + swiglustep_and_mul_triton(output, input)

365 +

366 # Activations without gated multiplication

367 elif activation == SILU_NO_MUL:

368 output.copy_(F.silu(input))

vllm/model_executor/models/registry.py

@@ -189,6 +189,7 @@ _TEXT_GENERATION_MODELS = {

189 "SeedOssForCausalLM": ("seed_oss", "SeedOssForCausalLM"),

190 "Step1ForCausalLM": ("step1", "Step1ForCausalLM"),

191 "Step3TextForCausalLM": ("step3_text", "Step3TextForCausalLM"),

192 + "Step3p5ForCausalLM": ("step3p5", "Step3p5ForCausalLM"),

193 "StableLMEpochForCausalLM": ("stablelm", "StablelmForCausalLM"),

194 "StableLmForCausalLM": ("stablelm", "StablelmForCausalLM"),

195 "Starcoder2ForCausalLM": ("starcoder2", "Starcoder2ForCausalLM"),

@@ -495,6 +496,7 @@ _SPECULATIVE_DECODING_MODELS = {

496 "MedusaModel": ("medusa", "Medusa"),

497 "OpenPanguMTPModel": ("openpangu_mtp", "OpenPanguMTP"),

498 "Qwen3NextMTP": ("qwen3_next_mtp", "Qwen3NextMTP"),

499 + "Step3p5MTP": ("step3p5_mtp", "Step3p5MTP"),

500 # Temporarily disabled.

501 # # TODO(woosuk): Re-enable this once the MLP Speculator is supported in V1.

502 # "MLPSpeculatorPreTrainedModel": ("mlp_speculator", "MLPSpeculator"),

vllm/model_executor/models/step3p5.py new

+894

@@ -0,0 +1,894 @@

1 +# SPDX-License-Identifier: Apache-2.0

2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3 +"""Inference-only Jurassic model."""

4 +

5 +from collections.abc import Iterable

6 +from typing import Any

7 +

8 +import torch

9 +from torch import nn

10 +from torch.nn.parameter import Parameter

11 +

12 +from vllm.compilation.decorators import support_torch_compile

13 +from vllm.config import CacheConfig, ModelConfig, VllmConfig

14 +from vllm.distributed import (

15 + get_dp_group,

16 + get_ep_group,

17 + get_pp_group,

18 + get_tensor_model_parallel_rank,

19 + get_tensor_model_parallel_world_size,

20 + get_tp_group,

21 +)

22 +from vllm.logger import init_logger

23 +from vllm.model_executor.layers.activation import SiluAndMul, SwigluStepAndMul

24 +from vllm.model_executor.layers.attention import Attention

25 +from vllm.model_executor.layers.fused_moe import FusedMoE

26 +from vllm.model_executor.layers.fused_moe.shared_fused_moe import SharedFusedMoE

27 +from vllm.model_executor.layers.layernorm import GemmaRMSNorm

28 +from vllm.model_executor.layers.linear import (

29 + ColumnParallelLinear,

30 + MergedColumnParallelLinear,

31 + QKVParallelLinear,

32 + ReplicatedLinear,

33 + RowParallelLinear,

34 +)

35 +from vllm.model_executor.layers.logits_processor import LogitsProcessor

36 +from vllm.model_executor.layers.quantization.base_config import QuantizationConfig

37 +from vllm.model_executor.layers.rotary_embedding import get_rope

38 +from vllm.model_executor.layers.vocab_parallel_embedding import (

39 + DEFAULT_VOCAB_PADDING_SIZE,

40 + ParallelLMHead,

41 + VocabParallelEmbedding,

42 +)

43 +from vllm.model_executor.model_loader.weight_utils import default_weight_loader

44 +from vllm.sequence import IntermediateTensors

45 +from vllm.v1.attention.backend import AttentionType

46 +

47 +from .interfaces import MixtureOfExperts, SupportsPP

48 +from .utils import (

49 + AutoWeightsLoader,

50 + PPMissingLayer,

51 + WeightsMapper,

52 + extract_layer_index,

53 + is_pp_missing_parameter,

54 + make_empty_intermediate_tensors_factory,

55 + make_layers,

56 + maybe_prefix,

57 +)

58 +

59 +logger = init_logger(__name__)

60 +

61 +

62 +class FP32ReplicatedLinear(ReplicatedLinear):

63 + """

64 + Use FP32 for higher precision.

65 + """

66 +

67 + def forward(

68 + self,

69 + x: torch.Tensor,

70 + ) -> torch.Tensor | tuple[torch.Tensor, Parameter | None]:

71 + assert self.params_dtype == torch.float32

72 + return super().forward(x.to(torch.float32))

73 +

74 +

75 +class Step3p5MLP(nn.Module):

76 + def __init__(

77 + self,

78 + config: ModelConfig,

79 + hidden_size: int,

80 + intermediate_size: int,

81 + hidden_act: str,

82 + quant_config: QuantizationConfig | None = None,

83 + reduce_results: bool = True,

84 + prefix: str = "",

85 + ) -> None:

86 + super().__init__()

87 + self.gate_up_proj = MergedColumnParallelLinear(

88 + hidden_size,

89 + [intermediate_size] * 2,

90 + bias=False,

91 + quant_config=quant_config,

92 + prefix=f"{prefix}.gate_up_proj",

93 + )

94 + self.down_proj = RowParallelLinear(

95 + intermediate_size,

96 + hidden_size,

97 + bias=False,

98 + quant_config=quant_config,

99 + reduce_results=reduce_results,

100 + prefix=f"{prefix}.down_proj",

101 + )

102 +

103 + if hidden_act != "silu":

104 + raise ValueError(

105 + f"Unsupported activation: {hidden_act}. Only silu is supported for now."

106 + )

107 + self.act_fn = SiluAndMul()

108 + self.prefix = prefix

109 + self.hidden_size = hidden_size

110 + self.limit = None

111 + layer_idx = extract_layer_index(prefix)

112 + if (

113 + config.swiglu_limits_shared

114 + and config.swiglu_limits_shared[layer_idx] is not None

115 + and config.swiglu_limits_shared[layer_idx] != 0

116 + ):

117 + self.limit = config.swiglu_limits_shared[layer_idx]

118 + self.act_fn = SwigluStepAndMul(limit=self.limit)

119 +

120 + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:

121 + gate_up, _ = self.gate_up_proj(hidden_states)

122 + intermediate_act = self.act_fn(gate_up)

123 + output, _ = self.down_proj(intermediate_act)

124 + return output

125 +

126 +

127 +class Step3p5Attention(nn.Module):

128 + def __init__(

129 + self,

130 + hidden_size: int,

131 + num_heads: int,

132 + num_kv_heads: int,

133 + max_position: int = 4096 * 32,

134 + head_dim: int | None = None,

135 + rms_norm_eps: float = 1e-06,

136 + qkv_bias: bool = False,

137 + rope_theta: float | list[float] | None = 10000,

138 + cache_config: CacheConfig | None = None,

139 + quant_config: QuantizationConfig | None = None,

140 + rope_scaling: dict[str, Any] | None = None,

141 + prefix: str = "",

142 + attn_type: str = AttentionType.DECODER,

143 + # Step3p5 specific args

144 + sliding_window: int | None = None,

145 + use_head_wise_attn_gate: bool = False,

146 + layer_types: list = None,

147 + use_rope_layers: list = None,

148 + yarn_only_types: list = None,

149 + swa_num_attention_heads: int | None = None,

150 + partial_rotary_factor: float = 1.0,

151 + ):

152 + super().__init__()

153 + self.hidden_size = hidden_size

154 + self.total_num_heads = num_heads

155 + tp_size = get_tensor_model_parallel_world_size()

156 + self.layer_idx = extract_layer_index(prefix)

157 + if layer_types:

158 + enable_sliding_window = layer_types[self.layer_idx] == "sliding_attention"

159 + else:

160 + enable_sliding_window = self.layer_idx % 2 == 0

161 + if yarn_only_types and layer_types[self.layer_idx] not in yarn_only_types:

162 + rope_scaling = None

163 +

164 + if sliding_window is not None and enable_sliding_window:

165 + sliding_window = sliding_window

166 + if swa_num_attention_heads is not None:

167 + num_heads = swa_num_attention_heads

168 + self.total_num_heads = swa_num_attention_heads

169 + else:

170 + sliding_window = None

171 +

172 + if isinstance(rope_theta, list):

173 + rope_theta = rope_theta[self.layer_idx]

174 +

175 + self.rank = get_tensor_model_parallel_rank()

176 + self.partial_rotary_factor = partial_rotary_factor

177 + assert self.total_num_heads % tp_size == 0

178 + self.num_heads = self.total_num_heads // tp_size

179 + self.total_num_kv_heads = num_kv_heads

180 + if self.total_num_kv_heads >= tp_size:

181 + # Number of KV heads is greater than TP size, so we partition

182 + # the KV heads across multiple tensor parallel GPUs.

183 + assert self.total_num_kv_heads % tp_size == 0

184 + else:

185 + # Number of KV heads is less than TP size, so we replicate

186 + # the KV heads across multiple tensor parallel GPUs.

187 + assert tp_size % self.total_num_kv_heads == 0

188 + self.num_kv_heads = max(1, self.total_num_kv_heads // tp_size)

189 + self.head_dim = head_dim or hidden_size // self.total_num_heads

190 + self.q_size = self.num_heads * self.head_dim

191 + self.kv_size = self.num_kv_heads * self.head_dim

192 + self.scaling = self.head_dim**-0.5

193 + self.rope_theta = rope_theta

194 + self.qkv_proj = QKVParallelLinear(

195 + hidden_size,

196 + self.head_dim,

197 + self.total_num_heads,

198 + self.total_num_kv_heads,

199 + bias=qkv_bias,

200 + quant_config=quant_config,

201 + prefix=f"{prefix}.qkv_proj",

202 + )

203 + self.o_proj = RowParallelLinear(

204 + self.total_num_heads * self.head_dim,

205 + hidden_size,

206 + bias=False,

207 + quant_config=quant_config,

208 + prefix=f"{prefix}.o_proj",

209 + )

210 +

211 + if rope_scaling is not None and not isinstance(rope_scaling, dict):

212 + raise ValueError("rope_scaling must be a dict for Step3p5Attention.")

213 +

214 + rope_parameters: dict[str, Any] = (

215 + dict(rope_scaling) if rope_scaling is not None else {}

216 + )

217 + rope_parameters.setdefault("rope_type", "default")

218 + rope_parameters["rope_theta"] = self.rope_theta

219 + rope_parameters["partial_rotary_factor"] = partial_rotary_factor

220 +

221 + self.rotary_emb = get_rope(

222 + head_size=self.head_dim,

223 + max_position=max_position,

224 + rope_parameters=rope_parameters,

225 + )

226 +

227 + self.q_norm = GemmaRMSNorm(self.head_dim, rms_norm_eps)

228 + self.k_norm = GemmaRMSNorm(self.head_dim, rms_norm_eps)

229 + self.use_head_wise_attn_gate = use_head_wise_attn_gate

230 + if use_head_wise_attn_gate:

231 + self.g_proj = ColumnParallelLinear(

232 + hidden_size,

233 + self.total_num_heads,

234 + bias=False,

235 + prefix=f"{prefix}.g_proj",

236 + )

237 +

238 + self.use_rope = True

239 + if use_rope_layers:

240 + self.use_rope = use_rope_layers[self.layer_idx]

241 +

242 + self.attn = Attention(

243 + self.num_heads,

244 + self.head_dim,

245 + self.scaling,

246 + num_kv_heads=self.num_kv_heads,

247 + cache_config=cache_config,

248 + quant_config=quant_config,

249 + prefix=f"{prefix}.attn",

250 + per_layer_sliding_window=sliding_window,

251 + attn_type=attn_type,

252 + )

253 +

254 + self.max_position_embeddings = max_position

255 + assert self.partial_rotary_factor == 1 or self.partial_rotary_factor == 0.5

256 + self.rotary_dim = (

257 + self.head_dim if self.partial_rotary_factor == 1 else self.head_dim // 2

258 + )

259 +

260 + def forward(

261 + self,

262 + positions: torch.Tensor,

263 + hidden_states: torch.Tensor,

264 + ) -> torch.Tensor:

265 + qkv, _ = self.qkv_proj(hidden_states)

266 + q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1)

267 + # Add qk-norm inline similar to Qwen3 MOE attention

268 + q_by_head = q.view(*q.shape[:-1], q.shape[-1] // self.head_dim, self.head_dim)

269 + q_by_head = self.q_norm(q_by_head.contiguous())

270 + q = q_by_head.view(q.shape)

271 +

272 + k_by_head = k.view(*k.shape[:-1], k.shape[-1] // self.head_dim, self.head_dim)

273 + k_by_head = self.k_norm(k_by_head.contiguous())

274 + k = k_by_head.view(k.shape)

275 + if self.use_rope:

276 + q, k = self.rotary_emb(positions, q, k)

277 + attn_output = self.attn(q, k, v)

278 + if self.use_head_wise_attn_gate:

279 + extra_dims, _ = self.g_proj(hidden_states)

280 + output = (

281 + attn_output.view(*attn_output.shape[:-1], self.num_heads, self.head_dim)

282 + * extra_dims.unsqueeze(-1).sigmoid()

283 + )

284 + attn_output = output.view(*attn_output.shape)

285 + output, _ = self.o_proj(attn_output)

286 + return output

287 +

288 +

289 +class FusedMoEBlock(nn.Module):

290 + def __init__(

291 + self,

292 + vllm_config: VllmConfig,

293 + prefix: str = "",

294 + ):

295 + super().__init__()

296 +

297 + self.tp_size = get_tensor_model_parallel_world_size()

298 + self.layer_idx = extract_layer_index(prefix)

299 +

300 + self.ep_size = get_ep_group().device_group.size()

301 + self.ep_rank = get_ep_group().device_group.rank()

302 + config = vllm_config.model_config.hf_config

303 + quant_config = vllm_config.quant_config

304 + parallel_config = vllm_config.parallel_config

305 +

306 + self.hidden_size = config.hidden_size

307 + self.enable_eplb = parallel_config.enable_eplb

308 + self.n_routed_experts = config.moe_num_experts

309 + self.n_logical_experts = self.n_routed_experts

310 + self.n_redundant_experts = parallel_config.eplb_config.num_redundant_experts

311 + self.n_physical_experts = self.n_logical_experts + self.n_redundant_experts

312 + self.n_local_physical_experts = self.n_physical_experts // self.ep_size

313 +

314 + self.physical_expert_start = self.ep_rank * self.n_local_physical_experts

315 + self.physical_expert_end = (

316 + self.physical_expert_start + self.n_local_physical_experts

317 + )

318 +

319 + if self.tp_size > config.moe_num_experts:

320 + raise ValueError(

321 + f"Tensor parallel size {self.tp_size} is greater than "

322 + f"the number of experts {config.moe_num_experts}."

323 + )

324 +

325 + self.gate = FP32ReplicatedLinear(

326 + config.hidden_size,

327 + config.moe_num_experts,

328 + bias=False,

329 + quant_config=None,

330 + params_dtype=torch.float32, # Use FP32 for higher precision.

331 + prefix=f"{prefix}.gate",

332 + )

333 + self.use_moe_router_bias = config.use_moe_router_bias

334 + assert self.use_moe_router_bias, "Only support use_moe_router_bias is true."

335 + self.routed_scaling_factor = config.moe_router_scaling_factor

336 + self.router_bias = nn.Parameter(

337 + torch.zeros(config.moe_num_experts, dtype=torch.float32),

338 + requires_grad=False,

339 + )

340 + self.need_fp32_gate = config.need_fp32_gate

341 + assert self.need_fp32_gate, (

342 + "Router logits must use FP32 precision for numerical stability."

343 + )

344 +

345 + activation = "silu"

346 + swiglu_limits = config.swiglu_limits or []

347 + swiglu_limit = (

348 + swiglu_limits[self.layer_idx]

349 + if self.layer_idx < len(swiglu_limits)

350 + else None

351 + )

352 + if swiglu_limit not in (None, 0):

353 + swiglu_limit = float(swiglu_limit)

354 + assert swiglu_limit == 7.0, (

355 + "Swiglu limit in fused moe block only suport 7.0 now."

356 + )

357 + activation = "swiglustep"

358 + logger.debug(

359 + "step3p5 layer_idx: %s, activation: %s, limit: %s",

360 + self.layer_idx,

361 + activation,

362 + swiglu_limit,

363 + )

364 +

365 + self.share_expert = Step3p5MLP(

366 + config=config,

367 + hidden_size=self.hidden_size,

368 + intermediate_size=config.share_expert_dim,

369 + hidden_act="silu",

370 + reduce_results=False,

371 + quant_config=quant_config,

372 + prefix=f"{prefix}.share_expert",

373 + )

374 + self.experts = SharedFusedMoE(

375 + shared_experts=self.share_expert,

376 + gate=self.gate,

377 + num_experts=config.moe_num_experts,

378 + top_k=config.moe_top_k,

379 + hidden_size=config.hidden_size,

380 + intermediate_size=config.moe_intermediate_size,

381 + reduce_results=False,

382 + renormalize=config.norm_expert_weight,

383 + quant_config=quant_config,

384 + activation=activation,

385 + prefix=f"{prefix}.experts",

386 + scoring_func=getattr(config, "moe_router_activation", "sigmoid"),

387 + e_score_correction_bias=self.router_bias,

388 + routed_scaling_factor=config.moe_router_scaling_factor,

389 + enable_eplb=self.enable_eplb,

390 + num_redundant_experts=self.n_redundant_experts,

391 + )

392 +

393 + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:

394 + num_tokens, hidden_dim = hidden_states.shape

395 + hidden_states = hidden_states.view(-1, hidden_dim)

396 +

397 + if self.experts.is_internal_router:

398 + # In this case, the gate/router runs inside the FusedMoE class

399 + fused_moe_out = self.experts(

400 + hidden_states=hidden_states, router_logits=hidden_states

401 + )

402 + else:

403 + # router_logits: (num_tokens, n_experts)

404 + router_logits, _ = self.gate(hidden_states)

405 + fused_moe_out = self.experts(

406 + hidden_states=hidden_states, router_logits=router_logits

407 + )

408 +

409 + shared_output, final_hidden_states = fused_moe_out

410 + if self.share_expert is None:

411 + assert shared_output is None

412 +

413 + if self.share_expert is not None:

414 + assert shared_output is not None

415 + final_hidden_states += shared_output

416 +

417 + if self.tp_size > 1:

418 + final_hidden_states = self.experts.maybe_all_reduce_tensor_model_parallel(

419 + final_hidden_states

420 + )

421 +

422 + return final_hidden_states.view(num_tokens, hidden_dim)

423 +

424 +

425 +class Step3p5DecoderLayer(nn.Module):

426 + def __init__(

427 + self,

428 + vllm_config: VllmConfig,

429 + prefix: str = "",

430 + ) -> None:

431 + super().__init__()

432 + config = vllm_config.model_config.hf_config

433 + self.hidden_size = config.hidden_size

434 + layer_idx = extract_layer_index(prefix)

435 + self.layer_idx = layer_idx

436 + cache_config = vllm_config.cache_config

437 + quant_config = vllm_config.quant_config

438 + if cache_config is not None:

439 + cache_config.sliding_window = None

440 + if config.att_impl_type == "GQA":

441 + num_attention_heads = None

442 + num_attention_groups = None

443 + head_dim = None

444 + if (

445 + getattr(config, "attention_other_setting", None)

446 + and getattr(config, "layer_types", [])

447 + and config.layer_types[layer_idx]

448 + == config.attention_other_setting["attention_type"]

449 + ):

450 + num_attention_heads = config.attention_other_setting[

451 + "num_attention_heads"

452 + ]

453 + num_attention_groups = config.attention_other_setting[

454 + "num_attention_groups"

455 + ]

456 + head_dim = config.attention_other_setting["head_dim"]

457 + partial_rotary_factors = getattr(config, "partial_rotary_factors", [])

458 + self.self_attn = Step3p5Attention(

459 + hidden_size=self.hidden_size,

460 + num_heads=num_attention_heads

461 + if num_attention_heads

462 + else config.num_attention_heads,

463 + max_position=config.max_position_embeddings,

464 + num_kv_heads=num_attention_groups

465 + if num_attention_groups

466 + else config.num_attention_groups,

467 + rope_theta=config.rope_theta,

468 + rms_norm_eps=config.rms_norm_eps,

469 + qkv_bias=getattr(config, "attention_bias", False),

470 + head_dim=head_dim if head_dim else getattr(config, "head_dim", None),

471 + cache_config=cache_config,

472 + quant_config=quant_config,

473 + rope_scaling=getattr(config, "rope_scaling", None),

474 + sliding_window=getattr(config, "sliding_window", None),

475 + use_head_wise_attn_gate=getattr(

476 + config, "use_head_wise_attn_gate", False

477 + ),

478 + layer_types=getattr(config, "layer_types", []),

479 + use_rope_layers=getattr(config, "use_rope_layers", []),

480 + yarn_only_types=getattr(config, "yarn_only_types", []),

481 + partial_rotary_factor=partial_rotary_factors[layer_idx]

482 + if partial_rotary_factors

483 + else 1.0,

484 + prefix=f"{prefix}.self_attn",

485 + )

486 + else:

487 + raise ValueError(

488 + f"Unsupported attention implementation: {config.att_impl_type}"

489 + )

490 + self.use_moe = False

491 + self.tp_group = get_tp_group()

492 + self.use_fused_all_reduce = (

493 + get_tensor_model_parallel_world_size() > 1

494 + and get_dp_group().world_size == 1

495 + )

496 + if self.use_fused_all_reduce:

497 + logger.warning_once("Enable custom fused all reduce...")

498 + else:

499 + logger.warning_once("Disable custom fused all reduce...")

500 +

501 + moe_layers_enum = getattr(config, "moe_layers_enum", None)

502 + if moe_layers_enum is not None:

503 + moe_layers_idx = [int(i) for i in moe_layers_enum.strip().split(",")]

504 + else:

505 + moe_layers_idx = [i for i in range(1, config.num_hidden_layers)]

506 + if layer_idx in moe_layers_idx:

507 + self.moe = FusedMoEBlock(

508 + vllm_config,

509 + prefix=f"{prefix}.moe",

510 + )

511 + self.use_moe = True

512 + else:

513 + self.mlp = Step3p5MLP(

514 + config=config,

515 + hidden_size=config.hidden_size,

516 + intermediate_size=config.intermediate_size,

517 + hidden_act="silu",

518 + quant_config=quant_config,

519 + reduce_results=True,

520 + prefix=f"{prefix}.mlp",

521 + )

522 + self.input_layernorm = GemmaRMSNorm(config.hidden_size, config.rms_norm_eps)

523 + self.post_attention_layernorm = GemmaRMSNorm(

524 + config.hidden_size, config.rms_norm_eps

525 + )

526 + self.prefix = prefix

527 +

528 + def add_and_maybe_inplace_all_reduce(

529 + self, in1: torch.Tensor, in2: torch.Tensor

530 + ) -> torch.Tensor:

531 + if not self.use_fused_all_reduce:

532 + return in1 + in2

533 + return self.tp_group.all_reduce(in1 + in2)

534 +

535 + def forward(

536 + self, positions: torch.Tensor, hidden_states: torch.Tensor

537 + ) -> torch.Tensor:

538 + residual = hidden_states

539 + hidden_states = self.input_layernorm(hidden_states)

540 +

541 + hidden_states = self.self_attn(

542 + positions=positions,

543 + hidden_states=hidden_states,

544 + )

545 + hidden_states += residual

546 + residual = hidden_states

547 + hidden_states = self.post_attention_layernorm(hidden_states)

548 +

549 + if self.use_moe:

550 + ffn_output = self.moe(hidden_states)

551 + else:

552 + ffn_output = self.mlp(hidden_states)

553 + hidden_states = ffn_output + residual

554 + return hidden_states

555 +

556 +

557 +@support_torch_compile

558 +class Step3p5Model(nn.Module):

559 + def __init__(self, vllm_config: VllmConfig, prefix: str = "") -> None:

560 + super().__init__()

561 +

562 + self.vllm_config = vllm_config

563 + config = vllm_config.model_config.hf_config

564 + self.vocab_size = config.vocab_size

565 + self.config = config

566 +

567 + self.moe_num_experts = config.moe_num_experts

568 +

569 + if get_pp_group().is_first_rank or (

570 + config.tie_word_embeddings and get_pp_group().is_last_rank

571 + ):

572 + self.embed_tokens = VocabParallelEmbedding(

573 + self.vocab_size,

574 + config.hidden_size,

575 + )

576 + else:

577 + self.embed_tokens = PPMissingLayer()

578 +

579 + self.start_layer, self.end_layer, self.layers = make_layers(

580 + config.num_hidden_layers,

581 + lambda prefix: Step3p5DecoderLayer(

582 + vllm_config,

583 + prefix=prefix,

584 + ),

585 + prefix=f"{prefix}.layers",

586 + )

587 + if get_pp_group().is_last_rank:

588 + self.norm = GemmaRMSNorm(config.hidden_size, config.rms_norm_eps)

589 + else:

590 + self.norm = PPMissingLayer()

591 +

592 + self.make_empty_intermediate_tensors = make_empty_intermediate_tensors_factory(

593 + ["hidden_states"], config.hidden_size

594 + )

595 +

596 + def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:

597 + return self.embed_tokens(input_ids)

598 +

599 + def forward(

600 + self,

601 + input_ids: torch.Tensor,

602 + positions: torch.Tensor,

603 + intermediate_tensors: IntermediateTensors | None = None,

604 + inputs_embeds: torch.Tensor | None = None,

605 + ) -> torch.Tensor:

606 + if get_pp_group().is_first_rank:

607 + if inputs_embeds is not None:

608 + hidden_states = inputs_embeds

609 + else:

610 + hidden_states = self.embed_input_ids(input_ids)

611 + else:

612 + assert intermediate_tensors is not None

613 + hidden_states = intermediate_tensors["hidden_states"]

614 + for i in range(self.start_layer, self.end_layer):

615 + layer = self.layers[i]

616 + hidden_states = layer(positions, hidden_states)

617 +

618 + if not get_pp_group().is_last_rank:

619 + return IntermediateTensors(

620 + {

621 + "hidden_states": hidden_states,

622 + }

623 + )

624 +

625 + return hidden_states

626 +

627 + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:

628 + config = self.config

629 + assert config.num_attention_groups > 1, "Only support GQA"

630 + qkv_params_mapping = []

631 + stacked_params_mapping = [

632 + # (param_name, shard_name, shard_id)

633 + ("qkv_proj", "q_proj", "q"),

634 + ("qkv_proj", "k_proj", "k"),

635 + ("qkv_proj", "v_proj", "v"),

636 + ("gate_up_proj", "gate_proj", 0),

637 + ("gate_up_proj", "up_proj", 1),

638 + ]

639 +

640 + params_dict = dict(self.named_parameters())

641 + loaded_params: set[str] = set()

642 +

643 + expert_params_mapping = [

644 + (".moe.experts.w13_weight", ".moe.gate_proj.weight", "w1"),

645 + (".moe.experts.w13_weight", ".moe.up_proj.weight", "w3"),

646 + (".moe.experts.w2_weight", ".moe.down_proj.weight", "w2"),

647 + ]

648 +

649 + disable_moe_stacked_params = [data[1] for data in expert_params_mapping]

650 +

651 + for name, loaded_weight in weights:

652 + if name.startswith("model."):

653 + local_name = name[len("model.") :]

654 + full_name = name

655 + else:

656 + local_name = name

657 + full_name = f"model.{name}" if name else "model"

658 +

659 + spec_layer = get_spec_layer_idx_from_weight_name(config, full_name)

660 + if spec_layer is not None:

661 + continue # skip spec decode layers for main model

662 +

663 + # Skip any layers beyond the main model's depth (e.g., MTP layers)

664 + if full_name.startswith("model.layers."):

665 + parts = full_name.split(".")

666 + if len(parts) > 2 and parts[2].isdigit():

667 + layer_idx = int(parts[2])

668 + if layer_idx >= config.num_hidden_layers:

669 + continue

670 +

671 + for param_name, weight_name, shard_id in stacked_params_mapping:

672 + if weight_name not in local_name:

673 + continue

674 + if any(

675 + disable_moe_stacked_param in local_name

676 + for disable_moe_stacked_param in disable_moe_stacked_params

677 + ):

678 + continue

679 + replaced_name = local_name.replace(weight_name, param_name)

680 + if is_pp_missing_parameter(replaced_name, self):

681 + continue

682 + if replaced_name not in params_dict:

683 + continue

684 + param = params_dict[replaced_name]

685 + weight_loader = param.weight_loader

686 + weight_loader(param, loaded_weight, shard_id)

687 + loaded_params.add(replaced_name)

688 + break

689 + else:

690 + for param_name, weight_name, shard_id in expert_params_mapping:

691 + if weight_name not in local_name:

692 + continue

693 + replaced_name = local_name.replace(weight_name, param_name)

694 + if is_pp_missing_parameter(replaced_name, self):

695 + continue

696 + if (

697 + replaced_name.endswith(".bias")

698 + or replaced_name.endswith("_bias")

699 + ) and replaced_name not in params_dict:

700 + continue

701 + if replaced_name not in params_dict:

702 + continue

703 + param = params_dict[replaced_name]

704 + weight_loader = param.weight_loader

705 + moe_expert_num = self.moe_num_experts

706 + assert loaded_weight.shape[0] == moe_expert_num

707 + for expert_id in range(moe_expert_num):

708 + loaded_weight_expert = loaded_weight[expert_id]

709 + weight_loader(

710 + param,

711 + loaded_weight_expert,

712 + replaced_name,

713 + shard_id=shard_id,

714 + expert_id=expert_id,

715 + )

716 + loaded_params.add(replaced_name)

717 + break

718 + else:

719 + for (

720 + param_name,

721 + weight_name,

722 + start_idx,

723 + end_idx,

724 + ) in qkv_params_mapping:

725 + if weight_name not in local_name:

726 + continue

727 + replaced_name = local_name.replace(weight_name, param_name)

728 + if is_pp_missing_parameter(replaced_name, self):

729 + continue

730 + if replaced_name not in params_dict:

731 + continue

732 + param = params_dict[replaced_name]

733 + dim = param.shape[param.output_dim]

734 + begin_idx = int(start_idx * dim)

735 + end_idx = int(end_idx * dim)

736 + param_slice = param.narrow(

737 + param.output_dim, begin_idx, end_idx - begin_idx

738 + )

739 + param_slice.copy_(loaded_weight)

740 + loaded_params.add(replaced_name)

741 + break

742 + else:

743 + if is_pp_missing_parameter(local_name, self):

744 + continue

745 + if "expert_bias" in local_name:

746 + logger.warning_once("ignore expert_bias")

747 + continue

748 + if local_name not in params_dict:

749 + continue

750 + param = params_dict[local_name]

751 + weight_loader = getattr(

752 + param, "weight_loader", default_weight_loader

753 + )

754 + weight_loader(param, loaded_weight)

755 + loaded_params.add(local_name)

756 + return loaded_params

757 +

758 +

759 +class Step3p5ForCausalLM(nn.Module, SupportsPP, MixtureOfExperts):

760 + hf_to_vllm_mapper = WeightsMapper(

761 + orig_to_new_substr={".share_expert.": ".moe.share_expert."}

762 + )

763 +

764 + def __init__(

765 + self,

766 + *,

767 + vllm_config: VllmConfig,

768 + prefix: str = "",

769 + ):

770 + super().__init__()

771 + config = vllm_config.model_config.hf_config

772 + lora_config = vllm_config.lora_config

773 + self.config = config

774 + self.vllm_config = vllm_config

775 +

776 + self.model = Step3p5Model(

777 + vllm_config=vllm_config, prefix=maybe_prefix(prefix, "model")

778 + )

779 +

780 + self.moe_layers: list[FusedMoEBlock] = []

781 + for layer in self.model.layers:

782 + if isinstance(layer, PPMissingLayer):

783 + continue

784 + assert isinstance(layer, Step3p5DecoderLayer)

785 + if hasattr(layer, "moe") and isinstance(layer.moe, FusedMoEBlock):

786 + self.moe_layers.append(layer.moe)

787 +

788 + if get_pp_group().is_last_rank:

789 + self.unpadded_vocab_size = config.vocab_size

790 + if lora_config:

791 + self.unpadded_vocab_size += lora_config.lora_extra_vocab_size

792 + self.lm_head = ParallelLMHead(

793 + self.unpadded_vocab_size,

794 + config.hidden_size,

795 + org_num_embeddings=config.vocab_size,

796 + padding_size=DEFAULT_VOCAB_PADDING_SIZE

797 + if not lora_config

798 + else lora_config.lora_vocab_padding_size,

799 + )

800 + self.logits_processor = LogitsProcessor(

801 + self.unpadded_vocab_size, config.vocab_size

802 + )

803 + else:

804 + self.lm_head = PPMissingLayer()

805 +

806 + self.make_empty_intermediate_tensors = (

807 + self.model.make_empty_intermediate_tensors

808 + )

809 +

810 + # Set MoE hyperparameters

811 + self.expert_weights = []

812 + assert len(self.moe_layers) > 0, "No MoE layers found in the model."

813 + example_layer = self.moe_layers[0]

814 + self.num_moe_layers = len(self.moe_layers)

815 + self.num_expert_groups = 1

816 + self.num_shared_experts = 0

817 + self.num_logical_experts = example_layer.n_logical_experts

818 + self.num_physical_experts = example_layer.n_physical_experts

819 + self.num_local_physical_experts = example_layer.n_local_physical_experts

820 + self.num_routed_experts = example_layer.n_routed_experts

821 + self.num_redundant_experts = example_layer.n_redundant_experts

822 +

823 + def forward(

824 + self,

825 + input_ids: torch.Tensor,

826 + positions: torch.Tensor,

827 + intermediate_tensors: IntermediateTensors | None = None,

828 + inputs_embeds: torch.Tensor | None = None,

829 + ):

830 + hidden_states = self.model(

831 + input_ids, positions, intermediate_tensors, inputs_embeds

832 + )

833 + return hidden_states

834 +

835 + def compute_logits(self, hidden_states: torch.Tensor) -> torch.Tensor:

836 + hidden_states = self.model.norm(hidden_states)

837 + logits = self.logits_processor(self.lm_head, hidden_states)

838 + return logits

839 +

840 + def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:

841 + return self.model.embed_tokens(input_ids)

842 +

843 + def set_eplb_state(

844 + self,

845 + expert_load_view: torch.Tensor,

846 + logical_to_physical_map: torch.Tensor,

847 + logical_replica_count: torch.Tensor,

848 + ) -> None:

849 + for layer_idx, layer in enumerate(self.moe_layers):

850 + experts = layer.experts

851 + assert isinstance(experts, FusedMoE)

852 + # Register the expert weights.

853 + self.expert_weights.append(experts.get_expert_weights())

854 + experts.set_eplb_state(

855 + moe_layer_idx=layer_idx,

856 + expert_load_view=expert_load_view,

857 + logical_to_physical_map=logical_to_physical_map,

858 + logical_replica_count=logical_replica_count,

859 + )

860 +

861 + def update_physical_experts_metadata(

862 + self,

863 + num_physical_experts: int,

864 + num_local_physical_experts: int,

865 + ) -> None:

866 + assert self.num_local_physical_experts == num_local_physical_experts

867 + self.num_physical_experts = num_physical_experts

868 + self.num_local_physical_experts = num_local_physical_experts

869 + self.num_redundant_experts = num_physical_experts - self.num_logical_experts

870 + for layer in self.moe_layers:

871 + assert isinstance(layer, FusedMoEBlock)

872 + layer.n_local_physical_experts = num_local_physical_experts

873 + layer.n_physical_experts = num_physical_experts

874 + layer.n_redundant_experts = self.num_redundant_experts

875 + layer.experts.update_expert_map()

876 +

877 + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:

878 + loader = AutoWeightsLoader(self)

879 + return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper)

880 +

881 +

882 +def get_spec_layer_idx_from_weight_name(

883 + config: ModelConfig, weight_name: str

884 +) -> int | None:

885 + if hasattr(config, "num_nextn_predict_layers") and (

886 + config.num_nextn_predict_layers > 0

887 + ):

888 + layer_idx = config.num_hidden_layers

889 + for i in range(config.num_nextn_predict_layers):

890 + if weight_name.startswith(

891 + f"layers.{layer_idx + i}." # Step3p5Model

892 + ) or weight_name.startswith(f"model.layers.{layer_idx + i}."): # Step3p5MTP

893 + return layer_idx + i

894 + return None

vllm/model_executor/models/step3p5_mtp.py new

+315

@@ -0,0 +1,315 @@

1 +# SPDX-License-Identifier: Apache-2.0

2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3 +from collections.abc import Iterable

4 +

5 +import torch

6 +import torch.nn as nn

7 +from transformers import PretrainedConfig

8 +

9 +from vllm.config import VllmConfig

10 +from vllm.logger import init_logger

11 +from vllm.model_executor.layers.layernorm import GemmaRMSNorm

12 +from vllm.model_executor.layers.logits_processor import LogitsProcessor

13 +from vllm.model_executor.layers.quantization import QuantizationConfig

14 +from vllm.model_executor.layers.vocab_parallel_embedding import (

15 + ParallelLMHead,

16 + VocabParallelEmbedding,

17 +)

18 +from vllm.model_executor.model_loader.weight_utils import default_weight_loader

19 +from vllm.sequence import IntermediateTensors

20 +

21 +from .step3p5 import Step3p5DecoderLayer, get_spec_layer_idx_from_weight_name

22 +from .utils import maybe_prefix

23 +

24 +logger = init_logger(__name__)

25 +

26 +

27 +class SharedHead(nn.Module):

28 + def __init__(

29 + self,

30 + config: PretrainedConfig,

31 + quant_config: QuantizationConfig | None = None,

32 + ) -> None:

33 + super().__init__()

34 + self.norm = GemmaRMSNorm(config.hidden_size, config.rms_norm_eps)

35 + self.head = ParallelLMHead(

36 + config.vocab_size, config.hidden_size, quant_config=quant_config

37 + )

38 +

39 + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:

40 + return self.norm(hidden_states)

41 +

42 +

43 +class Step3p5AMultiTokenPredictorLayer(nn.Module):

44 + def __init__(

45 + self,

46 + vllm_config: VllmConfig,

47 + prefix: str,

48 + ) -> None:

49 + super().__init__()

50 + config = vllm_config.model_config.hf_config

51 + quant_config = vllm_config.quant_config

52 + self.enorm = GemmaRMSNorm(config.hidden_size, config.rms_norm_eps)

53 + self.hnorm = GemmaRMSNorm(config.hidden_size, config.rms_norm_eps)

54 + self.eh_proj = nn.Linear(config.hidden_size * 2, config.hidden_size, bias=False)

55 + self.shared_head = SharedHead(config=config, quant_config=quant_config)

56 + self.mtp_block = Step3p5DecoderLayer(

57 + vllm_config,

58 + prefix=f"{prefix}.mtp_block",

59 + )

60 +

61 + def forward(

62 + self,

63 + input_ids: torch.Tensor,

64 + positions: torch.Tensor,

65 + previous_hidden_states: torch.Tensor,

66 + inputs_embeds: torch.Tensor | None = None,

67 + spec_step_index: int = 0,

68 + ) -> torch.Tensor:

69 + assert inputs_embeds is not None

70 + inputs_embeds = self.enorm(inputs_embeds)

71 + previous_hidden_states = self.hnorm(previous_hidden_states)

72 +

73 + hidden_states = self.eh_proj(

74 + torch.cat([inputs_embeds, previous_hidden_states], dim=-1)

75 + )

76 +

77 + hidden_states = self.mtp_block(positions=positions, hidden_states=hidden_states)

78 + return hidden_states

79 +

80 +

81 +class Step3p5AMultiTokenPredictor(nn.Module):

82 + def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):

83 + super().__init__()

84 + config = vllm_config.model_config.hf_config

85 + self.embed_tokens = VocabParallelEmbedding(

86 + config.vocab_size,

87 + config.hidden_size,

88 + )

89 + self.mtp_start_layer_idx = config.num_hidden_layers

90 + self.num_mtp_layers = config.num_nextn_predict_layers

91 + # to map the exact layer index from weights

92 + self.layers = torch.nn.ModuleDict(

93 + {

94 + str(idx): Step3p5AMultiTokenPredictorLayer(

95 + vllm_config,

96 + f"{prefix}.layers.{idx}",

97 + )

98 + for idx in range(

99 + self.mtp_start_layer_idx,

100 + self.mtp_start_layer_idx + self.num_mtp_layers,

101 + )

102 + }

103 + )

104 +

105 + self.logits_processor = LogitsProcessor(config.vocab_size)

106 +

107 + def forward(

108 + self,

109 + input_ids: torch.Tensor,

110 + positions: torch.Tensor,

111 + previous_hidden_states: torch.Tensor,

112 + inputs_embeds: torch.Tensor | None = None,

113 + spec_step_idx: int = 0,

114 + ) -> torch.Tensor:

115 + if inputs_embeds is None:

116 + inputs_embeds = self.embed_tokens(input_ids)

117 + current_step_idx = spec_step_idx % self.num_mtp_layers

118 + return self.layers[str(self.mtp_start_layer_idx + current_step_idx)](

119 + input_ids,

120 + positions,

121 + previous_hidden_states,

122 + inputs_embeds,

123 + current_step_idx,

124 + )

125 +

126 + def compute_logits(

127 + self,

128 + hidden_states: torch.Tensor,

129 + spec_step_idx: int = 0,

130 + ) -> torch.Tensor:

131 + current_step_idx = spec_step_idx % self.num_mtp_layers

132 + mtp_layer = self.layers[str(self.mtp_start_layer_idx + current_step_idx)]

133 + logits = self.logits_processor(

134 + mtp_layer.shared_head.head, mtp_layer.shared_head(hidden_states)

135 + )

136 + return logits

137 +

138 + def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:

139 + return self.embed_tokens(input_ids)

140 +

141 +

142 +class Step3p5MTP(nn.Module):

143 + def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):

144 + super().__init__()

145 + self.config = vllm_config.model_config.hf_config

146 + self.vllm_config = vllm_config

147 + self.model = Step3p5AMultiTokenPredictor(

148 + vllm_config=vllm_config, prefix=maybe_prefix(prefix, "model")

149 + )

150 +

151 + def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:

152 + return self.model.embed_input_ids(input_ids)

153 +

154 + def forward(

155 + self,

156 + input_ids: torch.Tensor,

157 + positions: torch.Tensor,

158 + hidden_states: torch.Tensor,

159 + intermediate_tensors: IntermediateTensors | None = None,

160 + inputs_embeds: torch.Tensor | None = None,

161 + spec_step_idx: int = 0,

162 + ) -> torch.Tensor:

163 + hidden_states = self.model(

164 + input_ids, positions, hidden_states, inputs_embeds, spec_step_idx

165 + )

166 + return hidden_states

167 +

168 + def compute_logits(

169 + self,

170 + hidden_states: torch.Tensor,

171 + spec_step_idx: int = 0,

172 + ) -> torch.Tensor | None:

173 + return self.model.compute_logits(hidden_states, spec_step_idx)

174 +

175 + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:

176 + stacked_params_mapping = [

177 + # (param_name, shard_name, shard_id)

178 + ("qkv_proj", "q_proj", "q"),

179 + ("qkv_proj", "k_proj", "k"),

180 + ("qkv_proj", "v_proj", "v"),

181 + ("gate_up_proj", "gate_proj", 0),

182 + ("gate_up_proj", "up_proj", 1),

183 + ]

184 +

185 + expert_params_mapping = [

186 + (".moe.experts.w13_weight", ".moe.gate_proj.weight", "w1"),

187 + (".moe.experts.w13_weight", ".moe.up_proj.weight", "w3"),

188 + (".moe.experts.w2_weight", ".moe.down_proj.weight", "w2"),

189 + ]

190 +

191 + params_dict = dict(self.named_parameters())

192 + loaded_params: set[str] = set()

193 + for name, loaded_weight in weights:

194 + if "rotary_emb.inv_freq" in name:

195 + continue

196 + spec_layer = get_spec_layer_idx_from_weight_name(self.config, name)

197 + if "embed_tokens" not in name and spec_layer is None:

198 + continue

199 + name = self._rewrite_spec_layer_name(spec_layer, name)

200 + for param_name, weight_name, shard_id in stacked_params_mapping:

201 + # Skip non-stacked layers and experts (experts handled below).

202 + if weight_name not in name:

203 + continue

204 + # We have mlp.experts[0].gate_proj in the checkpoint.

205 + # Since we handle the experts below in expert_params_mapping,

206 + # we need to skip here BEFORE we update the name, otherwise

207 + # name will be updated to mlp.experts[0].gate_up_proj, which

208 + # will then be updated below in expert_params_mapping

209 + # for mlp.experts[0].gate_gate_up_proj, which breaks load.

210 + if ("mlp.experts." in name) and name not in params_dict:

211 + continue

212 + if "experts" in name or "moe" in name:

213 + continue

214 + name = name.replace(weight_name, param_name)

215 + # Skip loading extra bias for GPTQ models.

216 + if name.endswith(".bias") and name not in params_dict:

217 + continue

218 +

219 + param = params_dict[name]

220 + weight_loader = param.weight_loader

221 + weight_loader(param, loaded_weight, shard_id)

222 + break

223 + else:

224 + for mapping in expert_params_mapping:

225 + param_name, weight_name, shard_id = mapping

226 + if weight_name not in name:

227 + continue

228 + name = name.replace(weight_name, param_name)

229 + # Skip loading extra bias for GPTQ models.

230 + if (

231 + name.endswith(".bias") or name.endswith("_bias")

232 + ) and name not in params_dict:

233 + continue

234 + param = params_dict[name]

235 + weight_loader = param.weight_loader

236 + for expert_id in range(loaded_weight.shape[0]):

237 + loaded_weight_expert = loaded_weight[expert_id]

238 + weight_loader(

239 + param,

240 + loaded_weight_expert,

241 + name,

242 + shard_id=shard_id,

243 + expert_id=expert_id,

244 + )

245 + loaded_params.add(name)

246 + break

247 + else:

248 + # Skip loading extra bias for GPTQ models.

249 + if (

250 + name.endswith(".bias")

251 + and name not in params_dict

252 + or "tok_embeddings" in name

253 + ):

254 + continue

255 +

256 + if spec_layer is not None and ".transformer." in name:

257 + name = name.replace(".transformer.", ".")

258 + if "shared_head" in name:

259 + name = name.replace("shared_head.output", "shared_head.head")

260 + if "embed_tokens" in name:

261 + assert (

262 + hasattr(self.config, "num_nextn_predict_layers")

263 + and self.config.num_nextn_predict_layers > 0

264 + )

265 + name = "model.embed_tokens.weight"

266 + param = params_dict[name]

267 + weight_loader = getattr(

268 + param, "weight_loader", default_weight_loader

269 + )

270 + weight_loader(param, loaded_weight)

271 + loaded_params.add(name)

272 + params_need_to_load = set(params_dict.keys())

273 + # Some KV cache scales are optional: checkpoints may omit them and vLLM

274 + # will fall back to default scales during initialization.

275 + optional_params = {

276 + name

277 + for name, param in params_dict.items()

278 + if name.endswith((".k_scale", ".v_scale", ".q_scale", ".prob_scale"))

279 + and getattr(param, "numel", lambda: 0)() == 1

280 + and getattr(param, "requires_grad", False) is False

281 + }

282 + params_need_to_load -= optional_params

283 + if params_need_to_load != loaded_params:

284 + missing_params = list(params_need_to_load - loaded_params)

285 + param_name_example = missing_params[0]

286 + raise RuntimeError(

287 + "Some parameters like "

288 + f"{param_name_example} are not in the checkpoint and will falsely "

289 + "use random initialization"

290 + )

291 + return loaded_params

292 +

293 + def _rewrite_spec_layer_name(self, spec_layer: int, name: str) -> str:

294 + """

295 + Rewrite the weight name to match the format of the original model.

296 + Add .mtp_block for modules in transformer layer block for spec layer

297 + """

298 + spec_layer_weight_names = [

299 + "embed_tokens",

300 + "enorm",

301 + "hnorm",

302 + "eh_proj",

303 + "shared_head",

304 + ]

305 + spec_layer_weight = False

306 + for weight_name in spec_layer_weight_names:

307 + if weight_name in name:

308 + spec_layer_weight = True

309 + break

310 + if not spec_layer_weight:

311 + # treat rest weights as weights for transformer layer block

312 + name = name.replace(

313 + f"model.layers.{spec_layer}.", f"model.layers.{spec_layer}.mtp_block."

314 + )

315 + return name

vllm/reasoning/__init__.py

@@ -84,6 +84,10 @@ _REASONING_PARSERS_TO_REGISTER = {

84 "step3_reasoning_parser",

85 "Step3ReasoningParser",

86 ),

87 + "step3p5": (

88 + "step3p5_reasoning_parser",

89 + "Step3p5ReasoningParser",

90 + ),

91 }

vllm/reasoning/step3p5_reasoning_parser.py new

+153

@@ -0,0 +1,153 @@

1 +# SPDX-License-Identifier: Apache-2.0

2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3 +

4 +from collections.abc import Sequence

5 +

6 +from vllm.entrypoints.openai.chat_completion.protocol import (

7 + ChatCompletionRequest,

8 +)

9 +from vllm.entrypoints.openai.engine.protocol import DeltaMessage

10 +from vllm.entrypoints.openai.responses.protocol import (

11 + ResponsesRequest,

12 +)

13 +from vllm.reasoning.basic_parsers import BaseThinkingReasoningParser

14 +from vllm.tokenizers import TokenizerLike

15 +

16 +

17 +class Step3p5ReasoningParser(BaseThinkingReasoningParser):

18 + """

19 + Reasoning parser for Step3p5 model.

20 +

21 + Step3p5 uses the <think>...</think> format, but it tends to emit an extra

22 + newline immediately before and/or after the </think> token. This parser trims:

23 + - the newline right before </think>

24 + - the newline right after </think>

25 + """

26 +

27 + @property

28 + def start_token(self) -> str:

29 + return "<think>"

30 +

31 + @property

32 + def end_token(self) -> str:

33 + return "</think>"

34 +

35 + def __init__(self, tokenizer: TokenizerLike, *args, **kwargs):

36 + super().__init__(tokenizer, *args, **kwargs)

37 +

38 + # Used to hold a trailing "\n" from reasoning content so we can decide

39 + # whether it is immediately before </think>.

40 + self._pending_reasoning_newline = False

41 +

42 + # Used to delay the reasoning end detection.

43 + # This is necessary to remove the newline appears immediately after </think>,

44 + # which may cause the end detection to be delayed by one round.

45 + self.end_offset = 1

46 +

47 + def is_reasoning_end(self, input_ids: Sequence[int]) -> bool:

48 + if self.end_token_id in input_ids and self.end_offset > 0:

49 + self.end_offset -= 1

50 + return False

51 + return self.end_offset < 1

52 +

53 + def is_reasoning_end_streaming(

54 + self, input_ids: Sequence[int], delta_ids: Sequence[int]

55 + ) -> bool:

56 + if self.end_token_id in input_ids and self.end_offset > 0:

57 + self.end_offset -= 1

58 + return False

59 + return self.end_offset < 1

60 +

61 + def extract_reasoning(

62 + self,

63 + model_output: str,

64 + request: ChatCompletionRequest | ResponsesRequest,

65 + ) -> tuple[str | None, str | None]:

66 + reasoning, content = super().extract_reasoning(model_output, request)

67 + if reasoning is not None:

68 + reasoning = reasoning.removesuffix("\n")

69 + if content is not None:

70 + content = content.removeprefix("\n")

71 + return reasoning or None, content or None

72 +

73 + def extract_reasoning_streaming(

74 + self,

75 + previous_text: str,

76 + current_text: str,

77 + delta_text: str,

78 + previous_token_ids: Sequence[int],

79 + current_token_ids: Sequence[int],

80 + delta_token_ids: Sequence[int],

81 + ) -> DeltaMessage | None:

82 + # Drop the immediate newline that models often emit after </think>.

83 + if previous_text.endswith(self.end_token) and delta_text:

84 + if delta_text == "\n":

85 + return None

86 + elif delta_text.startswith("\n"):

87 + remaining = delta_text.removeprefix("\n")

88 + return DeltaMessage(content=remaining) if remaining else None

89 +

90 + ret = super().extract_reasoning_streaming(

91 + previous_text,

92 + current_text,

93 + delta_text,

94 + previous_token_ids,

95 + current_token_ids,

96 + delta_token_ids,

97 + )

98 +

99 + if ret is None:

100 + return None

101 +

102 + # Compatibility path for models that don't generate the start token:

103 + # treat everything before </think> as reasoning and everything after

104 + # as content.

105 + if (

106 + self.start_token_id not in previous_token_ids

107 + and self.start_token_id not in delta_token_ids

108 + ):

109 + if self.end_token_id in delta_token_ids:

110 + end_index = delta_text.find(self.end_token)

111 + reasoning = delta_text[:end_index]

112 + content = delta_text[end_index + len(self.end_token) :]

113 + ret = DeltaMessage(reasoning=reasoning, content=content or None)

114 + elif self.end_token_id in previous_token_ids:

115 + ret = DeltaMessage(content=delta_text)

116 + else:

117 + ret = DeltaMessage(reasoning=delta_text)

118 +

119 + reasoning_to_output = ret.reasoning

120 + content_to_output = ret.content

121 +

122 + # Reasoning: handle the newline immediately before </think>.

123 + if reasoning_to_output is not None:

124 + if self._pending_reasoning_newline:

125 + reasoning_to_output = "\n" + reasoning_to_output

126 + self._pending_reasoning_newline = False

127 +

128 + if reasoning_to_output.endswith("\n"):

129 + reasoning_to_output = reasoning_to_output.removesuffix("\n")

130 + if self.end_token in delta_text:

131 + # Trailing "\n" is right before </think>, drop it.

132 + self._pending_reasoning_newline = False

133 + else:

134 + # Hold the trailing "\n" until we know whether </think> follows.

135 + self._pending_reasoning_newline = True

136 +

137 + # Content: handle the newline immediately after </think>.

138 + if content_to_output is not None:

139 + # No need to get into parser again to remove newline after </think>.

140 + self.end_offset -= 1

141 +

142 + # If we have content, reasoning must have ended.

143 + self._pending_reasoning_newline = False

144 +

145 + if self.end_token in delta_text and content_to_output.startswith("\n"):

146 + content_to_output = content_to_output.removeprefix("\n")

147 +

148 + reasoning_to_output = reasoning_to_output or None

149 + content_to_output = content_to_output or None

150 + if reasoning_to_output is None and content_to_output is None:

151 + return None

152 +

153 + return DeltaMessage(reasoning=reasoning_to_output, content=content_to_output)

vllm/tool_parsers/__init__.py

@@ -134,6 +134,10 @@ _TOOL_PARSERS_TO_REGISTER = {

134 "step3_tool_parser",

135 "Step3ToolParser",

136 ),

137 + "step3p5": (

138 + "step3p5_tool_parser",

139 + "Step3p5ToolParser",

140 + ),

141 "xlam": (

142 "xlam_tool_parser",

143 "xLAMToolParser",

vllm/tool_parsers/step3p5_tool_parser.py new

+1511

@@ -0,0 +1,1511 @@

1 +# SPDX-License-Identifier: Apache-2.0

2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3 +import ast

4 +import json

5 +from collections.abc import Sequence

6 +from typing import Any

7 +from xml.parsers.expat import ParserCreate

8 +

9 +import regex as re

10 +

11 +from vllm.entrypoints.chat_utils import make_tool_call_id

12 +from vllm.entrypoints.openai.chat_completion.protocol import (

13 + ChatCompletionRequest,

14 + ChatCompletionToolsParam,

15 +)

16 +from vllm.entrypoints.openai.engine.protocol import (

17 + DeltaFunctionCall,

18 + DeltaMessage,

19 + DeltaToolCall,

20 + ExtractedToolCallInformation,

21 + FunctionCall,

22 + ToolCall,

23 +)

24 +from vllm.logger import init_logger

25 +from vllm.tokenizers import TokenizerLike

26 +from vllm.tool_parsers.abstract_tool_parser import (

27 + ToolParser,

28 + ToolParserManager,

29 +)

30 +

31 +logger = init_logger(__name__)

32 +

33 +

34 +class StreamingXMLToolCallParser:

35 + """

36 + Simplified streaming XML tool call parser

37 + Supports streaming input, parsing, and output

38 + """

39 +

40 + def __init__(self):

41 + self.reset_streaming_state()

42 +

43 + # Tool configuration information

44 + self.tools: list[ChatCompletionToolsParam] | None = None

45 + self.tool_call_start_token: str = "<tool_call>"

46 + self.tool_call_end_token: str = "</tool_call>"

47 + self.function_start_token: str = "<function="

48 + self.function_end_token: str = "</function>"

49 + self.parameter_start_token: str = "<parameter="

50 + self.parameter_end_token: str = "</parameter>"

51 +

52 + def reset_streaming_state(self):

53 + """Reset streaming parsing state"""

54 +

55 + self.deltas = []

56 + # state for streaming

57 + self.tool_call_index = 0

58 + self.current_call_id = None

59 + self.last_completed_call_id = None

60 + self.current_function_name = None

61 + self.current_function_open = False

62 + self.parameters = {}

63 + self.current_param_name = None

64 + self.current_param_value = ""

65 + self.current_param_value_converted = ""

66 + self.current_param_is_first = False

67 + self.should_emit_end_newline = False

68 + self.start_quote_emitted = False

69 +

70 + self.streaming_buffer = ""

71 + self.last_processed_pos = 0

72 +

73 + self.text_content_buffer = ""

74 +

75 + # state for preprocessing and deferred parsing

76 + self._pre_inside_parameter = False

77 + self._pre_param_buffer = ""

78 + self._pre_current_param_name = None

79 + self.defer_current_parameter = False

80 + self.deferred_param_raw_value = ""

81 +

82 + # recreate parser

83 + self.parser = ParserCreate()

84 + self.setup_parser()

85 +

86 + def parse_single_streaming_chunks(self, xml_chunk: str) -> DeltaMessage:

87 + """

88 + Parse single streaming XML chunk and return Delta response

89 + This is the actual streaming interface that receives chunks

90 + one by one and maintains internal state

91 +

92 + Args:

93 + xml_chunk: Single XML chunk string

94 + Returns:

95 + DeltaMessage: Contains delta information generated by this chunk,

96 + returns empty response if no complete elements

97 + """

98 + # Record delta count before processing

99 + initial_delta_count = len(self.deltas)

100 +

101 + self.streaming_buffer += xml_chunk

102 +

103 + found_elements = self._process_complete_xml_elements()

104 +

105 + if found_elements:

106 + # If complete elements found, check if end events were missed

107 + # some tags may not have been triggered

108 + try:

109 + new_deltas = self.deltas[initial_delta_count:]

110 + # If this chunk contains </function>

111 + # but didn't generate '}', then complete it

112 + if (

113 + self.current_call_id is not None

114 + and self.function_end_token in xml_chunk

115 + ):

116 + # - Added '}' (non-empty parameter ending)

117 + # - Added '{}' (empty parameter function)

118 + has_function_close = any(

119 + (

120 + td.tool_calls

121 + and any(

122 + (

123 + tc.function

124 + and tc.id == self.current_call_id

125 + and isinstance(tc.function.arguments, str)

126 + and (tc.function.arguments in ("}", "{}"))

127 + )

128 + for tc in td.tool_calls

129 + )

130 + )

131 + for td in new_deltas

132 + )

133 + if not has_function_close:

134 + # Close potentially unclosed element

135 + if self.current_param_name:

136 + self._end_element("parameter")

137 + if self.current_function_name:

138 + self._end_element("function")

139 + # If this chunk contains </tool_call>

140 + # but didn't generate final empty delta, then complete it

141 + if (

142 + self.current_call_id is not None

143 + and self.tool_call_end_token in xml_chunk

144 + ):

145 + has_toolcall_close = any(

146 + (

147 + td.tool_calls

148 + and any(

149 + (

150 + tc.type == "function"

151 + and tc.function

152 + and tc.function.arguments == ""

153 + and tc.id == self.current_call_id

154 + )

155 + for tc in td.tool_calls

156 + )

157 + )

158 + for td in new_deltas

159 + )

160 + if not has_toolcall_close:

161 + # Close potentially unclosed element

162 + if self.current_param_name:

163 + self._end_element("parameter")

164 + if self.current_function_name:

165 + self._end_element("function")

166 + self._end_element("tool_call")

167 + except Exception as e:

168 + logger.warning("Error with fallback parsing: %s", e)

169 + # Merge newly generated deltas into single response

170 + result_delta = self._merge_new_deltas_to_single_response(

171 + initial_delta_count

172 + )

173 + return result_delta

174 + else:

175 + # No complete elements, check if there's unoutput text content

176 + if self.text_content_buffer and self.tool_call_index == 0:

177 + # Has text content but no tool_call yet, output text content

178 + text_delta = DeltaMessage(content=self.text_content_buffer)

179 + self._emit_delta(text_delta)

180 + # Clear buffer to avoid duplicate output

181 + self.text_content_buffer = ""

182 + return text_delta

183 +

184 + # If this chunk contains end tags but wasn't triggered by parser,

185 + # manually complete end events

186 + # Only execute when still on the same call as when entered,

187 + # to prevent accidentally closing new calls

188 + # in multi <tool_call> scenarios

189 + if self.current_call_id is not None and (

190 + self.function_end_token in xml_chunk

191 + or self.tool_call_end_token in xml_chunk

192 + ):

193 + # Close potentially unclosed element

194 + if self.current_param_name:

195 + self._end_element("parameter")

196 + if self.function_end_token in xml_chunk and self.current_function_name:

197 + self._end_element("function")

198 + if self.tool_call_end_token in xml_chunk:

199 + self._end_element("tool_call")

200 + # Return the merged delta result generated by this fallback

201 + result_delta = self._merge_new_deltas_to_single_response(

202 + initial_delta_count

203 + )

204 + return result_delta

205 +

206 + # No complete elements, return empty response

207 + return DeltaMessage(content=None)

208 +

209 + def _escape_xml_special_chars(self, text: str) -> str:

210 + """

211 + Escape XML special characters

212 + Args:

213 + text: Original text

214 + Returns:

215 + Escaped text

216 + """

217 + xml_escapes = {

218 + "&": "&",

219 + "<": "<",

220 + ">": ">",

221 + '"': """,

222 + "'": "'",

223 + }

224 +

225 + for char, escape in xml_escapes.items():

226 + text = text.replace(char, escape)

227 +

228 + return text

229 +

230 + def _process_complete_xml_elements(self) -> bool:

231 + """

232 + Process complete XML elements in buffer

233 +

234 + Returns:

235 + bool: Whether complete elements were found and processed

236 + """

237 + found_any = False

238 +

239 + while self.last_processed_pos < len(self.streaming_buffer):

240 + # Find next complete xml element

241 + element, end_pos = self._find_next_complete_element(self.last_processed_pos)

242 + if element is None:

243 + # No complete element found, wait for more data

244 + break

245 +

246 + # Check if this element should be skipped

247 + if self._should_skip_element(element):

248 + self.last_processed_pos = end_pos

249 + continue

250 +

251 + # Found complete XML element, process it

252 + try:

253 + preprocessed_element = self._preprocess_xml_chunk(element)

254 + # Check if this is the first tool_call start

255 + if (

256 + (

257 + preprocessed_element.strip().startswith("<tool_call>")

258 + or preprocessed_element.strip().startswith("<function name=")

259 + )

260 + and self.tool_call_index == 0

261 + ) and self.text_content_buffer:

262 + # First tool_call starts,

263 + # output previously collected text content first

264 + text_delta = DeltaMessage(content=self.text_content_buffer)

265 + self._emit_delta(text_delta)

266 + # Clear buffer for potential subsequent text content

267 + self.text_content_buffer = ""

268 +

269 + # If a new tool_call starts and

270 + # there are already completed tool_calls with function name

271 + if (

272 + preprocessed_element.strip().startswith("<tool_call>")

273 + and self.tool_call_index > 0

274 + and self.current_call_id

275 + and self.current_function_name

276 + ):

277 + # Reset parser state but preserve generated deltas

278 + if self.current_param_name:

279 + self._end_element("parameter")

280 + if self.current_function_open:

281 + self._end_element("function")

282 + # Output final tool_call tail delta

283 + final_delta = DeltaMessage(

284 + role=None,

285 + content=None,

286 + reasoning_content=None,

287 + tool_calls=[

288 + DeltaToolCall(

289 + index=self.tool_call_index - 1,

290 + id=self.current_call_id,

291 + type="function",

292 + function=DeltaFunctionCall(name=None, arguments=""),

293 + )

294 + ],

295 + )

296 + self._emit_delta(final_delta)

297 + # Reset XML parser and current call state

298 + self._reset_xml_parser_after_tool_call()

299 + # Parse preprocessed element

300 + self.parser.Parse(preprocessed_element, False)

301 + found_any = True

302 +

303 + except Exception as e:

304 + logger.warning("Error when parsing XML elements: %s", e)

305 +

306 + # Update processed position

307 + self.last_processed_pos = end_pos

308 +

309 + return found_any

310 +

311 + def _fix_incomplete_tag_in_chunk(self, chunk: str) -> str:

312 + """

313 + Fallback: fix incomplete <parameter=xxx or <function=xxx tags

314 + (missing >)

315 + Examples: <parameter=-C: -> <parameter=-C>, <parameter=parameter=-n:

316 + -> <parameter=-n>

317 + Also handles missing = cases: <function xxx> -> <function=xxx>,

318 + <functionxxx> -> <function=xxx>

319 + Only fixes tags that pass validation (parameter exists in tool definition)

320 + """

321 + # First, handle missing = cases for function tags

322 + chunk = self._fix_missing_equals_in_function_tag(chunk)

323 +

324 + for tag_type in ["parameter", "function"]:

325 + pattern = f"<{tag_type}="

326 + if pattern not in chunk:

327 + continue

328 +

329 + start_idx = chunk.find(pattern)

330 + after_tag = chunk[start_idx:]

331 + gt_pos = after_tag.find(">")

332 + lt_pos = after_tag.find("<", len(pattern))

333 +

334 + # Skip if already well-formed

335 + if (

336 + gt_pos != -1

337 + and (lt_pos == -1 or gt_pos < lt_pos)

338 + and pattern in after_tag[:gt_pos]

339 + ):

340 + continue

341 +

342 + # Extract tag name (stop at space, newline, or <)

343 + content = chunk[start_idx + len(pattern) :]

344 + end_pos = next(

345 + (i for i, ch in enumerate(content) if ch in (" ", "\n", "<")),

346 + len(content),

347 + )

348 + tag_name = content[:end_pos]

349 +

350 + if not tag_name:

351 + continue

352 +

353 + # Remove duplicate prefix: <parameter=parameter=xxx -> <parameter=xxx

354 + if tag_name.startswith(f"{tag_type}="):

355 + tag_name = tag_name[len(tag_type) + 1 :]

356 +

357 + # Remove trailing non-alphanumeric chars (keep - and _)

358 + while tag_name and not (

359 + tag_name[-1].isalnum() or tag_name[-1] in ("-", "_")

360 + ):

361 + tag_name = tag_name[:-1]

362 +

363 + if not tag_name:

364 + continue

365 +

366 + # Validate parameter exists in tool definition

367 + if tag_type == "parameter" and not self._validate_parameter_name(tag_name):

368 + continue

369 +

370 + # Apply fix

371 + chunk = chunk.replace(

372 + f"<{tag_type}={content[:end_pos]}", f"<{tag_type}={tag_name}>", 1

373 + )

374 +

375 + return chunk

376 +

377 + def _fix_missing_equals_in_function_tag(self, chunk: str) -> str:

378 + """

379 + Fix missing = in function tags: <function xxx> or <functionxxx>

380 + Examples:

381 + <function execute_bash> -> <function=execute_bash>

382 + <functionexecute_bash> -> <function=execute_bash>

383 + Only fixes if function name exists in tool definition

384 + """

385 + # already correct

386 + if "<function=" in chunk:

387 + return chunk

388 +

389 + # Pattern 1: <function xxx> (with space/newline but no =)

390 + pattern1 = r"<function\s+([a-zA-Z_][a-zA-Z0-9_]*)\s*>"

391 + match1 = re.search(pattern1, chunk)

392 + if match1:

393 + func_name = match1.group(1).strip()

394 + # must validate function name exists before fixing

395 + if func_name and self._validate_function_name(func_name):

396 + original = match1.group(0)

397 + fixed = f"<function={func_name}>"

398 + chunk = chunk.replace(original, fixed, 1)

399 + return chunk

400 +

401 + # Pattern 2: <functionxxx> (no space, no =)

402 + # only match <function followed by letters

403 + pattern2 = r"<function([a-zA-Z_][a-zA-Z0-9_]*)\s*>"

404 + match2 = re.search(pattern2, chunk)

405 + if match2:

406 + func_name = match2.group(1).strip()

407 + # must validate function name exists before fixing

408 + if func_name and self._validate_function_name(func_name):

409 + original = match2.group(0)

410 + fixed = f"<function={func_name}>"

411 + chunk = chunk.replace(original, fixed, 1)

412 + return chunk

413 +

414 + return chunk

415 +

416 + def _validate_function_name(self, func_name: str) -> bool:

417 + """Check if function name exists in tool definitions"""

418 + if not self.tools:

419 + return False

420 +

421 + for tool in self.tools:

422 + if (

423 + hasattr(tool, "type")

424 + and tool.type == "function"

425 + and hasattr(tool, "function")

426 + and hasattr(tool.function, "name")

427 + and tool.function.name == func_name

428 + ):

429 + return True

430 +

431 + return False

432 +

433 + def _validate_parameter_name(self, param_name: str) -> bool:

434 + """Check if parameter exists in current function's tool definition"""

435 + if not self.tools or not self.current_function_name:

436 + return True

437 +

438 + for tool in self.tools:

439 + if (

440 + hasattr(tool, "type")

441 + and tool.type == "function"

442 + and hasattr(tool, "function")

443 + and hasattr(tool.function, "name")

444 + and tool.function.name == self.current_function_name

445 + ):

446 + if not hasattr(tool.function, "parameters"):

447 + return True

448 + params = tool.function.parameters

449 + if isinstance(params, dict):

450 + properties = params.get("properties", params)

451 + return param_name in properties

452 + break

453 +

454 + return True

455 +

456 + def _should_skip_element(self, element: str) -> bool:

457 + """

458 + Determine whether an element should be skipped

459 +

460 + Args:

461 + element: Element to evaluate

462 +

463 + Returns:

464 + bool: True means should skip, False means should process

465 + """

466 +

467 + # If it's a tool_call XML tag, don't skip

468 + if (

469 + element.startswith(self.tool_call_start_token)

470 + or element.startswith(self.function_start_token)

471 + or element.startswith(self.parameter_start_token)

472 + ):

473 + return False

474 +

475 + # If currently not parsing tool calls and not blank,

476 + # collect this text instead of skipping

477 + # Only process other XML elements after tool_call appears,

478 + # otherwise treat as plain text

479 + if self.current_call_id is None and element:

480 + # Collect text content to buffer

481 + self.text_content_buffer += element

482 + return True # Still skip, but content has been collected

483 +

484 + # If currently parsing tool calls,

485 + # this might be parameter value, don't skip

486 + if self.current_call_id is not None:

487 + return False

488 +

489 + # Skip blank content

490 + return not element

491 +

492 + def _find_next_complete_element(self, start_pos: int) -> tuple[str | None, int]:

493 + """

494 + Find next complete XML element from specified position

495 +

496 + Args:

497 + start_pos: Position to start searching

498 +

499 + Returns:

500 + (Complete element string, element end position),

501 + returns (None, start_pos) if no complete element found

502 + """

503 + buffer = self.streaming_buffer[start_pos:]

504 +

505 + if not buffer:

506 + return None, start_pos

507 +

508 + if buffer.startswith("<"):

509 + # Check if this is an incomplete parameter/function tag

510 + # e.g., <parameter=-C: or <function=xxx

511 + is_incomplete_param = (

512 + buffer.startswith("<parameter=") and ">" not in buffer.split("\n")[0]

513 + )

514 + is_incomplete_func = (

515 + buffer.startswith("<function=") and ">" not in buffer.split("\n")[0]

516 + )

517 +

518 + if is_incomplete_param or is_incomplete_func:

519 + # Find the corresponding closing tag

520 + tag_type = "parameter" if is_incomplete_param else "function"

521 + closing_tag = f"</{tag_type}>"

522 + closing_pos = buffer.find(closing_tag)

523 +

524 + if closing_pos != -1:

525 + # Found closing tag, return complete element including closing tag

526 + complete_element = buffer[: closing_pos + len(closing_tag)]

527 + return complete_element, start_pos + closing_pos + len(closing_tag)

528 +

529 + # Need to ensure no new < appears,

530 + # find the nearest one between < and >

531 + tag_end = buffer.find("<", 1)

532 + tag_end2 = buffer.find(">", 1)

533 + if tag_end != -1 and tag_end2 != -1:

534 + # Next nearest is <

535 + if tag_end < tag_end2:

536 + return buffer[:tag_end], start_pos + tag_end

537 + # Next nearest is >, means found XML element

538 + else:

539 + return buffer[: tag_end2 + 1], start_pos + tag_end2 + 1

540 + elif tag_end != -1:

541 + return buffer[:tag_end], start_pos + tag_end

542 + elif tag_end2 != -1:

543 + return buffer[: tag_end2 + 1], start_pos + tag_end2 + 1

544 + else:

545 + # If currently not parsing tool calls (entering a tool_call),

546 + # check if starts with <tool_call> or <function=

547 + if self.current_call_id is None:

548 + # Check if might be start of <tool_call>

549 + if buffer == "<tool_call>"[: len(buffer)]:

550 + # Might be start of <tool_call>, wait for more data

551 + return None, start_pos

552 + elif (

553 + buffer.startswith("<function=")

554 + or buffer == "<function="[: len(buffer)]

555 + ):

556 + # Might be start of <function=, wait for more data

557 + # to get the complete function tag

558 + return None, start_pos

559 + else:

560 + # Not start of <tool_call> or <function=, treat as text

561 + return buffer, start_pos + len(buffer)

562 + else:

563 + # When parsing tool calls,

564 + # wait for more data to get complete tag

565 + return None, start_pos

566 + else:

567 + # Find text content (until next < or buffer end)

568 + next_tag_pos = buffer.find("<")

569 + if next_tag_pos != -1:

570 + # Found text content

571 + text_content = buffer[:next_tag_pos]

572 + return text_content, start_pos + next_tag_pos

573 + else:

574 + # Buffer end is all text, process

575 + # (no longer wait for more data)

576 + remaining = buffer

577 + return remaining, start_pos + len(remaining)

578 +

579 + def _merge_new_deltas_to_single_response(self, initial_count: int) -> DeltaMessage:

580 + """

581 + Merge newly generated deltas from this processing

582 + into a single DeltaMessage

583 +

584 + Args:

585 + initial_count: Delta count before processing

586 +

587 + Returns:

588 + Merged DeltaMessage containing all newly generated delta information

589 + """

590 + if len(self.deltas) <= initial_count:

591 + return DeltaMessage(content=None)

592 +

593 + # Get newly generated deltas

594 + new_deltas = self.deltas[initial_count:]

595 +

596 + if len(new_deltas) == 1:

597 + # Only one new delta, return directly

598 + return new_deltas[0]

599 +

600 + # Merge multiple new deltas

601 + merged_tool_calls: list[DeltaToolCall] = []

602 + merged_content: str = ""

603 +

604 + for delta in new_deltas:

605 + if delta.content:

606 + merged_content += delta.content

607 + if delta.tool_calls:

608 + # For tool_calls, we need to intelligently merge arguments

609 + for tool_call in delta.tool_calls:

610 + # Find if there's already a tool_call with the same call_id

611 + existing_call = None

612 + for existing in merged_tool_calls:

613 + if existing.id == tool_call.id:

614 + existing_call = existing

615 + break

616 +

617 + if existing_call and existing_call.function:

618 + # Merge to existing tool_call

619 + if tool_call.function and tool_call.function.name:

620 + existing_call.function.name = tool_call.function.name

621 + if (

622 + tool_call.function

623 + and tool_call.function.arguments is not None

624 + ):

625 + if existing_call.function.arguments is None:

626 + existing_call.function.arguments = ""

627 +

628 + # For streaming JSON parameters,

629 + # simply concatenate in order

630 + new_args = tool_call.function.arguments

631 + existing_call.function.arguments += new_args

632 + if tool_call.type:

633 + existing_call.type = tool_call.type

634 + else:

635 + # Add new tool_call

636 + merged_tool_calls.append(tool_call)

637 +

638 + return DeltaMessage(

639 + content=merged_content if merged_content else None,

640 + tool_calls=merged_tool_calls,

641 + )

642 +

643 + def _preprocess_xml_chunk(self, chunk: str) -> str:

644 + """

645 + Preprocess XML chunk, handle non-standard formats,

646 + and escape special characters

647 +

648 + Args:

649 + chunk: Original XML chunk

650 +

651 + Returns:

652 + Processed XML chunk

653 + """

654 +

655 + # Check if this is a tool_call related element

656 + is_tool_call = False

657 + if chunk.startswith(self.tool_call_start_token) or chunk.startswith(

658 + self.tool_call_end_token

659 + ):

660 + is_tool_call = True

661 + # Check for function tags (including malformed ones without =)

662 + # <function=xxx>, </function>, <function xxx>, <functionxxx>

663 + if (

664 + chunk.startswith(self.function_start_token)

665 + or chunk.startswith(self.function_end_token)

666 + or chunk.startswith("<function ")

667 + or re.match(r"^<function[a-zA-Z_]", chunk)

668 + ): # <functionXXX without space or =

669 + is_tool_call = True

670 + if chunk.startswith(self.parameter_start_token) or chunk.startswith(

671 + self.parameter_end_token

672 + ):

673 + is_tool_call = True

674 +

675 + # Fallback: fix incomplete <parameter= or <function= tags without

676 + # closing >

677 + # This handles cases like: <parameter=-C:\n or <parameter=-B 5\n

678 + # Apply when parsing tool calls OR when chunk looks like a function/

679 + # parameter tag

680 + if (

681 + self.current_call_id is not None

682 + or chunk.startswith("<function")

683 + or chunk.startswith("<parameter")

684 + ):

685 + chunk = self._fix_incomplete_tag_in_chunk(chunk)

686 +

687 + # Handle <function=name> format -> <function name="name">

688 + processed = re.sub(r"<function=([^>]+)>", r'<function name="\1">', chunk)

689 + # Handle <parameter=name> format -> <parameter name="name">

690 + processed = re.sub(r"<parameter=([^>]+)>", r'<parameter name="\1">', processed)

691 +

692 + original_chunk = chunk

693 + # If in parameter value accumulation mode

694 + if self._pre_inside_parameter:

695 + # Parameter end: output accumulated raw text

696 + # safely then return </parameter>

697 + if processed.startswith("</parameter>"):

698 + body_text = self._pre_param_buffer

699 + # Trigger deferred parsing mode

700 + # literal_eval+json output in end_element

701 + self.defer_current_parameter = True

702 + self.deferred_param_raw_value = body_text

703 + # Clean up state

704 + self._pre_inside_parameter = False

705 + self._pre_param_buffer = ""

706 + self._pre_current_param_name = None

707 + safe_text = self._escape_xml_special_chars(body_text)

708 + return f"{safe_text}</parameter>"

709 + else:

710 + # If this is the first block of content after entering parameter

711 + # evaluate if deferred parsing is needed;

712 + # If not needed, exit accumulation mode

713 + # and pass through directly

714 + if self._pre_param_buffer == "":

715 + # Get current parameter type

716 + param_type = (

717 + self._get_param_type(self._pre_current_param_name)

718 + if self._pre_current_param_name

719 + else "string"

720 + )

721 + # Only these types need deferred parsing to

722 + # handle Python literals containing single quotes

723 + is_object_type = param_type in ["object"]

724 + is_complex_type = (

725 + param_type in ["array", "arr", "sequence"]

726 + or param_type.startswith("dict")

727 + or param_type.startswith("list")

728 + )

729 +

730 + # Only delay when contains container symbols

731 + # and has single quotes and is complex type

732 + has_container_hint = (

733 + ("[" in original_chunk)

734 + or ("{" in original_chunk)

735 + or ("(" in original_chunk)

736 + )

737 +

738 + # Determine if deferred parsing is needed

739 + need_defer = False

740 + if is_complex_type:

741 + # Complex type, always need deferred parsing

742 + need_defer = True

743 + elif (

744 + is_object_type

745 + and has_container_hint

746 + and ("'" in original_chunk)

747 + ):

748 + # Object type with container symbols

749 + # and single quotes, need deferred parsing

750 + need_defer = True

751 +

752 + if not need_defer:

753 + # No need for deferred parsing,

754 + # exit parameter mode directly

755 + self._pre_inside_parameter = False

756 + return self._escape_xml_special_chars(original_chunk)

757 + self._pre_param_buffer += original_chunk

758 + return ""

759 +

760 + # Parameter start: enable accumulation

761 + if processed.startswith("<parameter name="):

762 + m = re.match(r'<parameter name="([^"]+)">', processed)

763 + if m:

764 + self._pre_current_param_name = m.group(1)

765 + self._pre_inside_parameter = True

766 + self._pre_param_buffer = ""

767 + return processed

768 +

769 + # If processed doesn't contain special_token, escape processed

770 + # This is because XML parsing encounters special characters

771 + # and reports errors, so escaping is needed

772 + if not is_tool_call:

773 + processed = self._escape_xml_special_chars(processed)

774 + return processed

775 +

776 + def _emit_delta(self, delta: DeltaMessage):

777 + """Emit Delta response (streaming output)"""

778 + self.deltas.append(delta)

779 +

780 + def _auto_close_open_parameter_if_needed(self, incoming_tag: str | None = None):

781 + """Before starting to process new elements,

782 + if there are unclosed tags from before,

783 + automatically complete their endings to the parser.

784 + - If there are unclosed parameters,

785 + it's equivalent to feeding `</parameter>`

786 + - When about to start a new function or tool_call,

787 + if there are unclosed functions, complete `</function>`.

788 + - When about to start a new tool_call,

789 + if there are unclosed tool_calls, complete `</tool_call>`.

790 + """

791 + # First close unclosed parameters

792 + if self.current_param_name:

793 + self._end_element("parameter")

794 +

795 + # If about to start new function or tool_call,

796 + # and there are unclosed functions, close function first

797 + if incoming_tag in ("function", "tool_call") and self.current_function_name:

798 + self._end_element("function")

799 +

800 + # If about to start new tool_call,

801 + # and there are unclosed tool_calls, close tool_call first

802 + if incoming_tag == "tool_call" and self.current_call_id:

803 + self._end_element("tool_call")

804 +

805 + def _start_element(self, name: str, attrs: dict[str, str]):

806 + """Handle XML start element events"""

807 +

808 + if name == "root":

809 + return

810 +

811 + if name == "tool_call":

812 + # Before opening new tool_call,

813 + # automatically complete previous unclosed tags

814 + self._auto_close_open_parameter_if_needed("tool_call")

815 +

816 + self.parameters = {}

817 + self.current_call_id = make_tool_call_id()

818 + self.current_param_is_first = True

819 + self.tool_call_index += 1

820 + elif name.startswith("function") or (name == "function"):

821 + # If missing tool_call, manually complete

822 + if not self.current_call_id:

823 + self._start_element("tool_call", {})

824 + # Before opening new function,

825 + # automatically complete previous unclosed tags (parameter/function)

826 + self._auto_close_open_parameter_if_needed("function")

827 + function_name = self._extract_function_name(name, attrs)

828 + self.current_function_name = function_name

829 + self.current_function_open = True

830 + if function_name:

831 + delta = DeltaMessage(

832 + tool_calls=[

833 + DeltaToolCall(

834 + index=self.tool_call_index - 1,

835 + id=self.current_call_id,

836 + type="function",

837 + function=DeltaFunctionCall(

838 + name=function_name, arguments=""

839 + ),

840 + )

841 + ]

842 + )

843 + self._emit_delta(delta)

844 + elif name.startswith("parameter") or (name == "parameter"):

845 + # If previous parameter hasn't ended normally,

846 + # complete its end first, then start new parameter

847 + self._auto_close_open_parameter_if_needed("parameter")

848 + param_name = self._extract_parameter_name(name, attrs)

849 + self.current_param_name = param_name

850 + self.current_param_value = ""

851 + self.current_param_value_converted = ""

852 + self.start_quote_emitted = False # Reset start quote flag

853 +

854 + # Only output parameter name and colon,

855 + # don't output quotes

856 + # decide after parameter value type is determined

857 + if param_name:

858 + if not self.parameters:

859 + # First parameter

860 + # start JSON, only output parameter name and colon

861 + json_start = f'{{"{param_name}": '

862 + delta = DeltaMessage(

863 + tool_calls=[

864 + DeltaToolCall(

865 + index=self.tool_call_index - 1,

866 + id=self.current_call_id,

867 + type="function",

868 + function=DeltaFunctionCall(

869 + name=None, arguments=json_start

870 + ),

871 + )

872 + ]

873 + )

874 + self._emit_delta(delta)

875 + self.current_param_is_first = True

876 + else:

877 + # Subsequent parameters

878 + # add comma and parameter name, no quotes

879 + json_continue = f', "{param_name}": '

880 + delta = DeltaMessage(

881 + tool_calls=[

882 + DeltaToolCall(

883 + index=self.tool_call_index - 1,

884 + id=self.current_call_id,

885 + type="function",

886 + function=DeltaFunctionCall(

887 + name=None, arguments=json_continue

888 + ),

889 + )

890 + ]

891 + )

892 + self._emit_delta(delta)

893 + self.current_param_is_first = False

894 +

895 + def _char_data(self, data: str):

896 + """Handle XML character data events"""

897 + if data and self.current_param_name:

898 + # If preprocessing stage determines deferred parsing is needed,

899 + # only cache character data, no streaming output

900 + if self.defer_current_parameter:

901 + original_data = data

902 + if self.should_emit_end_newline:

903 + original_data = "\n" + original_data

904 + self.should_emit_end_newline = False

905 + if original_data.endswith("\n"):

906 + self.should_emit_end_newline = True

907 + original_data = original_data[:-1]

908 + self.current_param_value += original_data

909 + return

910 +

911 + param_type = self._get_param_type(self.current_param_name)

912 +

913 + # Check if this is the first time receiving data for this parameter

914 + # If this is the first packet of data and starts with \n, remove \n

915 + if not self.current_param_value and data.startswith("\n"):

916 + data = data[1:]

917 +

918 + # Output start quote for string type (if not already output)

919 + if (

920 + param_type in ["string", "str", "text", "varchar", "char", "enum"]

921 + and not self.start_quote_emitted

922 + ):

923 + quote_delta = DeltaMessage(

924 + tool_calls=[

925 + DeltaToolCall(

926 + index=self.tool_call_index - 1,

927 + id=self.current_call_id,

928 + type="function",

929 + function=DeltaFunctionCall(name=None, arguments='"'),

930 + )

931 + ]

932 + )

933 + self._emit_delta(quote_delta)

934 + self.start_quote_emitted = True

935 +

936 + if not data:

937 + return

938 +

939 + original_data = data

940 + # Delay output of trailing newline

941 + if self.should_emit_end_newline:

942 + original_data = "\n" + original_data

943 + self.should_emit_end_newline = False

944 + if original_data.endswith("\n"):

945 + self.should_emit_end_newline = True

946 + original_data = original_data[:-1]

947 + self.current_param_value += original_data

948 +

949 + # convert parameter value by param_type

950 + converted_value = self._convert_param_value(

951 + self.current_param_value, param_type

952 + )

953 + output_data = self._convert_for_json_streaming(converted_value, param_type)

954 +

955 + delta_data = output_data[len(self.current_param_value_converted) :]

956 + self.current_param_value_converted = output_data

957 +

958 + delta = DeltaMessage(

959 + tool_calls=[

960 + DeltaToolCall(

961 + index=self.tool_call_index - 1,

962 + id=self.current_call_id,

963 + type="function",

964 + function=DeltaFunctionCall(name=None, arguments=delta_data),

965 + )

966 + ]

967 + )

968 + self._emit_delta(delta)

969 +

970 + def _end_element(self, name: str):

971 + """Handle XML end element events"""

972 +

973 + if name == "root":

974 + return

975 +

976 + # If function or tool_call ends and there are still unclosed parameters,

977 + # complete parameter end first

978 + if (

979 + name.startswith("function") or name == "function" or name == "tool_call"

980 + ) and self.current_param_name:

981 + self._auto_close_open_parameter_if_needed()

982 +

983 + if (

984 + name.startswith("parameter") or name == "parameter"

985 + ) and self.current_param_name:

986 + # End current parameter

987 + param_name = self.current_param_name

988 + param_value = self.current_param_value

989 +

990 + # If in deferred parsing mode,

991 + # perform overall parsing on raw content

992 + # accumulated in preprocessing stage and output once

993 + if self.defer_current_parameter:

994 + raw_text = (

995 + self.deferred_param_raw_value

996 + if self.deferred_param_raw_value

997 + else param_value

998 + )

999 + parsed_value = None

1000 + output_arguments = None

1001 + try:

1002 + # If previously delayed trailing newline,

1003 + # add it back before parsing

1004 + if self.should_emit_end_newline:

1005 + raw_for_parse = raw_text + "\n"

1006 + else:

1007 + raw_for_parse = raw_text

1008 + parsed_value = ast.literal_eval(raw_for_parse)

1009 + output_arguments = json.dumps(parsed_value, ensure_ascii=False)

1010 + except Exception:

1011 + # Fallback: output as string as-is

1012 + output_arguments = json.dumps(raw_text, ensure_ascii=False)

1013 + parsed_value = raw_text

1014 +

1015 + delta = DeltaMessage(

1016 + tool_calls=[

1017 + DeltaToolCall(

1018 + index=self.tool_call_index - 1,

1019 + id=self.current_call_id,

1020 + type="function",

1021 + function=DeltaFunctionCall(

1022 + name=None, arguments=output_arguments

1023 + ),

1024 + )

1025 + ]

1026 + )

1027 + self._emit_delta(delta)

1028 +

1029 + # Clean up and store

1030 + self.should_emit_end_newline = False

1031 + self.parameters[param_name] = parsed_value

1032 + self.current_param_name = None

1033 + self.current_param_value = ""

1034 + self.current_param_value_converted = ""

1035 + self.start_quote_emitted = False

1036 + self.defer_current_parameter = False

1037 + self.deferred_param_raw_value = ""

1038 + return

1039 +

1040 + param_type = self._get_param_type(param_name)

1041 +

1042 + # convert complete parameter value by param_type

1043 + converted_value = self._convert_param_value(param_value, param_type)

1044 +

1045 + # Decide whether to add end quote based on parameter type

1046 + if param_type in ["string", "str", "text", "varchar", "char", "enum"]:

1047 + # For empty string parameters, need special handling

1048 + if not param_value and not self.start_quote_emitted:

1049 + # No start quote output,

1050 + # directly output complete empty string

1051 + delta = DeltaMessage(

1052 + tool_calls=[

1053 + DeltaToolCall(

1054 + index=self.tool_call_index - 1,

1055 + id=self.current_call_id,

1056 + type="function",

1057 + function=DeltaFunctionCall(name=None, arguments='""'),

1058 + )

1059 + ]

1060 + )

1061 + self._emit_delta(delta)

1062 + else:

1063 + # Non-empty parameter value, output end quote

1064 + delta = DeltaMessage(

1065 + tool_calls=[

1066 + DeltaToolCall(

1067 + index=self.tool_call_index - 1,

1068 + id=self.current_call_id,

1069 + type="function",

1070 + function=DeltaFunctionCall(name=None, arguments='"'),

1071 + )

1072 + ]

1073 + )

1074 + self._emit_delta(delta)

1075 +

1076 + self.should_emit_end_newline = False

1077 + # Store converted value

1078 + self.parameters[param_name] = converted_value

1079 + self.current_param_name = None

1080 + self.current_param_value = ""

1081 + self.current_param_value_converted = ""

1082 + self.start_quote_emitted = False

1083 +

1084 + elif name.startswith("function") or name == "function":

1085 + # if there are parameters, close JSON object

1086 + if self.parameters:

1087 + delta = DeltaMessage(

1088 + tool_calls=[

1089 + DeltaToolCall(

1090 + index=self.tool_call_index - 1,

1091 + id=self.current_call_id,

1092 + type="function",

1093 + function=DeltaFunctionCall(name=None, arguments="}"),

1094 + )

1095 + ]

1096 + )

1097 + self._emit_delta(delta)

1098 + # return empty object

1099 + else:

1100 + delta = DeltaMessage(

1101 + tool_calls=[

1102 + DeltaToolCall(

1103 + index=self.tool_call_index - 1,

1104 + id=self.current_call_id,

1105 + type="function",

1106 + function=DeltaFunctionCall(name=None, arguments="{}"),

1107 + )

1108 + ]

1109 + )

1110 + self._emit_delta(delta)

1111 + self.current_function_open = False

1112 + self.current_function_name = (

1113 + None # Clear function name to prevent duplicate closing

1114 + )

1115 +

1116 + elif name == "tool_call":

1117 + # Before ending tool_call,

1118 + # ensure function is closed to complete missing right brace

1119 + if self.current_function_open:

1120 + # If there are still unclosed parameters, close them first

1121 + if self.current_param_name:

1122 + self._end_element("parameter")

1123 + # Close function, ensure output '}' or '{}'

1124 + self._end_element("function")

1125 + # Final Delta

1126 + delta = DeltaMessage(

1127 + tool_calls=[

1128 + DeltaToolCall(

1129 + index=self.tool_call_index - 1,

1130 + id=self.current_call_id,

1131 + type="function",

1132 + function=DeltaFunctionCall(name=None, arguments=""),

1133 + )

1134 + ]

1135 + )

1136 + self._emit_delta(delta)

1137 +

1138 + # Check if there's text content to output (between tool_calls)

1139 + if self.text_content_buffer.strip():

1140 + text_delta = DeltaMessage(content=self.text_content_buffer)

1141 + self._emit_delta(text_delta)

1142 +

1143 + self._reset_xml_parser_after_tool_call()

1144 +

1145 + def setup_parser(self):

1146 + """Set up XML parser event handlers"""

1147 + self.parser.buffer_text = True

1148 + self.parser.StartElementHandler = self._start_element

1149 + self.parser.EndElementHandler = self._end_element

1150 + self.parser.CharacterDataHandler = self._char_data

1151 +

1152 + def set_tools(self, tools: list[ChatCompletionToolsParam] | None):

1153 + """Set tool configuration information"""

1154 + self.tools = tools

1155 +

1156 + def _extract_function_name(self, name: str, attrs: dict[str, str]) -> str | None:

1157 + """Extract function name from various formats"""

1158 + if attrs and "name" in attrs:

1159 + return attrs["name"]

1160 +

1161 + if "=" in name:

1162 + parts = name.split("=", 1)

1163 + if len(parts) == 2 and parts[0] == "function":

1164 + return parts[1]

1165 +

1166 + return None

1167 +

1168 + def _extract_parameter_name(self, name: str, attrs: dict[str, str]) -> str | None:

1169 + """Extract parameter name from various formats"""

1170 + if attrs and "name" in attrs:

1171 + return attrs["name"]

1172 +

1173 + if "=" in name:

1174 + parts = name.split("=", 1)

1175 + if len(parts) == 2 and parts[0] == "parameter":

1176 + return parts[1]

1177 +

1178 + return None

1179 +

1180 + def _get_param_type(self, param_name: str) -> str:

1181 + """Get parameter type based on tool configuration, defaults to string

1182 + Args:

1183 + param_name: Parameter name

1184 +

1185 + Returns:

1186 + Parameter type

1187 + """

1188 + if not self.tools or not self.current_function_name:

1189 + return "string"

1190 +

1191 + for tool in self.tools:

1192 + if not hasattr(tool, "type") or not (

1193 + hasattr(tool, "function") and hasattr(tool.function, "name")

1194 + ):

1195 + continue

1196 + if (

1197 + tool.type == "function"

1198 + and tool.function.name == self.current_function_name

1199 + ):

1200 + if not hasattr(tool.function, "parameters"):

1201 + return "string"

1202 + params = tool.function.parameters

1203 + if isinstance(params, dict) and "properties" in params:

1204 + properties = params["properties"]

1205 + if param_name in properties and isinstance(

1206 + properties[param_name], dict

1207 + ):

1208 + return self.repair_param_type(

1209 + str(properties[param_name].get("type", "string"))

1210 + )

1211 + elif isinstance(params, dict) and param_name in params:

1212 + param_config = params[param_name]

1213 + if isinstance(param_config, dict):

1214 + return self.repair_param_type(

1215 + str(param_config.get("type", "string"))

1216 + )

1217 + break

1218 + return "string"

1219 +

1220 + def repair_param_type(self, param_type: str) -> str:

1221 + """Repair unknown parameter types by treating them as string

1222 + Args:

1223 + param_type: Parameter type

1224 +

1225 + Returns:

1226 + Repaired parameter type

1227 + """

1228 + if (

1229 + param_type in ["string", "str", "text", "varchar", "char", "enum"]

1230 + or param_type.startswith("int")

1231 + or param_type.startswith("uint")

1232 + or param_type.startswith("long")

1233 + or param_type.startswith("short")

1234 + or param_type.startswith("unsigned")

1235 + or param_type.startswith("num")

1236 + or param_type.startswith("float")

1237 + or param_type in ["boolean", "bool", "binary"]

1238 + or (

1239 + param_type in ["object", "array", "arr", "sequence"]

1240 + or param_type.startswith("dict")

1241 + or param_type.startswith("list")

1242 + )

1243 + ):

1244 + return param_type

1245 + else:

1246 + return "string"

1247 +

1248 + def _convert_param_value(self, param_value: str, param_type: str) -> Any:

1249 + """Convert value based on parameter type

1250 + Args:

1251 + param_value: Parameter value

1252 + param_type: Parameter type

1253 +

1254 + Returns:

1255 + Converted value

1256 + """

1257 + if param_value.lower() == "null":

1258 + return None

1259 +

1260 + param_type = param_type.strip().lower()

1261 + if param_type in ["string", "str", "text", "varchar", "char", "enum"]:

1262 + return param_value

1263 + elif (

1264 + param_type.startswith("int")

1265 + or param_type.startswith("uint")

1266 + or param_type.startswith("long")

1267 + or param_type.startswith("short")

1268 + or param_type.startswith("unsigned")

1269 + ):

1270 + try:

1271 + return int(param_value)

1272 + except (ValueError, TypeError):

1273 + logger.warning(

1274 + "Parsed value '%s' is not an integer, degenerating to string.",

1275 + param_value,

1276 + )

1277 + return param_value

1278 + elif param_type.startswith("num") or param_type.startswith("float"):

1279 + try:

1280 + float_param_value: float = float(param_value)

1281 + return (

1282 + float_param_value

1283 + if float_param_value - int(float_param_value) != 0

1284 + else int(float_param_value)

1285 + )

1286 + except (ValueError, TypeError):

1287 + logger.warning(

1288 + "Parsed value '%s' is not a float, degenerating to string.",

1289 + param_value,

1290 + )

1291 + return param_value

1292 + elif param_type in ["boolean", "bool", "binary"]:

1293 + param_value = param_value.lower()

1294 + return param_value == "true"

1295 + else:

1296 + return param_value

1297 +

1298 + def _convert_for_json_streaming(self, converted_value: Any, param_type: str) -> str:

1299 + """Convert converted_value based on

1300 + whether it's empty and if type is string

1301 + Args:

1302 + converted_value: Converted value

1303 + param_type: Parameter type

1304 +

1305 + Returns:

1306 + Converted string for streaming output

1307 + """

1308 + # Check if value is empty, but exclude numeric 0

1309 + if converted_value is None or converted_value == "":

1310 + return ""

1311 +

1312 + if param_type in ["string", "str", "text", "varchar", "char", "enum"]:

1313 + # String type, remove double quotes

1314 + return json.dumps(converted_value, ensure_ascii=False)[1:-1]

1315 + else:

1316 + # Non-string type, return complete JSON string

1317 + if not isinstance(converted_value, str):

1318 + return json.dumps(converted_value, ensure_ascii=False)

1319 + else:

1320 + return converted_value

1321 +

1322 + def _reset_xml_parser_after_tool_call(self):

1323 + """

1324 + Each tool_call is treated as a separate XML document,

1325 + so we need to reset the parser after each tool_call.

1326 + """

1327 +

1328 + # recreate XML parser

1329 + self.parser = ParserCreate()

1330 + self.setup_parser()

1331 +

1332 + # Reset current tool_call state

1333 + if self.current_call_id:

1334 + self.last_completed_call_id = self.current_call_id

1335 + self.current_call_id = None

1336 + self.current_function_name = None

1337 + self.current_function_open = False

1338 + self.parameters = {}

1339 + self.current_param_name = None

1340 + self.current_param_value = ""

1341 + self.current_param_value_converted = ""

1342 + self.current_param_is_first = False

1343 + self.should_emit_end_newline = False

1344 + self.start_quote_emitted = False

1345 + self.text_content_buffer = ""

1346 +

1347 + # Reset preprocessing and deferred parsing state

1348 + self._pre_inside_parameter = False

1349 + self._pre_param_buffer = ""

1350 + self._pre_current_param_name = None

1351 + self.defer_current_parameter = False

1352 + self.deferred_param_raw_value = ""

1353 +

1354 +

1355 +@ToolParserManager.register_module("step3p5")

1356 +class Step3p5ToolParser(ToolParser):

1357 + def __init__(self, tokenizer: TokenizerLike):

1358 + super().__init__(tokenizer)

1359 + self.parser = StreamingXMLToolCallParser()

1360 +

1361 + # Add missing attributes for compatibility with serving_chat.py

1362 + self.prev_tool_call_arr: list[dict] = []

1363 + self.streamed_args_for_tool: list[str] = []

1364 +

1365 + logger.info(

1366 + "vLLM Successfully import tool parser %s !", self.__class__.__name__

1367 + )

1368 +

1369 + def extract_tool_calls(

1370 + self,

1371 + model_output: str,

1372 + request: ChatCompletionRequest,

1373 + ) -> ExtractedToolCallInformation:

1374 + self.parser.reset_streaming_state()

1375 + # Reset tool call tracking arrays for new extraction

1376 + self.prev_tool_call_arr = []

1377 + self.streamed_args_for_tool = []

1378 + if request:

1379 + self.parser.set_tools(request.tools)

1380 + result = self.parser.parse_single_streaming_chunks(model_output)

1381 + if not result.tool_calls:

1382 + return ExtractedToolCallInformation(

1383 + tool_calls=[],

1384 + tools_called=False,

1385 + content=result.content,

1386 + )

1387 + else:

1388 + tool_calls = []

1389 + for tool_call in result.tool_calls:

1390 + if tool_call.function and tool_call.function.name:

1391 + tool_calls.append(

1392 + ToolCall(

1393 + id=tool_call.id,

1394 + type=tool_call.type,

1395 + function=FunctionCall(

1396 + name=tool_call.function.name,

1397 + arguments=tool_call.function.arguments,

1398 + ),

1399 + )

1400 + )

1401 +

1402 + # Update tool call tracking arrays for compatibility

1403 + tool_index = (

1404 + tool_call.index

1405 + if tool_call.index is not None

1406 + else len(self.prev_tool_call_arr) - 1

1407 + )

1408 +

1409 + # Ensure we have enough entries in our tracking arrays

1410 + while len(self.prev_tool_call_arr) <= tool_index:

1411 + self.prev_tool_call_arr.append({"name": "", "arguments": ""})

1412 + while len(self.streamed_args_for_tool) <= tool_index:

1413 + self.streamed_args_for_tool.append("")

1414 +

1415 + # Update tool call information

1416 + self.prev_tool_call_arr[tool_index]["name"] = (

1417 + tool_call.function.name

1418 + )

1419 + self.prev_tool_call_arr[tool_index]["arguments"] = (

1420 + tool_call.function.arguments

1421 + )

1422 +

1423 + # Update streamed arguments

1424 + if tool_call.function.arguments:

1425 + self.streamed_args_for_tool[tool_index] = (

1426 + tool_call.function.arguments

1427 + )

1428 +

1429 + return ExtractedToolCallInformation(

1430 + tool_calls=tool_calls,

1431 + tools_called=len(tool_calls) > 0,

1432 + content=result.content,

1433 + )

1434 +

1435 + def extract_tool_calls_streaming(

1436 + self,

1437 + previous_text: str,

1438 + current_text: str,

1439 + delta_text: str,

1440 + previous_token_ids: Sequence[int],

1441 + current_token_ids: Sequence[int],

1442 + delta_token_ids: Sequence[int],

1443 + request: ChatCompletionRequest,

1444 + ) -> DeltaMessage | None:

1445 + if not previous_text:

1446 + self.parser.reset_streaming_state()

1447 + # Reset tool call tracking arrays for new streaming session

1448 + self.prev_tool_call_arr = []

1449 + self.streamed_args_for_tool = []

1450 + if request:

1451 + self.parser.set_tools(request.tools)

1452 +

1453 + # Model sometimes outputs separately causing delta_text to be empty.

1454 + # If there were tool_calls before and all current tool_calls have ended,

1455 + # return an empty tool_call for outer streaming output

1456 + # to correctly output tool_call field

1457 + if not delta_text and delta_token_ids:

1458 + open_calls = current_text.count(

1459 + self.parser.tool_call_start_token

1460 + ) - current_text.count(self.parser.tool_call_end_token)

1461 + if (

1462 + open_calls == 0

1463 + and self.parser.tool_call_index > 0

1464 + or not self.parser.tool_call_index

1465 + and current_text

1466 + ):

1467 + return DeltaMessage(content="")

1468 + return None

1469 +

1470 + # Parse the delta text and get the result

1471 + result = self.parser.parse_single_streaming_chunks(delta_text)

1472 +

1473 + # Update tool call tracking arrays based on incremental parsing results

1474 + if result and result.tool_calls:

1475 + for tool_call in result.tool_calls:

1476 + if tool_call.function:

1477 + tool_index = (

1478 + tool_call.index

1479 + if tool_call.index is not None

1480 + else len(self.prev_tool_call_arr) - 1

1481 + )

1482 +

1483 + # Ensure we have enough entries in our tracking arrays

1484 + while len(self.prev_tool_call_arr) <= tool_index:

1485 + self.prev_tool_call_arr.append({"name": "", "arguments": ""})

1486 + while len(self.streamed_args_for_tool) <= tool_index:

1487 + self.streamed_args_for_tool.append("")

1488 +

1489 + # Update tool name if provided

1490 + if tool_call.function.name:

1491 + self.prev_tool_call_arr[tool_index]["name"] = (

1492 + tool_call.function.name

1493 + )

1494 +

1495 + # Update arguments incrementally

1496 + if tool_call.function.arguments is not None:

1497 + # Concatenate the incremental arguments

1498 + # to the existing streamed arguments

1499 + self.prev_tool_call_arr[tool_index]["arguments"] += (

1500 + tool_call.function.arguments

1501 + )

1502 + self.streamed_args_for_tool[tool_index] += (

1503 + tool_call.function.arguments

1504 + )

1505 + return result

1506 +

1507 + def parser_should_check_for_unstreamed_tool_arg_tokens(self) -> bool:

1508 + """

1509 + Skip the remaining_call calculation in serving_chat

1510 + """

1511 + return False

vllm/transformers_utils/config.py

@@ -97,6 +97,7 @@ _CONFIG_REGISTRY: dict[str, type[PretrainedConfig]] = LazyConfigDict(

97 ultravox="UltravoxConfig",

98 step3_vl="Step3VLConfig",

99 step3_text="Step3TextConfig",

100 + step3p5="Step3p5Config",

101 qwen3_asr="Qwen3ASRConfig",

102 qwen3_next="Qwen3NextConfig",

103 lfm2_moe="Lfm2MoeConfig",

vllm/transformers_utils/configs/__init__.py

@@ -52,6 +52,7 @@ _CLASS_TO_MODULE: dict[str, str] = {

52 "Step3VLConfig": "vllm.transformers_utils.configs.step3_vl",

53 "Step3VisionEncoderConfig": "vllm.transformers_utils.configs.step3_vl",

54 "Step3TextConfig": "vllm.transformers_utils.configs.step3_vl",

55 + "Step3p5Config": "vllm.transformers_utils.configs.step3p5",

56 "Qwen3ASRConfig": "vllm.transformers_utils.configs.qwen3_asr",

57 "Qwen3NextConfig": "vllm.transformers_utils.configs.qwen3_next",

58 "Tarsier2Config": "vllm.transformers_utils.configs.tarsier2",

@@ -95,6 +96,7 @@ __all__ = [

96 "Step3VLConfig",

97 "Step3VisionEncoderConfig",

98 "Step3TextConfig",

99 + "Step3p5Config",

100 "Qwen3ASRConfig",

101 "Qwen3NextConfig",

102 "Tarsier2Config",

vllm/transformers_utils/configs/step3p5.py new

+100

@@ -0,0 +1,100 @@

1 +# SPDX-License-Identifier: Apache-2.0

2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3 +from typing import Any

4 +

5 +from transformers.configuration_utils import PretrainedConfig

6 +

7 +

8 +class Step3p5Config(PretrainedConfig):

9 + model_type = "step3p5"

10 +

11 + def __init__(

12 + self,

13 + hidden_size: int = 5120,

14 + intermediate_size: int = 13312,

15 + num_attention_heads: int = 40,

16 + num_attention_groups: int = 8,

17 + num_hidden_layers: int = 48,

18 + max_seq_len: int = 4096,

19 + vocab_size: int = 65536,

20 + rms_norm_eps: float = 1e-5,

21 + moe_every_n_layer: int = 2,

22 + use_moe: bool = False,

23 + moe_intermediate_size: int = 10240,

24 + moe_num_experts: int = 16,

25 + moe_top_k: int = 4,

26 + moe_layer_offset: int = 0,

27 + rope_theta: float | list[float] | None = 500000,

28 + rope_scaling: dict[str, Any] | None = None,

29 + head_dim: int | None = None,

30 + share_expert_dim: int | None = None,

31 + norm_expert_weight: bool = True,

32 + bos_token_id: list[int] | int | None = None,

33 + eos_token_id: list[int] | int | None = None,

34 + moe_router_activation: str = "softmax",

35 + moe_router_scaling_factor: float = 1.0,

36 + att_impl_type: str = "GQA",

37 + use_head_wise_attn_gate: bool = False,

38 + use_moe_router_bias: bool = True,

39 + need_fp32_gate: bool = True,

40 + layer_types: list[str] | None = None,

41 + use_rope_layers: list[bool] | None = None,

42 + yarn_only_types: list[str] | None = None,

43 + attention_other_setting: dict[str, Any] | None = None,

44 + num_nextn_predict_layers: int = 0,

45 + swiglu_limits: list[float] | None = None,

46 + swiglu_limits_shared: list[float] | None = None,

47 + max_position_embeddings: int | None = None,

48 + **kwargs,

49 + ):

50 + self.hidden_size = hidden_size

51 + self.intermediate_size = intermediate_size

52 + self.num_attention_heads = num_attention_heads

53 + self.num_attention_groups = num_attention_groups

54 + self.num_hidden_layers = num_hidden_layers

55 + self.max_seq_len = max_seq_len

56 + self.vocab_size = vocab_size

57 + self.rms_norm_eps = rms_norm_eps

58 + self.use_moe = use_moe

59 + self.moe_intermediate_size = moe_intermediate_size

60 + self.moe_every_n_layer = moe_every_n_layer

61 + self.moe_num_experts = moe_num_experts

62 + self.num_experts_per_tok = moe_top_k

63 + self.moe_top_k = moe_top_k

64 + self.moe_layer_offset = moe_layer_offset

65 +

66 + self.rope_theta = rope_theta

67 + self.rope_scaling = rope_scaling

68 + self.head_dim = head_dim

69 + if share_expert_dim is None:

70 + self.share_expert_dim = self.moe_intermediate_size * self.moe_top_k

71 + else:

72 + self.share_expert_dim = share_expert_dim

73 + self.norm_expert_weight = norm_expert_weight

74 +

75 + self.max_position_embeddings = max_position_embeddings

76 + self.moe_router_activation = moe_router_activation

77 + self.moe_router_scaling_factor = moe_router_scaling_factor

78 + self.use_moe_router_bias = use_moe_router_bias

79 + self.need_fp32_gate = need_fp32_gate

80 +

81 + self.att_impl_type = att_impl_type

82 + self.use_head_wise_attn_gate = use_head_wise_attn_gate

83 + self.layer_types = layer_types

84 + self.use_rope_layers = use_rope_layers

85 + self.yarn_only_types = yarn_only_types

86 + self.attention_other_setting = attention_other_setting

87 + self.num_nextn_predict_layers = num_nextn_predict_layers

88 + self.swiglu_limits = swiglu_limits

89 + self.swiglu_limits_shared = swiglu_limits_shared

90 +

91 + resolved_bos_token_id = 1 if bos_token_id is None else bos_token_id

92 + resolved_eos_token_id = [2, 3] if eos_token_id is None else eos_token_id

93 + self.bos_token_id = resolved_bos_token_id

94 + self.eos_token_id = resolved_eos_token_id

95 +

96 + super().__init__(

97 + bos_token_id=resolved_bos_token_id,

98 + eos_token_id=resolved_eos_token_id,

99 + **kwargs,

100 + )