@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 63ed2409e8
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, 06:30 UTC

Add K-EXAONE-236B-A23B (#31621)

Signed-off-by: lkm2835 <lkm2835@gmail.com> Signed-off-by: Cyrus Leung <cyrus.tl.leung@gmail.com> Co-authored-by: lgai-exaone <exaonemodels@lgresearch.ai> Co-authored-by: Cyrus Leung <cyrus.tl.leung@gmail.com>

Kyungmin Lee committed Jan 13, 2026 at 01:30 UTC 63ed2409e8cf55ed01df866a38c9d5dad34a393e
7 files changed +856
docs/models/supported_models.md
+1
@@ -375,6 +375,7 @@ th {
375 | `Ernie4_5ForCausalLM` | Ernie4.5 | `baidu/ERNIE-4.5-0.3B-PT`, etc. | ✅︎ | ✅︎ |
376 | `Ernie4_5_MoeForCausalLM` | Ernie4.5MoE | `baidu/ERNIE-4.5-21B-A3B-PT`, `baidu/ERNIE-4.5-300B-A47B-PT`, etc. |✅︎| ✅︎ |
377 | `ExaoneForCausalLM` | EXAONE-3 | `LGAI-EXAONE/EXAONE-3.0-7.8B-Instruct`, etc. | ✅︎ | ✅︎ |
378 +| `ExaoneMoeCausalLM` | K-EXAONE | `LGAI-EXAONE/K-EXAONE-236B-A23B`, etc. | | |
379 | `Exaone4ForCausalLM` | EXAONE-4 | `LGAI-EXAONE/EXAONE-4.0-32B`, etc. | ✅︎ | ✅︎ |
380 | `Fairseq2LlamaForCausalLM` | Llama (fairseq2 format) | `mgleize/fairseq2-dummy-Llama-3.2-1B`, etc. | ✅︎ | ✅︎ |
381 | `FalconForCausalLM` | Falcon | `tiiuae/falcon-7b`, `tiiuae/falcon-40b`, `tiiuae/falcon-rw-7b`, etc. | | ✅︎ |
tests/models/registry.py
+8
@@ -250,6 +250,9 @@ _TEXT_GENERATION_EXAMPLE_MODELS = {
250 "LGAI-EXAONE/EXAONE-3.0-7.8B-Instruct", trust_remote_code=True
251 ),
252 "Exaone4ForCausalLM": _HfExamplesInfo("LGAI-EXAONE/EXAONE-4.0-32B"),
253 + "ExaoneMoEForCausalLM": _HfExamplesInfo(
254 + "LGAI-EXAONE/K-EXAONE-236B-A23B", min_transformers_version="5.0.0"
255 + ),
256 "Fairseq2LlamaForCausalLM": _HfExamplesInfo("mgleize/fairseq2-dummy-Llama-3.2-1B"),
257 "FalconForCausalLM": _HfExamplesInfo("tiiuae/falcon-7b"),
258 "FalconH1ForCausalLM": _HfExamplesInfo("tiiuae/Falcon-H1-0.5B-Base"),
@@ -1005,6 +1008,11 @@ _SPECULATIVE_DECODING_EXAMPLE_MODELS = {
1008 trust_remote_code=True,
1009 speculative_model="baidu/ERNIE-4.5-21B-A3B-PT",
1010 ),
1011 + "ExaoneMoeMTP": _HfExamplesInfo(
1012 + "LGAI-EXAONE/K-EXAONE-236B-A23B",
1013 + speculative_model="LGAI-EXAONE/K-EXAONE-236B-A23B",
1014 + min_transformers_version="5.0.0",
1015 + ),
1016 "Glm4MoeMTPModel": _HfExamplesInfo(
1017 "zai-org/GLM-4.5",
1018 speculative_model="zai-org/GLM-4.5",
vllm/config/speculative.py
+10
@@ -33,6 +33,7 @@ MTPModelTypes = Literal[
33 "mimo_mtp",
34 "glm4_moe_mtp",
35 "ernie_mtp",
36 + "exaone_moe_mtp",
37 "qwen3_next_mtp",
38 "longcat_flash_mtp",
39 "mtp",
@@ -219,6 +220,15 @@ class SpeculativeConfig:
220 hf_config.update(
221 {"n_predict": n_predict, "architectures": ["Qwen3NextMTP"]}
222 )
223 +
224 + if hf_config.model_type == "exaone_moe":
225 + hf_config.model_type = "exaone_moe_mtp"
226 + if hf_config.model_type == "exaone_moe_mtp":
227 + n_predict = getattr(hf_config, "num_nextn_predict_layers", None)
228 + hf_config.update(
229 + {"n_predict": n_predict, "architectures": ["ExaoneMoeMTP"]}
230 + )
231 +
232 if hf_config.model_type == "longcat_flash":
233 hf_config.model_type = "longcat_flash_mtp"
234 n_predict = getattr(hf_config, "num_nextn_predict_layers", 1)
vllm/model_executor/models/exaone4.py
+2
@@ -72,6 +72,7 @@ class Exaone4GatedMLP(nn.Module):
72 intermediate_size: int,
73 hidden_act: str,
74 quant_config: QuantizationConfig | None = None,
75 + reduce_results: bool = True,
76 bias: bool = False,
77 prefix: str = "",
78 ) -> None:
@@ -88,6 +89,7 @@ class Exaone4GatedMLP(nn.Module):
89 output_size=hidden_size,
90 bias=bias,
91 quant_config=quant_config,
92 + reduce_results=reduce_results,
93 prefix=f"{prefix}.down_proj",
94 )
95 if hidden_act != "silu":
vllm/model_executor/models/exaone_moe.py new
+578
@@ -0,0 +1,578 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +# ruff: noqa: E501
4 +
5 +# Licensed under the Apache License, Version 2.0 (the "License");
6 +# you may not use this file except in compliance with the License.
7 +# You may obtain a copy of the License at
8 +#
9 +# http://www.apache.org/licenses/LICENSE-2.0
10 +#
11 +# Unless required by applicable law or agreed to in writing, software
12 +# distributed under the License is distributed on an "AS IS" BASIS,
13 +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
14 +# See the License for the specific language governing permissions and
15 +# limitations under the License.
16 +"""Inference-only K-EXAONE-236B-A22B model compatible with HuggingFace weights."""
17 +
18 +import typing
19 +from collections.abc import Callable, Iterable
20 +from itertools import islice
21 +
22 +import torch
23 +from torch import nn
24 +from transformers import PretrainedConfig
25 +
26 +from vllm.compilation.decorators import support_torch_compile
27 +from vllm.config import CacheConfig, VllmConfig, get_current_vllm_config
28 +from vllm.distributed import (
29 + get_ep_group,
30 + get_pp_group,
31 + get_tensor_model_parallel_world_size,
32 +)
33 +from vllm.model_executor.layers.fused_moe import FusedMoE
34 +from vllm.model_executor.layers.layernorm import RMSNorm
35 +from vllm.model_executor.layers.linear import ReplicatedLinear
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 + DEFAULT_VOCAB_PADDING_SIZE,
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.sequence import IntermediateTensors
48 +
49 +from .exaone4 import Exaone4Attention as ExaoneMoeAttention
50 +from .exaone4 import Exaone4GatedMLP as ExaoneMoeGatedMLP
51 +from .interfaces import SupportsLoRA, SupportsPP
52 +from .utils import (
53 + AutoWeightsLoader,
54 + PPMissingLayer,
55 + extract_layer_index,
56 + is_pp_missing_parameter,
57 + make_empty_intermediate_tensors_factory,
58 + make_layers,
59 + maybe_prefix,
60 +)
61 +
62 +
63 +class ExaoneMoe(nn.Module):
64 + def __init__(
65 + self,
66 + config: PretrainedConfig,
67 + quant_config: QuantizationConfig | None = None,
68 + prefix: str = "",
69 + enable_eplb: bool = False,
70 + ):
71 + super().__init__()
72 + self.tp_size = get_tensor_model_parallel_world_size()
73 +
74 + self.routed_scaling_factor = config.routed_scaling_factor
75 +
76 + self.ep_group = get_ep_group().device_group
77 + self.ep_rank = self.ep_group.rank()
78 + self.ep_size = self.ep_group.size()
79 + self.n_routed_experts = config.num_experts
80 +
81 + if self.tp_size > config.num_experts:
82 + raise ValueError(
83 + f"Tensor parallel size {self.tp_size} is greater than "
84 + f"the number of experts {config.num_experts}."
85 + )
86 +
87 + self.gate = ReplicatedLinear(
88 + config.hidden_size,
89 + config.num_experts,
90 + bias=False,
91 + quant_config=None,
92 + prefix=f"{prefix}.gate",
93 + )
94 +
95 + self.e_score_correction_bias = nn.Parameter(
96 + torch.empty(config.num_experts, dtype=torch.float32)
97 + )
98 +
99 + # Load balancing settings.
100 + vllm_config = get_current_vllm_config()
101 + eplb_config = vllm_config.parallel_config.eplb_config
102 + self.enable_eplb = enable_eplb
103 +
104 + self.n_logical_experts = self.n_routed_experts
105 + eplb_config.num_redundant_experts = (
106 + eplb_config.num_redundant_experts
107 + if eplb_config.num_redundant_experts is not None
108 + else 0
109 + )
110 + self.n_redundant_experts = eplb_config.num_redundant_experts
111 + self.n_physical_experts = self.n_logical_experts + self.n_redundant_experts
112 + self.n_local_physical_experts = self.n_physical_experts // self.ep_size
113 +
114 + self.physical_expert_start = self.ep_rank * self.n_local_physical_experts
115 + self.physical_expert_end = (
116 + self.physical_expert_start + self.n_local_physical_experts
117 + )
118 +
119 + self.experts = FusedMoE(
120 + num_experts=self.n_routed_experts,
121 + top_k=config.num_experts_per_tok,
122 + hidden_size=config.hidden_size,
123 + intermediate_size=config.moe_intermediate_size,
124 + reduce_results=False,
125 + renormalize=config.norm_topk_prob,
126 + quant_config=quant_config,
127 + use_grouped_topk=True,
128 + num_expert_group=config.n_group,
129 + topk_group=config.topk_group,
130 + prefix=f"{prefix}.experts",
131 + scoring_func="sigmoid",
132 + routed_scaling_factor=self.routed_scaling_factor,
133 + e_score_correction_bias=self.e_score_correction_bias,
134 + enable_eplb=self.enable_eplb,
135 + num_redundant_experts=self.n_redundant_experts,
136 + )
137 +
138 + if getattr(config, "num_shared_experts", 0) > 0:
139 + intermediate_size = config.moe_intermediate_size * config.num_shared_experts
140 + self.shared_experts = ExaoneMoeGatedMLP(
141 + hidden_size=config.hidden_size,
142 + intermediate_size=intermediate_size,
143 + hidden_act=config.hidden_act,
144 + quant_config=quant_config,
145 + reduce_results=self.experts.must_reduce_shared_expert_outputs(),
146 + prefix=f"{prefix}.shared_experts",
147 + )
148 + else:
149 + self.shared_experts = None
150 +
151 + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
152 + # NOTE: hidden_states can have either 1D or 2D shape.
153 + orig_shape = hidden_states.shape
154 + hidden_dim = hidden_states.shape[-1]
155 + hidden_states = hidden_states.view(-1, hidden_dim)
156 +
157 + # router_logits: (num_tokens, n_experts)
158 + router_logits, _ = self.gate(hidden_states)
159 +
160 + final_hidden_states = self.experts(
161 + hidden_states=hidden_states, router_logits=router_logits
162 + )
163 +
164 + if self.shared_experts is not None:
165 + shared_output = self.shared_experts(hidden_states)
166 + final_hidden_states = final_hidden_states + shared_output
167 +
168 + if self.tp_size > 1:
169 + final_hidden_states = self.experts.maybe_all_reduce_tensor_model_parallel( # noqa E501
170 + final_hidden_states
171 + )
172 +
173 + return final_hidden_states.view(orig_shape)
174 +
175 +
176 +class ExaoneMoeDecoderLayer(nn.Module):
177 + def __init__(
178 + self,
179 + config: PretrainedConfig,
180 + cache_config: CacheConfig | None = None,
181 + quant_config: QuantizationConfig | None = None,
182 + mtp_layer: bool = None,
183 + prefix: str = "",
184 + ) -> None:
185 + super().__init__()
186 + layer_idx = extract_layer_index(prefix)
187 + self.hidden_size = config.hidden_size
188 + max_position_embeddings = getattr(config, "max_position_embeddings", 8192)
189 + # Support abacusai/Smaug-72B-v0.1 with attention_bias
190 + # Support internlm/internlm-7b with bias
191 + attention_bias = getattr(config, "attention_bias", False) or getattr(
192 + config, "bias", False
193 + )
194 +
195 + self.self_attn = ExaoneMoeAttention(
196 + config=config,
197 + hidden_size=self.hidden_size,
198 + num_heads=config.num_attention_heads,
199 + num_kv_heads=getattr(
200 + config, "num_key_value_heads", config.num_attention_heads
201 + ),
202 + max_position_embeddings=max_position_embeddings,
203 + quant_config=quant_config,
204 + bias=attention_bias,
205 + cache_config=cache_config,
206 + prefix=f"{prefix}.self_attn",
207 + )
208 +
209 + if config.is_moe_layer[layer_idx] and not mtp_layer:
210 + self.mlp = ExaoneMoe(
211 + config=config, quant_config=quant_config, prefix=f"{prefix}.mlp"
212 + )
213 + else:
214 + self.mlp = ExaoneMoeGatedMLP(
215 + hidden_size=self.hidden_size,
216 + intermediate_size=config.intermediate_size,
217 + hidden_act=config.hidden_act,
218 + quant_config=quant_config,
219 + bias=getattr(config, "mlp_bias", False),
220 + prefix=f"{prefix}.mlp",
221 + )
222 + self.input_layernorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
223 + self.post_attention_layernorm = RMSNorm(
224 + config.hidden_size, eps=config.rms_norm_eps
225 + )
226 +
227 + def forward(
228 + self,
229 + positions: torch.Tensor,
230 + hidden_states: torch.Tensor,
231 + residual: torch.Tensor | None,
232 + ) -> tuple[torch.Tensor, torch.Tensor]:
233 + if residual is None:
234 + residual = hidden_states
235 + hidden_states = self.input_layernorm(hidden_states)
236 + else:
237 + hidden_states, residual = self.input_layernorm(hidden_states, residual)
238 +
239 + # Self Attention
240 + hidden_states = self.self_attn(
241 + positions=positions,
242 + hidden_states=hidden_states,
243 + )
244 +
245 + # Fully Connected
246 + hidden_states, residual = self.post_attention_layernorm(hidden_states, residual)
247 + hidden_states = self.mlp(hidden_states)
248 +
249 + return hidden_states, residual
250 +
251 +
252 +@support_torch_compile
253 +class ExaoneMoeModel(nn.Module):
254 + def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
255 + super().__init__()
256 +
257 + config = vllm_config.model_config.hf_config
258 + cache_config = vllm_config.cache_config
259 + quant_config = vllm_config.quant_config
260 + lora_config = vllm_config.lora_config
261 + self.num_redundant_experts = (
262 + vllm_config.parallel_config.eplb_config.num_redundant_experts
263 + )
264 +
265 + self.config = config
266 + self.quant_config = quant_config
267 + lora_vocab = (
268 + (lora_config.lora_extra_vocab_size * (lora_config.max_loras or 1))
269 + if lora_config
270 + else 0
271 + )
272 + self.vocab_size = config.vocab_size + lora_vocab
273 + if get_pp_group().is_first_rank or (
274 + config.tie_word_embeddings and get_pp_group().is_last_rank
275 + ):
276 + self.embed_tokens = VocabParallelEmbedding(
277 + self.vocab_size,
278 + config.hidden_size,
279 + org_num_embeddings=config.vocab_size,
280 + quant_config=quant_config,
281 + )
282 + else:
283 + self.embed_tokens = PPMissingLayer()
284 + self.start_layer, self.end_layer, self.layers = make_layers(
285 + config.num_hidden_layers,
286 + lambda prefix: ExaoneMoeDecoderLayer(
287 + config=config,
288 + cache_config=cache_config,
289 + quant_config=quant_config,
290 + prefix=prefix,
291 + ),
292 + prefix=f"{prefix}.layers",
293 + )
294 + if get_pp_group().is_last_rank:
295 + self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
296 + else:
297 + self.norm = PPMissingLayer()
298 +
299 + self.make_empty_intermediate_tensors = make_empty_intermediate_tensors_factory(
300 + ["hidden_states", "residual"], config.hidden_size
301 + )
302 +
303 + def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:
304 + return self.embed_tokens(input_ids)
305 +
306 + def forward(
307 + self,
308 + input_ids: torch.Tensor | None,
309 + positions: torch.Tensor,
310 + intermediate_tensors: IntermediateTensors | None,
311 + inputs_embeds: torch.Tensor | None = None,
312 + ) -> torch.Tensor | IntermediateTensors:
313 + if get_pp_group().is_first_rank:
314 + if inputs_embeds is not None:
315 + hidden_states = inputs_embeds
316 + else:
317 + hidden_states = self.embed_input_ids(input_ids)
318 + residual = None
319 + else:
320 + assert intermediate_tensors is not None
321 + hidden_states = intermediate_tensors["hidden_states"]
322 + residual = intermediate_tensors["residual"]
323 + for layer in islice(self.layers, self.start_layer, self.end_layer):
324 + hidden_states, residual = layer(
325 + positions,
326 + hidden_states,
327 + residual,
328 + )
329 + if not get_pp_group().is_last_rank:
330 + return IntermediateTensors(
331 + {"hidden_states": hidden_states, "residual": residual}
332 + )
333 +
334 + hidden_states, _ = self.norm(hidden_states, residual)
335 + return hidden_states
336 +
337 + def get_expert_mapping(self) -> list[tuple[str, str, int, str]]:
338 + # Params for weights, fp8 weight scales, fp8 activation scales
339 + # (param_name, weight_name, expert_id, shard_id)
340 + return FusedMoE.make_expert_params_mapping(
341 + ckpt_gate_proj_name="gate_proj",
342 + ckpt_down_proj_name="down_proj",
343 + ckpt_up_proj_name="up_proj",
344 + num_experts=self.config.num_experts,
345 + num_redundant_experts=self.num_redundant_experts,
346 + )
347 +
348 + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
349 + stacked_params_mapping = [
350 + # (param_name, shard_name, shard_id)
351 + (".qkv_proj", ".q_proj", "q"),
352 + (".qkv_proj", ".k_proj", "k"),
353 + (".qkv_proj", ".v_proj", "v"),
354 + (".gate_up_proj", ".gate_proj", 0),
355 + (".gate_up_proj", ".up_proj", 1),
356 + ]
357 +
358 + # Skip loading extra parameters for GPTQ/modelopt models.
359 + ignore_suffixes = (
360 + ".bias",
361 + "_bias",
362 + ".k_scale",
363 + "_k_scale",
364 + ".v_scale",
365 + "_v_scale",
366 + ".weight_scale",
367 + "_weight_scale",
368 + ".input_scale",
369 + "_input_scale",
370 + )
371 +
372 + params_dict = dict(self.named_parameters())
373 + loaded_params: set[str] = set()
374 + expert_params_mapping = self.get_expert_mapping()
375 + for name, loaded_weight in weights:
376 + if name.startswith("mtp."):
377 + continue
378 + if "rotary_emb.inv_freq" in name:
379 + continue
380 + if "rotary_emb.cos_cached" in name or "rotary_emb.sin_cached" in name:
381 + # Models trained using ColossalAI may include these tensors in
382 + # the checkpoint. Skip them.
383 + continue
384 + if self.quant_config is not None and (
385 + scale_name := self.quant_config.get_cache_scale(name)
386 + ):
387 + # Loading kv cache quantization scales
388 + param = params_dict[scale_name]
389 + weight_loader = getattr(param, "weight_loader", default_weight_loader)
390 + loaded_weight = (
391 + loaded_weight if loaded_weight.dim() == 0 else loaded_weight[0]
392 + )
393 + weight_loader(param, loaded_weight)
394 + loaded_params.add(scale_name)
395 + continue
396 + for param_name, weight_name, shard_id in stacked_params_mapping:
397 + if weight_name not in name:
398 + continue
399 + if "mlp.experts" in name:
400 + continue
401 + name = name.replace(weight_name, param_name)
402 + # Skip loading extra bias for GPTQ models.
403 + if name.endswith(".bias") and name not in params_dict:
404 + continue
405 + if is_pp_missing_parameter(name, self):
406 + continue
407 + if name not in params_dict:
408 + continue
409 +
410 + param = params_dict[name]
411 + weight_loader = param.weight_loader
412 + weight_loader(param, loaded_weight, shard_id)
413 + break
414 + else:
415 + is_expert_weight = False
416 + for mapping in expert_params_mapping:
417 + param_name, weight_name, expert_id, shard_id = mapping
418 + if weight_name not in name:
419 + continue
420 +
421 + # Anyway, this is an expert weight and should not be
422 + # attempted to load as other weights later
423 + is_expert_weight = True
424 +
425 + # Do not modify `name` since the loop may continue here
426 + # Instead, create a new variable
427 + name_mapped = name.replace(weight_name, param_name)
428 +
429 + if is_pp_missing_parameter(name_mapped, self):
430 + continue
431 +
432 + # Skip loading extra parameters for GPTQ/modelopt models.
433 + if (
434 + name_mapped.endswith(ignore_suffixes)
435 + and name_mapped not in params_dict
436 + ):
437 + continue
438 +
439 + param = params_dict[name_mapped]
440 + # We should ask the weight loader to return success or not
441 + # here since otherwise we may skip experts with other
442 + # available replicas.
443 + weight_loader = typing.cast(
444 + Callable[..., bool], param.weight_loader
445 + )
446 + success = weight_loader(
447 + param,
448 + loaded_weight,
449 + name_mapped,
450 + shard_id=shard_id,
451 + expert_id=expert_id,
452 + return_success=True,
453 + )
454 + if success:
455 + name = name_mapped
456 + break
457 + else:
458 + if is_expert_weight:
459 + continue
460 + # Skip loading extra bias for GPTQ models.
461 + if name.endswith(".bias") and name not in params_dict:
462 + continue
463 + # Skip loading extra parameters for GPTQ/modelopt models.
464 + if name.endswith(ignore_suffixes) and name not in params_dict:
465 + continue
466 + # Remapping the name of FP8 kv-scale.
467 + name = maybe_remap_kv_scale_name(name, params_dict)
468 + if name is None:
469 + continue
470 +
471 + if is_pp_missing_parameter(name, self):
472 + continue
473 +
474 + param = params_dict[name]
475 + weight_loader = getattr(
476 + param, "weight_loader", default_weight_loader
477 + )
478 + weight_loader(param, loaded_weight)
479 + loaded_params.add(name)
480 + return loaded_params
481 +
482 +
483 +class ExaoneMoeForCausalLM(nn.Module, SupportsLoRA, SupportsPP):
484 + packed_modules_mapping = {
485 + "qkv_proj": [
486 + "q_proj",
487 + "k_proj",
488 + "v_proj",
489 + ],
490 + "gate_up_proj": [
491 + "gate_proj",
492 + "up_proj",
493 + ],
494 + }
495 +
496 + # LoRA specific attributes
497 + embedding_modules = {
498 + "embed_tokens": "input_embeddings",
499 + "lm_head": "output_embeddings",
500 + }
501 + embedding_padding_modules = ["lm_head"]
502 +
503 + def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
504 + super().__init__()
505 + config = vllm_config.model_config.hf_config.get_text_config()
506 + quant_config = vllm_config.quant_config
507 + lora_config = vllm_config.lora_config
508 +
509 + self.config = config
510 + self.lora_config = lora_config
511 + self.quant_config = quant_config
512 +
513 + self.model = ExaoneMoeModel(
514 + vllm_config=vllm_config,
515 + prefix=maybe_prefix(prefix, "model"),
516 + )
517 + if get_pp_group().is_last_rank:
518 + self.unpadded_vocab_size = config.vocab_size
519 + if lora_config:
520 + self.unpadded_vocab_size += lora_config.lora_extra_vocab_size
521 + self.lm_head = ParallelLMHead(
522 + self.unpadded_vocab_size,
523 + config.hidden_size,
524 + org_num_embeddings=config.vocab_size,
525 + padding_size=DEFAULT_VOCAB_PADDING_SIZE
526 + # We need bigger padding if using lora for kernel
527 + # compatibility
528 + if not lora_config
529 + else lora_config.lora_vocab_padding_size,
530 + quant_config=quant_config,
531 + )
532 + if config.tie_word_embeddings:
533 + self.lm_head.weight = self.model.embed_tokens.weight
534 +
535 + logit_scale = getattr(config, "logit_scale", 1.0)
536 + self.logits_processor = LogitsProcessor(
537 + self.unpadded_vocab_size, config.vocab_size, logit_scale
538 + )
539 + else:
540 + self.lm_head = PPMissingLayer()
541 +
542 + self.make_empty_intermediate_tensors = (
543 + self.model.make_empty_intermediate_tensors
544 + )
545 +
546 + def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:
547 + return self.model.embed_input_ids(input_ids)
548 +
549 + def forward(
550 + self,
551 + input_ids: torch.Tensor,
552 + positions: torch.Tensor,
553 + intermediate_tensors: IntermediateTensors | None = None,
554 + inputs_embeds: torch.Tensor | None = None,
555 + ) -> torch.Tensor | IntermediateTensors:
556 + model_output = self.model(
557 + input_ids, positions, intermediate_tensors, inputs_embeds
558 + )
559 + return model_output
560 +
561 + def compute_logits(
562 + self,
563 + hidden_states: torch.Tensor,
564 + ) -> torch.Tensor | None:
565 + logits = self.logits_processor(self.lm_head, hidden_states)
566 + return logits
567 +
568 + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
569 + loader = AutoWeightsLoader(
570 + self,
571 + # With tie_word_embeddings, we can skip lm_head.weight
572 + # The weight might appear unnecessarily in the files if the model is
573 + # processed with quantization, LoRA, fine-tuning, etc.
574 + skip_prefixes=(
575 + ["lm_head.", "mtp."] if self.config.tie_word_embeddings else ["mtp."]
576 + ),
577 + )
578 + return loader.load_weights(weights)
vllm/model_executor/models/exaone_moe_mtp.py new
+255
@@ -0,0 +1,255 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""Inference-only ExaoneMoe MTP model."""
4 +
5 +from collections.abc import Iterable
6 +
7 +import torch
8 +from torch import nn
9 +
10 +from vllm.compilation.decorators import support_torch_compile
11 +from vllm.config import VllmConfig
12 +from vllm.distributed.parallel_state import get_pp_group
13 +from vllm.logger import init_logger
14 +from vllm.model_executor.layers.layernorm import RMSNorm
15 +from vllm.model_executor.layers.linear import ColumnParallelLinear
16 +from vllm.model_executor.layers.logits_processor import LogitsProcessor
17 +from vllm.model_executor.layers.vocab_parallel_embedding import (
18 + ParallelLMHead,
19 + VocabParallelEmbedding,
20 +)
21 +from vllm.model_executor.model_loader.weight_utils import default_weight_loader
22 +from vllm.model_executor.models.exaone_moe import ExaoneMoeDecoderLayer
23 +from vllm.sequence import IntermediateTensors
24 +
25 +from .utils import (
26 + AutoWeightsLoader,
27 + is_pp_missing_parameter,
28 + maybe_prefix,
29 +)
30 +
31 +logger = init_logger(__name__)
32 +
33 +KVCache = tuple[torch.Tensor, torch.Tensor]
34 +
35 +
36 +@support_torch_compile
37 +class ExaoneMoeMultiTokenPredictor(nn.Module):
38 + def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
39 + super().__init__()
40 +
41 + model_config = vllm_config.model_config
42 + quant_config = vllm_config.quant_config
43 + lora_config = vllm_config.lora_config
44 + config = model_config.hf_config
45 +
46 + self.config = config
47 + lora_vocab = (
48 + (lora_config.lora_extra_vocab_size * (lora_config.max_loras or 1))
49 + if lora_config
50 + else 0
51 + )
52 + self.vocab_size = config.vocab_size + lora_vocab
53 + self.org_vocab_size = config.vocab_size
54 +
55 + self.mtp_start_layer_idx = config.num_hidden_layers
56 + self.num_mtp_layers = getattr(config, "num_nextn_predict_layers", 1)
57 +
58 + self.embed_tokens = VocabParallelEmbedding(
59 + self.vocab_size,
60 + config.hidden_size,
61 + org_num_embeddings=config.vocab_size,
62 + )
63 +
64 + self.fc = ColumnParallelLinear(
65 + self.config.hidden_size * 2,
66 + self.config.hidden_size,
67 + gather_output=True,
68 + bias=False,
69 + return_bias=False,
70 + quant_config=quant_config,
71 + prefix=f"{prefix}.fc",
72 + )
73 + self.layers = nn.ModuleList(
74 + ExaoneMoeDecoderLayer(
75 + vllm_config.model_config.hf_config,
76 + quant_config=quant_config,
77 + prefix=f"{prefix}.layers.{idx}",
78 + mtp_layer=True,
79 + )
80 + for idx in range(self.num_mtp_layers)
81 + )
82 +
83 + self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
84 + self.pre_fc_norm_hidden = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
85 + self.pre_fc_norm_embedding = RMSNorm(
86 + config.hidden_size, eps=config.rms_norm_eps
87 + )
88 +
89 + def get_input_embeddings(self, input_ids: torch.Tensor) -> torch.Tensor:
90 + return self.embed_tokens(input_ids)
91 +
92 + def forward(
93 + self,
94 + input_ids: torch.Tensor,
95 + positions: torch.Tensor,
96 + hidden_states: torch.Tensor,
97 + intermediate_tensors: IntermediateTensors | None = None,
98 + inputs_embeds: torch.Tensor | None = None,
99 + spec_step_idx: int = 0,
100 + ) -> torch.Tensor:
101 + if get_pp_group().is_first_rank:
102 + if inputs_embeds is None:
103 + inputs_embeds = self.get_input_embeddings(input_ids)
104 + assert hidden_states.shape[-1] == inputs_embeds.shape[-1]
105 + inputs_embeds = self.pre_fc_norm_embedding(inputs_embeds)
106 + hidden_states = self.pre_fc_norm_hidden(hidden_states)
107 + hidden_states = torch.cat([inputs_embeds, hidden_states], dim=-1)
108 + hidden_states = self.fc(hidden_states)
109 + residual = None
110 + else:
111 + assert intermediate_tensors is not None
112 + hidden_states = intermediate_tensors["hidden_states"]
113 + residual = intermediate_tensors["residual"]
114 +
115 + current_step_idx = spec_step_idx % self.num_mtp_layers
116 + hidden_states, residual = self.layers[current_step_idx](
117 + positions=positions,
118 + hidden_states=hidden_states,
119 + residual=residual,
120 + )
121 +
122 + if not get_pp_group().is_last_rank:
123 + return IntermediateTensors(
124 + {"hidden_states": hidden_states, "residual": residual}
125 + )
126 +
127 + hidden_states, _ = self.norm(hidden_states, residual)
128 + return hidden_states
129 +
130 + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
131 + stacked_params_mapping = [
132 + # (param_name, shard_name, shard_id)
133 + ("qkv_proj", "q_proj", "q"),
134 + ("qkv_proj", "k_proj", "k"),
135 + ("qkv_proj", "v_proj", "v"),
136 + ("gate_up_proj", "gate_proj", 0),
137 + ("gate_up_proj", "up_proj", 1),
138 + ]
139 +
140 + # Params for weights, fp8 weight scales, fp8 activation scales
141 + # (param_name, weight_name, expert_id, shard_id)
142 + params_dict = dict(self.named_parameters())
143 + loaded_params: set[str] = set()
144 + for name, loaded_weight in weights:
145 + if "rotary_emb.inv_freq" in name:
146 + continue
147 +
148 + for param_name, weight_name, shard_id in stacked_params_mapping:
149 + if weight_name not in name:
150 + continue
151 +
152 + if "mlp.experts" in name:
153 + continue
154 +
155 + name = name.replace(weight_name, param_name)
156 + # Skip loading extra bias for GPTQ models.
157 + if name.endswith(".bias") and name not in params_dict:
158 + continue
159 + # Skip layers on other devices.
160 + if is_pp_missing_parameter(name, self):
161 + continue
162 + if name not in params_dict:
163 + continue
164 + param = params_dict[name]
165 + weight_loader = param.weight_loader
166 + weight_loader(param, loaded_weight, shard_id)
167 + break
168 + else:
169 + # Skip loading extra bias for GPTQ models.
170 + if name.endswith(".bias") and name not in params_dict:
171 + continue
172 + if is_pp_missing_parameter(name, self):
173 + continue
174 +
175 + param = params_dict[name]
176 + weight_loader = getattr(param, "weight_loader", default_weight_loader)
177 + weight_loader(param, loaded_weight)
178 + loaded_params.add(name)
179 + return loaded_params
180 +
181 +
182 +@support_torch_compile
183 +class ExaoneMoeMTP(nn.Module):
184 + def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
185 + config = vllm_config.model_config.hf_config
186 + self.vllm_config = vllm_config
187 + cache_config = vllm_config.cache_config
188 + assert not cache_config.enable_prefix_caching, (
189 + "ExaoneMoeMTP currently does not support prefix caching"
190 + )
191 +
192 + self.quant_config = vllm_config.quant_config
193 +
194 + super().__init__()
195 + self.config = config
196 + self.model = ExaoneMoeMultiTokenPredictor(
197 + vllm_config=vllm_config, prefix=maybe_prefix(prefix, "mtp")
198 + )
199 + self.unpadded_vocab_size = config.vocab_size
200 + self.lm_head = ParallelLMHead(
201 + self.unpadded_vocab_size,
202 + config.hidden_size,
203 + org_num_embeddings=config.vocab_size,
204 + # padding_size=DEFAULT_VOCAB_PADDING_SIZE,
205 + prefix=maybe_prefix(prefix, "lm_head"),
206 + )
207 + if config.tie_word_embeddings:
208 + self.lm_head.weight = self.model.embed_tokens.weight
209 + self.logits_processor = LogitsProcessor(
210 + self.unpadded_vocab_size, config.vocab_size
211 + )
212 +
213 + def get_input_embeddings(self, input_ids: torch.Tensor) -> torch.Tensor:
214 + return self.model.get_input_embeddings(input_ids)
215 +
216 + def forward(
217 + self,
218 + input_ids: torch.Tensor,
219 + positions: torch.Tensor,
220 + hidden_states: torch.Tensor,
221 + intermediate_tensors: IntermediateTensors | None = None,
222 + inputs_embeds: torch.Tensor | None = None,
223 + spec_step_idx: int = 0,
224 + **kwargs: object,
225 + ):
226 + hidden_states = self.model(
227 + input_ids,
228 + positions,
229 + hidden_states,
230 + intermediate_tensors,
231 + inputs_embeds,
232 + spec_step_idx,
233 + )
234 + return hidden_states
235 +
236 + def compute_logits(
237 + self,
238 + hidden_states: torch.Tensor,
239 + spec_step_idx: int = 0,
240 + ) -> torch.Tensor | None:
241 + return self.logits_processor(self.lm_head, hidden_states)
242 +
243 + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
244 + shared_weight_names = ["embed_tokens", "lm_head"]
245 +
246 + def remap_weight_names(weights):
247 + for name, weight in weights:
248 + if name.startswith("mtp."):
249 + name = name.replace("mtp.", "model.")
250 + elif not any(key in name for key in shared_weight_names):
251 + continue
252 + yield name, weight
253 +
254 + loader = AutoWeightsLoader(self)
255 + return loader.load_weights(remap_weight_names(weights))
vllm/model_executor/models/registry.py
+2
@@ -98,6 +98,7 @@ _TEXT_GENERATION_MODELS = {
98 "Ernie4_5_MoeForCausalLM": ("ernie45_moe", "Ernie4_5_MoeForCausalLM"),
99 "ExaoneForCausalLM": ("exaone", "ExaoneForCausalLM"),
100 "Exaone4ForCausalLM": ("exaone4", "Exaone4ForCausalLM"),
101 + "ExaoneMoEForCausalLM": ("exaone_moe", "ExaoneMoeForCausalLM"),
102 "Fairseq2LlamaForCausalLM": ("fairseq2_llama", "Fairseq2LlamaForCausalLM"),
103 "FalconForCausalLM": ("falcon", "FalconForCausalLM"),
104 "FalconMambaForCausalLM": ("mamba", "MambaForCausalLM"),
@@ -457,6 +458,7 @@ _SPECULATIVE_DECODING_MODELS = {
458 "EagleDeepSeekMTPModel": ("deepseek_eagle", "EagleDeepseekV3ForCausalLM"),
459 "DeepSeekMTPModel": ("deepseek_mtp", "DeepSeekMTP"),
460 "ErnieMTPModel": ("ernie_mtp", "ErnieMTP"),
461 + "ExaoneMoeMTP": ("exaone_moe_mtp", "ExaoneMoeMTP"),
462 "LongCatFlashMTPModel": ("longcat_flash_mtp", "LongCatFlashMTP"),
463 "Glm4MoeMTPModel": ("glm4_moe_mtp", "Glm4MoeMTP"),
464 "MedusaModel": ("medusa", "Medusa"),