@tduncan / Nvidia-Optimization-Stack-VLLM / commits / 9ae11a6b89

[Model] Add K-EXAONE-2.0-750B-A37B (#50524)

Signed-off-by: lkm2835 <lkm2835@gmail.com> Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>

Kyungmin Lee committed Aug 3, 2026 at 22:33 UTC 9ae11a6b895dea5b2275e0b6ad96327af67d3b0c

6 files changed +161 -11

docs/models/supported_models.md

+1 -1

@@ -364,7 +364,7 @@ th {

364 | `Ernie4_5ForCausalLM` | Ernie4.5 | `baidu/ERNIE-4.5-0.3B-PT`, etc. | ✅︎ | ✅︎ |

365 | `Ernie4_5_MoeForCausalLM` | Ernie4.5MoE | `baidu/ERNIE-4.5-21B-A3B-PT`, `baidu/ERNIE-4.5-300B-A47B-PT`, etc. | ✅︎ | ✅︎ |

366 | `ExaoneForCausalLM` | EXAONE-3 | `LGAI-EXAONE/EXAONE-3.0-7.8B-Instruct`, etc. | ✅︎ | ✅︎ |

368 | `Exaone4ForCausalLM` | EXAONE-4 | `LGAI-EXAONE/EXAONE-4.0-32B`, etc. | ✅︎ | ✅︎ |

369 | `Fairseq2LlamaForCausalLM` | Llama (fairseq2 format) | `mgleize/fairseq2-dummy-Llama-3.2-1B`, etc. | ✅︎ | ✅︎ |

tests/models/registry.py

+1 -1

@@ -267,7 +267,7 @@ _TEXT_GENERATION_EXAMPLE_MODELS = {

267 "LGAI-EXAONE/EXAONE-3.0-7.8B-Instruct", trust_remote_code=True

268 ),

269 "Exaone4ForCausalLM": _HfExamplesInfo("LGAI-EXAONE/EXAONE-4.0-32B"),

270 - "ExaoneMoEForCausalLM": _HfExamplesInfo(

270 + "ExaoneMoeForCausalLM": _HfExamplesInfo(

271 "LGAI-EXAONE/K-EXAONE-236B-A23B", min_transformers_version="5.1.0"

272 ),

273 "Fairseq2LlamaForCausalLM": _HfExamplesInfo("mgleize/fairseq2-dummy-Llama-3.2-1B"),

vllm/model_executor/models/exaone4.py

+5 -2

@@ -31,7 +31,7 @@ from transformers import Exaone4Config

31 from vllm.compilation.decorators import support_torch_compile

32 from vllm.config import CacheConfig, VllmConfig

33 from vllm.distributed import get_pp_group, get_tensor_model_parallel_world_size

34 -from vllm.model_executor.layers.activation import SiluAndMul

34 +from vllm.model_executor.layers.activation import SiluAndMul, SiluAndMulWithClamp

35 from vllm.model_executor.layers.attention import Attention

36 from vllm.model_executor.layers.layernorm import RMSNorm

37 from vllm.model_executor.layers.linear import (

@@ -70,6 +70,7 @@ class Exaone4GatedMLP(nn.Module):

70 quant_config: QuantizationConfig | None = None,

71 reduce_results: bool = True,

72 bias: bool = False,

73 + swiglu_limit: float | None = None,

74 prefix: str = "",

75 use_data_parallel: bool = False,

76 ) -> None:

@@ -95,7 +96,9 @@ class Exaone4GatedMLP(nn.Module):

96 raise ValueError(

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

98 )

98 - self.act_fn = SiluAndMul()

99 + self.act_fn = (

100 + SiluAndMul() if swiglu_limit is None else SiluAndMulWithClamp(swiglu_limit)

101 + )

102

103 def forward(self, x):

104 gate_up, _ = self.gate_up_proj(x)

vllm/model_executor/models/exaone_moe.py

+152 -5

@@ -29,19 +29,25 @@ from vllm.distributed import (

29 get_pp_group,

30 get_tensor_model_parallel_world_size,

31 )

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

33 from vllm.model_executor.layers.fused_moe import FusedMoEFactory

34 from vllm.model_executor.layers.layernorm import RMSNorm

34 -from vllm.model_executor.layers.linear import ReplicatedLinear

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

36 + QKVParallelLinear,

37 + ReplicatedLinear,

38 + RowParallelLinear,

39 +)

40 from vllm.model_executor.layers.logits_processor import LogitsProcessor

41 from vllm.model_executor.layers.quantization import QuantizationConfig

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

43 from vllm.model_executor.layers.vocab_parallel_embedding import (

44 DEFAULT_VOCAB_PADDING_SIZE,

45 ParallelLMHead,

46 VocabParallelEmbedding,

47 )

48 from vllm.sequence import IntermediateTensors

49 +from vllm.transformers_utils.config import set_default_rope_theta

44 -from .exaone4 import Exaone4Attention as ExaoneMoeAttention

51 from .exaone4 import Exaone4GatedMLP as ExaoneMoeGatedMLP

52 from .interfaces import SupportsLoRA, SupportsPP

53 from .utils import (

@@ -59,6 +65,7 @@ class ExaoneMoe(nn.Module):

65 def __init__(

66 self,

67 config: PretrainedConfig,

68 + swiglu_limit: float | None = None,

69 quant_config: QuantizationConfig | None = None,

70 prefix: str = "",

71 enable_eplb: bool = False,

@@ -119,6 +126,7 @@ class ExaoneMoe(nn.Module):

126 hidden_act=config.hidden_act,

127 quant_config=quant_config,

128 reduce_results=False,

129 + swiglu_limit=swiglu_limit,

130 prefix=f"{prefix}.shared_experts",

131 )

132 else:

@@ -133,6 +141,8 @@ class ExaoneMoe(nn.Module):

141 intermediate_size=config.moe_intermediate_size,

142 renormalize=config.norm_topk_prob,

143 quant_config=quant_config,

144 + activation="silu",

145 + swiglu_limit=swiglu_limit,

146 use_grouped_topk=True,

147 num_expert_group=config.n_group,

148 topk_group=config.topk_group,

@@ -157,13 +167,138 @@ class ExaoneMoe(nn.Module):

167 return final_hidden_states.view(orig_shape)

168

169

170 +class ExaoneMoeAttention(nn.Module):

171 + def __init__(

172 + self,

173 + config,

174 + hidden_size: int,

175 + num_heads: int,

176 + num_kv_heads: int,

177 + max_position_embeddings: int = 8192,

178 + quant_config: QuantizationConfig | None = None,

179 + bias: bool = False,

180 + cache_config: CacheConfig | None = None,

181 + prefix: str = "",

182 + is_mtp: bool = False,

183 + ) -> None:

184 + super().__init__()

185 + self.hidden_size = hidden_size

186 + tp_size = get_tensor_model_parallel_world_size()

187 + self.total_num_heads = num_heads

188 + assert self.total_num_heads % tp_size == 0

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

190 + self.total_num_kv_heads = num_kv_heads

191 + if self.total_num_kv_heads >= tp_size:

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

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

194 + assert self.total_num_kv_heads % tp_size == 0

195 + else:

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

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

198 + assert tp_size % self.total_num_kv_heads == 0

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

200 + # MistralConfig has an optional head_dim introduced by Mistral-Nemo

201 + self.head_dim = getattr(config, "head_dim", None)

202 + if self.head_dim is None:

203 + self.head_dim = self.hidden_size // self.total_num_heads

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

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

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

207 + self.max_position_embeddings = max_position_embeddings

208 +

209 + self.qkv_proj = QKVParallelLinear(

210 + hidden_size=hidden_size,

211 + head_size=self.head_dim,

212 + total_num_heads=self.total_num_heads,

213 + total_num_kv_heads=self.total_num_kv_heads,

214 + bias=bias,

215 + quant_config=quant_config,

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

217 + )

218 +

219 + self.o_proj = RowParallelLinear(

220 + input_size=self.total_num_heads * self.head_dim,

221 + output_size=hidden_size,

222 + bias=bias,

223 + quant_config=quant_config,

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

225 + )

226 +

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

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

229 +

230 + is_neox_style = True

231 + if quant_config is not None and quant_config.get_name() == "gguf":

232 + is_neox_style = False

233 +

234 + layer_idx = extract_layer_index(prefix)

235 +

236 + if config.sliding_windows is not None:

237 + self.sliding_window_size = config.sliding_windows[layer_idx]

238 +

239 + if config.layer_types[layer_idx] == "full_attention":

240 + self.sliding_window_size = None

241 +

242 + if is_mtp:

243 + self.sliding_window = (

244 + config.mtp_layer_types[layer_idx] == "sliding_attention"

245 + )

246 + if config.mtp_sliding_windows is not None:

247 + self.sliding_window_size = config.mtp_sliding_windows[layer_idx]

248 +

249 + if config.mtp_layer_types[layer_idx] == "full_attention":

250 + self.sliding_window_size = None

251 +

252 + # apply rotary embeddings to every layer in full attention models

253 + self.apply_rope_all_layers = "sliding_attention" not in config.layer_types

254 +

255 + set_default_rope_theta(config, default_theta=1000000)

256 + self.rotary_emb = get_rope(

257 + self.head_dim,

258 + max_position=max_position_embeddings,

259 + rope_parameters=config.rope_parameters,

260 + is_neox_style=is_neox_style,

261 + )

262 + self.attn = Attention(

263 + self.num_heads,

264 + self.head_dim,

265 + self.scaling,

266 + num_kv_heads=self.num_kv_heads,

267 + cache_config=cache_config,

268 + quant_config=quant_config,

269 + per_layer_sliding_window=self.sliding_window_size,

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

271 + )

272 +

273 + def forward(

274 + self,

275 + positions: torch.Tensor,

276 + hidden_states: torch.Tensor,

277 + ) -> torch.Tensor:

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

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

280 +

281 + q = q.unflatten(-1, (self.num_heads, self.head_dim))

282 + q = self.q_norm(q)

283 + q = q.flatten(-2, -1)

284 + k = k.unflatten(-1, (self.num_kv_heads, self.head_dim))

285 + k = self.k_norm(k)

286 + k = k.flatten(-2, -1)

287 +

288 + if self.sliding_window_size or self.apply_rope_all_layers:

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

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

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

292 + return output

293 +

294 +

295 class ExaoneMoeDecoderLayer(nn.Module):

296 def __init__(

297 self,

298 config: PretrainedConfig,

299 cache_config: CacheConfig | None = None,

300 quant_config: QuantizationConfig | None = None,

166 - mtp_layer: bool = None,

301 + is_mtp: bool = None,

302 prefix: str = "",

303 ) -> None:

304 super().__init__()

@@ -187,17 +322,29 @@ class ExaoneMoeDecoderLayer(nn.Module):

322 bias=attention_bias,

323 cache_config=cache_config,

324 prefix=f"{prefix}.self_attn",

325 + is_mtp=is_mtp,

326 )

327

192 - if config.is_moe_layer[layer_idx] and not mtp_layer:

328 + swiglu_limits = getattr(config, "swiglu_limits", None)

329 + if swiglu_limits is not None:

330 + swiglu_limit = swiglu_limits[layer_idx]

331 + swiglu_limit = swiglu_limit if swiglu_limit > 0 else None

332 + else:

333 + swiglu_limit = None

334 +

335 + if config.mlp_layer_types[layer_idx] == "sparse" and not is_mtp:

336 self.mlp = ExaoneMoe(

194 - config=config, quant_config=quant_config, prefix=f"{prefix}.mlp"

337 + config=config,

338 + swiglu_limit=swiglu_limit,

339 + quant_config=quant_config,

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

341 )

342 else:

343 self.mlp = ExaoneMoeGatedMLP(

344 hidden_size=self.hidden_size,

345 intermediate_size=config.intermediate_size,

346 hidden_act=config.hidden_act,

347 + swiglu_limit=swiglu_limit,

348 quant_config=quant_config,

349 bias=getattr(config, "mlp_bias", False),

350 prefix=f"{prefix}.mlp",

vllm/model_executor/models/exaone_moe_mtp.py

+1 -1

@@ -81,7 +81,7 @@ class ExaoneMoeMultiTokenPredictor(nn.Module):

81 vllm_config.model_config.hf_config,

82 quant_config=quant_config,

83 prefix=f"{prefix}.layers.{idx}",

84 - mtp_layer=True,

84 + is_mtp=True,

85 )

86 for idx in range(self.num_mtp_layers)

87 )

vllm/model_executor/models/registry.py

+1 -1

@@ -97,7 +97,7 @@ _TEXT_GENERATION_MODELS = {

97 "Ernie4_5_MoeForCausalLM": ("ernie45_moe", "Ernie4_5_MoeForCausalLM"),

98 "ExaoneForCausalLM": ("exaone", "ExaoneForCausalLM"),

99 "Exaone4ForCausalLM": ("exaone4", "Exaone4ForCausalLM"),

100 - "ExaoneMoEForCausalLM": ("exaone_moe", "ExaoneMoeForCausalLM"),

100 + "ExaoneMoeForCausalLM": ("exaone_moe", "ExaoneMoeForCausalLM"),

101 "Fairseq2LlamaForCausalLM": ("fairseq2_llama", "Fairseq2LlamaForCausalLM"),

102 "FalconForCausalLM": ("falcon", "FalconForCausalLM"),

103 "FalconMambaForCausalLM": ("mamba", "MambaForCausalLM"),