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

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 15, 01:00 UTC

GLM-4 Update (#20736)

Signed-off-by: zRzRzRzRzRzRzR <2448370773@qq.com> Signed-off-by: Isotr0py <mozf@mail2.sysu.edu.cn> Signed-off-by: Lu Fang <fanglu@fb.com> Co-authored-by: Isotr0py <mozf@mail2.sysu.edu.cn> Co-authored-by: Lu Fang <fanglu@fb.com>

Yuxuan Zhang committed Jul 20, 2025 at 06:40 UTC 10eb24cc91315481414fba0e0134209e6a9d7c94

14 files changed +2006 -11

benchmarks/kernels/benchmark_moe.py

+5 -1

@@ -576,7 +576,11 @@ def main(args: argparse.Namespace):

576 topk = config.num_experts_per_tok

577 intermediate_size = config.intermediate_size

578 shard_intermediate_size = 2 * intermediate_size // args.tp_size

579 - elif config.architectures[0] in ("DeepseekV3ForCausalLM", "DeepseekV2ForCausalLM"):

579 + elif config.architectures[0] in (

580 + "DeepseekV3ForCausalLM",

581 + "DeepseekV2ForCausalLM",

582 + "Glm4MoeForCausalLM",

583 + ):

584 E = config.n_routed_experts

585 topk = config.num_experts_per_tok

586 intermediate_size = config.moe_intermediate_size

benchmarks/kernels/benchmark_moe_permute_unpermute.py

@@ -318,6 +318,7 @@ def main(args: argparse.Namespace):

318 elif (

319 config.architectures[0] == "DeepseekV3ForCausalLM"

320 or config.architectures[0] == "DeepseekV2ForCausalLM"

321 + or config.architectures[0] == "Glm4MoeForCausalLM"

322 ):

323 E = config.n_routed_experts

324 topk = config.num_experts_per_tok

docs/models/supported_models.md

@@ -576,6 +576,7 @@ Specified using `--task generate`.

576 | `Gemma3ForConditionalGeneration` | Gemma 3 | T + I+ | `google/gemma-3-4b-it`, `google/gemma-3-27b-it`, etc. | ✅︎ | ✅︎ | ⚠️ |

577 | `GLM4VForCausalLM`^ | GLM-4V | T + I | `THUDM/glm-4v-9b`, `THUDM/cogagent-9b-20241220`, etc. | ✅︎ | ✅︎ | ✅︎ |

578 | `Glm4vForConditionalGeneration` | GLM-4.1V-Thinking | T + IE+ + VE+ | `THUDM/GLM-4.1V-9B-Thinking`, etc. | ✅︎ | ✅︎ | ✅︎ |

579 +| `Glm4MoeForCausalLM` | GLM-4.5 | T + IE+ + VE+ | `THUDM/GLM-4.5`, etc. | ✅︎ | ✅︎ | ✅︎ |

580 | `GraniteSpeechForConditionalGeneration` | Granite Speech | T + A | `ibm-granite/granite-speech-3.3-8b` | ✅︎ | ✅︎ | ✅︎ |

581 | `H2OVLChatModel` | H2OVL | T + IE+ | `h2oai/h2ovl-mississippi-800m`, `h2oai/h2ovl-mississippi-2b`, etc. | | ✅︎ | ✅︎ |

582 | `Idefics3ForConditionalGeneration` | Idefics3 | T + I | `HuggingFaceM4/Idefics3-8B-Llama3`, etc. | ✅︎ | | ✅︎ |

tests/models/registry.py

@@ -360,6 +360,9 @@ _MULTIMODAL_EXAMPLE_MODELS = {

360 trust_remote_code=True,

361 hf_overrides={"architectures": ["GLM4VForCausalLM"]}), # noqa: E501

362 "Glm4vForConditionalGeneration": _HfExamplesInfo("THUDM/GLM-4.1V-9B-Thinking", min_transformers_version="4.53"), # noqa: E501

363 + "Glm4MoeForCausalLM": _HfExamplesInfo("THUDM/GLM-4.5",

364 + min_transformers_version="4.54",

365 + is_available_online=False), # noqa: E501

366 "H2OVLChatModel": _HfExamplesInfo("h2oai/h2ovl-mississippi-800m",

367 extras={"2b": "h2oai/h2ovl-mississippi-2b"}, # noqa: E501

368 max_transformers_version="4.48", # noqa: E501

@@ -485,6 +488,10 @@ _SPECULATIVE_DECODING_EXAMPLE_MODELS = {

488 is_available_online=False,

489 speculative_model="openbmb/MiniCPM-2B-sft-bf16",

490 tokenizer="openbmb/MiniCPM-2B-sft-bf16"),

491 + "Glm4MoeMTPModel": _HfExamplesInfo("THUDM/GLM-4.5",

492 + speculative_model="THUDM/GLM-4.5",

493 + min_transformers_version="4.54",

494 + is_available_online=False),

495 "MiMoMTPModel": _HfExamplesInfo("XiaomiMiMo/MiMo-7B-RL",

496 trust_remote_code=True,

497 speculative_model="XiaomiMiMo/MiMo-7B-RL")

tests/tool_use/test_glm4_moe_tool_parser.py new

+410

@@ -0,0 +1,410 @@

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

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

3 +# ruff: noqa: E501

4 +

5 +import json

6 +

7 +import pytest

8 +

9 +from vllm.entrypoints.openai.protocol import FunctionCall, ToolCall

10 +from vllm.entrypoints.openai.tool_parsers import Glm4MoeModelToolParser

11 +from vllm.transformers_utils.tokenizer import get_tokenizer

12 +

13 +pytest.skip("skip glm4_moe parser test", allow_module_level=True)

14 +# Use a common model that is likely to be available

15 +MODEL = "THUDM/GLM-4.5"

16 +

17 +

18 +@pytest.fixture(scope="module")

19 +def glm4_moe_tokenizer():

20 + return get_tokenizer(tokenizer_name=MODEL)

21 +

22 +

23 +@pytest.fixture

24 +def glm4_moe_tool_parser(glm4_moe_tokenizer):

25 + return Glm4MoeModelToolParser(glm4_moe_tokenizer)

26 +

27 +

28 +def assert_tool_calls(actual_tool_calls: list[ToolCall],

29 + expected_tool_calls: list[ToolCall]):

30 + assert len(actual_tool_calls) == len(expected_tool_calls)

31 +

32 + for actual_tool_call, expected_tool_call in zip(actual_tool_calls,

33 + expected_tool_calls):

34 + assert isinstance(actual_tool_call.id, str)

35 + assert len(actual_tool_call.id) > 0

36 +

37 + assert actual_tool_call.type == "function"

38 + assert actual_tool_call.function.name == expected_tool_call.function.name

39 + # Compare arguments as JSON objects to handle formatting differences

40 + actual_args = json.loads(actual_tool_call.function.arguments)

41 + expected_args = json.loads(expected_tool_call.function.arguments)

42 + assert actual_args == expected_args

43 +

44 +

45 +def test_extract_tool_calls_no_tools(glm4_moe_tool_parser):

46 + model_output = "This is a test"

47 + extracted_tool_calls = glm4_moe_tool_parser.extract_tool_calls(

48 + model_output, request=None) # type: ignore[arg-type]

49 + assert not extracted_tool_calls.tools_called

50 + assert extracted_tool_calls.tool_calls == []

51 + assert extracted_tool_calls.content == model_output

52 +

53 +

54 +@pytest.mark.parametrize(

55 + ids=[

56 + "single_tool_call",

57 + "multiple_tool_calls",

58 + "tool_call_with_content_before",

59 + "tool_call_with_mixed_args",

60 + "tool_call_with_chinese_content",

61 + ],

62 + argnames=["model_output", "expected_tool_calls", "expected_content"],

63 + argvalues=[

64 + (

65 + """<tool_call>get_current_weather

66 + <arg_key>city</arg_key>

67 + <arg_value>Dallas</arg_value>

68 + <arg_key>state</arg_key>

69 + <arg_value>TX</arg_value>

70 + <arg_key>unit</arg_key>

71 + <arg_value>fahrenheit</arg_value>

72 + </tool_call>""",

73 + [

74 + ToolCall(function=FunctionCall(

75 + name="get_current_weather",

76 + arguments=json.dumps({

77 + "city": "Dallas",

78 + "state": "TX",

79 + "unit": "fahrenheit",

80 + }),

81 + ))

82 + ],

83 + None,

84 + ),

85 + (

86 + """<tool_call>get_current_weather

87 + <arg_key>city</arg_key>

88 + <arg_value>Dallas</arg_value>

89 + <arg_key>state</arg_key>

90 + <arg_value>TX</arg_value>

91 + <arg_key>unit</arg_key>

92 + <arg_value>fahrenheit</arg_value>

93 + </tool_call>

94 + <tool_call>get_current_weather

95 + <arg_key>city</arg_key>

96 + <arg_value>Orlando</arg_value>

97 + <arg_key>state</arg_key>

98 + <arg_value>FL</arg_value>

99 + <arg_key>unit</arg_key>

100 + <arg_value>fahrenheit</arg_value>

101 + </tool_call>""",

102 + [

103 + ToolCall(function=FunctionCall(

104 + name="get_current_weather",

105 + arguments=json.dumps({

106 + "city": "Dallas",

107 + "state": "TX",

108 + "unit": "fahrenheit",

109 + }),

110 + )),

111 + ToolCall(function=FunctionCall(

112 + name="get_current_weather",

113 + arguments=json.dumps({

114 + "city": "Orlando",

115 + "state": "FL",

116 + "unit": "fahrenheit",

117 + }),

118 + )),

119 + ],

120 + None,

121 + ),

122 + (

123 + """I'll help you check the weather. <tool_call>get_current_weather

124 + <arg_key>city</arg_key>

125 + <arg_value>Seattle</arg_value>

126 + <arg_key>state</arg_key>

127 + <arg_value>WA</arg_value>

128 + <arg_key>unit</arg_key>

129 + <arg_value>celsius</arg_value>

130 + </tool_call>""",

131 + [

132 + ToolCall(function=FunctionCall(

133 + name="get_current_weather",

134 + arguments=json.dumps({

135 + "city": "Seattle",

136 + "state": "WA",

137 + "unit": "celsius",

138 + }),

139 + ))

140 + ],

141 + "I'll help you check the weather.",

142 + ),

143 + (

144 + """<tool_call>get_current_weather

145 + <arg_key>city</arg_key>

146 + <arg_value>New York</arg_value>

147 + <arg_key>state</arg_key>

148 + <arg_value>NY</arg_value>

149 + <arg_key>unit</arg_key>

150 + <arg_value>celsius</arg_value>

151 + </tool_call>""",

152 + [

153 + ToolCall(function=FunctionCall(

154 + name="get_current_weather",

155 + arguments=json.dumps({

156 + "city": "New York",

157 + "state": "NY",

158 + "unit": "celsius",

159 + }),

160 + ))

161 + ],

162 + None,

163 + ),

164 + ("""I will help you get the weather.<tool_call>get_weather

165 + <arg_key>city</arg_key>

166 + <arg_value>Beijing</arg_value>

167 + <arg_key>date</arg_key>

168 + <arg_value>2025-08-01</arg_value>

169 + </tool_call>""", [

170 + ToolCall(function=FunctionCall(

171 + name="get_weather",

172 + arguments=json.dumps({

173 + "city": "Beijing",

174 + "date": "2025-08-01",

175 + }),

176 + ))

177 + ], "I will help you get the weather."),

178 + ],

179 +)

180 +def test_extract_tool_calls(glm4_moe_tool_parser, model_output,

181 + expected_tool_calls, expected_content):

182 + extracted_tool_calls = glm4_moe_tool_parser.extract_tool_calls(

183 + model_output, request=None) # type: ignore[arg-type]

184 + assert extracted_tool_calls.tools_called

185 + assert_tool_calls(extracted_tool_calls.tool_calls, expected_tool_calls)

186 +

187 + assert extracted_tool_calls.content == expected_content

188 +

189 +

190 +def test_extract_tool_calls_with_thinking_tags(glm4_moe_tool_parser):

191 + """Test tool extraction when thinking tags are present."""

192 + model_output = """<think>I want to get the weather.</think>

193 +

194 +I will help you get the weather.

195 +<tool_call>get_weather

196 +<arg_key>city</arg_key>

197 +<arg_value>Beijing</arg_value>

198 +<arg_key>date</arg_key>

199 +<arg_value>2025-08-01</arg_value>

200 +</tool_call>"""

201 +

202 + extracted_tool_calls = glm4_moe_tool_parser.extract_tool_calls(

203 + model_output, request=None) # type: ignore[arg-type]

204 +

205 + assert extracted_tool_calls.tools_called

206 + assert len(extracted_tool_calls.tool_calls) == 1

207 + assert extracted_tool_calls.tool_calls[0].function.name == "get_weather"

208 +

209 + expected_content = """<think>I want to get the weather.</think>

210 +

211 +I will help you get the weather."""

212 + assert extracted_tool_calls.content == expected_content

213 +

214 +

215 +def test_extract_tool_calls_malformed_xml(glm4_moe_tool_parser):

216 + """Test that malformed XML is handled gracefully."""

217 + model_output = """<tool_call>get_weather

218 +<arg_key>city</arg_key>

219 +<arg_value>Seattle</arg_value>

220 +<arg_key>incomplete_arg

221 +<arg_value>value</arg_value>

222 +</tool_call>"""

223 +

224 + extracted_tool_calls = glm4_moe_tool_parser.extract_tool_calls(

225 + model_output, request=None) # type: ignore[arg-type]

226 +

227 + # Should handle malformed XML gracefully

228 + # The parser should either extract what it can or return no tool calls

229 + # depending on how robust we want the parsing to be

230 + assert isinstance(extracted_tool_calls.tools_called, bool)

231 + assert isinstance(extracted_tool_calls.tool_calls, list)

232 +

233 +

234 +def test_extract_tool_calls_empty_arguments(glm4_moe_tool_parser):

235 + """Test tool calls with no arguments."""

236 + model_output = """<tool_call>get_current_time

237 +</tool_call>"""

238 +

239 + extracted_tool_calls = glm4_moe_tool_parser.extract_tool_calls(

240 + model_output, request=None) # type: ignore[arg-type]

241 +

242 + assert extracted_tool_calls.tools_called

243 + assert len(extracted_tool_calls.tool_calls) == 1

244 + assert extracted_tool_calls.tool_calls[

245 + 0].function.name == "get_current_time"

246 + # Empty arguments should result in empty JSON object

247 + assert extracted_tool_calls.tool_calls[0].function.arguments == "{}"

248 +

249 +

250 +def test_extract_tool_calls_mixed_content(glm4_moe_tool_parser):

251 + """Test extraction with mixed content and multiple tool calls."""

252 + model_output = """I will help you get the weather info.

253 +

254 +<tool_call>get_weather

255 +<arg_key>city</arg_key>

256 +<arg_value>Beijing</arg_value>

257 +<arg_key>date</arg_key>

258 +<arg_value>2025-08-01</arg_value>

259 +</tool_call>

260 +

261 +meaningwhile, I will also check the weather in Shanghai.

262 +

263 +<tool_call>get_weather

264 +<arg_key>city</arg_key>

265 +<arg_value>Shanghai</arg_value>

266 +<arg_key>date</arg_key>

267 +<arg_value>2025-08-01</arg_value>

268 +</tool_call>"""

269 +

270 + extracted_tool_calls = glm4_moe_tool_parser.extract_tool_calls(

271 + model_output, request=None) # type: ignore[arg-type]

272 +

273 + assert extracted_tool_calls.tools_called

274 + assert len(extracted_tool_calls.tool_calls) == 2

275 +

276 + # Check first tool call

277 + assert extracted_tool_calls.tool_calls[0].function.name == "get_weather"

278 + args1 = json.loads(extracted_tool_calls.tool_calls[0].function.arguments)

279 + assert args1["city"] == "Beijing"

280 + assert args1["date"] == "2025-08-01"

281 +

282 + # Check second tool call

283 + assert extracted_tool_calls.tool_calls[1].function.name == "get_weather"

284 + args2 = json.loads(extracted_tool_calls.tool_calls[1].function.arguments)

285 + assert args2["city"] == "Shanghai"

286 + assert args2["date"] == "2025-08-01"

287 +

288 + # Content should be everything before the first tool call

289 + assert extracted_tool_calls.content == "I will help you get the weather info."

290 +

291 +

292 +def test_streaming_basic_functionality(glm4_moe_tool_parser):

293 + """Test basic streaming functionality."""

294 + # Reset streaming state

295 + glm4_moe_tool_parser.current_tool_name_sent = False

296 + glm4_moe_tool_parser.prev_tool_call_arr = []

297 + glm4_moe_tool_parser.current_tool_id = -1

298 + glm4_moe_tool_parser.streamed_args_for_tool = []

299 +

300 + # Test with a simple tool call

301 + current_text = """<tool_call>get_weather

302 +<arg_key>city</arg_key>

303 +<arg_value>Beijing</arg_value>

304 +</tool_call>"""

305 +

306 + # Mock token IDs for testing

307 + tool_call_start_id = glm4_moe_tool_parser.tool_call_start_token_id or 12345

308 + tool_call_end_id = glm4_moe_tool_parser.tool_call_end_token_id or 12346

309 +

310 + result = glm4_moe_tool_parser.extract_tool_calls_streaming(

311 + previous_text="",

312 + current_text=current_text,

313 + delta_text="</tool_call>",

314 + previous_token_ids=[],

315 + current_token_ids=[tool_call_start_id, tool_call_end_id],

316 + delta_token_ids=[tool_call_end_id],

317 + request=None,

318 + )

319 +

320 + # The result behavior depends on the streaming state

321 + # This test mainly ensures no exceptions are thrown

322 + assert result is None or hasattr(result, 'tool_calls') or hasattr(

323 + result, 'content')

324 +

325 +

326 +def test_streaming_no_tool_calls(glm4_moe_tool_parser):

327 + """Test streaming when there are no tool calls."""

328 + current_text = "This is just regular text without any tool calls."

329 +

330 + result = glm4_moe_tool_parser.extract_tool_calls_streaming(

331 + previous_text="This is just regular text",

332 + current_text=current_text,

333 + delta_text=" without any tool calls.",

334 + previous_token_ids=[],

335 + current_token_ids=[],

336 + delta_token_ids=[],

337 + request=None,

338 + )

339 +

340 + # Should return the delta text as content

341 + assert result is not None

342 + assert hasattr(result, 'content')

343 + assert result.content == " without any tool calls."

344 +

345 +

346 +def test_streaming_with_content_before_tool_calls(glm4_moe_tool_parser):

347 + """Test streaming when there's content before tool calls."""

348 + # Reset streaming state

349 + glm4_moe_tool_parser.current_tool_name_sent = False

350 + glm4_moe_tool_parser.prev_tool_call_arr = []

351 + glm4_moe_tool_parser.current_tool_id = -1

352 + glm4_moe_tool_parser.streamed_args_for_tool = []

353 +

354 + current_text = "I will help you get the weather<tool_call>"

355 +

356 + result = glm4_moe_tool_parser.extract_tool_calls_streaming(

357 + previous_text="I will help you",

358 + current_text=current_text,

359 + delta_text="get the weather.<tool_call>",

360 + previous_token_ids=[],

361 + current_token_ids=[],

362 + delta_token_ids=[],

363 + request=None,

364 + )

365 +

366 + # Should return content when no tool call tokens are detected

367 + assert result is not None

368 + assert hasattr(result, 'content')

369 + assert result.content == "get the weather.<tool_call>"

370 +

371 +

372 +def test_extract_tool_calls_special_characters(glm4_moe_tool_parser):

373 + """Test tool calls with special characters and unicode."""

374 + model_output = """<tool_call>send_message

375 +<arg_key>recipient</arg_key>

376 +<arg_value>Amy</arg_value>

377 +<arg_key>message</arg_key>

378 +<arg_value>It is a nice day</arg_value>

379 +<arg_key>priority</arg_key>

380 +<arg_value>high</arg_value>

381 +</tool_call>"""

382 +

383 + extracted_tool_calls = glm4_moe_tool_parser.extract_tool_calls(

384 + model_output, request=None) # type: ignore[arg-type]

385 +

386 + assert extracted_tool_calls.tools_called

387 + assert len(extracted_tool_calls.tool_calls) == 1

388 + assert extracted_tool_calls.tool_calls[0].function.name == "send_message"

389 +

390 + args = json.loads(extracted_tool_calls.tool_calls[0].function.arguments)

391 + assert args["recipient"] == "Amy"

392 + assert args["message"] == "It is a nice day"

393 + assert args["priority"] == "high"

394 +

395 +

396 +def test_extract_tool_calls_incomplete_tool_call(glm4_moe_tool_parser):

397 + """Test incomplete tool calls (missing closing tag)."""

398 + model_output = """<tool_call>get_weather

399 +<arg_key>city</arg_key>

400 +<arg_value>Beijing</arg_value>

401 +<arg_key>date</arg_key>

402 +<arg_value>2025-08-01</arg_value>"""

403 +

404 + extracted_tool_calls = glm4_moe_tool_parser.extract_tool_calls(

405 + model_output, request=None) # type: ignore[arg-type]

406 +

407 + # Incomplete tool calls should not be extracted

408 + assert not extracted_tool_calls.tools_called

409 + assert extracted_tool_calls.tool_calls == []

410 + assert extracted_tool_calls.content == model_output

vllm/config.py

+12 -3

@@ -1333,7 +1333,8 @@ class ModelConfig:

1333 self, parallel_config: "ParallelConfig") -> tuple[int, int]:

1334 from vllm.distributed.utils import get_pp_indices

1335 if (self.hf_text_config.model_type == "deepseek_mtp"

1336 - or self.hf_config.model_type == "mimo_mtp"):

1336 + or self.hf_config.model_type == "mimo_mtp"

1337 + or self.hf_config.model_type == "glm4_moe_mtp"):

1338 total_num_hidden_layers = getattr(self.hf_text_config,

1339 "num_nextn_predict_layers", 0)

1340 else:

@@ -2663,7 +2664,15 @@ class SpeculativeConfig:

2664 "n_predict": n_predict,

2665 "architectures": ["MiMoMTPModel"]

2666 })

2666 - return hf_config

2667 +

2668 + if hf_config.architectures[0] == "Glm4MoeForCausalLM":

2669 + hf_config.model_type = "glm4_moe_mtp"

2670 + n_predict = getattr(hf_config, "num_nextn_predict_layers", None)

2671 + hf_config.update({

2672 + "num_hidden_layers": 0,

2673 + "n_predict": n_predict,

2674 + "architectures": ["Glm4MoeMTPModel"]

2675 + })

2676

2677 return hf_config

2678

@@ -2774,7 +2783,7 @@ class SpeculativeConfig:

2783 "mlp_speculator"):

2784 self.method = "mlp_speculator"

2785 elif (self.draft_model_config.hf_config.model_type

2777 - in ("deepseek_mtp", "mimo_mtp")):

2786 + in ("deepseek_mtp", "mimo_mtp", "glm4_moe_mtp")):

2787 self.method = "deepseek_mtp"

2788 if self.num_speculative_tokens > 1:

2789 logger.warning(

vllm/entrypoints/openai/tool_parsers/__init__.py

+19 -6

@@ -3,6 +3,7 @@

4 from .abstract_tool_parser import ToolParser, ToolParserManager

5 from .deepseekv3_tool_parser import DeepSeekV3ToolParser

6 +from .glm4_moe_tool_parser import Glm4MoeModelToolParser

7 from .granite_20b_fc_tool_parser import Granite20bFCToolParser

8 from .granite_tool_parser import GraniteToolParser

9 from .hermes_tool_parser import Hermes2ProToolParser

@@ -19,10 +20,22 @@ from .pythonic_tool_parser import PythonicToolParser

20 from .xlam_tool_parser import xLAMToolParser

22 __all__ = [

22 - "ToolParser", "ToolParserManager", "Granite20bFCToolParser",

23 - "GraniteToolParser", "Hermes2ProToolParser", "MistralToolParser",

24 - "Internlm2ToolParser", "Llama3JsonToolParser", "JambaToolParser",

25 - "Llama4PythonicToolParser", "PythonicToolParser", "Phi4MiniJsonToolParser",

26 - "DeepSeekV3ToolParser", "xLAMToolParser", "MinimaxToolParser",

27 - "KimiK2ToolParser", "HunyuanA13BToolParser"

23 + "ToolParser",

24 + "ToolParserManager",

25 + "Granite20bFCToolParser",

26 + "GraniteToolParser",

27 + "Hermes2ProToolParser",

28 + "MistralToolParser",

29 + "Internlm2ToolParser",

30 + "Llama3JsonToolParser",

31 + "JambaToolParser",

32 + "Llama4PythonicToolParser",

33 + "PythonicToolParser",

34 + "Phi4MiniJsonToolParser",

35 + "DeepSeekV3ToolParser",

36 + "xLAMToolParser",

37 + "MinimaxToolParser",

38 + "KimiK2ToolParser",

39 + "HunyuanA13BToolParser",

40 + "Glm4MoeModelToolParser",

41 ]

vllm/entrypoints/openai/tool_parsers/glm4_moe_tool_parser.py new

+402

@@ -0,0 +1,402 @@

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

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

3 +# code modified from deepseekv3_tool_parser.py

4 +

5 +from collections.abc import Sequence

6 +from typing import Union

7 +

8 +import regex as re

9 +

10 +from vllm.entrypoints.openai.protocol import (ChatCompletionRequest,

11 + DeltaFunctionCall, DeltaMessage,

12 + DeltaToolCall,

13 + ExtractedToolCallInformation,

14 + FunctionCall, ToolCall)

15 +from vllm.entrypoints.openai.tool_parsers.abstract_tool_parser import (

16 + ToolParser, ToolParserManager)

17 +from vllm.logger import init_logger

18 +from vllm.transformers_utils.tokenizer import AnyTokenizer

19 +

20 +logger = init_logger(__name__)

21 +

22 +

23 +@ToolParserManager.register_module("glm4_moe")

24 +class Glm4MoeModelToolParser(ToolParser):

25 +

26 + def __init__(self, tokenizer: AnyTokenizer):

27 + super().__init__(tokenizer)

28 + self.current_tool_name_sent = False

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

30 + self.current_tool_id = -1

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

32 + self.tool_call_start_token = "<tool_call>"

33 + self.tool_call_end_token = "</tool_call>"

34 +

35 + self.tool_calls_start_token = self.tool_call_start_token

36 +

37 + # Updated regex for the XML-based format

38 + self.tool_call_regex = re.compile(

39 + r"<tool_call>\s*"

40 + r"(?P<function_name>[^\n<]+)\s*" # 函数名（到换行或 <）

41 + r"(?P<arguments>(?:\s*<arg_key>[^<]+</arg_key>\s*"

42 + r"<arg_value>[^<]*</arg_value>\s*)*)\s*"

43 + r"</tool_call>",

44 + re.DOTALL,

45 + )

46 +

47 + # Regex for parsing individual arguments

48 + self.arg_regex = re.compile(

49 + r"<arg_key>(?P<key>[^<]+)</arg_key>\s*<arg_value>(?P<value>[^<]*)</arg_value>",

50 + re.DOTALL,

51 + )

52 +

53 + # Streaming regex

54 + self.stream_tool_call_portion_regex = re.compile(

55 + r"(?P<function_name>[^\n<]+)\s*"

56 + r"(?P<arguments>(?:\s*<arg_key>[^<]+</arg_key>\s*"

57 + r"<arg_value>[^<]*</arg_value>\s*)*)",

58 + re.DOTALL,

59 + )

60 +

61 + # For streaming, we also need a regex to match just the function name

62 + self.stream_tool_call_name_regex = re.compile(

63 + r"(?P<function_name>[^\n<]+)",

64 + re.DOTALL,

65 + )

66 +

67 + if not self.model_tokenizer:

68 + raise ValueError(

69 + "The model tokenizer must be passed to the ToolParser "

70 + "constructor during construction.")

71 +

72 + self.tool_call_start_token_id = self.vocab.get(

73 + self.tool_call_start_token)

74 + self.tool_call_end_token_id = self.vocab.get(self.tool_call_end_token)

75 +

76 + def _parse_arguments(self, args_text: str) -> str:

77 + """Parse XML-based arguments into JSON format."""

78 + if not args_text or not args_text.strip():

79 + return "{}"

80 +

81 + args_dict = {}

82 + matches = self.arg_regex.findall(args_text)

83 +

84 + for key, value in matches:

85 + args_dict[key.strip()] = value.strip()

86 +

87 + import json

88 + return json.dumps(args_dict, ensure_ascii=False)

89 +

90 + def extract_tool_calls(

91 + self,

92 + model_output: str,

93 + request: ChatCompletionRequest,

94 + ) -> ExtractedToolCallInformation:

95 +

96 + # sanity check; avoid unnecessary processing

97 + if self.tool_calls_start_token not in model_output:

98 + return ExtractedToolCallInformation(tools_called=False,

99 + tool_calls=[],

100 + content=model_output)

101 +

102 + try:

103 + # Find all tool calls in the output

104 + function_call_matches = self.tool_call_regex.findall(model_output)

105 +

106 + logger.debug("function_call_matches: %s", function_call_matches)

107 +

108 + if not function_call_matches:

109 + return ExtractedToolCallInformation(

110 + tools_called=False,

111 + tool_calls=[],

112 + content=model_output,

113 + )

114 +

115 + tool_calls = []

116 + for i, match in enumerate(function_call_matches):

117 + function_name, function_args_xml = match

118 + function_name = function_name.strip()

119 +

120 + # Parse XML arguments to JSON

121 + function_args_json = self._parse_arguments(function_args_xml)

122 +

123 + tool_calls.append(

124 + ToolCall(

125 + id=f"call_{i}",

126 + type='function',

127 + function=FunctionCall(name=function_name,

128 + arguments=function_args_json),

129 + ))

130 +

131 + # Extract content before the first tool call

132 + content = model_output[:model_output.find(self.

133 + tool_calls_start_token)]

134 + return ExtractedToolCallInformation(

135 + tools_called=bool(tool_calls),

136 + tool_calls=tool_calls,

137 + content=content.strip() if content.strip() else None,

138 + )

139 +

140 + except Exception:

141 + logger.exception("Error in extracting tool call from response.")

142 + return ExtractedToolCallInformation(tools_called=False,

143 + tool_calls=[],

144 + content=model_output)

145 +

146 + def extract_tool_calls_streaming(

147 + self,

148 + previous_text: str,

149 + current_text: str,

150 + delta_text: str,

151 + previous_token_ids: Sequence[int],

152 + current_token_ids: Sequence[int],

153 + delta_token_ids: Sequence[int],

154 + request: ChatCompletionRequest,

155 + ) -> Union[DeltaMessage, None]:

156 +

157 + logger.debug("delta_text: %s", delta_text)

158 + logger.debug("delta_token_ids: %s", delta_token_ids)

159 + # check to see if we should be streaming a tool call - is there a

160 + if self.tool_call_start_token_id not in current_token_ids:

161 + logger.debug("No tool call tokens found!")

162 + return DeltaMessage(content=delta_text)

163 + delta_text = delta_text.replace(self.tool_calls_start_token,

164 + "").replace(self.tool_call_end_token,

165 + "")

166 + try:

167 +

168 + # figure out where we are in the parsing by counting tool call

169 + # start & end tags

170 + prev_tool_start_count = previous_token_ids.count(

171 + self.tool_call_start_token_id)

172 + prev_tool_end_count = previous_token_ids.count(

173 + self.tool_call_end_token_id)

174 + cur_tool_start_count = current_token_ids.count(

175 + self.tool_call_start_token_id)

176 + cur_tool_end_count = current_token_ids.count(

177 + self.tool_call_end_token_id)

178 + tool_call_portion = None

179 + text_portion = None

180 +

181 + # case: if we're generating text, OR rounding out a tool call

182 + if (cur_tool_start_count == cur_tool_end_count

183 + and prev_tool_end_count == cur_tool_end_count

184 + and self.tool_call_end_token not in delta_text):

185 + logger.debug("Generating text content! skipping tool parsing.")

186 + return DeltaMessage(content=delta_text)

187 +

188 + if self.tool_call_end_token in delta_text:

189 + logger.debug("tool_call_end_token in delta_text")

190 + full_text = current_text + delta_text

191 + tool_call_portion = full_text.split(

192 + self.tool_call_start_token)[-1].split(

193 + self.tool_call_end_token)[0].rstrip()

194 + delta_text = delta_text.split(

195 + self.tool_call_end_token)[0].rstrip()

196 + text_portion = delta_text.split(

197 + self.tool_call_end_token)[-1].lstrip()

198 +

199 + # case -- we're starting a new tool call

200 + if (cur_tool_start_count > cur_tool_end_count

201 + and cur_tool_start_count > prev_tool_start_count):

202 + if len(delta_token_ids) > 1:

203 + tool_call_portion = current_text.split(

204 + self.tool_call_start_token)[-1]

205 + else:

206 + tool_call_portion = None

207 + delta = None

208 +

209 + text_portion = None

210 +

211 + # set cursors and state appropriately

212 + self.current_tool_id += 1

213 + self.current_tool_name_sent = False

214 + self.streamed_args_for_tool.append("")

215 + logger.debug("Starting on a new tool %s", self.current_tool_id)

216 +

217 + # case -- we're updating an existing tool call

218 + elif (cur_tool_start_count > cur_tool_end_count

219 + and cur_tool_start_count == prev_tool_start_count):

220 +

221 + # get the portion of the text that's the tool call

222 + tool_call_portion = current_text.split(

223 + self.tool_call_start_token)[-1]

224 + text_portion = None

225 +

226 + # case -- the current tool call is being closed.

227 + elif (cur_tool_start_count == cur_tool_end_count

228 + and cur_tool_end_count >= prev_tool_end_count):

229 + if self.prev_tool_call_arr is None or len(

230 + self.prev_tool_call_arr) == 0:

231 + logger.debug(

232 + "attempting to close tool call, but no tool call")

233 + return None

234 + diff = self.prev_tool_call_arr[self.current_tool_id].get(

235 + "arguments")

236 + if diff:

237 + diff = (diff.encode("utf-8").decode("unicode_escape")

238 + if diff is str else diff)

239 + if '"}' not in delta_text:

240 + return None

241 + end_loc = delta_text.rindex('"}')

242 + diff = delta_text[:end_loc] + '"}'

243 + logger.debug(

244 + "Finishing tool and found diff that had not "

245 + "been streamed yet: %s",

246 + diff,

247 + )

248 + self.streamed_args_for_tool[self.current_tool_id] += diff

249 + return DeltaMessage(tool_calls=[

250 + DeltaToolCall(

251 + index=self.current_tool_id,

252 + function=DeltaFunctionCall(

253 + arguments=diff).model_dump(exclude_none=True),

254 + )

255 + ])

256 +

257 + # case -- otherwise we're just generating text

258 + else:

259 + text = delta_text.replace(self.tool_call_start_token, "")

260 + text = text.replace(self.tool_call_end_token, "")

261 + delta = DeltaMessage(tool_calls=[], content=text)

262 + return delta

263 +

264 + current_tool_call = dict()

265 + if tool_call_portion:

266 + current_tool_call_matches = (

267 + self.stream_tool_call_portion_regex.match(

268 + tool_call_portion))

269 + if current_tool_call_matches:

270 + tool_id, tool_args = (current_tool_call_matches.groups())

271 + tool_name = tool_id.split('.')[1].split(':')[0]

272 + current_tool_call['id'] = tool_id

273 + current_tool_call["name"] = tool_name

274 + current_tool_call["arguments"] = tool_args

275 + else:

276 + current_tool_call_name_matches = (

277 + self.stream_tool_call_name_regex.match(

278 + tool_call_portion))

279 + if current_tool_call_name_matches:

280 + tool_id_str, = current_tool_call_name_matches.groups()

281 + tool_name = tool_id_str.split('.')[1].split(':')[0]

282 + current_tool_call['id'] = tool_id_str

283 + current_tool_call["name"] = tool_name

284 + current_tool_call["arguments"] = ""

285 + else:

286 + logger.debug("Not enough token")

287 + return None

288 +

289 + # case - we haven't sent the tool name yet. If it's available, send

290 + # it. otherwise, wait until it's available.

291 + if not self.current_tool_name_sent:

292 + if current_tool_call is None:

293 + return None

294 + function_name: Union[str, None] = current_tool_call.get("name")

295 + tool_id = current_tool_call.get("id")

296 + if function_name:

297 + self.current_tool_name_sent = True

298 + return DeltaMessage(tool_calls=[

299 + DeltaToolCall(

300 + index=self.current_tool_id,

301 + type="function",

302 + id=tool_id,

303 + function=DeltaFunctionCall(

304 + name=function_name).model_dump(

305 + exclude_none=True),

306 + )

307 + ])

308 + else:

309 + return None

310 +

311 + # case -- otherwise, send the tool call delta

312 +

313 + # if the tool call portion is None, send the delta as text

314 + if tool_call_portion is None:

315 + # if there's text but not tool calls, send that -

316 + # otherwise None to skip chunk

317 + delta = (DeltaMessage(

318 + content=delta_text) if text_portion is not None else None)

319 + return delta

320 +

321 + # now, the nitty-gritty of tool calls

322 + # now we have the portion to parse as tool call.

323 +

324 + logger.debug("Trying to parse current tool call with ID %s",

325 + self.current_tool_id)

326 +

327 + # if we're starting a new tool call, push an empty object in as

328 + # a placeholder for the arguments

329 + if len(self.prev_tool_call_arr) <= self.current_tool_id:

330 + self.prev_tool_call_arr.append({})

331 +

332 + # main logic for tool parsing here - compare prev. partially-parsed

333 + # JSON to the current partially-parsed JSON

334 + prev_arguments = self.prev_tool_call_arr[self.current_tool_id].get(

335 + "arguments")

336 + cur_arguments = current_tool_call.get("arguments")

337 +

338 + logger.debug("diffing old arguments: %s", prev_arguments)

339 + logger.debug("against new ones: %s", cur_arguments)

340 +

341 + # case -- no arguments have been created yet. skip sending a delta.

342 + if not cur_arguments and not prev_arguments:

343 + logger.debug("Skipping text %s - no arguments", delta_text)

344 + delta = None

345 +

346 + # case -- prev arguments are defined, but non are now.

347 + # probably impossible, but not a fatal error - just keep going

348 + elif not cur_arguments and prev_arguments:

349 + logger.error("should be impossible to have arguments reset "

350 + "mid-call. skipping streaming anything.")

351 + delta = None

352 +

353 + # case -- we now have the first info about arguments available from

354 + # autocompleting the JSON

355 + elif cur_arguments and not prev_arguments:

356 +

357 + delta = DeltaMessage(tool_calls=[

358 + DeltaToolCall(

359 + index=self.current_tool_id,

360 + function=DeltaFunctionCall(

361 + arguments=cur_arguments).model_dump(

362 + exclude_none=True),

363 + )

364 + ])

365 + self.streamed_args_for_tool[

366 + self.current_tool_id] = cur_arguments

367 +

368 + # last case -- we have an update to existing arguments.

369 + elif cur_arguments and prev_arguments:

370 + if (isinstance(delta_text, str)

371 + and cur_arguments != prev_arguments

372 + and len(cur_arguments) > len(prev_arguments)

373 + and cur_arguments.startswith(prev_arguments)):

374 + delta_arguments = cur_arguments[len(prev_arguments):]

375 + logger.debug("got diff %s", delta_text)

376 +

377 + delta = DeltaMessage(tool_calls=[

378 + DeltaToolCall(

379 + index=self.current_tool_id,

380 + function=DeltaFunctionCall(

381 + arguments=delta_arguments).model_dump(

382 + exclude_none=True),

383 + )

384 + ])

385 + self.streamed_args_for_tool[

386 + self.current_tool_id] = cur_arguments

387 + else:

388 + delta = None

389 +

390 + # handle saving the state for the current tool into

391 + # the "prev" list for use in diffing for the next iteration

392 + if self.current_tool_id == len(self.prev_tool_call_arr) - 1:

393 + self.prev_tool_call_arr[

394 + self.current_tool_id] = current_tool_call

395 + else:

396 + self.prev_tool_call_arr.append(current_tool_call)

397 +

398 + return delta

399 +

400 + except Exception:

401 + logger.exception("Error trying to handle streaming tool call.")

402 + return None # do not stream a delta. skip this token ID.

vllm/model_executor/models/glm4_moe.py new

+685

@@ -0,0 +1,685 @@

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

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

3 +

7 +#

8 +# This code is based on EleutherAI's GPT-NeoX library and the GPT-NeoX

9 +# and OPT implementations in this library. It has been modified from its

10 +# original forms to accommodate minor architectural differences compared

11 +# to GPT-NeoX and OPT used by the Meta AI team that trained the model.

12 +#

13 +# Licensed under the Apache License, Version 2.0 (the "License");

14 +# you may not use this file except in compliance with the License.

15 +# You may obtain a copy of the License at

16 +#

17 +# http://www.apache.org/licenses/LICENSE-2.0

18 +#

19 +# Unless required by applicable law or agreed to in writing, software

20 +# distributed under the License is distributed on an "AS IS" BASIS,

21 +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.

22 +# See the License for the specific language governing permissions and

23 +# limitations under the License.

24 +"""Inference-only GLM-4.5 model compatible with HuggingFace weights."""

25 +import typing

26 +from collections.abc import Callable, Iterable

27 +from typing import Any, Optional, Union

28 +

29 +import torch

30 +from torch import nn

31 +from transformers import PretrainedConfig

32 +

33 +from vllm.attention import Attention

34 +from vllm.compilation.decorators import support_torch_compile

35 +from vllm.config import CacheConfig, VllmConfig, get_current_vllm_config

36 +from vllm.distributed import (get_ep_group, get_pp_group,

37 + get_tensor_model_parallel_world_size)

38 +from vllm.logger import init_logger

39 +from vllm.model_executor.layers.activation import SiluAndMul

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

41 +from vllm.model_executor.layers.layernorm import RMSNorm

42 +from vllm.model_executor.layers.linear import (MergedColumnParallelLinear,

43 + QKVParallelLinear,

44 + ReplicatedLinear,

45 + RowParallelLinear)

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

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

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

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

50 + ParallelLMHead, VocabParallelEmbedding)

51 +from vllm.model_executor.model_loader.weight_utils import (

52 + default_weight_loader, maybe_remap_kv_scale_name)

53 +from vllm.model_executor.sampling_metadata import SamplingMetadata

54 +from vllm.sequence import IntermediateTensors

55 +

56 +from .interfaces import SupportsPP

57 +from .utils import (AutoWeightsLoader, PPMissingLayer, is_pp_missing_parameter,

58 + make_empty_intermediate_tensors_factory, make_layers,

59 + maybe_prefix)

60 +

61 +logger = init_logger(__name__)

62 +

63 +

64 +class Glm4MoeMLP(nn.Module):

65 +

66 + def __init__(

67 + self,

68 + hidden_size: int,

69 + intermediate_size: int,

70 + hidden_act: str,

71 + quant_config: Optional[QuantizationConfig] = None,

72 + reduce_results: bool = True,

73 + prefix: str = "",

74 + ) -> None:

75 + super().__init__()

76 + self.gate_up_proj = MergedColumnParallelLinear(

77 + hidden_size, [intermediate_size] * 2,

78 + bias=False,

79 + quant_config=quant_config,

80 + prefix=f"{prefix}.gate_up_proj")

81 + self.down_proj = RowParallelLinear(intermediate_size,

82 + hidden_size,

83 + bias=False,

84 + quant_config=quant_config,

85 + reduce_results=reduce_results,

86 + prefix=f"{prefix}.down_proj")

87 + if hidden_act != "silu":

88 + raise ValueError(f"Unsupported activation: {hidden_act}. "

89 + "Only silu is supported for now.")

90 + self.act_fn = SiluAndMul()

91 +

92 + def forward(self, x):

93 + gate_up, _ = self.gate_up_proj(x)

94 + x = self.act_fn(gate_up)

95 + x, _ = self.down_proj(x)

96 + return x

97 +

98 +

99 +class Glm4MoE(nn.Module):

100 +

101 + def __init__(

102 + self,

103 + config: PretrainedConfig,

104 + quant_config: Optional[QuantizationConfig] = None,

105 + prefix: str = "",

106 + enable_eplb: bool = False,

107 + ):

108 + super().__init__()

109 + self.tp_size = get_tensor_model_parallel_world_size()

110 + self.routed_scaling_factor = config.routed_scaling_factor

111 +

112 + self.ep_group = get_ep_group().device_group

113 + self.ep_rank = self.ep_group.rank()

114 + self.ep_size = self.ep_group.size()

115 + self.n_routed_experts: int = config.n_routed_experts

116 + self.n_shared_experts: int = config.n_shared_experts

117 +

118 + if config.hidden_act != "silu":

119 + raise ValueError(f"Unsupported activation: {config.hidden_act}. "

120 + "Only silu is supported for now.")

121 +

122 + self.gate = ReplicatedLinear(config.hidden_size,

123 + config.n_routed_experts,

124 + bias=False,

125 + quant_config=None,

126 + prefix=f"{prefix}.gate")

127 +

128 + # noaux_tc is not set in transformers new config now

129 + self.gate.e_score_correction_bias = (nn.Parameter(

130 + torch.empty(config.n_routed_experts)))

131 +

132 + # Load balancing settings.

133 + vllm_config = get_current_vllm_config()

134 + parallel_config = vllm_config.parallel_config

135 + self.enable_eplb = enable_eplb

136 +

137 + self.n_redundant_experts = parallel_config.num_redundant_experts

138 + self.n_logical_experts = self.n_routed_experts

139 + self.n_physical_experts = (self.n_logical_experts +

140 + self.n_redundant_experts)

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

142 +

143 + self.physical_expert_start = (self.ep_rank *

144 + self.n_local_physical_experts)

145 + self.physical_expert_end = (self.physical_expert_start +

146 + self.n_local_physical_experts)

147 +

148 + self.experts = FusedMoE(

149 + num_experts=config.n_routed_experts,

150 + top_k=config.num_experts_per_tok,

151 + hidden_size=config.hidden_size,

152 + intermediate_size=config.moe_intermediate_size,

153 + reduce_results=False,

154 + renormalize=config.norm_topk_prob,

155 + quant_config=quant_config,

156 + use_grouped_topk=True,

157 + num_expert_group=config.n_group,

158 + topk_group=config.topk_group,

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

160 + scoring_func="sigmoid",

161 + e_score_correction_bias=self.gate.e_score_correction_bias,

162 + enable_eplb=self.enable_eplb,

163 + num_redundant_experts=self.n_redundant_experts)

164 +

165 + if config.n_shared_experts is not None:

166 + intermediate_size = (config.moe_intermediate_size *

167 + config.n_shared_experts)

168 + self.shared_experts = Glm4MoeMLP(

169 + hidden_size=config.hidden_size,

170 + intermediate_size=intermediate_size,

171 + hidden_act=config.hidden_act,

172 + quant_config=quant_config,

173 + reduce_results=self.experts.must_reduce_shared_expert_outputs(

174 + ),

175 + prefix=f"{prefix}.shared_experts",

176 + )

177 +

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

179 + num_tokens, hidden_dim = hidden_states.shape

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

181 +

182 + if self.n_shared_experts is not None:

183 + shared_output = self.shared_experts(hidden_states)

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

185 + final_hidden_states = self.experts(

186 + hidden_states=hidden_states,

187 + router_logits=router_logits) * self.routed_scaling_factor

188 + if shared_output is not None:

189 + final_hidden_states = final_hidden_states + shared_output

190 + if self.tp_size > 1:

191 + final_hidden_states = (

192 + self.experts.maybe_all_reduce_tensor_model_parallel(

193 + final_hidden_states))

194 + return final_hidden_states.view(num_tokens, hidden_dim)

195 +

196 +

197 +class Glm4MoeAttention(nn.Module):

198 +

199 + def __init__(

200 + self,

201 + config: PretrainedConfig,

202 + hidden_size: int,

203 + num_heads: int,

204 + num_kv_heads: int,

205 + rope_theta: float = 10000,

206 + rope_scaling: Optional[dict[str, Any]] = None,

207 + max_position_embeddings: int = 131072,

208 + head_dim: Optional[int] = None,

209 + rms_norm_eps: float = 1e-05,

210 + qkv_bias: bool = False,

211 + use_qk_norm: bool = False,

212 + cache_config: Optional[CacheConfig] = None,

213 + quant_config: Optional[QuantizationConfig] = None,

214 + prefix: str = "",

215 + ) -> None:

216 + super().__init__()

217 + self.hidden_size = hidden_size

218 + tp_size = get_tensor_model_parallel_world_size()

219 + self.total_num_heads = num_heads

220 + assert self.total_num_heads % tp_size == 0

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

222 + self.total_num_kv_heads = num_kv_heads

223 + if self.total_num_kv_heads >= tp_size:

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

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

226 + assert self.total_num_kv_heads % tp_size == 0

227 + else:

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

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

230 + assert tp_size % self.total_num_kv_heads == 0

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

232 + self.head_dim = head_dim or (hidden_size // self.total_num_heads)

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

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

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

236 + self.rope_theta = rope_theta

237 + self.max_position_embeddings = max_position_embeddings

238 + self.use_qk_norm = use_qk_norm

239 +

240 + self.qkv_proj = QKVParallelLinear(hidden_size,

241 + self.head_dim,

242 + self.total_num_heads,

243 + self.total_num_kv_heads,

244 + bias=qkv_bias,

245 + quant_config=quant_config,

246 + prefix=f"{prefix}.qkv_proj")

247 +

248 + self.o_proj = RowParallelLinear(self.total_num_heads * self.head_dim,

249 + hidden_size,

250 + bias=False,

251 + quant_config=quant_config,

252 + prefix=f"{prefix}.o_proj")

253 +

254 + partial_rotary_factor = getattr(config, "partial_rotary_factor", 0.5)

255 + self.rotary_emb = get_rope(

256 + self.head_dim,

257 + rotary_dim=self.head_dim,

258 + max_position=max_position_embeddings,

259 + base=rope_theta,

260 + rope_scaling=rope_scaling,

261 + partial_rotary_factor=partial_rotary_factor,

262 + )

263 + self.attn = Attention(

264 + self.num_heads,

265 + self.head_dim,

266 + self.scaling,

267 + num_kv_heads=self.num_kv_heads,

268 + cache_config=cache_config,

269 + quant_config=quant_config,

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

271 + )

272 +

273 + if self.use_qk_norm:

274 + self.q_norm = RMSNorm(self.head_dim, eps=rms_norm_eps)

275 + self.k_norm = RMSNorm(self.head_dim, eps=rms_norm_eps)

276 +

277 + def forward(

278 + self,

279 + positions: torch.Tensor,

280 + hidden_states: torch.Tensor,

281 + ) -> torch.Tensor:

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

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

284 + if self.use_qk_norm:

285 + q = self.q_norm(q.reshape(-1, self.num_heads,

286 + self.head_dim)).reshape(q.shape)

287 + k = self.k_norm(k.reshape(-1, self.num_kv_heads,

288 + self.head_dim)).reshape(k.shape)

289 +

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

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

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

293 + return output

294 +

295 +

296 +class Glm4MoeDecoderLayer(nn.Module):

297 +

298 + def __init__(

299 + self,

300 + config: PretrainedConfig,

301 + cache_config: Optional[CacheConfig] = None,

302 + quant_config: Optional[QuantizationConfig] = None,

303 + prefix: str = "",

304 + enable_eplb: bool = False,

305 + ) -> None:

306 + super().__init__()

307 + self.hidden_size = config.hidden_size

308 + rope_theta = getattr(config, "rope_theta", 10000)

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

310 + max_position_embeddings = getattr(config, "max_position_embeddings",

311 + 131072)

312 + # DecoderLayers are created with `make_layers` which passes the prefix

313 + # with the layer's index.

314 + layer_idx = int(prefix.split(sep='.')[-1])

315 + self.layer_idx = layer_idx

316 +

317 + self.self_attn = Glm4MoeAttention(

318 + config=config,

319 + hidden_size=self.hidden_size,

320 + num_heads=config.num_attention_heads,

321 + num_kv_heads=config.num_key_value_heads,

322 + rope_theta=rope_theta,

323 + rope_scaling=rope_scaling,

324 + max_position_embeddings=max_position_embeddings,

325 + head_dim=config.head_dim,

326 + rms_norm_eps=config.rms_norm_eps,

327 + qkv_bias=config.attention_bias,

328 + cache_config=cache_config,

329 + quant_config=quant_config,

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

331 + use_qk_norm=config.use_qk_norm,

332 + )

333 +

334 + if (config.n_routed_experts is not None

335 + and layer_idx >= config.first_k_dense_replace):

336 + self.mlp = Glm4MoE(

337 + config=config,

338 + quant_config=quant_config,

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

340 + enable_eplb=enable_eplb,

341 + )

342 + else:

343 + self.mlp = Glm4MoeMLP(hidden_size=config.hidden_size,

344 + intermediate_size=config.intermediate_size,

345 + hidden_act=config.hidden_act,

346 + quant_config=quant_config,

347 + prefix=f"{prefix}.mlp")

348 +

349 + self.input_layernorm = RMSNorm(config.hidden_size,

350 + eps=config.rms_norm_eps)

351 + self.post_attention_layernorm = RMSNorm(config.hidden_size,

352 + eps=config.rms_norm_eps)

353 + self.routed_scaling_factor = config.routed_scaling_factor

354 +

355 + def forward(

356 + self,

357 + positions: torch.Tensor,

358 + hidden_states: torch.Tensor,

359 + residual: Optional[torch.Tensor],

360 + ) -> tuple[torch.Tensor, torch.Tensor]:

361 + if residual is None:

362 + residual = hidden_states

363 + hidden_states = self.input_layernorm(hidden_states)

364 + else:

365 + hidden_states, residual = self.input_layernorm(

366 + hidden_states, residual)

367 + hidden_states = self.self_attn(positions=positions,

368 + hidden_states=hidden_states)

369 + hidden_states, residual = self.post_attention_layernorm(

370 + hidden_states, residual)

371 + hidden_states = self.mlp(hidden_states)

372 + return hidden_states, residual

373 +

374 +

375 +@support_torch_compile

376 +class Glm4MoeModel(nn.Module):

377 +

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

379 + super().__init__()

380 +

381 + config = vllm_config.model_config.hf_config

382 + cache_config = vllm_config.cache_config

383 + quant_config = vllm_config.quant_config

384 + enable_eplb = vllm_config.parallel_config.enable_eplb

385 + self.config = config

386 +

387 + self.vocab_size = config.vocab_size

388 +

389 + if get_pp_group().is_first_rank:

390 + self.embed_tokens = VocabParallelEmbedding(

391 + config.vocab_size,

392 + config.hidden_size,

393 + quant_config=quant_config,

394 + prefix=f"{prefix}.embed_tokens")

395 + else:

396 + self.embed_tokens = PPMissingLayer()

397 +

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

399 + config.num_hidden_layers,

400 + lambda prefix: Glm4MoeDecoderLayer(

401 + config=config,

402 + cache_config=cache_config,

403 + quant_config=quant_config,

404 + prefix=prefix,

405 + enable_eplb=enable_eplb,

406 + ),

407 + prefix=f"{prefix}.layers")

408 +

409 + if get_pp_group().is_last_rank:

410 + self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)

411 + else:

412 + self.norm = PPMissingLayer()

413 + self.make_empty_intermediate_tensors = (

414 + make_empty_intermediate_tensors_factory(

415 + ["hidden_states", "residual"], config.hidden_size))

416 +

417 + def get_input_embeddings(self, input_ids: torch.Tensor) -> torch.Tensor:

418 + return self.embed_tokens(input_ids)

419 +

420 + def forward(

421 + self,

422 + input_ids: torch.Tensor,

423 + positions: torch.Tensor,

424 + intermediate_tensors: Optional[IntermediateTensors] = None,

425 + inputs_embeds: Optional[torch.Tensor] = None,

426 + ) -> Union[torch.Tensor, IntermediateTensors]:

427 + if get_pp_group().is_first_rank:

428 + if inputs_embeds is not None:

429 + hidden_states = inputs_embeds

430 + else:

431 + hidden_states = self.get_input_embeddings(input_ids)

432 + residual = None

433 + else:

434 + assert intermediate_tensors is not None

435 + hidden_states = intermediate_tensors["hidden_states"]

436 + residual = intermediate_tensors["residual"]

437 +

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

439 + layer = self.layers[i]

440 + hidden_states, residual = layer(positions, hidden_states, residual)

441 +

442 + if not get_pp_group().is_last_rank:

443 + return IntermediateTensors({

444 + "hidden_states": hidden_states,

445 + "residual": residual

446 + })

447 +

448 + hidden_states, _ = self.norm(hidden_states, residual)

449 + return hidden_states

450 +

451 + def make_empty_intermediate_tensors(

452 + self, batch_size: int, dtype: torch.dtype,

453 + device: torch.device) -> IntermediateTensors:

454 + return IntermediateTensors({

455 + "hidden_states":

456 + torch.zeros((batch_size, self.config.hidden_size),

457 + dtype=dtype,

458 + device=device),

459 + "residual":

460 + torch.zeros((batch_size, self.config.hidden_size),

461 + dtype=dtype,

462 + device=device),

463 + })

464 +

465 + def load_weights(self, weights: Iterable[tuple[str,

466 + torch.Tensor]]) -> set[str]:

467 + stacked_params_mapping = [

468 + # (param_name, shard_name, shard_id)

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

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

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

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

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

474 + ]

475 +

476 + # Params for weights, fp8 weight scales, fp8 activation scales

477 + # (param_name, weight_name, expert_id, shard_id)

478 + expert_params_mapping = FusedMoE.make_expert_params_mapping(

479 + ckpt_gate_proj_name="gate_proj",

480 + ckpt_down_proj_name="down_proj",

481 + ckpt_up_proj_name="up_proj",

482 + num_experts=self.config.n_routed_experts)

483 +

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

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

486 + for name, loaded_weight in weights:

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

488 + if spec_layer is not None:

489 + continue

490 + for (param_name, weight_name, shard_id) in stacked_params_mapping:

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

492 + if weight_name not in name:

493 + continue

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

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

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

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

498 + # will then be updated below in expert_params_mapping

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

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

501 + continue

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

503 + # Skip loading extra bias for GPTQ models.

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

505 + continue

506 + if is_pp_missing_parameter(name, self):

507 + continue

508 +

509 + param = params_dict[name]

510 + weight_loader = param.weight_loader

511 + weight_loader(param, loaded_weight, shard_id)

512 + break

513 + else:

514 + is_expert_weight = False

515 + for mapping in expert_params_mapping:

516 + param_name, weight_name, expert_id, shard_id = mapping

517 + if weight_name not in name:

518 + continue

519 +

520 + # Anyway, this is an expert weight and should not be

521 + # attempted to load as other weights later

522 + is_expert_weight = True

523 +

524 + # Do not modify `name` since the loop may continue here

525 + # Instead, create a new variable

526 + name_mapped = name.replace(weight_name, param_name)

527 +

528 + if is_pp_missing_parameter(name_mapped, self):

529 + continue

530 +

531 + param = params_dict[name_mapped]

532 + # We should ask the weight loader to return success or not

533 + # here since otherwise we may skip experts with other

534 + # available replicas.

535 + weight_loader = typing.cast(Callable[..., bool],

536 + param.weight_loader)

537 + success = weight_loader(param,

538 + loaded_weight,

539 + name_mapped,

540 + shard_id=shard_id,

541 + expert_id=expert_id,

542 + return_success=True)

543 + if success:

544 + name = name_mapped

545 + break

546 + else:

547 + if is_expert_weight:

548 + # We've checked that this is an expert weight

549 + # However it's not mapped locally to this rank

550 + # So we simply skip it

551 + continue

552 +

553 + # Skip loading extra bias for GPTQ models.

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

555 + continue

556 +

557 + # Remapping the name of FP8 kv-scale.

558 + name = maybe_remap_kv_scale_name(name, params_dict)

559 + if name is None:

560 + continue

561 +

562 + if is_pp_missing_parameter(name, self):

563 + continue

564 +

565 + param = params_dict[name]

566 + weight_loader = getattr(param, "weight_loader",

567 + default_weight_loader)

568 + weight_loader(param, loaded_weight)

569 + loaded_params.add(name)

570 +

571 + return loaded_params

572 +

573 +

574 +class Glm4MoeForCausalLM(nn.Module, SupportsPP):

575 + packed_modules_mapping = {

576 + "qkv_proj": [

577 + "q_proj",

578 + "k_proj",

579 + "v_proj",

580 + ],

581 + "gate_up_proj": [

582 + "gate_proj",

583 + "up_proj",

584 + ],

585 + }

586 +

587 + fall_back_to_pt_during_load = False

588 +

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

590 + super().__init__()

591 + config = vllm_config.model_config.hf_config

592 + quant_config = vllm_config.quant_config

593 + self.config = config

594 + self.quant_config = quant_config

595 + self.model = Glm4MoeModel(vllm_config=vllm_config,

596 + prefix=maybe_prefix(prefix, "model"))

597 + if get_pp_group().is_last_rank:

598 + self.lm_head = ParallelLMHead(config.vocab_size,

599 + config.hidden_size,

600 + quant_config=quant_config)

601 + else:

602 + self.lm_head = PPMissingLayer()

603 + if self.config.tie_word_embeddings:

604 + self.lm_head.weight = self.model.embed_tokens.weight

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

606 + self.make_empty_intermediate_tensors = (

607 + self.model.make_empty_intermediate_tensors)

608 + self.expert_weights = []

609 +

610 + # Set MoE hyperparameters

611 + self.num_moe_layers = (config.num_hidden_layers -

612 + config.first_k_dense_replace)

613 + self.num_expert_groups = config.n_group

614 +

615 + self.moe_layers: list[FusedMoE] = []

616 + for layer in self.model.layers:

617 + assert isinstance(layer, Glm4MoeDecoderLayer)

618 + if isinstance(layer.mlp, Glm4MoE):

619 + self.moe_layers.append(layer.mlp.experts)

620 +

621 + # Pick last one layer since the first ones may be dense layers.

622 + example_moe = typing.cast(

623 + Glm4MoE, self.model.layers[config.num_hidden_layers - 1].mlp)

624 + self.num_logical_experts = example_moe.n_logical_experts

625 + self.num_physical_experts = example_moe.n_physical_experts

626 + self.num_local_physical_experts = example_moe.n_local_physical_experts

627 + self.num_routed_experts = example_moe.n_routed_experts

628 + self.num_shared_experts = example_moe.n_shared_experts

629 + self.num_redundant_experts = example_moe.n_redundant_experts

630 +

631 + def set_eplb_state(

632 + self,

633 + expert_load_view: torch.Tensor,

634 + logical_to_physical_map: torch.Tensor,

635 + logical_replica_count: torch.Tensor,

636 + ) -> None:

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

638 + # Register the expert weights.

639 + self.expert_weights.append(layer.get_expert_weights())

640 + layer.set_eplb_state(

641 + moe_layer_idx=layer_idx,

642 + expert_load_view=expert_load_view,

643 + logical_to_physical_map=logical_to_physical_map,

644 + logical_replica_count=logical_replica_count,

645 + )

646 +

647 + def get_input_embeddings(self, input_ids: torch.Tensor) -> torch.Tensor:

648 + return self.model.get_input_embeddings(input_ids)

649 +

650 + def forward(

651 + self,

652 + input_ids: torch.Tensor,

653 + positions: torch.Tensor,

654 + intermediate_tensors: Optional[IntermediateTensors] = None,

655 + inputs_embeds: Optional[torch.Tensor] = None,

656 + ) -> Union[torch.Tensor, IntermediateTensors]:

657 + hidden_states = self.model(input_ids, positions, intermediate_tensors,

658 + inputs_embeds)

659 + return hidden_states

660 +

661 + def compute_logits(

662 + self,

663 + hidden_states: torch.Tensor,

664 + sampling_metadata: SamplingMetadata,

665 + ) -> Optional[torch.Tensor]:

666 + logits = self.logits_processor(self.lm_head, hidden_states,

667 + sampling_metadata)

668 + return logits

669 +

670 + def load_weights(self, weights: Iterable[tuple[str,

671 + torch.Tensor]]) -> set[str]:

672 + loader = AutoWeightsLoader(self)

673 + return loader.load_weights(weights)

674 +

675 +

676 +def get_spec_layer_idx_from_weight_name(config: PretrainedConfig,

677 + weight_name: str) -> Optional[int]:

678 + if hasattr(config,

679 + "num_nextn_predict_layers") and (config.num_nextn_predict_layers

680 + > 0):

681 + layer_idx = config.num_hidden_layers

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

683 + if f"layers.{layer_idx+i}." in weight_name:

684 + return layer_idx + i

685 + return None

vllm/model_executor/models/glm4_moe_mtp.py new

+307

@@ -0,0 +1,307 @@

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

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

3 +

7 +#

8 +# This code is based on EleutherAI's GPT-NeoX library and the GPT-NeoX

9 +# and OPT implementations in this library. It has been modified from its

10 +# original forms to accommodate minor architectural differences compared

11 +# to GPT-NeoX and OPT used by the Meta AI team that trained the model.

12 +#

13 +# Licensed under the Apache License, Version 2.0 (the "License");

14 +# you may not use this file except in compliance with the License.

15 +# You may obtain a copy of the License at

16 +#

17 +# http://www.apache.org/licenses/LICENSE-2.0

18 +#

19 +# Unless required by applicable law or agreed to in writing, software

20 +# distributed under the License is distributed on an "AS IS" BASIS,

21 +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.

22 +# See the License for the specific language governing permissions and

23 +# limitations under the License.

24 +"""Inference-only GLM-4.5 MTP model compatible with HuggingFace weights."""

25 +

26 +from collections.abc import Iterable

27 +from typing import Optional

28 +

29 +import torch

30 +import torch.nn as nn

31 +from transformers import PretrainedConfig

32 +

33 +from vllm.config import CacheConfig, VllmConfig

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

35 +from vllm.model_executor.layers.layernorm import RMSNorm

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

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

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

39 + ParallelLMHead, VocabParallelEmbedding)

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

41 +from vllm.model_executor.sampling_metadata import SamplingMetadata

42 +from vllm.sequence import IntermediateTensors

43 +

44 +from .glm4_moe import Glm4MoeDecoderLayer, get_spec_layer_idx_from_weight_name

45 +from .interfaces import SupportsPP

46 +from .utils import maybe_prefix

47 +

48 +

49 +class SharedHead(nn.Module):

50 +

51 + def __init__(

52 + self,

53 + config: PretrainedConfig,

54 + quant_config: Optional[QuantizationConfig] = None,

55 + ) -> None:

56 + super().__init__()

57 + self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)

58 + self.head = ParallelLMHead(config.vocab_size,

59 + config.hidden_size,

60 + quant_config=quant_config)

61 +

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

63 + return self.norm(hidden_states)

64 +

65 +

66 +class Glm4MoeMultiTokenPredictorLayer(nn.Module):

67 +

68 + def __init__(

69 + self,

70 + config: PretrainedConfig,

71 + prefix: str,

72 + cache_config: Optional[CacheConfig] = None,

73 + quant_config: Optional[QuantizationConfig] = None,

74 + ) -> None:

75 + super().__init__()

76 + self.enorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)

77 + self.hnorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)

78 + self.eh_proj = nn.Linear(config.hidden_size * 2,

79 + config.hidden_size,

80 + bias=False)

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

82 + self.mtp_block = Glm4MoeDecoderLayer(config=config,

83 + cache_config=cache_config,

84 + quant_config=quant_config,

85 + prefix=prefix)

86 +

87 + def forward(

88 + self,

89 + input_ids: torch.Tensor,

90 + positions: torch.Tensor,

91 + previous_hidden_states: torch.Tensor,

92 + inputs_embeds: Optional[torch.Tensor] = None,

93 + spec_step_index: int = 0,

94 + ) -> torch.Tensor:

95 + assert inputs_embeds is not None

96 + # masking inputs at position 0, as not needed by MTP

97 + inputs_embeds[positions == 0] = 0

98 + inputs_embeds = self.enorm(inputs_embeds)

99 + previous_hidden_states = self.hnorm(previous_hidden_states)

100 +

101 + hidden_states = self.eh_proj(

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

103 +

104 + hidden_states, residual = self.mtp_block(positions=positions,

105 + hidden_states=hidden_states,

106 + residual=None)

107 + hidden_states = residual + hidden_states

108 + return hidden_states

109 +

110 +

111 +class Glm4MoeMultiTokenPredictor(nn.Module):

112 +

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

114 + super().__init__()

115 + config = vllm_config.model_config.hf_config

116 + self.mtp_start_layer_idx = config.num_hidden_layers

117 + self.num_mtp_layers = config.num_nextn_predict_layers

118 + # to map the exact layer index from weights

119 + self.layers = torch.nn.ModuleDict({

120 + str(idx):

121 + Glm4MoeMultiTokenPredictorLayer(

122 + config,

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

124 + cache_config=vllm_config.cache_config,

125 + quant_config=vllm_config.quant_config,

126 + )

127 + for idx in range(self.mtp_start_layer_idx,

128 + self.mtp_start_layer_idx + self.num_mtp_layers)

129 + })

130 + self.embed_tokens = VocabParallelEmbedding(

131 + config.vocab_size,

132 + config.hidden_size,

133 + )

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

135 +

136 + def forward(

137 + self,

138 + input_ids: torch.Tensor,

139 + positions: torch.Tensor,

140 + previous_hidden_states: torch.Tensor,

141 + inputs_embeds: Optional[torch.Tensor] = None,

142 + spec_step_idx: int = 0,

143 + ) -> torch.Tensor:

144 + if inputs_embeds is None:

145 + inputs_embeds = self.embed_tokens(input_ids)

146 + current_step_idx = (spec_step_idx % self.num_mtp_layers)

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

148 + input_ids,

149 + positions,

150 + previous_hidden_states,

151 + inputs_embeds,

152 + current_step_idx,

153 + )

154 +

155 + def compute_logits(

156 + self,

157 + hidden_states: torch.Tensor,

158 + sampling_metadata: SamplingMetadata,

159 + spec_step_idx: int = 0,

160 + ) -> torch.Tensor:

161 + current_step_idx = (spec_step_idx % self.num_mtp_layers)

162 + mtp_layer = self.layers[str(self.mtp_start_layer_idx +

163 + current_step_idx)]

164 + logits = self.logits_processor(mtp_layer.shared_head.head,

165 + mtp_layer.shared_head(hidden_states),

166 + sampling_metadata)

167 + return logits

168 +

169 +

170 +class Glm4MoeMTP(nn.Module, SupportsPP):

171 +

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

173 + super().__init__()

174 + self.config = vllm_config.model_config.hf_config

175 + self.model = Glm4MoeMultiTokenPredictor(vllm_config=vllm_config,

176 + prefix=maybe_prefix(

177 + prefix, "model"))

178 +

179 + def forward(

180 + self,

181 + input_ids: torch.Tensor,

182 + positions: torch.Tensor,

183 + previous_hidden_states: torch.Tensor,

184 + intermediate_tensors: Optional[IntermediateTensors] = None,

185 + inputs_embeds: Optional[torch.Tensor] = None,

186 + spec_step_idx: int = 0,

187 + ) -> torch.Tensor:

188 + hidden_states = self.model(input_ids, positions,

189 + previous_hidden_states, inputs_embeds,

190 + spec_step_idx)

191 + return hidden_states

192 +

193 + def compute_logits(

194 + self,

195 + hidden_states: torch.Tensor,

196 + sampling_metadata: SamplingMetadata,

197 + spec_step_idx: int = 0,

198 + ) -> Optional[torch.Tensor]:

199 + return self.model.compute_logits(hidden_states, sampling_metadata,

200 + spec_step_idx)

201 +

202 + def load_weights(self, weights: Iterable[tuple[str,

203 + torch.Tensor]]) -> set[str]:

204 + stacked_params_mapping = [

205 + # (param_name, shard_name, shard_id)

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

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

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

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

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

211 + ]

212 +

213 + # Params for weights, fp8 weight scales, fp8 activation scales

214 + # (param_name, weight_name, expert_id, shard_id)

215 + expert_params_mapping = FusedMoE.make_expert_params_mapping(

216 + ckpt_gate_proj_name="gate_proj",

217 + ckpt_down_proj_name="down_proj",

218 + ckpt_up_proj_name="up_proj",

219 + num_experts=self.config.n_routed_experts)

220 +

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

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

223 + for name, loaded_weight in weights:

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

225 + if spec_layer is None:

226 + continue

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

228 + for (param_name, weight_name, shard_id) in stacked_params_mapping:

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

230 + if weight_name not in name:

231 + continue

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

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

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

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

236 + # will then be updated below in expert_params_mapping

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

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

239 + continue

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

241 + # Skip loading extra bias for GPTQ models.

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

243 + continue

244 +

245 + param = params_dict[name]

246 + weight_loader = param.weight_loader

247 + weight_loader(param, loaded_weight, shard_id)

248 + break

249 + else:

250 + for mapping in expert_params_mapping:

251 + param_name, weight_name, expert_id, shard_id = mapping

252 + if weight_name not in name:

253 + continue

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

255 +

256 + param = params_dict[name]

257 + weight_loader = param.weight_loader

258 + weight_loader(param,

259 + loaded_weight,

260 + name,

261 + shard_id=shard_id,

262 + expert_id=expert_id)

263 + break

264 + else:

265 + # Skip loading extra bias for GPTQ models.

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

267 + continue

268 +

269 + # According to DeepSeek-V3 Technical Report, MTP modules

270 + # shares embedding layer. We only load the first weights.

271 + if (spec_layer != self.model.mtp_start_layer_idx

272 + and ".layers" not in name):

273 + continue

274 +

275 + param = params_dict[name]

276 + weight_loader = getattr(param, "weight_loader",

277 + default_weight_loader)

278 + weight_loader(param, loaded_weight)

279 + loaded_params.add(name)

280 + return loaded_params

281 +

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

283 + """

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

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

286 + and rename shared layer weights to be top level.

287 + """

288 + spec_layer_weight_names = [

289 + "embed_tokens", "enorm", "hnorm", "eh_proj", "shared_head"

290 + ]

291 + shared_weight_names = ["embed_tokens"]

292 + spec_layer_weight = False

293 + shared_weight = False

294 + for weight_name in spec_layer_weight_names:

295 + if weight_name in name:

296 + spec_layer_weight = True

297 + if weight_name in shared_weight_names:

298 + shared_weight = True

299 + break

300 + if not spec_layer_weight:

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

302 + name = name.replace(f"model.layers.{spec_layer}.",

303 + f"model.layers.{spec_layer}.mtp_block.")

304 + elif shared_weight:

305 + # treat shared weights as top level weights

306 + name = name.replace(f"model.layers.{spec_layer}.", "model.")

307 + return name

vllm/model_executor/models/registry.py

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

67 "Gemma3nForConditionalGeneration": ("gemma3n", "Gemma3nForConditionalGeneration"), # noqa: E501

68 "GlmForCausalLM": ("glm", "GlmForCausalLM"),

69 "Glm4ForCausalLM": ("glm4", "Glm4ForCausalLM"),

70 + "Glm4MoeForCausalLM": ("glm4_moe", "Glm4MoeForCausalLM"),

71 "GPT2LMHeadModel": ("gpt2", "GPT2LMHeadModel"),

72 "GPTBigCodeForCausalLM": ("gpt_bigcode", "GPTBigCodeForCausalLM"),

73 "GPTJForCausalLM": ("gpt_j", "GPTJForCausalLM"),

@@ -244,6 +245,7 @@ _SPECULATIVE_DECODING_MODELS = {

245 "EagleMiniCPMForCausalLM": ("minicpm_eagle", "EagleMiniCPMForCausalLM"),

246 "Eagle3LlamaForCausalLM": ("llama_eagle3", "Eagle3LlamaForCausalLM"),

247 "DeepSeekMTPModel": ("deepseek_mtp", "DeepSeekMTP"),

248 + "Glm4MoeMTPModel": ("glm4_moe_mtp", "Glm4MoeMTP"),

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

250 # Temporarily disabled.

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

vllm/reasoning/__init__.py

@@ -3,6 +3,7 @@

4 from .abs_reasoning_parsers import ReasoningParser, ReasoningParserManager

5 from .deepseek_r1_reasoning_parser import DeepSeekR1ReasoningParser

6 +from .glm4_moe_reasoning_parser import Glm4MoeModelReasoningParser

7 from .granite_reasoning_parser import GraniteReasoningParser

8 from .hunyuan_a13b_reasoning_parser import HunyuanA13BReasoningParser

9 from .qwen3_reasoning_parser import Qwen3ReasoningParser

@@ -14,4 +15,5 @@ __all__ = [

15 "GraniteReasoningParser",

16 "HunyuanA13BReasoningParser",

17 "Qwen3ReasoningParser",

18 + "Glm4MoeModelReasoningParser",

19 ]

vllm/reasoning/glm4_moe_reasoning_parser.py new

+151

@@ -0,0 +1,151 @@

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

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

3 +

4 +from collections.abc import Sequence

5 +from typing import Optional, Union

6 +

7 +from transformers import PreTrainedTokenizerBase

8 +

9 +from vllm.entrypoints.openai.protocol import (ChatCompletionRequest,

10 + DeltaMessage)

11 +from vllm.logger import init_logger

12 +from vllm.reasoning import ReasoningParser, ReasoningParserManager

13 +

14 +logger = init_logger(__name__)

15 +

16 +

17 +@ReasoningParserManager.register_module("glm4_moe")

18 +class Glm4MoeModelReasoningParser(ReasoningParser):

19 + """

20 + Reasoning parser for the Glm4MoeModel model.

21 +

22 + The Glm4MoeModel model uses <think>...</think> tokens to denote reasoning

23 + text within its output. The model provides a strict switch to disable

24 + reasoning output via the 'enable_thinking=False' parameter. This parser

25 + extracts the reasoning content enclosed by <think> and </think> tokens

26 + from the model's output.

27 + """

28 +

29 + def __init__(self, tokenizer: PreTrainedTokenizerBase):

30 + super().__init__(tokenizer)

31 + self.think_start_token = "<think>"

32 + self.think_end_token = "</think>"

33 +

34 + if not self.model_tokenizer:

35 + raise ValueError(

36 + "The model tokenizer must be passed to the ReasoningParser "

37 + "constructor during construction.")

38 +

39 + self.think_start_token_id = self.vocab.get(self.think_start_token)

40 + self.think_end_token_id = self.vocab.get(self.think_end_token)

41 + if (self.think_start_token_id is None

42 + or self.think_end_token_id is None):

43 + raise RuntimeError(

44 + "Glm4MoeModel reasoning parser could not locate "

45 + "think start/end tokens in the tokenizer!")

46 +

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

48 + return self.think_end_token_id in input_ids

49 +

50 + def extract_content_ids(self, input_ids: list[int]) -> list[int]:

51 + """

52 + Extract the content after the end tokens

53 + """

54 + if self.think_end_token_id not in input_ids[:-1]:

55 + return []

56 + else:

57 + return input_ids[input_ids.index(self.think_end_token_id) + 1:]

58 +

59 + def extract_reasoning_content_streaming(

60 + self,

61 + previous_text: str,

62 + current_text: str,

63 + delta_text: str,

64 + previous_token_ids: Sequence[int],

65 + current_token_ids: Sequence[int],

66 + delta_token_ids: Sequence[int],

67 + ) -> Union[DeltaMessage, None]:

68 + """

69 + Extract reasoning content from a delta message.

70 + Handles streaming output where previous + delta = current.

71 + Uses token IDs for faster processing.

72 + For text <think>abc</think>xyz:

73 + - 'abc' goes to reasoning_content

74 + - 'xyz' goes to content

75 + """

76 + # Skip single special tokens

77 + if len(delta_token_ids) == 1 and (delta_token_ids[0] in [

78 + self.think_start_token_id, self.think_end_token_id

79 + ]):

80 + return None

81 +

82 + if self.think_start_token_id in previous_token_ids:

83 + if self.think_end_token_id in delta_token_ids:

84 + # <think> in previous, </think> in delta,

85 + # extract reasoning content

86 + end_index = delta_text.find(self.think_end_token)

87 + reasoning_content = delta_text[:end_index]

88 + content = delta_text[end_index + len(self.think_end_token):]

89 + return DeltaMessage(reasoning_content=reasoning_content,

90 + content=content if content else None)

91 + elif self.think_end_token_id in previous_token_ids:

92 + # <think> in previous, </think> in previous,

93 + # reasoning content continues

94 + return DeltaMessage(content=delta_text)

95 + else:

96 + # <think> in previous, no </think> in previous or delta,

97 + # reasoning content continues

98 + return DeltaMessage(reasoning_content=delta_text)

99 + elif self.think_start_token_id in delta_token_ids:

100 + if self.think_end_token_id in delta_token_ids:

101 + # <think> in delta, </think> in delta, extract reasoning content

102 + start_index = delta_text.find(self.think_start_token)

103 + end_index = delta_text.find(self.think_end_token)

104 + reasoning_content = delta_text[start_index +

105 + len(self.think_start_token

106 + ):end_index]

107 + content = delta_text[end_index + len(self.think_end_token):]

108 + return DeltaMessage(reasoning_content=reasoning_content,

109 + content=content if content else None)

110 + else:

111 + # <think> in delta, no </think> in delta,

112 + # reasoning content continues

113 + return DeltaMessage(reasoning_content=delta_text)

114 + else:

115 + # thinking is disabled, just content

116 + return DeltaMessage(content=delta_text)

117 +

118 + def extract_reasoning_content(

119 + self, model_output: str, request: ChatCompletionRequest

120 + ) -> tuple[Optional[str], Optional[str]]:

121 + """

122 + Extract reasoning content from the model output.

123 +

124 + For text <think>abc</think>xyz:

125 + - 'abc' goes to reasoning_content

126 + - 'xyz' goes to content

127 +

128 + Returns:

129 + tuple[Optional[str], Optional[str]]: reasoning content and content

130 + """

131 +

132 + # Check if the model output contains the <think> and </think> tokens.

133 + if (self.think_start_token not in model_output

134 + or self.think_end_token not in model_output):

135 + return None, model_output

136 + # Check if the <think> is present in the model output, remove it

137 + # if it is present.

138 + model_output_parts = model_output.partition(self.think_start_token)

139 + model_output = model_output_parts[2] if model_output_parts[

140 + 1] else model_output_parts[0]

141 + # Check if the model output contains the </think> tokens.

142 + # If the end token is not found, return the model output as is.

143 + if self.think_end_token not in model_output:

144 + return None, model_output

145 +

146 + # Extract reasoning content from the model output.

147 + reasoning_content, _, content = model_output.partition(

148 + self.think_end_token)

149 +

150 + final_content = content or None

151 + return reasoning_content, final_content

vllm/worker/worker.py

+2 -1

@@ -77,7 +77,8 @@ class Worker(LocalOrDistributedWorkerBase):

77 "mlp_speculator",

78 "eagle",

79 "deepseek_mtp",

80 - "mimo_mtp")) \

80 + "glm4_moe_mtp",

81 + "mimo_mtp")) \

82 else {"return_hidden_states": True}

84 ModelRunnerClass: Type[GPUModelRunnerBase] = ModelRunner