@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 71832ba71e
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 10, 00:15 UTC

[GLM-4.7] GLM Model support for GLM-Lite (#31386)

Signed-off-by: zRzRzRzRzRzRzR <2448370773@qq.com> Signed-off-by: Yuxuan Zhang <2448370773@qq.com>

Yuxuan Zhang committed Jan 19, 2026 at 17:18 UTC 71832ba71e77c01a1d249e014384522b46813749
9 files changed +1135 -1
benchmarks/kernels/benchmark_moe.py
+1
@@ -643,6 +643,7 @@ def main(args: argparse.Namespace):
643 "DeepseekV3ForCausalLM",
644 "DeepseekV32ForCausalLM",
645 "Glm4MoeForCausalLM",
646 + "Glm4MoeLiteForCausalLM",
647 "NemotronHForCausalLM",
648 ):
649 E = config.n_routed_experts
benchmarks/kernels/benchmark_moe_permute_unpermute.py
+1
@@ -330,6 +330,7 @@ def main(args: argparse.Namespace):
330 config.architectures[0] == "DeepseekV3ForCausalLM"
331 or config.architectures[0] == "DeepseekV2ForCausalLM"
332 or config.architectures[0] == "Glm4MoeForCausalLM"
333 + or config.architectures[0] == "Glm4MoeLiteForCausalLM"
334 ):
335 E = config.n_routed_experts
336 topk = config.num_experts_per_tok
docs/features/tool_calling.md
+1
@@ -369,6 +369,7 @@ Flags: `--tool-call-parser glm45`
369 Supported models:
370
371 * `zai-org/GLM-4.7`
372 +* `zai-org/GLM-4.7-Flash`
373
374 Flags: `--tool-call-parser glm47`
375
tests/models/registry.py
+10
@@ -271,6 +271,11 @@ _TEXT_GENERATION_EXAMPLE_MODELS = {
271 "GlmForCausalLM": _HfExamplesInfo("zai-org/glm-4-9b-chat-hf"),
272 "Glm4ForCausalLM": _HfExamplesInfo("zai-org/GLM-4-9B-0414"),
273 "Glm4MoeForCausalLM": _HfExamplesInfo("zai-org/GLM-4.5"),
274 + "Glm4MoeLiteForCausalLM": _HfExamplesInfo(
275 + "zai-org/GLM-4.7-Flash",
276 + min_transformers_version="5.0.0.dev",
277 + is_available_online=False,
278 + ),
279 "GPT2LMHeadModel": _HfExamplesInfo("openai-community/gpt2", {"alias": "gpt2"}),
280 "GPTBigCodeForCausalLM": _HfExamplesInfo(
281 "bigcode/starcoder",
@@ -1040,6 +1045,11 @@ _SPECULATIVE_DECODING_EXAMPLE_MODELS = {
1045 "zai-org/GLM-4.5",
1046 speculative_model="zai-org/GLM-4.5",
1047 ),
1048 + "Glm4MoeLiteMTPModel": _HfExamplesInfo(
1049 + "zai-org/GLM-4.7-Flash",
1050 + speculative_model="zai-org/GLM-4.7-Flash",
1051 + is_available_online=False,
1052 + ),
1053 "LongCatFlashMTPModel": _HfExamplesInfo(
1054 "meituan-longcat/LongCat-Flash-Chat",
1055 trust_remote_code=True,
vllm/config/speculative.py
+12
@@ -33,6 +33,7 @@ MTPModelTypes = Literal[
33 "deepseek_mtp",
34 "mimo_mtp",
35 "glm4_moe_mtp",
36 + "glm4_moe_lite_mtp",
37 "ernie_mtp",
38 "exaone_moe_mtp",
39 "qwen3_next_mtp",
@@ -206,6 +207,17 @@ class SpeculativeConfig:
207 }
208 )
209
210 + if hf_config.architectures[0] == "Glm4MoeLiteForCausalLM":
211 + hf_config.model_type = "glm4_moe_lite_mtp"
212 + n_predict = getattr(hf_config, "num_nextn_predict_layers", None)
213 + hf_config.update(
214 + {
215 + "num_hidden_layers": 0,
216 + "n_predict": n_predict,
217 + "architectures": ["Glm4MoeLiteMTPModel"],
218 + }
219 + )
220 +
221 if hf_config.model_type == "ernie4_5_moe":
222 hf_config.model_type = "ernie_mtp"
223 if hf_config.model_type == "ernie_mtp":
vllm/model_executor/models/glm4_moe_lite.py new
+642
@@ -0,0 +1,642 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +
4 +# Copyright 2025 The ZhipuAI Team.
5 +# Copyright 2023 The vLLM team.
6 +# Copyright 2022 EleutherAI and the HuggingFace Inc. team. All rights reserved.
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.7-Flash model compatible with HuggingFace weights."""
25 +
26 +import typing
27 +from collections.abc import Callable, Iterable
28 +from itertools import islice
29 +from typing import TYPE_CHECKING
30 +
31 +import torch
32 +from torch import nn
33 +
34 +if TYPE_CHECKING:
35 + from transformers.models.glm4_moe_lite import Glm4MoeLiteConfig
36 +
37 +from vllm._aiter_ops import rocm_aiter_ops
38 +from vllm.compilation.decorators import support_torch_compile
39 +from vllm.config import VllmConfig
40 +from vllm.distributed import (
41 + get_pp_group,
42 +)
43 +from vllm.logger import init_logger
44 +from vllm.model_executor.layers.fused_moe import SharedFusedMoE
45 +from vllm.model_executor.layers.layernorm import RMSNorm
46 +from vllm.model_executor.layers.logits_processor import LogitsProcessor
47 +from vllm.model_executor.layers.vocab_parallel_embedding import (
48 + ParallelLMHead,
49 + VocabParallelEmbedding,
50 +)
51 +from vllm.model_executor.model_loader.weight_utils import (
52 + default_weight_loader,
53 + maybe_remap_kv_scale_name,
54 +)
55 +from vllm.model_executor.models.deepseek_v2 import (
56 + DeepseekV2Attention,
57 + DeepseekV2MLAAttention,
58 +)
59 +from vllm.model_executor.models.glm4_moe import (
60 + Glm4MixtureOfExperts,
61 + Glm4MoE,
62 + Glm4MoeMLP,
63 +)
64 +from vllm.platforms import current_platform
65 +from vllm.sequence import IntermediateTensors
66 +
67 +from .interfaces import SupportsLoRA, SupportsPP
68 +from .utils import (
69 + AutoWeightsLoader,
70 + PPMissingLayer,
71 + is_pp_missing_parameter,
72 + make_empty_intermediate_tensors_factory,
73 + make_layers,
74 + maybe_prefix,
75 +)
76 +
77 +logger = init_logger(__name__)
78 +
79 +
80 +class Glm4MoeLiteMLP(Glm4MoeMLP):
81 + pass
82 +
83 +
84 +class Glm4MoeLite(Glm4MoE):
85 + pass
86 +
87 +
88 +class Glm4LiteMixtureOfExperts(Glm4MixtureOfExperts):
89 + pass
90 +
91 +
92 +class Glm4MoeLiteAttention(DeepseekV2Attention):
93 + pass
94 +
95 +
96 +class Glm4MoeLiteMLAAttention(DeepseekV2MLAAttention):
97 + pass
98 +
99 +
100 +class Glm4MoeLiteDecoderLayer(nn.Module):
101 + def __init__(
102 + self,
103 + vllm_config: VllmConfig,
104 + prefix: str,
105 + config: "Glm4MoeLiteConfig | None" = None,
106 + topk_indices_buffer: torch.Tensor | None = None,
107 + ) -> None:
108 + super().__init__()
109 +
110 + if config is None:
111 + config = vllm_config.model_config.hf_config
112 + model_config = vllm_config.model_config
113 + cache_config = vllm_config.cache_config
114 + quant_config = vllm_config.quant_config
115 +
116 + self.hidden_size = config.hidden_size
117 + max_position_embeddings = getattr(config, "max_position_embeddings", 8192)
118 + moe_layer_freq = getattr(config, "moe_layer_freq", 1)
119 + # DecoderLayers are created with `make_layers` which passes the prefix
120 + # with the layer's index.
121 + layer_idx = int(prefix.split(sep=".")[-1])
122 + self.layer_idx = layer_idx
123 +
124 + # verify MLA attention specific fields
125 + qk_nope_head_dim = getattr(config, "qk_nope_head_dim", 0)
126 + qk_rope_head_dim = getattr(config, "qk_rope_head_dim", 0)
127 + v_head_dim = getattr(config, "v_head_dim", 0)
128 + kv_lora_rank = getattr(config, "kv_lora_rank", 0)
129 +
130 + if model_config.use_mla:
131 + attn_cls = Glm4MoeLiteMLAAttention
132 + else:
133 + attn_cls = Glm4MoeLiteAttention
134 +
135 + self.self_attn = attn_cls(
136 + vllm_config=vllm_config,
137 + config=config,
138 + hidden_size=self.hidden_size,
139 + num_heads=config.num_attention_heads,
140 + qk_nope_head_dim=qk_nope_head_dim,
141 + qk_rope_head_dim=qk_rope_head_dim,
142 + v_head_dim=v_head_dim,
143 + q_lora_rank=config.q_lora_rank if hasattr(config, "q_lora_rank") else None,
144 + kv_lora_rank=kv_lora_rank,
145 + max_position_embeddings=max_position_embeddings,
146 + cache_config=cache_config,
147 + quant_config=quant_config,
148 + prefix=f"{prefix}.self_attn",
149 + topk_indices_buffer=topk_indices_buffer,
150 + )
151 +
152 + if (
153 + config.n_routed_experts is not None
154 + and layer_idx >= config.first_k_dense_replace
155 + and layer_idx % moe_layer_freq == 0
156 + ):
157 + self.mlp = Glm4MoeLite(
158 + config=config,
159 + quant_config=quant_config,
160 + prefix=f"{prefix}.mlp",
161 + )
162 + else:
163 + self.mlp = Glm4MoeLiteMLP(
164 + hidden_size=config.hidden_size,
165 + intermediate_size=config.intermediate_size,
166 + hidden_act=config.hidden_act,
167 + quant_config=quant_config,
168 + prefix=f"{prefix}.mlp",
169 + )
170 + self.input_layernorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
171 + self.post_attention_layernorm = RMSNorm(
172 + config.hidden_size, eps=config.rms_norm_eps
173 + )
174 + self.routed_scaling_factor = getattr(config, "routed_scaling_factor", 1.0)
175 +
176 + def forward(
177 + self,
178 + positions: torch.Tensor,
179 + hidden_states: torch.Tensor,
180 + residual: torch.Tensor | None,
181 + llama_4_scaling: torch.Tensor | None = None,
182 + ) -> torch.Tensor:
183 + # Self Attention
184 + if residual is None:
185 + residual = hidden_states.clone()
186 + hidden_states = self.input_layernorm(hidden_states)
187 + else:
188 + hidden_states, residual = self.input_layernorm(hidden_states, residual)
189 +
190 + attn_kwargs = {
191 + "positions": positions,
192 + "hidden_states": hidden_states,
193 + }
194 + attn_kwargs["llama_4_scaling"] = llama_4_scaling
195 + hidden_states = self.self_attn(**attn_kwargs)
196 +
197 + hidden_states, residual = self.post_attention_layernorm(hidden_states, residual)
198 + hidden_states = self.mlp(hidden_states)
199 +
200 + return hidden_states, residual
201 +
202 +
203 +@support_torch_compile(
204 + dynamic_arg_dims={
205 + "input_ids": 0,
206 + "positions": -1,
207 + "intermediate_tensors": 0,
208 + "inputs_embeds": 0,
209 + }
210 +)
211 +class Glm4MoeLiteModel(nn.Module):
212 + def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
213 + super().__init__()
214 +
215 + config = vllm_config.model_config.hf_config
216 + quant_config = vllm_config.quant_config
217 + self.config = config
218 + self.device = current_platform.device_type
219 +
220 + self.vocab_size = config.vocab_size
221 + self.is_v32 = hasattr(config, "index_topk")
222 + if self.is_v32:
223 + topk_tokens = config.index_topk
224 + topk_indices_buffer = torch.empty(
225 + vllm_config.scheduler_config.max_num_batched_tokens,
226 + topk_tokens,
227 + dtype=torch.int32,
228 + device=self.device,
229 + )
230 + else:
231 + topk_indices_buffer = None
232 +
233 + if get_pp_group().is_first_rank:
234 + self.embed_tokens = VocabParallelEmbedding(
235 + config.vocab_size,
236 + config.hidden_size,
237 + quant_config=quant_config,
238 + prefix=f"{prefix}.embed_tokens",
239 + )
240 + else:
241 + self.embed_tokens = PPMissingLayer()
242 +
243 + self.start_layer, self.end_layer, self.layers = make_layers(
244 + config.num_hidden_layers,
245 + lambda prefix: Glm4MoeLiteDecoderLayer(
246 + vllm_config=vllm_config,
247 + config=config,
248 + prefix=prefix,
249 + topk_indices_buffer=topk_indices_buffer,
250 + ),
251 + prefix=f"{prefix}.layers",
252 + )
253 +
254 + if get_pp_group().is_last_rank:
255 + self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
256 + else:
257 + self.norm = PPMissingLayer()
258 + self.make_empty_intermediate_tensors = make_empty_intermediate_tensors_factory(
259 + ["hidden_states", "residual"], config.hidden_size
260 + )
261 +
262 + def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:
263 + return self.embed_tokens(input_ids)
264 +
265 + def forward(
266 + self,
267 + input_ids: torch.Tensor,
268 + positions: torch.Tensor,
269 + intermediate_tensors: IntermediateTensors | None = None,
270 + inputs_embeds: torch.Tensor | None = None,
271 + ) -> torch.Tensor | IntermediateTensors:
272 + if get_pp_group().is_first_rank:
273 + if inputs_embeds is not None:
274 + hidden_states = inputs_embeds
275 + else:
276 + hidden_states = self.embed_input_ids(input_ids)
277 + residual = None
278 + else:
279 + assert intermediate_tensors is not None
280 + hidden_states = intermediate_tensors["hidden_states"]
281 + residual = intermediate_tensors["residual"]
282 +
283 + for layer in islice(self.layers, self.start_layer, self.end_layer):
284 + hidden_states, residual = layer(positions, hidden_states, residual)
285 +
286 + if not get_pp_group().is_last_rank:
287 + return IntermediateTensors(
288 + {"hidden_states": hidden_states, "residual": residual}
289 + )
290 +
291 + hidden_states, _ = self.norm(hidden_states, residual)
292 + return hidden_states
293 +
294 + def make_empty_intermediate_tensors(
295 + self, batch_size: int, dtype: torch.dtype, device: torch.device
296 + ) -> IntermediateTensors:
297 + return IntermediateTensors(
298 + {
299 + "hidden_states": torch.zeros(
300 + (batch_size, self.config.hidden_size), dtype=dtype, device=device
301 + ),
302 + "residual": torch.zeros(
303 + (batch_size, self.config.hidden_size), dtype=dtype, device=device
304 + ),
305 + }
306 + )
307 +
308 + def get_expert_mapping(self) -> list[tuple[str, str, int, str]]:
309 + # Params for weights, fp8 weight scales, fp8 activation scales
310 + # (param_name, weight_name, expert_id, shard_id)
311 + return SharedFusedMoE.make_expert_params_mapping(
312 + self,
313 + ckpt_gate_proj_name="gate_proj",
314 + ckpt_down_proj_name="down_proj",
315 + ckpt_up_proj_name="up_proj",
316 + num_experts=self.config.n_routed_experts,
317 + )
318 +
319 + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
320 + rocm_aiter_moe_shared_expert_enabled = (
321 + rocm_aiter_ops.is_fusion_moe_shared_experts_enabled()
322 + )
323 + stacked_params_mapping = [
324 + # (param_name, shard_name, shard_id)
325 + ("gate_up_proj", "gate_proj", 0),
326 + ("gate_up_proj", "up_proj", 1),
327 + ]
328 + mla_params_mapping = [
329 + ("fused_qkv_a_proj", "q_a_proj", 0),
330 + ("fused_qkv_a_proj", "kv_a_proj_with_mqa", 1),
331 + ]
332 +
333 + stacked_params_mapping.extend(mla_params_mapping)
334 +
335 + # Params for weights, fp8 weight scales, fp8 activation scales
336 + # (param_name, weight_name, expert_id, shard_id)
337 + expert_params_mapping = SharedFusedMoE.make_expert_params_mapping(
338 + self,
339 + ckpt_gate_proj_name="gate_proj",
340 + ckpt_down_proj_name="down_proj",
341 + ckpt_up_proj_name="up_proj",
342 + num_experts=self.config.n_routed_experts
343 + + (
344 + self.config.n_shared_experts
345 + if rocm_aiter_moe_shared_expert_enabled
346 + else 0
347 + ),
348 + )
349 +
350 + params_dict = dict(self.named_parameters())
351 + loaded_params: set[str] = set()
352 + for name, loaded_weight in weights:
353 + if "rotary_emb.inv_freq" in name:
354 + continue
355 +
356 + spec_layer = get_spec_layer_idx_from_weight_name(self.config, name)
357 + if spec_layer is not None:
358 + continue # skip spec decode layers for main model
359 +
360 + is_fusion_moe_shared_experts_layer = (
361 + rocm_aiter_moe_shared_expert_enabled and ("mlp.shared_experts" in name)
362 + )
363 +
364 + for param_name, weight_name, shard_id in stacked_params_mapping:
365 + # Skip non-stacked layers and experts (experts handled below).
366 + if weight_name not in name:
367 + continue
368 + # We have mlp.experts[0].gate_proj in the checkpoint.
369 + # Since we handle the experts below in expert_params_mapping,
370 + # we need to skip here BEFORE we update the name, otherwise
371 + # name will be updated to mlp.experts[0].gate_up_proj, which
372 + # will then be updated below in expert_params_mapping
373 + # for mlp.experts[0].gate_gate_up_proj, which breaks load.
374 + if ("mlp.experts." in name) and name not in params_dict:
375 + continue
376 + if is_fusion_moe_shared_experts_layer:
377 + continue
378 + name_mapped = name.replace(weight_name, param_name)
379 +
380 + # QKV fusion is optional, fall back to normal
381 + # weight loading if it's not enabled
382 + # if go with fusion option, then update name
383 + if (
384 + param_name == "fused_qkv_a_proj"
385 + ) and name_mapped not in params_dict:
386 + continue
387 + else:
388 + name = name_mapped
389 + # Skip loading extra bias for GPTQ models.
390 + if name.endswith(".bias") and name not in params_dict:
391 + continue
392 +
393 + if is_pp_missing_parameter(name, self):
394 + continue
395 +
396 + param = params_dict[name]
397 + weight_loader = param.weight_loader
398 + weight_loader(param, loaded_weight, shard_id)
399 + break
400 + else:
401 + is_expert_weight = False
402 +
403 + # Special handling: when AITER fusion_shared_experts is enabled,
404 + # checkpoints may provide a single widened shared_experts tensor
405 + # without explicit expert indices
406 + # (e.g. ...mlp.shared_experts.gate_proj.weight).
407 + # For models with multiple shared experts, split that tensor
408 + # evenly into per-shared-expert slices and load them into
409 + # appended expert slots mlp.experts.{n_routed_experts + j}.*
410 + # accordingly.
411 + num_chunks = 1
412 + if is_fusion_moe_shared_experts_layer:
413 + num_chunks = getattr(self.config, "n_shared_experts", 1) or 1
414 + # Determine split axis based on op type
415 + # gate/up: ColumnParallel → split along dim 0
416 + # down: RowParallel → split along dim 1
417 + split_dim = 1 if "down_proj.weight" in name else 0
418 + total = loaded_weight.shape[split_dim]
419 + assert total % num_chunks == 0, (
420 + f"Shared expert weight dim {total} "
421 + f"not divisible by num_chunks {num_chunks}"
422 + )
423 + chunk_size = total // num_chunks
424 +
425 + for j in range(num_chunks):
426 + chunk_name = name
427 + weight_to_load = loaded_weight
428 +
429 + if is_fusion_moe_shared_experts_layer:
430 + if split_dim == 0:
431 + weight_to_load = loaded_weight[
432 + j * chunk_size : (j + 1) * chunk_size, :
433 + ]
434 + else:
435 + weight_to_load = loaded_weight[
436 + :, j * chunk_size : (j + 1) * chunk_size
437 + ]
438 + # Synthesize an expert-style name so expert mapping
439 + # can route it
440 + chunk_name = name.replace(
441 + "mlp.shared_experts",
442 + f"mlp.experts.{self.config.n_routed_experts + j}",
443 + )
444 +
445 + # Use expert_params_mapping to locate the destination
446 + # param and delegate to its expert-aware weight_loader
447 + # with expert_id.
448 + for mapping in expert_params_mapping:
449 + param_name, weight_name, expert_id, shard_id = mapping
450 + if weight_name not in chunk_name:
451 + continue
452 +
453 + # Anyway, this is an expert weight and should not be
454 + # attempted to load as other weights later
455 + is_expert_weight = True
456 +
457 + # Do not modify `name` since the loop may continue here
458 + # Instead, create a new variable
459 + name_mapped = chunk_name.replace(weight_name, param_name)
460 +
461 + if is_pp_missing_parameter(name_mapped, self):
462 + continue
463 +
464 + param = params_dict[name_mapped]
465 + # We should ask the weight loader to return success or
466 + # not here since otherwise we may skip experts with
467 + # other available replicas.
468 + weight_loader = typing.cast(
469 + Callable[..., bool], param.weight_loader
470 + )
471 + success = weight_loader(
472 + param,
473 + weight_to_load,
474 + name_mapped,
475 + shard_id=shard_id,
476 + expert_id=expert_id,
477 + return_success=True,
478 + )
479 + if success:
480 + if not is_fusion_moe_shared_experts_layer:
481 + name = name_mapped
482 + else:
483 + loaded_params.add(name_mapped)
484 + break
485 + else:
486 + if is_expert_weight:
487 + # We've checked that this is an expert weight
488 + # However it's not mapped locally to this rank
489 + # So we simply skip it
490 + continue
491 +
492 + # Skip loading extra bias for GPTQ models.
493 + if name.endswith(".bias") and name not in params_dict:
494 + continue
495 +
496 + # Remapping the name of FP8 kv-scale.
497 + name = maybe_remap_kv_scale_name(name, params_dict)
498 + if name is None:
499 + continue
500 +
501 + if is_pp_missing_parameter(name, self):
502 + continue
503 +
504 + param = params_dict[name]
505 + weight_loader = getattr(
506 + param, "weight_loader", default_weight_loader
507 + )
508 + weight_loader(param, loaded_weight)
509 + if not is_fusion_moe_shared_experts_layer:
510 + loaded_params.add(name)
511 +
512 + return loaded_params
513 +
514 +
515 +class Glm4MoeLiteForCausalLM(
516 + nn.Module, SupportsPP, SupportsLoRA, Glm4LiteMixtureOfExperts
517 +):
518 + packed_modules_mapping = {
519 + "gate_up_proj": ["gate_proj", "up_proj"],
520 + }
521 +
522 + def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
523 + super().__init__()
524 + config = vllm_config.model_config.hf_config
525 + quant_config = vllm_config.quant_config
526 + self.config = config
527 + self.quant_config = quant_config
528 +
529 + qk_nope_head_dim = getattr(config, "qk_nope_head_dim", 0)
530 + qk_rope_head_dim = getattr(config, "qk_rope_head_dim", 0)
531 + self.use_mha = config.model_type == "deepseek" or all(
532 + dim == 0 for dim in (qk_nope_head_dim, qk_rope_head_dim)
533 + )
534 +
535 + if self.use_mha:
536 + self.packed_modules_mapping["qkv_proj"] = ["q_proj", "k_proj", "v_proj"]
537 +
538 + # `packed_modules_mapping` needs to be modified before
539 + # initializing DeepseekV2Model, as it is passed inplace to
540 + # quantization config init and may be used to select the
541 + # quant_method for relevant layers during initialization.
542 + self.fuse_qkv_a_proj = (
543 + hasattr(config, "q_lora_rank") and config.q_lora_rank is not None
544 + )
545 + if self.fuse_qkv_a_proj:
546 + self.packed_modules_mapping["fused_qkv_a_proj"] = [
547 + "q_a_proj",
548 + "kv_a_proj_with_mqa",
549 + ]
550 +
551 + self.model = Glm4MoeLiteModel(
552 + vllm_config=vllm_config, prefix=maybe_prefix(prefix, "model")
553 + )
554 + if get_pp_group().is_last_rank:
555 + self.lm_head = ParallelLMHead(
556 + config.vocab_size,
557 + config.hidden_size,
558 + quant_config=quant_config,
559 + prefix=maybe_prefix(prefix, "lm_head"),
560 + )
561 + else:
562 + self.lm_head = PPMissingLayer()
563 + self.logits_processor = LogitsProcessor(config.vocab_size)
564 + self.make_empty_intermediate_tensors = (
565 + self.model.make_empty_intermediate_tensors
566 + )
567 + # Set MoE hyperparameters
568 + self.num_moe_layers = (
569 + self.config.num_hidden_layers - self.config.first_k_dense_replace
570 + )
571 + self.set_moe_parameters()
572 +
573 + def set_moe_parameters(self):
574 + self.expert_weights = []
575 +
576 + self.num_expert_groups = getattr(self.config, "n_group", 1)
577 +
578 + self.moe_layers = []
579 + self.moe_mlp_layers = []
580 + example_moe = None
581 + for layer in self.model.layers:
582 + if isinstance(layer, PPMissingLayer):
583 + continue
584 +
585 + assert isinstance(layer, Glm4MoeLiteDecoderLayer)
586 + if isinstance(layer.mlp, Glm4MoeLite):
587 + # Pick last one layer since the first ones may be dense layers.
588 + example_moe = layer.mlp
589 + self.moe_mlp_layers.append(layer.mlp)
590 + self.moe_layers.append(layer.mlp.experts)
591 +
592 + self.extract_moe_parameters(example_moe)
593 +
594 + def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:
595 + return self.model.embed_input_ids(input_ids)
596 +
597 + def forward(
598 + self,
599 + input_ids: torch.Tensor,
600 + positions: torch.Tensor,
601 + intermediate_tensors: IntermediateTensors | None = None,
602 + inputs_embeds: torch.Tensor | None = None,
603 + ) -> torch.Tensor | IntermediateTensors:
604 + hidden_states = self.model(
605 + input_ids, positions, intermediate_tensors, inputs_embeds
606 + )
607 + return hidden_states
608 +
609 + def compute_logits(
610 + self,
611 + hidden_states: torch.Tensor,
612 + ) -> torch.Tensor | None:
613 + logits = self.logits_processor(self.lm_head, hidden_states)
614 + return logits
615 +
616 + def get_expert_mapping(self) -> list[tuple[str, str, int, str]]:
617 + # Params for weights, fp8 weight scales, fp8 activation scales
618 + # (param_name, weight_name, expert_id, shard_id)
619 + return SharedFusedMoE.make_expert_params_mapping(
620 + ckpt_gate_proj_name="gate_proj",
621 + ckpt_down_proj_name="down_proj",
622 + ckpt_up_proj_name="up_proj",
623 + num_experts=self.config.n_routed_experts,
624 + num_redundant_experts=0,
625 + )
626 +
627 + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
628 + loader = AutoWeightsLoader(self)
629 + return loader.load_weights(weights)
630 +
631 +
632 +def get_spec_layer_idx_from_weight_name(
633 + config: "Glm4MoeLiteConfig", weight_name: str
634 +) -> int | None:
635 + if hasattr(config, "num_nextn_predict_layers") and (
636 + config.num_nextn_predict_layers > 0
637 + ):
638 + layer_idx = config.num_hidden_layers
639 + for i in range(config.num_nextn_predict_layers):
640 + if f"layers.{layer_idx + i}." in weight_name:
641 + return layer_idx + i
642 + return None
vllm/model_executor/models/glm4_moe_lite_mtp.py new
+464
@@ -0,0 +1,464 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +
4 +# Copyright 2025 The ZhipuAI Team.
5 +# Copyright 2023 The vLLM team.
6 +# Copyright 2022 EleutherAI and the HuggingFace Inc. team. All rights reserved.
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.7-Flash MTP model compatible with HuggingFace weights."""
25 +
26 +import typing
27 +from collections.abc import Callable, Iterable
28 +
29 +import torch
30 +import torch.nn as nn
31 +from transformers import PretrainedConfig
32 +
33 +from vllm._aiter_ops import rocm_aiter_ops
34 +from vllm.config import VllmConfig
35 +from vllm.model_executor.layers.fused_moe import FusedMoE, SharedFusedMoE
36 +from vllm.model_executor.layers.layernorm import RMSNorm
37 +from vllm.model_executor.layers.logits_processor import LogitsProcessor
38 +from vllm.model_executor.layers.quantization import QuantizationConfig
39 +from vllm.model_executor.layers.vocab_parallel_embedding import (
40 + ParallelLMHead,
41 + VocabParallelEmbedding,
42 +)
43 +from vllm.model_executor.model_loader.weight_utils import (
44 + default_weight_loader,
45 + maybe_remap_kv_scale_name,
46 +)
47 +from vllm.platforms import current_platform
48 +from vllm.sequence import IntermediateTensors
49 +
50 +from .glm4_moe_lite import (
51 + Glm4MixtureOfExperts,
52 + Glm4MoeLite,
53 + Glm4MoeLiteDecoderLayer,
54 + get_spec_layer_idx_from_weight_name,
55 +)
56 +from .interfaces import SupportsPP
57 +from .utils import maybe_prefix
58 +
59 +
60 +class SharedHead(nn.Module):
61 + def __init__(
62 + self,
63 + config: PretrainedConfig,
64 + prefix: str,
65 + quant_config: QuantizationConfig | None = None,
66 + ) -> None:
67 + super().__init__()
68 + self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
69 + self.head = ParallelLMHead(
70 + config.vocab_size,
71 + config.hidden_size,
72 + quant_config=quant_config,
73 + prefix=maybe_prefix(prefix, "head"),
74 + )
75 +
76 + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
77 + return self.norm(hidden_states)
78 +
79 +
80 +class Glm4MoeLiteMultiTokenPredictorLayer(nn.Module):
81 + def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
82 + super().__init__()
83 +
84 + config = vllm_config.speculative_config.draft_model_config.hf_config
85 + self.config = config
86 + quant_config = vllm_config.quant_config
87 +
88 + self.enorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
89 + self.hnorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
90 + self.eh_proj = nn.Linear(config.hidden_size * 2, config.hidden_size, bias=False)
91 +
92 + self.device = current_platform.device_type
93 +
94 + self.is_v32 = hasattr(config, "index_topk")
95 + if self.is_v32:
96 + topk_tokens = config.index_topk
97 + topk_indices_buffer = torch.empty(
98 + vllm_config.scheduler_config.max_num_batched_tokens,
99 + topk_tokens,
100 + dtype=torch.int32,
101 + device=self.device,
102 + )
103 + else:
104 + topk_indices_buffer = None
105 +
106 + self.shared_head = SharedHead(
107 + config=config, prefix=prefix, quant_config=quant_config
108 + )
109 + self.mtp_block = Glm4MoeLiteDecoderLayer(
110 + vllm_config=vllm_config,
111 + prefix=prefix,
112 + config=self.config,
113 + topk_indices_buffer=topk_indices_buffer,
114 + )
115 +
116 + def forward(
117 + self,
118 + input_ids: torch.Tensor,
119 + positions: torch.Tensor,
120 + previous_hidden_states: torch.Tensor,
121 + inputs_embeds: torch.Tensor | None = None,
122 + spec_step_index: int = 0,
123 + ) -> torch.Tensor:
124 + assert inputs_embeds is not None
125 + # masking inputs at position 0, as not needed by MTP
126 + inputs_embeds[positions == 0] = 0
127 + inputs_embeds = self.enorm(inputs_embeds)
128 + previous_hidden_states = self.hnorm(previous_hidden_states)
129 +
130 + hidden_states = self.eh_proj(
131 + torch.cat([inputs_embeds, previous_hidden_states], dim=-1)
132 + )
133 +
134 + hidden_states, residual = self.mtp_block(
135 + positions=positions, hidden_states=hidden_states, residual=None
136 + )
137 + hidden_states = residual + hidden_states
138 + return hidden_states
139 +
140 +
141 +class Glm4MoeLiteMultiTokenPredictor(nn.Module):
142 + def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
143 + super().__init__()
144 + config = vllm_config.model_config.hf_config
145 + self.mtp_start_layer_idx = config.num_hidden_layers
146 + self.num_mtp_layers = config.num_nextn_predict_layers
147 + # to map the exact layer index from weights
148 + self.layers = torch.nn.ModuleDict(
149 + {
150 + str(idx): Glm4MoeLiteMultiTokenPredictorLayer(
151 + vllm_config=vllm_config,
152 + prefix=f"{prefix}.layers.{idx}",
153 + )
154 + for idx in range(
155 + self.mtp_start_layer_idx,
156 + self.mtp_start_layer_idx + self.num_mtp_layers,
157 + )
158 + }
159 + )
160 + self.embed_tokens = VocabParallelEmbedding(
161 + config.vocab_size,
162 + config.hidden_size,
163 + )
164 + self.logits_processor = LogitsProcessor(config.vocab_size)
165 +
166 + def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:
167 + return self.embed_tokens(input_ids)
168 +
169 + def forward(
170 + self,
171 + input_ids: torch.Tensor,
172 + positions: torch.Tensor,
173 + previous_hidden_states: torch.Tensor,
174 + inputs_embeds: torch.Tensor | None = None,
175 + spec_step_idx: int = 0,
176 + ) -> torch.Tensor:
177 + if inputs_embeds is None:
178 + inputs_embeds = self.embed_tokens(input_ids)
179 + current_step_idx = spec_step_idx % self.num_mtp_layers
180 + return self.layers[str(self.mtp_start_layer_idx + current_step_idx)](
181 + input_ids,
182 + positions,
183 + previous_hidden_states,
184 + inputs_embeds,
185 + current_step_idx,
186 + )
187 +
188 + def compute_logits(
189 + self,
190 + hidden_states: torch.Tensor,
191 + spec_step_idx: int = 0,
192 + ) -> torch.Tensor:
193 + current_step_idx = spec_step_idx % self.num_mtp_layers
194 + mtp_layer = self.layers[str(self.mtp_start_layer_idx + current_step_idx)]
195 + logits = self.logits_processor(
196 + mtp_layer.shared_head.head, mtp_layer.shared_head(hidden_states)
197 + )
198 + return logits
199 +
200 +
201 +class Glm4MoeLiteMTP(nn.Module, SupportsPP, Glm4MixtureOfExperts):
202 + def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
203 + super().__init__()
204 + self.config = vllm_config.model_config.hf_config
205 + self.model = Glm4MoeLiteMultiTokenPredictor(
206 + vllm_config=vllm_config, prefix=maybe_prefix(prefix, "model")
207 + )
208 +
209 + self.expert_weights = []
210 +
211 + # Set MoE hyperparameters
212 + self.num_moe_layers = self.config.num_nextn_predict_layers
213 + self.num_expert_groups = self.config.n_group
214 +
215 + self.moe_layers: list[FusedMoE] = []
216 + self.moe_mlp_layers: list[Glm4MoeLite] = []
217 + example_moe = None
218 + for layer in self.model.layers.values():
219 + assert isinstance(layer, Glm4MoeLiteMultiTokenPredictorLayer)
220 + layer = layer.mtp_block
221 + assert isinstance(layer, Glm4MoeLiteDecoderLayer)
222 + if isinstance(layer.mlp, Glm4MoeLite):
223 + example_moe = layer.mlp
224 + self.moe_mlp_layers.append(layer.mlp)
225 + self.moe_layers.append(layer.mlp.experts)
226 + self.extract_moe_parameters(example_moe)
227 +
228 + def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:
229 + return self.model.embed_input_ids(input_ids)
230 +
231 + def forward(
232 + self,
233 + input_ids: torch.Tensor,
234 + positions: torch.Tensor,
235 + hidden_states: torch.Tensor,
236 + intermediate_tensors: IntermediateTensors | None = None,
237 + inputs_embeds: torch.Tensor | None = None,
238 + spec_step_idx: int = 0,
239 + ) -> torch.Tensor:
240 + hidden_states = self.model(
241 + input_ids, positions, hidden_states, inputs_embeds, spec_step_idx
242 + )
243 + return hidden_states
244 +
245 + def compute_logits(
246 + self,
247 + hidden_states: torch.Tensor,
248 + spec_step_idx: int = 0,
249 + ) -> torch.Tensor | None:
250 + return self.model.compute_logits(hidden_states, spec_step_idx)
251 +
252 + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
253 + rocm_aiter_moe_shared_expert_enabled = (
254 + rocm_aiter_ops.is_fusion_moe_shared_experts_enabled()
255 + )
256 + stacked_params_mapping = [
257 + ("gate_up_proj", "gate_proj", 0),
258 + ("gate_up_proj", "up_proj", 1),
259 + ("fused_qkv_a_proj", "q_a_proj", 0),
260 + ("fused_qkv_a_proj", "kv_a_proj_with_mqa", 1),
261 + ]
262 +
263 + expert_params_mapping = SharedFusedMoE.make_expert_params_mapping(
264 + self,
265 + ckpt_gate_proj_name="gate_proj",
266 + ckpt_down_proj_name="down_proj",
267 + ckpt_up_proj_name="up_proj",
268 + num_experts=self.config.n_routed_experts
269 + + (
270 + self.config.n_shared_experts
271 + if rocm_aiter_moe_shared_expert_enabled
272 + else 0
273 + ),
274 + )
275 +
276 + params_dict = dict(self.named_parameters())
277 + loaded_params: set[str] = set()
278 + for name, loaded_weight in weights:
279 + if "rotary_emb.inv_freq" in name:
280 + continue
281 + spec_layer = get_spec_layer_idx_from_weight_name(self.config, name)
282 + if spec_layer is None:
283 + continue
284 + is_fusion_moe_shared_experts_layer = (
285 + rocm_aiter_moe_shared_expert_enabled and ("mlp.shared_experts" in name)
286 + )
287 + name = self._rewrite_spec_layer_name(spec_layer, name)
288 + for param_name, weight_name, shard_id in stacked_params_mapping:
289 + # Skip non-stacked layers and experts (experts handled below).
290 + if weight_name not in name:
291 + continue
292 + # We have mlp.experts[0].gate_proj in the checkpoint.
293 + # Since we handle the experts below in expert_params_mapping,
294 + # we need to skip here BEFORE we update the name, otherwise
295 + # name will be updated to mlp.experts[0].gate_up_proj, which
296 + # will then be updated below in expert_params_mapping
297 + # for mlp.experts[0].gate_gate_up_proj, which breaks load.
298 + if ("mlp.experts." in name) and name not in params_dict:
299 + continue
300 + if is_fusion_moe_shared_experts_layer:
301 + continue
302 + name_mapped = name.replace(weight_name, param_name)
303 +
304 + # QKV fusion is optional, fall back to normal
305 + # weight loading if it's not enabled
306 + if (
307 + param_name == "fused_qkv_a_proj"
308 + ) and name_mapped not in params_dict:
309 + continue
310 + else:
311 + name = name_mapped
312 +
313 + # Skip loading extra bias for GPTQ models.
314 + if name.endswith(".bias") and name not in params_dict:
315 + continue
316 +
317 + param = params_dict[name]
318 + weight_loader = param.weight_loader
319 + weight_loader(param, loaded_weight, shard_id)
320 + break
321 + else:
322 + # Special handling: when AITER fusion_shared_experts is enabled,
323 + # checkpoints may provide a single widened shared_experts tensor
324 + # without explicit expert indices
325 + # (e.g. ...mlp.shared_experts.gate_proj.weight).
326 + # For models with multiple shared experts, split that tensor
327 + # evenly into per-shared-expert slices and load them into
328 + # appended expert slots mlp.experts.{n_routed_experts + j}.*
329 + # accordingly.
330 + num_chunks = 1
331 + if is_fusion_moe_shared_experts_layer:
332 + num_chunks = getattr(self.config, "n_shared_experts", 1) or 1
333 + # Determine split axis based on op type
334 + # gate/up: ColumnParallel → split along dim 0
335 + # down: RowParallel → split along dim 1
336 + split_dim = 1 if "down_proj.weight" in name else 0
337 + total = loaded_weight.shape[split_dim]
338 + assert total % num_chunks == 0, (
339 + f"Shared expert weight dim {total} "
340 + f"not divisible by num_chunks {num_chunks}"
341 + )
342 + chunk_size = total // num_chunks
343 +
344 + for j in range(num_chunks):
345 + chunk_name = name
346 + weight_to_load = loaded_weight
347 +
348 + if is_fusion_moe_shared_experts_layer:
349 + if split_dim == 0:
350 + weight_to_load = loaded_weight[
351 + j * chunk_size : (j + 1) * chunk_size, :
352 + ]
353 + else:
354 + weight_to_load = loaded_weight[
355 + :, j * chunk_size : (j + 1) * chunk_size
356 + ]
357 + # Synthesize an expert-style name so expert mapping
358 + # can route it
359 + chunk_name = name.replace(
360 + "mlp.shared_experts",
361 + f"mlp.experts.{self.config.n_routed_experts + j}",
362 + )
363 +
364 + # Use expert_params_mapping to locate the destination
365 + # param and delegate to its expert-aware weight_loader
366 + # with expert_id.
367 + is_expert_weight = False
368 + for mapping in expert_params_mapping:
369 + param_name, weight_name, expert_id, shard_id = mapping
370 + if weight_name not in chunk_name:
371 + continue
372 +
373 + # Anyway, this is an expert weight and should not be
374 + # attempted to load as other weights later
375 + is_expert_weight = True
376 +
377 + # Do not modify `name` since the loop may continue here
378 + # Instead, create a new variable
379 + name_mapped = chunk_name.replace(weight_name, param_name)
380 +
381 + param = params_dict[name_mapped]
382 + # We should ask the weight loader to return success or
383 + # not here since otherwise we may skip experts with
384 + # other available replicas.
385 + weight_loader = typing.cast(
386 + Callable[..., bool], param.weight_loader
387 + )
388 + success = weight_loader(
389 + param,
390 + weight_to_load,
391 + name_mapped,
392 + shard_id=shard_id,
393 + expert_id=expert_id,
394 + return_success=True,
395 + )
396 + if success:
397 + if not is_fusion_moe_shared_experts_layer:
398 + name = name_mapped
399 + else:
400 + loaded_params.add(name_mapped)
401 + break
402 + else:
403 + if is_expert_weight:
404 + # We've checked that this is an expert weight
405 + # However it's not mapped locally to this rank
406 + # So we simply skip it
407 + continue
408 +
409 + # Skip loading extra bias for GPTQ models.
410 + if name.endswith(".bias") and name not in params_dict:
411 + continue
412 +
413 + name = maybe_remap_kv_scale_name(name, params_dict)
414 + if name is None:
415 + continue
416 +
417 + # According to DeepSeek-V3 Technical Report, MTP modules
418 + # shares embedding layer. We only load the first weights.
419 + if (
420 + spec_layer != self.model.mtp_start_layer_idx
421 + and ".layers" not in name
422 + ):
423 + continue
424 +
425 + param = params_dict[name]
426 + weight_loader = getattr(
427 + param, "weight_loader", default_weight_loader
428 + )
429 + weight_loader(param, loaded_weight)
430 + if not is_fusion_moe_shared_experts_layer:
431 + loaded_params.add(name)
432 + return loaded_params
433 +
434 + def _rewrite_spec_layer_name(self, spec_layer: int, name: str) -> str:
435 + """
436 + Rewrite the weight name to match the format of the original model.
437 + Add .mtp_block for modules in transformer layer block for spec layer
438 + and rename shared layer weights to be top level.
439 + """
440 + spec_layer_weight_names = [
441 + "embed_tokens",
442 + "enorm",
443 + "hnorm",
444 + "eh_proj",
445 + "shared_head",
446 + ]
447 + shared_weight_names = ["embed_tokens"]
448 + spec_layer_weight = False
449 + shared_weight = False
450 + for weight_name in spec_layer_weight_names:
451 + if weight_name in name:
452 + spec_layer_weight = True
453 + if weight_name in shared_weight_names:
454 + shared_weight = True
455 + break
456 + if not spec_layer_weight:
457 + # treat rest weights as weights for transformer layer block
458 + name = name.replace(
459 + f"model.layers.{spec_layer}.", f"model.layers.{spec_layer}.mtp_block."
460 + )
461 + elif shared_weight:
462 + # treat shared weights as top level weights
463 + name = name.replace(f"model.layers.{spec_layer}.", "model.")
464 + return name
vllm/model_executor/models/glm4_moe_mtp.py
+2 -1
@@ -21,7 +21,8 @@
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."""
24 +"""Inference-only GLM-4.5, GLM-4.6, GLM-4.7 MTP
25 +model compatible with HuggingFace weights."""
26
27 from collections.abc import Iterable
28
vllm/model_executor/models/registry.py
+2
@@ -112,6 +112,7 @@ _TEXT_GENERATION_MODELS = {
112 "GlmForCausalLM": ("glm", "GlmForCausalLM"),
113 "Glm4ForCausalLM": ("glm4", "Glm4ForCausalLM"),
114 "Glm4MoeForCausalLM": ("glm4_moe", "Glm4MoeForCausalLM"),
115 + "Glm4MoeLiteForCausalLM": ("glm4_moe_lite", "Glm4MoeLiteForCausalLM"),
116 "GptOssForCausalLM": ("gpt_oss", "GptOssForCausalLM"),
117 "GPT2LMHeadModel": ("gpt2", "GPT2LMHeadModel"),
118 "GPTBigCodeForCausalLM": ("gpt_bigcode", "GPTBigCodeForCausalLM"),
@@ -465,6 +466,7 @@ _SPECULATIVE_DECODING_MODELS = {
466 "ExaoneMoeMTP": ("exaone_moe_mtp", "ExaoneMoeMTP"),
467 "LongCatFlashMTPModel": ("longcat_flash_mtp", "LongCatFlashMTP"),
468 "Glm4MoeMTPModel": ("glm4_moe_mtp", "Glm4MoeMTP"),
469 + "Glm4MoeLiteMTPModel": ("glm4_moe_lite_mtp", "Glm4MoeLiteMTP"),
470 "MedusaModel": ("medusa", "Medusa"),
471 "OpenPanguMTPModel": ("openpangu_mtp", "OpenPanguMTP"),
472 "Qwen3NextMTP": ("qwen3_next_mtp", "Qwen3NextMTP"),