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

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 13, 20:45 UTC

[model] Add support for openPangu_Ultra_MoE (#27521)

Signed-off-by: yuantao <2422264527@qq.com> Signed-off-by: yt0428 <51468697+yt0428@users.noreply.github.com> Co-authored-by: Jee Jee Li <pandaleefree@gmail.com>

yt0428 committed Nov 5, 2025 at 00:17 UTC 05cae69f0f4f0237779c45a3907aa13ee5a07667

8 files changed +1379 -1

docs/models/supported_models.md

@@ -404,6 +404,8 @@ th {

404 | `OPTForCausalLM` | OPT, OPT-IML | `facebook/opt-66b`, `facebook/opt-iml-max-30b`, etc. | ✅︎ | ✅︎ |

407 +| `PanguEmbeddedForCausalLM` |openPangu-Embedded-7B | `FreedomIntelligence/openPangu-Embedded-7B-V1.1` | ✅︎ | ✅︎ |

408 +| `PanguUltraMoEForCausalLM` |openpangu-ultra-moe-718b-model | `FreedomIntelligence/openPangu-Ultra-MoE-718B-V1.1` | ✅︎ | ✅︎ |

409 | `PhiForCausalLM` | Phi | `microsoft/phi-1_5`, `microsoft/phi-2`, etc. | ✅︎ | ✅︎ |

410 | `Phi3ForCausalLM` | Phi-4, Phi-3 | `microsoft/Phi-4-mini-instruct`, `microsoft/Phi-4`, `microsoft/Phi-3-mini-4k-instruct`, `microsoft/Phi-3-mini-128k-instruct`, `microsoft/Phi-3-medium-128k-instruct`, etc. | ✅︎ | ✅︎ |

411 | `PhiMoEForCausalLM` | Phi-3.5-MoE | `microsoft/Phi-3.5-MoE-instruct`, etc. | ✅︎ | ✅︎ |

tests/models/registry.py

+13

@@ -363,6 +363,11 @@ _TEXT_GENERATION_EXAMPLE_MODELS = {

363 "Olmo2ForCausalLM": _HfExamplesInfo("allenai/OLMo-2-0425-1B"),

364 "Olmo3ForCausalLM": _HfExamplesInfo("shanearora/2025-sep-a-base-model"),

365 "OlmoeForCausalLM": _HfExamplesInfo("allenai/OLMoE-1B-7B-0924-Instruct"),

366 + "OpenPanguMTPModel": _HfExamplesInfo(

367 + "FreedomIntelligence/openPangu-Ultra-MoE-718B-V1.1",

368 + trust_remote_code=True,

369 + is_available_online=False,

370 + ),

371 "OPTForCausalLM": _HfExamplesInfo(

372 "facebook/opt-125m", {"1b": "facebook/opt-iml-max-1.3b"}

373 ),

@@ -370,6 +375,14 @@ _TEXT_GENERATION_EXAMPLE_MODELS = {

375 "OrionStarAI/Orion-14B-Chat", trust_remote_code=True

376 ),

377 "OuroForCausalLM": _HfExamplesInfo("ByteDance/Ouro-1.4B", trust_remote_code=True),

378 + "PanguEmbeddedForCausalLM": _HfExamplesInfo(

379 + "FreedomIntelligence/openPangu-Embedded-7B-V1.1", trust_remote_code=True

380 + ),

381 + "PanguUltraMoEForCausalLM": _HfExamplesInfo(

382 + "FreedomIntelligence/openPangu-Ultra-MoE-718B-V1.1",

383 + trust_remote_code=True,

384 + is_available_online=False,

385 + ),

386 "PersimmonForCausalLM": _HfExamplesInfo("adept/persimmon-8b-chat"),

387 "PhiForCausalLM": _HfExamplesInfo("microsoft/phi-2"),

388 "Phi3ForCausalLM": _HfExamplesInfo("microsoft/Phi-3-mini-4k-instruct"),

vllm/config/model.py

@@ -1231,6 +1231,8 @@ class ModelConfig:

1231 "kimi_k2",

1232 "kimi_linear",

1233 "longcat_flash",

1234 + "pangu_ultra_moe",

1235 + "pangu_ultra_moe_mtp",

1236 ):

1237 return self.hf_text_config.kv_lora_rank is not None

1238 elif self.hf_text_config.model_type == "eagle":

@@ -1379,6 +1381,7 @@ class ModelConfig:

1381 or self.hf_config.model_type == "glm4_moe_mtp"

1382 or self.hf_config.model_type == "ernie_mtp"

1383 or self.hf_config.model_type == "qwen3_next_mtp"

1384 + or self.hf_config.model_type == "pangu_ultra_moe_mtp"

1385 ):

1386 total_num_hidden_layers = getattr(

1387 self.hf_text_config, "num_nextn_predict_layers", 0

vllm/config/speculative.py

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

41 "qwen3_next_mtp",

42 "mimo_mtp",

43 "longcat_flash_mtp",

44 + "pangu_ultra_moe_mtp",

45 "mtp",

46 "suffix",

47 ]

@@ -51,6 +52,7 @@ MTP_MODEL_TYPES = (

52 "ernie_mtp",

53 "qwen3_next_mtp",

54 "longcat_flash_mtp",

55 + "pangu_ultra_moe_mtp",

56 )

@@ -179,6 +181,13 @@ class SpeculativeConfig:

181 hf_config.update(

182 {"n_predict": n_predict, "architectures": ["DeepSeekMTPModel"]}

183 )

184 + if hf_config.model_type in ("pangu_ultra_moe"):

185 + hf_config.model_type = "pangu_ultra_moe_mtp"

186 + if hf_config.model_type == "pangu_ultra_moe_mtp":

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

188 + hf_config.update(

189 + {"n_predict": n_predict, "architectures": ["OpenPanguMTPModel"]}

190 + )

191

192 if hf_config.architectures[0] == "MiMoForCausalLM":

193 hf_config.model_type = "mimo_mtp"

vllm/model_executor/models/openpangu.py new

+1078

@@ -0,0 +1,1078 @@

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

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

6 +#

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

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

9 +# original forms to accommodate minor architectural differences compared

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

11 +#

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

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

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

15 +#

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

17 +#

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

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

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

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

22 +# limitations under the License.

23 +

24 +import typing

25 +from collections.abc import Callable, Iterable

26 +from typing import Any

27 +

28 +import torch

29 +from torch import nn

30 +from transformers import PretrainedConfig

31 +

32 +from vllm.attention import Attention, AttentionType

33 +from vllm.compilation.decorators import support_torch_compile

34 +from vllm.config import CacheConfig, ParallelConfig, VllmConfig

35 +from vllm.distributed import (

36 + get_ep_group,

37 + get_pp_group,

38 + get_tensor_model_parallel_world_size,

39 + get_tp_group,

40 + tensor_model_parallel_all_gather,

41 +)

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

43 +from vllm.model_executor.layers.fused_moe import SharedFusedMoE

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

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

46 + ColumnParallelLinear,

47 + MergedColumnParallelLinear,

48 + QKVParallelLinear,

49 + ReplicatedLinear,

50 + RowParallelLinear,

51 +)

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

53 +from vllm.model_executor.layers.mla import MLAModules, MultiHeadLatentAttentionWrapper

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

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

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

57 + ParallelLMHead,

58 + VocabParallelEmbedding,

59 +)

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

61 + default_weight_loader,

62 + maybe_remap_kv_scale_name,

63 +)

64 +from vllm.model_executor.models.interfaces import (

65 + MixtureOfExperts,

66 + SupportsLoRA,

67 + SupportsPP,

68 +)

69 +from vllm.model_executor.models.utils import (

70 + AutoWeightsLoader,

71 + PPMissingLayer,

72 + extract_layer_index,

73 + is_pp_missing_parameter,

74 + make_empty_intermediate_tensors_factory,

75 + make_layers,

76 + maybe_prefix,

77 + sequence_parallel_chunk,

78 +)

79 +from vllm.sequence import IntermediateTensors

80 +

81 +

82 +def check_ffn_act_fn(act_fn: str):

83 + if act_fn != "silu":

84 + raise ValueError(

85 + f"Unsupported activation: {act_fn}. Only silu is supported for now."

86 + )

87 +

88 +

89 +class OpenPanguMLP(nn.Module):

90 + def __init__(

91 + self,

92 + hidden_size: int,

93 + intermediate_size: int,

94 + hidden_act: str,

95 + quant_config: QuantizationConfig | None = None,

96 + bias: bool = False,

97 + reduce_results: bool = True,

98 + is_sequence_parallel=False,

99 + prefix: str = "",

100 + ) -> None:

101 + super().__init__()

102 + self.gate_up_proj = MergedColumnParallelLinear(

103 + hidden_size,

104 + [intermediate_size] * 2,

105 + bias=bias,

106 + quant_config=quant_config,

107 + disable_tp=is_sequence_parallel,

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

109 + )

110 + self.down_proj = RowParallelLinear(

111 + intermediate_size,

112 + hidden_size,

113 + bias=bias,

114 + quant_config=quant_config,

115 + reduce_results=reduce_results,

116 + disable_tp=is_sequence_parallel,

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

118 + )

119 +

120 + check_ffn_act_fn(hidden_act)

121 + self.act_fn = SiluAndMul()

122 +

123 + def forward(self, x: torch.Tensor) -> torch.Tensor:

124 + return self.down_proj(self.act_fn(self.gate_up_proj(x)[0]))[0]

125 +

126 +

127 +class OpenPanguMoE(nn.Module):

128 + def __init__(

129 + self,

130 + config: PretrainedConfig,

131 + parallel_config: ParallelConfig,

132 + quant_config: QuantizationConfig | None = None,

133 + prefix: str = "",

134 + ):

135 + super().__init__()

136 + self.tp_size = get_tensor_model_parallel_world_size()

137 + self.tp_rank = get_tp_group().rank_in_group

138 +

139 + self.routed_scaling_factor = config.routed_scaling_factor

140 + self.ep_group = get_ep_group().device_group

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

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

143 + self.n_routed_experts: int = config.n_routed_experts

144 + self.n_shared_experts: int = config.n_shared_experts

145 +

146 + self.is_sequence_parallel = parallel_config.use_sequence_parallel_moe

147 + check_ffn_act_fn(config.hidden_act)

148 +

149 + self.gate = ReplicatedLinear(

150 + config.hidden_size,

151 + config.n_routed_experts,

152 + bias=False,

153 + quant_config=None,

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

155 + )

156 + self.gate.e_score_correction_bias = None

157 +

158 + # Load balancing settings.

159 + eplb_config = parallel_config.eplb_config

160 + self.enable_eplb = parallel_config.enable_eplb

161 +

162 + self.n_redundant_experts = eplb_config.num_redundant_experts

163 + self.n_logical_experts = self.n_routed_experts

164 + self.n_physical_experts = self.n_logical_experts + self.n_redundant_experts

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

166 +

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

168 + self.physical_expert_end = (

169 + self.physical_expert_start + self.n_local_physical_experts

170 + )

171 +

172 + if config.n_shared_experts is not None:

173 + intermediate_size = config.moe_intermediate_size * config.n_shared_experts

174 + self.shared_experts = OpenPanguMLP(

175 + hidden_size=config.hidden_size,

176 + intermediate_size=intermediate_size,

177 + hidden_act=config.hidden_act,

178 + quant_config=quant_config,

179 + is_sequence_parallel=self.is_sequence_parallel,

180 + reduce_results=False,

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

182 + )

183 + else:

184 + self.shared_experts = None

185 +

186 + self.experts = SharedFusedMoE(

187 + shared_experts=self.shared_experts,

188 + num_experts=config.n_routed_experts,

189 + top_k=config.num_experts_per_tok,

190 + hidden_size=config.hidden_size,

191 + intermediate_size=config.moe_intermediate_size,

192 + reduce_results=False,

193 + renormalize=config.norm_topk_prob,

194 + quant_config=quant_config,

195 + use_grouped_topk=True,

196 + num_expert_group=1,

197 + topk_group=1,

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

199 + scoring_func="sigmoid",

200 + # we do scaling outside, set factor to 1.0 to avoid double mul

201 + routed_scaling_factor=1.0,

202 + e_score_correction_bias=self.gate.e_score_correction_bias,

203 + enable_eplb=self.enable_eplb,

204 + num_redundant_experts=self.n_redundant_experts,

205 + is_sequence_parallel=self.is_sequence_parallel,

206 + )

207 +

208 + def forward(

209 + self,

210 + hidden_states: torch.Tensor,

211 + ) -> torch.Tensor:

212 + num_tokens, hidden_dim = hidden_states.shape

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

214 +

215 + if self.is_sequence_parallel:

216 + hidden_states = sequence_parallel_chunk(hidden_states)

217 +

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

219 +

220 + fused_moe_out = self.experts(

221 + hidden_states=hidden_states, router_logits=router_logits

222 + )

223 +

224 + shared_output, final_hidden_states = fused_moe_out

225 + if self.shared_experts is None:

226 + assert shared_output is None

227 +

228 + if hidden_states.dtype != torch.float16:

229 + final_hidden_states *= self.routed_scaling_factor

230 + elif self.shared_experts is not None:

231 + assert shared_output is not None

232 + shared_output *= 1.0 / self.routed_scaling_factor

233 +

234 + if self.shared_experts is not None:

235 + assert shared_output is not None

236 + final_hidden_states += shared_output

237 +

238 + if self.is_sequence_parallel:

239 + final_hidden_states = tensor_model_parallel_all_gather(

240 + final_hidden_states, 0

241 + )

242 + final_hidden_states = final_hidden_states[:num_tokens]

243 + elif self.tp_size > 1:

244 + final_hidden_states = self.experts.maybe_all_reduce_tensor_model_parallel(

245 + final_hidden_states

246 + )

247 +

248 + return final_hidden_states.view(num_tokens, hidden_dim)

249 +

250 +

251 +class OpenPanguMLAAttention(nn.Module):

252 + def __init__(

253 + self,

254 + config: PretrainedConfig,

255 + hidden_size: int,

256 + num_heads: int,

257 + qk_nope_head_dim: int,

258 + qk_rope_head_dim: int,

259 + v_head_dim: int,

260 + q_lora_rank: int | None,

261 + kv_lora_rank: int,

262 + rope_theta: float = 10000,

263 + max_position_embeddings: int = 8192,

264 + cache_config: CacheConfig | None = None,

265 + quant_config: QuantizationConfig | None = None,

266 + prefix: str = "",

267 + ) -> None:

268 + super().__init__()

269 + self.hidden_size = hidden_size

270 + self.num_heads = num_heads

271 + self.qk_nope_head_dim = qk_nope_head_dim

272 + self.qk_rope_head_dim = qk_rope_head_dim

273 + self.qk_head_dim = qk_nope_head_dim + qk_rope_head_dim

274 + self.v_head_dim = v_head_dim

275 + self.q_lora_rank = q_lora_rank

276 + self.kv_lora_rank = kv_lora_rank

277 + self.rope_theta = rope_theta

278 +

279 + self.tp_size = get_tensor_model_parallel_world_size()

280 + if num_heads % self.tp_size != 0:

281 + raise ValueError(

282 + f"num_heads {num_heads} is not divisible by tp_size {self.tp_size}."

283 + )

284 + self.num_local_heads = num_heads // self.tp_size

285 +

286 + self.scaling = self.qk_head_dim**-0.5

287 + self.max_position_embeddings = max_position_embeddings

288 +

289 + self.prefix = prefix

290 +

291 + if self.q_lora_rank is not None:

292 + self.fused_qkv_a_proj = MergedColumnParallelLinear(

293 + self.hidden_size,

294 + [self.q_lora_rank, self.kv_lora_rank + self.qk_rope_head_dim],

295 + bias=False,

296 + quant_config=quant_config,

297 + prefix=f"{prefix}.fused_qkv_a_proj",

298 + disable_tp=True,

299 + )

300 + self.q_a_layernorm = RMSNorm(self.q_lora_rank, eps=config.rms_norm_eps)

301 + self.q_b_proj = ColumnParallelLinear(

302 + q_lora_rank,

303 + self.num_heads * self.qk_head_dim,

304 + bias=False,

305 + quant_config=quant_config,

306 + prefix=f"{prefix}.q_b_proj",

307 + )

308 + else:

309 + self.q_proj = ColumnParallelLinear(

310 + self.hidden_size,

311 + self.num_heads * self.qk_head_dim,

312 + bias=False,

313 + quant_config=quant_config,

314 + prefix=f"{prefix}.q_proj",

315 + )

316 + self.kv_a_proj_with_mqa = ReplicatedLinear(

317 + self.hidden_size,

318 + self.kv_lora_rank + self.qk_rope_head_dim,

319 + bias=False,

320 + quant_config=quant_config,

321 + prefix=f"{prefix}.kv_a_proj_with_mqa",

322 + )

323 +

324 + self.kv_a_layernorm = RMSNorm(self.kv_lora_rank, eps=config.rms_norm_eps)

325 + self.kv_b_proj = ColumnParallelLinear(

326 + self.kv_lora_rank,

327 + self.num_heads * (self.qk_nope_head_dim + self.v_head_dim),

328 + bias=False,

329 + quant_config=quant_config,

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

331 + )

332 +

333 + self.o_proj = RowParallelLinear(

334 + self.num_heads * self.v_head_dim,

335 + self.hidden_size,

336 + bias=False,

337 + quant_config=quant_config,

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

339 + )

340 +

341 + # TODO: remove hard coding

342 + rope_scaling = {

343 + "beta_fast": 32,

344 + "beta_slow": 1,

345 + "factor": 1,

346 + "mscale": 1.0,

347 + "mscale_all_dim": 1.0,

348 + "original_max_position_embeddings": max_position_embeddings,

349 + "type": "yarn",

350 + "rope_type": "deepseek_yarn",

351 + }

352 + self.rotary_emb = get_rope(

353 + qk_rope_head_dim,

354 + rotary_dim=qk_rope_head_dim,

355 + max_position=max_position_embeddings,

356 + base=rope_theta,

357 + rope_scaling=rope_scaling,

358 + is_neox_style=False,

359 + )

360 +

361 + mla_modules = MLAModules(

362 + kv_a_layernorm=self.kv_a_layernorm,

363 + kv_b_proj=self.kv_b_proj,

364 + rotary_emb=self.rotary_emb,

365 + o_proj=self.o_proj,

366 + fused_qkv_a_proj=self.fused_qkv_a_proj

367 + if self.q_lora_rank is not None

368 + else None,

369 + kv_a_proj_with_mqa=self.kv_a_proj_with_mqa

370 + if self.q_lora_rank is None

371 + else None,

372 + q_a_layernorm=self.q_a_layernorm if self.q_lora_rank is not None else None,

373 + q_b_proj=self.q_b_proj if self.q_lora_rank is not None else None,

374 + q_proj=self.q_proj if self.q_lora_rank is None else None,

375 + indexer=None,

376 + is_sparse=False,

377 + topk_indices_buffer=None,

378 + )

379 +

380 + self.mla_attn = MultiHeadLatentAttentionWrapper(

381 + self.hidden_size,

382 + self.num_local_heads,

383 + self.scaling,

384 + self.qk_nope_head_dim,

385 + self.qk_rope_head_dim,

386 + self.v_head_dim,

387 + self.q_lora_rank,

388 + self.kv_lora_rank,

389 + mla_modules,

390 + cache_config,

391 + quant_config,

392 + prefix,

393 + )

394 +

395 + def forward(

396 + self,

397 + positions: torch.Tensor,

398 + hidden_states: torch.Tensor,

399 + ) -> torch.Tensor:

400 + return self.mla_attn(positions, hidden_states)

401 +

402 +

403 +class OpenPanguEmbeddedAttention(nn.Module):

404 + def __init__(

405 + self,

406 + config: PretrainedConfig,

407 + hidden_size: int,

408 + num_heads: int,

409 + num_kv_heads: int,

410 + rope_theta: float = 10000,

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

412 + max_position_embeddings: int = 8192,

413 + quant_config: QuantizationConfig | None = None,

414 + bias: bool = False,

415 + bias_o_proj: bool = False,

416 + cache_config: CacheConfig | None = None,

417 + prefix: str = "",

418 + attn_type: str = AttentionType.DECODER,

419 + ) -> None:

420 + super().__init__()

421 + layer_idx = extract_layer_index(prefix)

422 + self.hidden_size = hidden_size

423 + tp_size = get_tensor_model_parallel_world_size()

424 + self.total_num_heads = num_heads

425 + if self.total_num_heads % tp_size != 0:

426 + raise ValueError(

427 + f"total_num_heads {self.total_num_heads} "

428 + f"is not divisible by tp_size {tp_size}."

429 + )

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

431 + self.total_num_kv_heads = num_kv_heads

432 + if self.total_num_kv_heads > tp_size and self.total_num_kv_heads % tp_size != 0:

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

434 + # the KV heads across multiple tensor parallel ranks.

435 + raise ValueError(

436 + "Number of KV heads is greater than TP size, "

437 + f"but total_num_kv_heads {self.total_num_kv_heads} "

438 + f"is not divisible by tp_size {tp_size}."

439 + )

440 + elif (

441 + self.total_num_kv_heads < tp_size and tp_size % self.total_num_kv_heads != 0

442 + ):

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

444 + # the KV heads across multiple tensor parallel ranks.

445 + raise ValueError(

446 + f"Number of KV heads is less than TP size, but tp_size {tp_size} "

447 + f"is not divisible by total_num_kv_heads {self.total_num_kv_heads}."

448 + )

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

450 + head_dim = getattr(config, "head_dim", None)

451 + if head_dim is None:

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

453 + self.head_dim = head_dim

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

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

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

457 + self.rope_theta = rope_theta

458 + self.max_position_embeddings = max_position_embeddings

459 +

460 + self.qkv_proj = QKVParallelLinear(

461 + hidden_size=hidden_size,

462 + head_size=self.head_dim,

463 + total_num_heads=self.total_num_heads,

464 + total_num_kv_heads=self.total_num_kv_heads,

465 + bias=bias,

466 + quant_config=quant_config,

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

468 + )

469 +

470 + self.o_proj = RowParallelLinear(

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

472 + output_size=hidden_size,

473 + bias=bias_o_proj,

474 + quant_config=quant_config,

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

476 + )

477 +

478 + self._init_rotary_emb(

479 + config, rope_scaling=rope_scaling, quant_config=quant_config

480 + )

481 +

482 + if hasattr(config, "interleaved_sliding_window"):

483 + interleaved_sliding_window = config.interleaved_sliding_window

484 + if isinstance(interleaved_sliding_window, int):

485 + sliding_window = interleaved_sliding_window

486 + elif isinstance(interleaved_sliding_window, list):

487 + sw_idx = layer_idx % len(interleaved_sliding_window)

488 + sliding_window = interleaved_sliding_window[sw_idx]

489 + else:

490 + raise ValueError(

491 + f"{type(interleaved_sliding_window)} "

492 + "for interleaved_sliding_window is not supported."

493 + )

494 + else:

495 + sliding_window = None

496 +

497 + self.attn = Attention(

498 + self.num_heads,

499 + self.head_dim,

500 + self.scaling,

501 + num_kv_heads=self.num_kv_heads,

502 + cache_config=cache_config,

503 + quant_config=quant_config,

504 + per_layer_sliding_window=sliding_window,

505 + attn_type=attn_type,

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

507 + )

508 +

509 + def forward(

510 + self,

511 + positions: torch.Tensor,

512 + hidden_states: torch.Tensor,

513 + ) -> torch.Tensor:

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

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

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

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

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

519 + return output

520 +

521 + def _init_rotary_emb(

522 + self,

523 + config: PretrainedConfig,

524 + rope_scaling: dict[str, Any] | None,

525 + quant_config: QuantizationConfig | None,

526 + ) -> None:

527 + is_neox_style = True

528 + is_gguf = quant_config and quant_config.get_name() == "gguf"

529 + if is_gguf and config.model_type == "PanguEmbedded":

530 + is_neox_style = False

531 +

532 + self.rotary_emb = get_rope(

533 + self.head_dim,

534 + rotary_dim=self.head_dim,

535 + max_position=self.max_position_embeddings,

536 + base=self.rope_theta,

537 + rope_scaling=rope_scaling,

538 + is_neox_style=is_neox_style,

539 + )

540 +

541 +

542 +class OpenPanguDecoderLayer(nn.Module):

543 + def __init__(

544 + self,

545 + config: PretrainedConfig,

546 + prefix: str,

547 + vllm_config: VllmConfig,

548 + ) -> None:

549 + super().__init__()

550 +

551 + if config is None:

552 + config = vllm_config.model_config.hf_config

553 + cache_config = vllm_config.cache_config

554 + quant_config = vllm_config.quant_config

555 + parallel_config = vllm_config.parallel_config

556 +

557 + self.hidden_size = config.hidden_size

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

559 + max_position_embeddings = getattr(config, "max_position_embeddings", 8192)

560 +

561 + layer_idx = int(prefix.split(sep=".")[-1])

562 + self.layer_idx = layer_idx

563 +

564 + self.use_mla = (

565 + hasattr(config, "qk_nope_head_dim")

566 + and hasattr(config, "qk_rope_head_dim")

567 + and hasattr(config, "v_head_dim")

568 + and hasattr(config, "kv_lora_rank")

569 + )

570 + if self.use_mla:

571 + self.self_attn = OpenPanguMLAAttention(

572 + config=config,

573 + hidden_size=self.hidden_size,

574 + num_heads=config.num_attention_heads,

575 + qk_nope_head_dim=config.qk_nope_head_dim,

576 + qk_rope_head_dim=config.qk_rope_head_dim,

577 + v_head_dim=config.v_head_dim,

578 + q_lora_rank=(

579 + config.q_lora_rank if hasattr(config, "q_lora_rank") else None

580 + ),

581 + kv_lora_rank=config.kv_lora_rank,

582 + rope_theta=rope_theta,

583 + max_position_embeddings=max_position_embeddings,

584 + cache_config=cache_config,

585 + quant_config=quant_config,

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

587 + )

588 + else:

589 + attention_bias = getattr(config, "attention_bias", False) or getattr(

590 + config, "bias", False

591 + )

592 + bias_o_proj = attention_bias

593 + if hasattr(config, "qkv_bias"):

594 + attention_bias = config.qkv_bias

595 + # By default, PanguEmbedded uses causal attention

596 + # as it is a decoder-only model.

597 + # You can override the HF config with `is_causal=False` to enable

598 + # bidirectional attention, which is used in some embedding models

599 + if getattr(config, "is_causal", True):

600 + attn_type = AttentionType.DECODER

601 + else:

602 + attn_type = AttentionType.ENCODER_ONLY

603 + self.self_attn = OpenPanguEmbeddedAttention(

604 + config=config,

605 + hidden_size=self.hidden_size,

606 + num_heads=config.num_attention_heads,

607 + num_kv_heads=getattr(

608 + config, "num_key_value_heads", config.num_attention_heads

609 + ),

610 + rope_theta=rope_theta,

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

612 + max_position_embeddings=max_position_embeddings,

613 + quant_config=quant_config,

614 + bias=attention_bias,

615 + bias_o_proj=bias_o_proj,

616 + cache_config=cache_config,

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

618 + attn_type=attn_type,

619 + )

620 +

621 + if (

622 + getattr(config, "n_routed_experts", None) is not None

623 + and layer_idx >= config.first_k_dense_replace

624 + ):

625 + self.mlp = OpenPanguMoE(

626 + config=config,

627 + parallel_config=parallel_config,

628 + quant_config=quant_config,

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

630 + )

631 + else:

632 + self.mlp = OpenPanguMLP(

633 + hidden_size=self.hidden_size,

634 + intermediate_size=config.intermediate_size,

635 + hidden_act=config.hidden_act,

636 + quant_config=quant_config,

637 + bias=getattr(config, "mlp_bias", False),

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

639 + )

640 + self.routed_scaling_factor = getattr(config, "routed_scaling_factor", None)

641 + self.num_hidden_layers = config.num_hidden_layers

642 + self.first_k_dense_replace = getattr(

643 + config, "first_k_dense_replace", self.num_hidden_layers

644 + )

645 +

646 + self.input_layernorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)

647 + self.post_attention_layernorm = RMSNorm(

648 + config.hidden_size, eps=config.rms_norm_eps

649 + )

650 + self.tp_group = get_tp_group().device_group

651 + self.sandwich_norm = getattr(config, "sandwich_norm", False)

652 + if self.sandwich_norm:

653 + self.pre_mlp_layernorm = RMSNorm(

654 + config.hidden_size, eps=config.rms_norm_eps

655 + )

656 + self.post_mlp_layernorm = RMSNorm(

657 + config.hidden_size, eps=config.rms_norm_eps

658 + )

659 +

660 + def forward(

661 + self,

662 + positions: torch.Tensor,

663 + hidden_states: torch.Tensor,

664 + residual: torch.Tensor | None,

665 + ) -> torch.Tensor:

666 + if residual is None:

667 + residual = hidden_states.clone()

668 + hidden_states = self.input_layernorm(hidden_states)

669 + else:

670 + hidden_states, residual = self.input_layernorm(hidden_states, residual)

671 +

672 + hidden_states = self.self_attn(

673 + positions=positions,

674 + hidden_states=hidden_states,

675 + )

676 +

677 + if (

678 + self.routed_scaling_factor is not None

679 + and hidden_states.dtype == torch.float16

680 + ):

681 + # Fix FP16 overflow

682 + # We scale both hidden_states and residual before

683 + # rmsnorm, and rmsnorm result would not affect by scale.

684 + hidden_states *= 1.0 / self.routed_scaling_factor

685 + if self.layer_idx == 0:

686 + # The residual is shared by all layers, we only scale it on

687 + # first layer.

688 + residual *= 1.0 / self.routed_scaling_factor

689 +

690 + if self.sandwich_norm:

691 + hidden_states = self.post_attention_layernorm(hidden_states)

692 + hidden_states, residual = self.pre_mlp_layernorm(hidden_states, residual)

693 + else:

694 + hidden_states, residual = self.post_attention_layernorm(

695 + hidden_states, residual

696 + )

697 +

698 + # Fully Connected

699 + hidden_states = self.mlp(hidden_states)

700 +

701 + if (

702 + self.routed_scaling_factor is not None

703 + and isinstance(self.mlp, OpenPanguMLP)

704 + and hidden_states.dtype == torch.float16

705 + ):

706 + hidden_states *= 1.0 / self.routed_scaling_factor

707 +

708 + if self.sandwich_norm:

709 + hidden_states = self.post_mlp_layernorm(hidden_states)

710 +

711 + return hidden_states, residual

712 +

713 +

714 +@support_torch_compile

715 +class OpenPanguModel(nn.Module):

716 + fall_back_to_pt_during_load = False

717 +

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

719 + super().__init__()

720 +

721 + config = vllm_config.model_config.hf_config

722 + quant_config = vllm_config.quant_config

723 + eplb_config = vllm_config.parallel_config.eplb_config

724 + self.config = config

725 + self.num_redundant_experts = eplb_config.num_redundant_experts

726 +

727 + self.padding_idx = config.pad_token_id

728 + self.vocab_size = config.vocab_size

729 +

730 + if get_pp_group().is_first_rank or (

731 + config.tie_word_embeddings and get_pp_group().is_last_rank

732 + ):

733 + self.embed_tokens = VocabParallelEmbedding(

734 + config.vocab_size,

735 + config.hidden_size,

736 + quant_config=quant_config,

737 + prefix=f"{prefix}.embed_tokens",

738 + )

739 + else:

740 + self.embed_tokens = PPMissingLayer()

741 +

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

743 + config.num_hidden_layers,

744 + lambda prefix: OpenPanguDecoderLayer(config, prefix, vllm_config),

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

746 + )

747 +

748 + if get_pp_group().is_last_rank:

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

750 + else:

751 + self.norm = PPMissingLayer()

752 + self.make_empty_intermediate_tensors = make_empty_intermediate_tensors_factory(

753 + ["hidden_states", "residual"], config.hidden_size

754 + )

755 +

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

757 + return self.embed_tokens(input_ids)

758 +

759 + def forward(

760 + self,

761 + input_ids: torch.Tensor,

762 + positions: torch.Tensor,

763 + intermediate_tensors: IntermediateTensors | None,

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

765 + ) -> torch.Tensor | IntermediateTensors:

766 + if get_pp_group().is_first_rank:

767 + if inputs_embeds is not None:

768 + hidden_states = inputs_embeds

769 + else:

770 + hidden_states = self.get_input_embeddings(input_ids)

771 + residual = None

772 + else:

773 + assert intermediate_tensors is not None

774 + hidden_states = intermediate_tensors["hidden_states"]

775 + residual = intermediate_tensors["residual"]

776 +

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

778 + layer = self.layers[i]

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

780 +

781 + if not get_pp_group().is_last_rank:

782 + return IntermediateTensors(

783 + {"hidden_states": hidden_states, "residual": residual}

784 + )

785 +

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

787 + return hidden_states

788 +

789 + def load_attn_mlp_weight(

790 + self,

791 + attn_mlp_replace_mapping: list[tuple[str, str, int]],

792 + params_dict: dict[str, Any],

793 + weight_name: str,

794 + loaded_weight: torch.Tensor,

795 + loaded_params: set[str],

796 + ) -> bool:

797 + for param_name, origin_name, shard_id in attn_mlp_replace_mapping:

798 + if origin_name not in weight_name or (

799 + ("mlp.experts." in weight_name) and weight_name not in params_dict

800 + ):

801 + continue

802 + weight_name_mapped = weight_name.replace(origin_name, param_name)

803 + if (

804 + param_name == "fused_qkv_a_proj"

805 + and weight_name_mapped not in params_dict

806 + ):

807 + continue

808 + else:

809 + weight_name = weight_name_mapped

810 + if weight_name.endswith(".bias") and weight_name not in params_dict:

811 + continue

812 + if is_pp_missing_parameter(weight_name, self):

813 + continue

814 +

815 + param = params_dict[weight_name]

816 + weight_loader = param.weight_loader

817 + weight_loader(param, loaded_weight, shard_id)

818 + loaded_params.add(weight_name)

819 + return True

820 + return False

821 +

822 + def load_expert_weight(

823 + self,

824 + expert_merge_mapping: list[tuple[str, str, int, str]],

825 + params_dict: dict[str, Any],

826 + weight_name: str,

827 + loaded_weight: torch.Tensor,

828 + loaded_params: set[str],

829 + flag_dict: dict[str, bool],

830 + ) -> bool:

831 + for mapping in expert_merge_mapping:

832 + param_name, origin_name, expert_id, shard_id = mapping

833 + if origin_name not in weight_name:

834 + continue

835 + flag_dict["is_expert_weight"] = True

836 + weight_name_mapped = weight_name.replace(origin_name, param_name)

837 + if is_pp_missing_parameter(weight_name_mapped, self):

838 + continue

839 + param = params_dict[weight_name_mapped]

840 + weight_loader = typing.cast(Callable[..., bool], param.weight_loader)

841 + success = weight_loader(

842 + param,

843 + loaded_weight,

844 + weight_name_mapped,

845 + shard_id=shard_id,

846 + expert_id=expert_id,

847 + return_success=True,

848 + )

849 + if success:

850 + weight_name = weight_name_mapped

851 + loaded_params.add(weight_name_mapped)

852 + return True

853 + return False

854 +

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

856 + attn_mlp_replace_mapping = [

857 + (".qkv_proj", ".q_proj", "q"),

858 + (".qkv_proj", ".k_proj", "k"),

859 + (".qkv_proj", ".v_proj", "v"),

860 + (".fused_qkv_a_proj", ".q_a_proj", 0),

861 + (".fused_qkv_a_proj", ".kv_a_proj_with_mqa", 1),

862 + (".gate_up_proj", ".gate_proj", 0),

863 + (".gate_up_proj", ".up_proj", 1),

864 + ]

865 + has_experts = hasattr(self.config, "n_routed_experts")

866 + if has_experts:

867 + expert_merge_mapping = SharedFusedMoE.make_expert_params_mapping(

868 + ckpt_gate_proj_name="gate_proj",

869 + ckpt_down_proj_name="down_proj",

870 + ckpt_up_proj_name="up_proj",

871 + num_experts=self.config.n_routed_experts,

872 + num_redundant_experts=self.num_redundant_experts,

873 + )

874 +

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

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

877 + for name, loaded_weight in weights:

878 + if "rotary_emb.inv_freq" in name:

879 + continue

880 + if self.config.tie_word_embeddings and "lm_head.weight" in name:

881 + continue

882 +

883 + if (

884 + "layers" in name

885 + and hasattr(self.config, "num_nextn_predict_layers")

886 + and (self.config.num_nextn_predict_layers > 0)

887 + ):

888 + layer_idx = int(name.split("layers.")[-1].split(".")[0])

889 + mtp_idx = layer_idx - self.config.num_hidden_layers

890 + if mtp_idx >= 0 and mtp_idx < self.config.num_nextn_predict_layers:

891 + continue # skip spec decode layers for main model

892 +

893 + flag_dict = {"is_expert_weight": False}

894 + if (

895 + self.load_attn_mlp_weight(

896 + attn_mlp_replace_mapping,

897 + params_dict,

898 + name,

899 + loaded_weight,

900 + loaded_params,

901 + )

902 + or has_experts

903 + and self.load_expert_weight(

904 + expert_merge_mapping,

905 + params_dict,

906 + name,

907 + loaded_weight,

908 + loaded_params,

909 + flag_dict,

910 + )

911 + ):

912 + continue

913 + else:

914 + if flag_dict["is_expert_weight"]:

915 + continue

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

917 + continue

918 + name = maybe_remap_kv_scale_name(name, params_dict)

919 + if name is None:

920 + continue

921 + if is_pp_missing_parameter(name, self):

922 + continue

923 +

924 + param = params_dict[name]

925 + weight_loader = getattr(param, "weight_loader", default_weight_loader)

926 + weight_loader(param, loaded_weight)

927 + loaded_params.add(name)

928 + return loaded_params

929 +

930 +

931 +class OpenPanguModelBase(nn.Module, SupportsPP, SupportsLoRA):

932 + packed_modules_mapping = {

933 + "qkv_proj": ["q_proj", "k_proj", "v_proj"],

934 + "gate_up_proj": ["gate_proj", "up_proj"],

935 + }

936 +

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

938 + super().__init__()

939 + config = vllm_config.model_config.hf_config

940 + quant_config = vllm_config.quant_config

941 + self.config = config

942 + self.quant_config = quant_config

943 +

944 + self.fuse_qkv_a_proj = (

945 + hasattr(config, "q_lora_rank") and config.q_lora_rank is not None

946 + )

947 + if self.fuse_qkv_a_proj:

948 + self.packed_modules_mapping["fused_qkv_a_proj"] = [

949 + "q_a_proj",

950 + "kv_a_proj_with_mqa",

951 + ]

952 +

953 + self.model = OpenPanguModel(

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

955 + )

956 + if get_pp_group().is_last_rank:

957 + self.lm_head = ParallelLMHead(

958 + config.vocab_size,

959 + config.hidden_size,

960 + quant_config=quant_config,

961 + prefix=maybe_prefix(prefix, "lm_head"),

962 + )

963 + if config.tie_word_embeddings:

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

965 + else:

966 + self.lm_head = PPMissingLayer()

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

968 + self.make_empty_intermediate_tensors = (

969 + self.model.make_empty_intermediate_tensors

970 + )

971 +

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

973 + return self.model.get_input_embeddings(input_ids)

974 +

975 + def forward(

976 + self,

977 + input_ids: torch.Tensor,

978 + positions: torch.Tensor,

979 + intermediate_tensors: IntermediateTensors | None = None,

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

981 + ) -> torch.Tensor | IntermediateTensors:

982 + hidden_states = self.model(

983 + input_ids, positions, intermediate_tensors, inputs_embeds

984 + )

985 + return hidden_states

986 +

987 + def compute_logits(

988 + self,

989 + hidden_states: torch.Tensor,

990 + ) -> torch.Tensor | None:

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

992 + return logits

993 +

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

995 + loader = AutoWeightsLoader(

996 + self,

997 + skip_prefixes=(["lm_head."] if self.config.tie_word_embeddings else None),

998 + )

999 + return loader.load_weights(weights)

1000 +

1001 +

1002 +class OpenPanguMoEModel(OpenPanguModelBase, MixtureOfExperts):

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

1004 + super().__init__(vllm_config=vllm_config, prefix=prefix)

1005 + config = vllm_config.model_config.hf_config

1006 +

1007 + # Set MoE hyperparameters

1008 + self.expert_weights = []

1009 + self.num_moe_layers = config.num_hidden_layers - config.first_k_dense_replace

1010 + self.num_expert_groups = 1

1011 +

1012 + self.moe_layers: list[SharedFusedMoE] = []

1013 + example_moe = None

1014 + for layer in self.model.layers:

1015 + if isinstance(layer, PPMissingLayer):

1016 + continue

1017 +

1018 + assert isinstance(layer, OpenPanguDecoderLayer)

1019 + if isinstance(layer.mlp, OpenPanguMoE):

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

1021 + example_moe = layer.mlp

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

1023 +

1024 + if example_moe is None:

1025 + raise RuntimeError("No MOE layer found in model.layers.")

1026 +

1027 + self.num_logical_experts = example_moe.n_logical_experts

1028 + self.num_physical_experts = example_moe.n_physical_experts

1029 + self.num_local_physical_experts = example_moe.n_local_physical_experts

1030 + self.n_routed_experts = example_moe.n_routed_experts

1031 + self.n_shared_experts = example_moe.n_shared_experts

1032 + self.num_redundant_experts = example_moe.n_redundant_experts

1033 +

1034 + def set_eplb_state(

1035 + self,

1036 + expert_load_view: torch.Tensor,

1037 + logical_to_physical_map: torch.Tensor,

1038 + logical_replica_count: torch.Tensor,

1039 + ) -> None:

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

1041 + # Register the expert weights.

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

1043 + layer.set_eplb_state(

1044 + moe_layer_idx=layer_idx,

1045 + expert_load_view=expert_load_view,

1046 + logical_to_physical_map=logical_to_physical_map,

1047 + logical_replica_count=logical_replica_count,

1048 + )

1049 +

1050 + def update_physical_experts_metadata(

1051 + self,

1052 + num_physical_experts: int,

1053 + num_local_physical_experts: int,

1054 + ) -> None:

1055 + assert self.num_local_physical_experts == num_local_physical_experts

1056 + self.num_physical_experts = num_physical_experts

1057 + self.num_local_physical_experts = num_local_physical_experts

1058 + self.num_redundant_experts = num_physical_experts - self.num_logical_experts

1059 + for layer in self.model.layers:

1060 + if isinstance(layer.mlp, OpenPanguMoE):

1061 + moe = layer.mlp

1062 + moe.n_local_physical_experts = num_local_physical_experts

1063 + moe.n_physical_experts = num_physical_experts

1064 + moe.n_redundant_experts = self.num_redundant_experts

1065 + moe.experts.update_expert_map()

1066 +

1067 +

1068 +class OpenPanguEmbeddedModel(OpenPanguModelBase):

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

1070 + super().__init__(vllm_config=vllm_config, prefix=prefix)

1071 +

1072 +

1073 +class PanguEmbeddedForCausalLM(OpenPanguEmbeddedModel):

1074 + pass

1075 +

1076 +

1077 +class PanguUltraMoEForCausalLM(OpenPanguMoEModel):

1078 + pass

vllm/model_executor/models/openpangu_mtp.py new

+265

@@ -0,0 +1,265 @@

1 +#

4 +#

5 +# This file is a part of the vllm-ascend project.

6 +#

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

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

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

10 +#

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

12 +#

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

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

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

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

17 +# limitations under the License.

18 +

19 +# Adapted from

20 +# https://github.com/vllm-project/vllm/blob/v0.7.3/vllm/model_executor/models/deepseek_mtp.py

21 +from collections.abc import Iterable

22 +

23 +import torch

24 +import torch.nn as nn

25 +

26 +from vllm.compilation.decorators import support_torch_compile

27 +from vllm.config import VllmConfig

28 +

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

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

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

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

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

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

35 + VocabParallelEmbedding,

36 +)

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

38 +from vllm.model_executor.models.deepseek_mtp import (

39 + DeepSeekMultiTokenPredictor,

40 + DeepSeekMultiTokenPredictorLayer,

41 + SharedHead,

42 +)

43 +from vllm.model_executor.models.utils import maybe_prefix

44 +from vllm.sequence import IntermediateTensors

45 +

46 +from .interfaces import SupportsPP

47 +from .openpangu import OpenPanguDecoderLayer

48 +

49 +

50 +class OpenPanguMultiTokenPredictorLayer(DeepSeekMultiTokenPredictorLayer):

51 + def __init__(self, vllm_config: VllmConfig, prefix: str) -> None:

52 + nn.Module.__init__(self)

53 +

54 + config = vllm_config.speculative_config.draft_model_config.hf_config

55 + self.config = config

56 + quant_config = vllm_config.quant_config

57 +

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

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

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

61 + self.shared_head = SharedHead(

62 + config=config,

63 + quant_config=quant_config,

64 + prefix=maybe_prefix(prefix, "shared_head"),

65 + )

66 + self.mtp_block = OpenPanguDecoderLayer(config, prefix, vllm_config)

67 +

68 +

69 +class OpenPanguMultiTokenPredictor(DeepSeekMultiTokenPredictor):

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

71 + nn.Module.__init__(self)

72 + config = vllm_config.model_config.hf_config

73 + self.mtp_start_layer_idx = config.num_hidden_layers

74 + self.num_mtp_layers = config.num_nextn_predict_layers

75 + # to map the exact layer index from weights

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

77 + {

78 + str(idx): OpenPanguMultiTokenPredictorLayer(

79 + vllm_config, f"{prefix}.layers.{idx}"

80 + )

81 + for idx in range(

82 + self.mtp_start_layer_idx,

83 + self.mtp_start_layer_idx + self.num_mtp_layers,

84 + )

85 + }

86 + )

87 + self.embed_tokens = VocabParallelEmbedding(

88 + config.vocab_size,

89 + config.hidden_size,

90 + )

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

92 +

93 +

94 +@support_torch_compile

95 +class OpenPanguMTP(nn.Module, SupportsPP):

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

97 + super().__init__()

98 + self.config = vllm_config.model_config.hf_config

99 + self.model = OpenPanguMultiTokenPredictor(

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

101 + )

102 +

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

104 + return self.model.get_input_embeddings(input_ids)

105 +

106 + def forward(

107 + self,

108 + input_ids: torch.Tensor,

109 + positions: torch.Tensor,

110 + hidden_states: torch.Tensor,

111 + intermediate_tensors: IntermediateTensors | None = None,

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

113 + spec_step_idx: int = 0,

114 + ) -> torch.Tensor:

115 + hidden_states = self.model(

116 + input_ids,

117 + positions,

118 + hidden_states,

119 + inputs_embeds,

120 + spec_step_idx,

121 + )

122 + return hidden_states

123 +

124 + def compute_logits(

125 + self,

126 + hidden_states: torch.Tensor,

127 + spec_step_idx: int = 0,

128 + ) -> torch.Tensor | None:

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

130 +

131 + def get_spec_layer(self, name):

132 + if (

133 + "layers" in name

134 + and hasattr(self.config, "num_nextn_predict_layers")

135 + and self.config.num_nextn_predict_layers > 0

136 + ):

137 + layer_idx = int(name.split("layers.")[-1].split(".")[0])

138 + mtp_idx = layer_idx - self.config.num_hidden_layers

139 + if mtp_idx >= 0 and mtp_idx < self.config.num_nextn_predict_layers:

140 + return layer_idx

141 + return None

142 +

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

144 + stacked_params_mapping = [

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

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

147 + ("fused_qkv_a_proj", "q_a_proj", 0),

148 + ("fused_qkv_a_proj", "kv_a_proj_with_mqa", 1),

149 + ]

150 +

151 + expert_params_mapping = FusedMoE.make_expert_params_mapping(

152 + ckpt_gate_proj_name="gate_proj",

153 + ckpt_down_proj_name="down_proj",

154 + ckpt_up_proj_name="up_proj",

155 + num_experts=self.config.n_routed_experts,

156 + )

157 +

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

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

160 + for name, loaded_weight in weights:

161 + if "rotary_emb.inv_freq" in name:

162 + continue

163 + spec_layer = self.get_spec_layer(name)

164 + if spec_layer is None:

165 + continue

166 +

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

168 + for param_name, weight_name, shard_id in stacked_params_mapping:

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

170 + if weight_name not in name:

171 + continue

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

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

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

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

176 + # will then be updated below in expert_params_mapping

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

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

179 + continue

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

181 +

182 + # QKV fusion is optional, fall back to normal

183 + # weight loading if it's not enabled

184 + if (

185 + param_name == "fused_qkv_a_proj"

186 + ) and name_mapped not in params_dict:

187 + continue

188 + else:

189 + name = name_mapped

190 +

191 + # Skip loading extra bias for GPTQ models.

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

193 + continue

194 +

195 + param = params_dict[name]

196 + weight_loader = param.weight_loader

197 + weight_loader(param, loaded_weight, shard_id)

198 + break

199 + else:

200 + for mapping in expert_params_mapping:

201 + param_name, weight_name, expert_id, shard_id = mapping

202 + if weight_name not in name:

203 + continue

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

205 +

206 + param = params_dict[name]

207 + weight_loader = param.weight_loader

208 + weight_loader(

209 + param,

210 + loaded_weight,

211 + name,

212 + shard_id=shard_id,

213 + expert_id=expert_id,

214 + )

215 + break

216 + else:

217 + # Skip loading extra bias for GPTQ models.

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

219 + continue

220 +

221 + if (

222 + spec_layer != self.model.mtp_start_layer_idx

223 + and ".layers" not in name

224 + ):

225 + continue

226 +

227 + param = params_dict[name]

228 + weight_loader = getattr(

229 + param, "weight_loader", default_weight_loader

230 + )

231 + weight_loader(param, loaded_weight)

232 + loaded_params.add(name)

233 + return loaded_params

234 +

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

236 + """

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

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

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

240 + """

241 + spec_layer_weight_names = [

242 + "embed_tokens",

243 + "enorm",

244 + "hnorm",

245 + "eh_proj",

246 + "shared_head",

247 + ]

248 + shared_weight_names = ["embed_tokens"]

249 + spec_layer_weight = False

250 + shared_weight = False

251 + for weight_name in spec_layer_weight_names:

252 + if weight_name in name:

253 + spec_layer_weight = True

254 + if weight_name in shared_weight_names:

255 + shared_weight = True

256 + break

257 + if not spec_layer_weight:

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

259 + name = name.replace(

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

261 + )

262 + elif shared_weight:

263 + # treat shared weights as top level weights

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

265 + return name

vllm/model_executor/models/registry.py

@@ -149,6 +149,8 @@ _TEXT_GENERATION_MODELS = {

149 "OPTForCausalLM": ("opt", "OPTForCausalLM"),

150 "OrionForCausalLM": ("orion", "OrionForCausalLM"),

151 "OuroForCausalLM": ("ouro", "OuroForCausalLM"),

152 + "PanguEmbeddedForCausalLM": ("openpangu", "PanguEmbeddedForCausalLM"),

153 + "PanguUltraMoEForCausalLM": ("openpangu", "PanguUltraMoEForCausalLM"),

154 "PersimmonForCausalLM": ("persimmon", "PersimmonForCausalLM"),

155 "PhiForCausalLM": ("phi", "PhiForCausalLM"),

156 "Phi3ForCausalLM": ("phi3", "Phi3ForCausalLM"),

@@ -406,6 +408,7 @@ _SPECULATIVE_DECODING_MODELS = {

408 "LongCatFlashMTPModel": ("longcat_flash_mtp", "LongCatFlashMTP"),

409 "Glm4MoeMTPModel": ("glm4_moe_mtp", "Glm4MoeMTP"),

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

411 + "OpenPanguMTPModel": ("openpangu_mtp", "OpenPanguMTP"),

412 "Qwen3NextMTP": ("qwen3_next_mtp", "Qwen3NextMTP"),

413 # Temporarily disabled.

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

vllm/v1/spec_decode/eagle.py

+6 -1

@@ -316,7 +316,12 @@ class EagleProposer:

316 positions = target_positions[:, last_token_indices]

317 else:

318 positions = target_positions[last_token_indices]

319 - if self.method in ("deepseek_mtp", "ernie_mtp", "longcat_flash_mtp"):

319 + if self.method in (

320 + "deepseek_mtp",

321 + "ernie_mtp",

322 + "longcat_flash_mtp",

323 + "pangu_ultra_moe_mtp",

324 + ):

325 hidden_states = self.hidden_states[last_token_indices]

326 else:

327 hidden_states = hidden_states[last_token_indices]