@tduncan / Nvidia-Optimization-Stack-VLLM / commits / 0ff0477ff9

[Perf][ROCm] Add a ROCm path for Hy4 and compile the backbone (#57526)

Signed-off-by: Aakif Nawaz <aakif.nawaz@amd.com>

akii96 committed Sep 21, 2026 at 11:20 UTC 0ff0477ff955154f37db48791260cdf27fa35e60
4 files changed +782 -7
tests/models/test_hyv4_rocm.py
+48
@@ -51,6 +51,18 @@ def _num_captured_graphs(_worker):
51 return compilation_counter.num_cudagraph_captured
52
53
54 +def _piecewise_compile_state(worker):
55 + from vllm.compilation.counter import compilation_counter
56 +
57 + # `cudagraph_mode` is returned by name: the enum does not survive the
58 + # collective_rpc round trip, it arrives as its bare value.
59 + return (
60 + compilation_counter.num_models_seen,
61 + compilation_counter.num_cudagraph_captured,
62 + worker.vllm_config.compilation_config.cudagraph_mode.name,
63 + )
64 +
65 +
66 @pytest.mark.parametrize(
67 "tp,enforce_eager,mtp,dtype,kv_cache_dtype",
68 [
@@ -135,3 +147,39 @@ def test_hyv4_dummy_generation(
147 and metric.value > 0
148 for metric in model.llm.get_metrics()
149 ), "MTP drafter did not run"
150 +
151 +
152 +def test_hyv4_backbone_reaches_the_compiler(vllm_runner, enable_pickle) -> None:
153 + """Guard against the backbone losing `@support_torch_compile`.
154 +
155 + An undecorated backbone is not a hard failure: vLLM downgrades
156 + `cudagraph_mode` to NONE and serves the whole model eager, which on TP8
157 + gfx942 cost ~4x decode latency.
158 + """
159 + with vllm_runner(
160 + "tencent/Hy4-preview",
161 + trust_remote_code=True,
162 + skip_tokenizer_init=True,
163 + load_format="dummy",
164 + dtype="bfloat16",
165 + hf_overrides=_small_hyv4_config,
166 + max_model_len=256,
167 + max_num_batched_tokens=128,
168 + max_num_seqs=8,
169 + enable_chunked_prefill=True,
170 + enable_prefix_caching=False,
171 + kv_cache_memory_bytes=128 * 1024 * 1024,
172 + enforce_eager=False,
173 + seed=17,
174 + compilation_config=CompilationConfig(
175 + mode=CompilationMode.VLLM_COMPILE,
176 + cudagraph_mode=CUDAGraphMode.PIECEWISE,
177 + cudagraph_capture_sizes=[1, 2, 4],
178 + ),
179 + ) as model:
180 + for models_seen, captured, mode in model.collective_rpc(
181 + _piecewise_compile_state
182 + ):
183 + assert models_seen > 0, "backbone was never handed to the compiler"
184 + assert mode == CUDAGraphMode.PIECEWISE.name
185 + assert captured > 0
vllm/models/hy_v4/__init__.py
+14 -7
@@ -14,20 +14,27 @@ The package is organized like `vllm.models.deepseek_v32`: this module is the
14 only public entry point and dispatches on the current platform, so registry
15 entries never reach into a platform subpackage.
16
17 -NVIDIA and ROCm are supported. The port also drops the reference
18 -implementation's HPC/TPCP fusion paths, which depend on infrastructure that
19 -does not exist in this tree.
17 +NVIDIA and ROCm are supported, the latter through `amd/`. The port also drops
18 +the reference implementation's HPC/TPCP fusion paths, which depend on
19 +infrastructure that does not exist in this tree.
20 """
21
22 +from typing import TYPE_CHECKING
23 +
24 from vllm.platforms import current_platform
25
26 if current_platform.is_xpu():
27 raise NotImplementedError("hy_v4 does not yet support XPU.")
26 -else:
27 - # The implementation is shared by CUDA and ROCm; platform-specific
28 - # attention behavior is selected by the normal attention backend router.
28 +
29 +# The NVIDIA branch is the static default that type-checkers see; the ROCm
30 +# branch overrides it at runtime (kept type-compatible via type: ignore).
31 +if TYPE_CHECKING or not current_platform.is_rocm():
32 from .nvidia.model import HYV4ForCausalLM
30 - from .nvidia.mtp import HYV4MTP
33 +else:
34 + from .amd.model import HYV4ForCausalLM # type: ignore[assignment]
35 +
36 +# The draft head has no ROCm-specific implementation yet.
37 +from .nvidia.mtp import HYV4MTP # noqa: E402
38
39 __all__ = [
40 "HYV4ForCausalLM",
vllm/models/hy_v4/amd/__init__.py new
+2
@@ -0,0 +1,2 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
vllm/models/hy_v4/amd/model.py new
+718
@@ -0,0 +1,718 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +
4 +# Adapted from
5 +# Copyright 2023 The vLLM team.
6 +# Copyright 2022 EleutherAI and the HuggingFace Inc. team. All rights reserved.
7 +#
8 +# This code is based on EleutherAI's GPT-NeoX library and the GPT-NeoX
9 +# and OPT implementations in this library. It has been modified from its
10 +# original forms to accommodate minor architectural differences compared
11 +# to GPT-NeoX and OPT used by the Meta AI team that trained the model.
12 +#
13 +# Licensed under the Apache License, Version 2.0 (the "License");
14 +# you may not use this file except in compliance with the License.
15 +# You may obtain a copy of the License at
16 +#
17 +# http://www.apache.org/licenses/LICENSE-2.0
18 +#
19 +# Unless required by applicable law or agreed to in writing, software
20 +# distributed under the License is distributed on an "AS IS" BASIS,
21 +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
22 +# See the License for the specific language governing permissions and
23 +# limitations under the License.
24 +"""Inference-only HY V4 model compatible with HuggingFace weights (ROCm)."""
25 +
26 +import typing
27 +from collections.abc import Callable, Iterable, MutableSequence, Sequence
28 +from itertools import islice
29 +
30 +import regex as re
31 +import torch
32 +from torch import nn
33 +from transformers import PretrainedConfig
34 +
35 +from vllm.compilation.decorators import support_torch_compile
36 +from vllm.config import CacheConfig, VllmConfig
37 +from vllm.distributed import (
38 + get_pp_group,
39 + get_tensor_model_parallel_rank,
40 + get_tensor_model_parallel_world_size,
41 +)
42 +from vllm.logger import init_logger
43 +from vllm.model_executor.layers.fused_moe import fused_moe_make_expert_params_mapping
44 +from vllm.model_executor.layers.layernorm import RMSNorm
45 +from vllm.model_executor.layers.logits_processor import LogitsProcessor
46 +from vllm.model_executor.layers.quantization import QuantizationConfig
47 +from vllm.model_executor.layers.vocab_parallel_embedding import (
48 + ParallelLMHead,
49 + VocabParallelEmbedding,
50 +)
51 +from vllm.model_executor.model_loader.weight_utils import default_weight_loader
52 +from vllm.model_executor.models.deepseek_v2 import _try_load_fp8_indexer_wk
53 +from vllm.model_executor.models.interfaces import SupportsLoRA, SupportsPP
54 +from vllm.model_executor.models.utils import (
55 + AutoWeightsLoader,
56 + PPMissingLayer,
57 + WeightsMapper,
58 + get_pp_missing_layer_names,
59 + is_pp_missing_parameter,
60 + make_empty_intermediate_tensors_factory,
61 + make_layers,
62 + maybe_prefix,
63 +)
64 +from vllm.platforms import current_platform
65 +from vllm.sequence import IntermediateTensors
66 +
67 +from ..nvidia.attention import (
68 + HYV4MLAAttention,
69 + compute_skip_topk_layers,
70 + is_skip_topk_indexer_weight,
71 +)
72 +from ..nvidia.hc import HYV4HCHeadLayer, HYV4HCLayer
73 +from ..nvidia.moe import HYV4FeedForward, HYV4MoEFused
74 +
75 +logger = init_logger(__name__)
76 +
77 +
78 +def _normalize_hyv4_config(config: PretrainedConfig) -> PretrainedConfig:
79 + """Populate the aliases consumed by the shared MoE implementation."""
80 + config.router_scaling_factor = config.routed_scaling_factor
81 + config.num_experts = config.n_routed_experts
82 + config.expert_hidden_dim = config.moe_intermediate_size
83 + config.num_shared_experts = config.n_shared_experts
84 + config.route_norm = config.norm_topk_prob
85 + return config
86 +
87 +
88 +class HYV4DecoderLayer(nn.Module):
89 + """One HY V4 decoder layer: MLA attention plus a dense or MoE MLP.
90 +
91 + When``config.enable_ihc`` is set the layer runs on ``hc_mult`` residual
92 + channels and each sub-block is wrapped by an `HYV4HCLayer` boundary;
93 + otherwise it uses the standard single-stream residual.
94 + """
95 +
96 + def __init__(
97 + self,
98 + config: PretrainedConfig,
99 + vllm_config: VllmConfig,
100 + cache_config: CacheConfig | None = None,
101 + quant_config: QuantizationConfig | None = None,
102 + prefix: str = "",
103 + topk_indices_buffer: torch.Tensor | None = None,
104 + ) -> None:
105 + super().__init__()
106 + self.hidden_size = config.hidden_size
107 + self.config = config
108 + self.enable_ihc = getattr(config, "enable_ihc", False)
109 + layer_idx = int(prefix.split(".")[-1])
110 + self.layer_idx = layer_idx
111 +
112 + max_position_embeddings = getattr(config, "max_position_embeddings", 8192)
113 + self.self_attn = HYV4MLAAttention(
114 + vllm_config=vllm_config,
115 + config=config,
116 + hidden_size=self.hidden_size,
117 + num_heads=config.num_attention_heads,
118 + qk_nope_head_dim=config.qk_nope_head_dim,
119 + qk_rope_head_dim=config.qk_rope_head_dim,
120 + v_head_dim=config.v_head_dim,
121 + q_lora_rank=config.q_lora_rank,
122 + kv_lora_rank=config.kv_lora_rank,
123 + max_position_embeddings=max_position_embeddings,
124 + cache_config=cache_config,
125 + quant_config=quant_config,
126 + prefix=f"{prefix}.self_attn",
127 + layer_idx=layer_idx,
128 + topk_indices_buffer=topk_indices_buffer,
129 + )
130 + self.input_layernorm = RMSNorm(config.hidden_size, config.rms_norm_eps)
131 + if config.mlp_layer_types[layer_idx] == "dense":
132 + self.mlp = HYV4FeedForward(
133 + hidden_size=config.hidden_size,
134 + intermediate_size=config.intermediate_size,
135 + hidden_act=config.hidden_act,
136 + quant_config=quant_config,
137 + prefix=f"{prefix}.mlp",
138 + )
139 + self.block_type = "feedforward"
140 + else:
141 + self.mlp = HYV4MoEFused(
142 + config=config, quant_config=quant_config, prefix=f"{prefix}.mlp"
143 + )
144 + self.block_type = "moe"
145 + self.post_attention_layernorm = RMSNorm(config.hidden_size, config.rms_norm_eps)
146 + self.hc_attn_layer = HYV4HCLayer(
147 + config, layer_idx, prefix=f"{prefix}.hc_attn_layer"
148 + )
149 + self.hc_mlp_layer = HYV4HCLayer(
150 + config, layer_idx, prefix=f"{prefix}.hc_mlp_layer"
151 + )
152 +
153 + def forward(
154 + self,
155 + positions: torch.Tensor,
156 + hidden_states: torch.Tensor,
157 + residual: torch.Tensor | None,
158 + ) -> tuple[torch.Tensor, torch.Tensor | None]:
159 + if self.enable_ihc:
160 + return self._forward_ihc(positions, hidden_states)
161 + return self._forward_normal(positions, hidden_states, residual)
162 +
163 + def _forward_normal(
164 + self,
165 + positions: torch.Tensor,
166 + hidden_states: torch.Tensor,
167 + residual: torch.Tensor | None,
168 + ) -> tuple[torch.Tensor, torch.Tensor]:
169 + """Standard single-residual-stream forward (iHC disabled)."""
170 + if residual is not None:
171 + hidden_states = hidden_states + residual
172 + residual = hidden_states
173 + hidden_states = self.input_layernorm(hidden_states)
174 +
175 + hidden_states = self.self_attn(
176 + positions=positions,
177 + hidden_states=hidden_states,
178 + )
179 +
180 + hidden_states = hidden_states + residual
181 + residual = hidden_states
182 + hidden_states = self.post_attention_layernorm(hidden_states)
183 +
184 + hidden_states = self.mlp(hidden_states)
185 +
186 + return hidden_states, residual
187 +
188 + def _forward_ihc(
189 + self,
190 + positions: torch.Tensor,
191 + hidden_states: torch.Tensor,
192 + ) -> tuple[torch.Tensor, None]:
193 + """IHC forward: each sub-block reduces and re-scatters the channels."""
194 + hidden_states = self.hc_attn_layer.prepare_input(hidden_states)
195 + hidden_states, post_gates, residual = self.hc_attn_layer.pre(hidden_states)
196 + hidden_states = self.input_layernorm(hidden_states)
197 + hidden_states = self.self_attn(
198 + positions=positions,
199 + hidden_states=hidden_states,
200 + )
201 + hidden_states = self.hc_attn_layer.post(hidden_states, residual, post_gates)
202 +
203 + hidden_states = self.hc_mlp_layer.prepare_input(hidden_states)
204 + hidden_states, post_gates, residual = self.hc_mlp_layer.pre(hidden_states)
205 + hidden_states = self.post_attention_layernorm(hidden_states)
206 + hidden_states = self.mlp(hidden_states)
207 + hidden_states = self.hc_mlp_layer.post(hidden_states, residual, post_gates)
208 +
209 + # Under iHC the residual is carried inside hidden_states.
210 + return hidden_states, None
211 +
212 +
213 +@support_torch_compile
214 +class HYV4Model(nn.Module):
215 + """HY V4 backbone."""
216 +
217 + def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
218 + super().__init__()
219 +
220 + config = _normalize_hyv4_config(vllm_config.model_config.hf_config)
221 + cache_config = vllm_config.cache_config
222 + quant_config = vllm_config.quant_config
223 + parallel_config = vllm_config.parallel_config
224 + eplb_config = parallel_config.eplb_config
225 + self.num_redundant_experts = eplb_config.num_redundant_experts
226 + self.device = current_platform.device_type
227 + self.vocab_size = config.vocab_size
228 + self.is_sparse = hasattr(config, "index_topk")
229 + if self.is_sparse:
230 + self.topk_indices_buffer = torch.empty(
231 + vllm_config.scheduler_config.max_num_batched_tokens,
232 + config.index_topk,
233 + dtype=torch.int32,
234 + device=self.device,
235 + )
236 + else:
237 + self.topk_indices_buffer = None
238 + self.config = config
239 + self.quant_config = quant_config
240 + self.enable_ihc = getattr(config, "enable_ihc", False)
241 +
242 + if get_pp_group().is_first_rank or (
243 + self.config.tie_word_embeddings and get_pp_group().is_last_rank
244 + ):
245 + self.embed_tokens = VocabParallelEmbedding(
246 + config.vocab_size,
247 + config.hidden_size,
248 + )
249 + else:
250 + self.embed_tokens = PPMissingLayer()
251 + self.start_layer, self.end_layer, self.layers = make_layers(
252 + config.num_hidden_layers,
253 + lambda prefix: HYV4DecoderLayer(
254 + config=config,
255 + vllm_config=vllm_config,
256 + cache_config=cache_config,
257 + quant_config=quant_config,
258 + prefix=prefix,
259 + topk_indices_buffer=self.topk_indices_buffer,
260 + ),
261 + prefix=f"{prefix}.layers",
262 + )
263 + if get_pp_group().is_last_rank:
264 + self.norm = RMSNorm(config.hidden_size, config.rms_norm_eps)
265 + else:
266 + self.norm = PPMissingLayer()
267 + if self.enable_ihc:
268 + # iHC head layer: merge the residual channels into one stream.
269 + if get_pp_group().is_last_rank:
270 + self.hc_head = HYV4HCHeadLayer(
271 + config,
272 + hidden_size=config.hidden_size,
273 + hc_mult=config.hc_mult,
274 + hc_eps=config.hc_eps,
275 + prefix=f"{prefix}.hc_head",
276 + )
277 + else:
278 + self.hc_head = PPMissingLayer()
279 + self.make_empty_intermediate_tensors = (
280 + make_empty_intermediate_tensors_factory(
281 + ["hidden_states"], config.hc_mult * config.hidden_size
282 + )
283 + )
284 + else:
285 + self.make_empty_intermediate_tensors = (
286 + make_empty_intermediate_tensors_factory(
287 + ["hidden_states", "residual"], config.hidden_size
288 + )
289 + )
290 +
291 + # MoE hyperparameters (consumed by EPLB).
292 + self.expert_weights: MutableSequence[Sequence[torch.Tensor]] = []
293 + self.num_expert_groups = 1
294 + self.moe_layers: list[nn.Module] = []
295 + example_layer: HYV4MoEFused | None = None
296 + for layer in self.layers:
297 + if isinstance(layer, PPMissingLayer):
298 + continue
299 +
300 + assert isinstance(layer, HYV4DecoderLayer)
301 + if layer.block_type == "moe":
302 + assert isinstance(layer.mlp, HYV4MoEFused)
303 + example_layer = layer.mlp
304 + self.moe_layers.append(layer.mlp.experts)
305 +
306 + if example_layer is None:
307 + self.num_moe_layers = 0
308 + self.num_logical_experts = 0
309 + self.num_physical_experts = 0
310 + self.num_local_physical_experts = 0
311 + self.num_routed_experts = 0
312 + self.num_redundant_experts = 0
313 + return
314 +
315 + self.num_moe_layers = len(self.moe_layers)
316 + self.num_logical_experts = example_layer.n_logical_experts
317 + self.num_physical_experts = example_layer.n_physical_experts
318 + self.num_local_physical_experts = example_layer.n_local_physical_experts
319 + self.num_routed_experts = example_layer.n_routed_experts
320 + self.num_redundant_experts = example_layer.n_redundant_experts
321 +
322 + def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:
323 + return self.embed_tokens(input_ids)
324 +
325 + def update_physical_experts_metadata(
326 + self,
327 + num_physical_experts: int,
328 + num_local_physical_experts: int,
329 + ) -> None:
330 + assert self.num_local_physical_experts == num_local_physical_experts
331 + self.num_physical_experts = num_physical_experts
332 + self.num_local_physical_experts = num_local_physical_experts
333 + self.num_redundant_experts = num_physical_experts - self.num_logical_experts
334 + for layer in self.layers:
335 + if isinstance(layer, PPMissingLayer):
336 + continue
337 + if isinstance(layer.mlp, HYV4MoEFused):
338 + moe = layer.mlp
339 + moe.n_local_physical_experts = num_local_physical_experts
340 + moe.n_physical_experts = num_physical_experts
341 + moe.n_redundant_experts = self.num_redundant_experts
342 + moe.experts.update_expert_map()
343 +
344 + def get_expert_mapping(self) -> list[tuple[str, str, int, str]]:
345 + # (param_name, weight_name, expert_id, shard_id) for weights, fp8
346 + # weight scales and fp8 activation scales.
347 + if not hasattr(self, "_cached_expert_params_mapping"):
348 + self._cached_expert_params_mapping = fused_moe_make_expert_params_mapping(
349 + self,
350 + ckpt_gate_proj_name="gate_proj",
351 + ckpt_down_proj_name="down_proj",
352 + ckpt_up_proj_name="up_proj",
353 + num_experts=self.config.num_experts,
354 + num_redundant_experts=self.num_redundant_experts,
355 + )
356 + return self._cached_expert_params_mapping
357 +
358 + def forward(
359 + self,
360 + input_ids: torch.Tensor,
361 + positions: torch.Tensor,
362 + intermediate_tensors: IntermediateTensors | None = None,
363 + inputs_embeds: torch.Tensor | None = None,
364 + ) -> torch.Tensor | IntermediateTensors:
365 + if get_pp_group().is_first_rank:
366 + if inputs_embeds is not None:
367 + hidden_states = inputs_embeds
368 + else:
369 + hidden_states = self.embed_input_ids(input_ids)
370 + residual = None
371 + else:
372 + assert intermediate_tensors is not None
373 + hidden_states = intermediate_tensors["hidden_states"]
374 + # In iHC mode the flattened [num_tokens, hc*h] tensor from the
375 + # previous PP stage is reshaped back to 3D by the first layer's
376 + # prepare_input, and residual is unused.
377 + residual = None if self.enable_ihc else intermediate_tensors["residual"]
378 +
379 + for layer in islice(self.layers, self.start_layer, self.end_layer):
380 + hidden_states, residual = layer(positions, hidden_states, residual)
381 +
382 + if not get_pp_group().is_last_rank:
383 + if self.enable_ihc:
384 + # hidden_states is [num_tokens, hc, h]; flatten the channel dim
385 + # for PP transfer (matches the 2D receive buffer).
386 + return IntermediateTensors({"hidden_states": hidden_states.flatten(1)})
387 + return IntermediateTensors(
388 + {"hidden_states": hidden_states, "residual": residual}
389 + )
390 + if self.enable_ihc:
391 + hidden_states = self.hc_head(hidden_states)
392 + else:
393 + hidden_states = hidden_states + residual
394 +
395 + return self.norm(hidden_states)
396 +
397 + def load_fused_expert_weights(
398 + self,
399 + name: str,
400 + params_dict: dict,
401 + loaded_weight: torch.Tensor,
402 + shard_id: str,
403 + num_experts: int,
404 + ) -> bool:
405 + param = params_dict[name]
406 + weight_loader = typing.cast(Callable[..., bool], param.weight_loader)
407 + loaded_local_expert = False
408 + for expert_id in range(num_experts):
409 + curr_expert_weight = loaded_weight[expert_id]
410 + success = weight_loader(
411 + param,
412 + curr_expert_weight,
413 + name,
414 + shard_id,
415 + expert_id,
416 + return_success=True,
417 + )
418 + if success:
419 + loaded_local_expert = True
420 +
421 + return loaded_local_expert
422 +
423 + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
424 + # FP8 indexer wk dequant buffer (weight and scale arrive separately).
425 + pending_wk_fp8: dict[str, dict[str, torch.Tensor]] = {}
426 + pp_missing_layer_names = get_pp_missing_layer_names(self)
427 + skip_topk_layers = compute_skip_topk_layers(self.config)
428 +
429 + # Must not be cached on `self`: `process_weights_after_loading` swaps in
430 + # kernel-specific expert weights via `replace_parameter`, and a cache
431 + # outliving this call would pin the pre-shuffle storage (OOM on large
432 + # MoE) and make a later reload target orphaned tensors.
433 + params_dict = dict(self.named_parameters())
434 + # Split per-expert mapping (V3 style): experts.0.gate_proj.weight
435 + split_expert_params_mapping = self.get_expert_mapping()
436 + loaded_params: set[str] = set()
437 +
438 + # Sink weights are sharded like the q/k/v linears.
439 + sink_tp_size = get_tensor_model_parallel_world_size()
440 + sink_tp_rank = get_tensor_model_parallel_rank()
441 + n_local_head = self.config.num_attention_heads // sink_tp_size
442 + head_rank_start = n_local_head * sink_tp_rank
443 + head_rank_end = n_local_head * (sink_tp_rank + 1)
444 + base_layer = (
445 + "base_layer." if any(".base_layer." in name for name in params_dict) else ""
446 + )
447 + # Fused expert mapping: experts.gate_up_proj (all experts in one tensor).
448 + # The packed weights are owned by the RoutedExperts submodule, so the
449 + # targets are experts.routed_experts.[base_layer.]w{13,2}_weight.
450 + fused_expert_prefix = f".experts.routed_experts.{base_layer}"
451 + fused_expert_params_mapping = [
452 + (f"{fused_expert_prefix}w13_weight", ".experts.gate_up_proj", 0, "w1"),
453 + (f"{fused_expert_prefix}w2_weight", ".experts.down_proj", 0, "w2"),
454 + ]
455 + num_experts = getattr(self.config, "num_experts", 0)
456 +
457 + def _should_skip_missing_param(param_name: str) -> bool:
458 + # Sparse checkpoints may contain indexer weights for layers that
459 + # fall back to dense MLA on unsupported GPUs.
460 + return ".indexer." in param_name
461 +
462 + def _is_split_expert_weight(weight_name: str) -> bool:
463 + """Whether this weight is in split (per-expert) format."""
464 + # Split format: mlp.experts.<id>.gate_proj.weight
465 + return bool(re.search(r"\.experts\.\d+\.", weight_name))
466 +
467 + def _is_fused_expert_weight(weight_name: str) -> bool:
468 + """Whether this weight is in fused (all-experts-packed) format."""
469 + return ".experts.gate_up_proj" in weight_name or (
470 + ".experts.down_proj" in weight_name
471 + and not _is_split_expert_weight(weight_name)
472 + )
473 +
474 + for name, loaded_weight in weights:
475 + if is_skip_topk_indexer_weight(name, skip_topk_layers):
476 + continue
477 + if _try_load_fp8_indexer_wk(
478 + name,
479 + loaded_weight,
480 + pending_wk_fp8,
481 + params_dict,
482 + loaded_params,
483 + pp_missing_layer_names,
484 + ):
485 + continue
486 + if self.config.tie_word_embeddings and "lm_head.weight" in name:
487 + continue
488 + # Names and stacked shard ids are normalized upstream by
489 + # AutoWeightsLoader, through quant_config.get_cache_scale_mapper()
490 + # and hf_to_vllm_mapper.
491 + stacked_shard_id = getattr(loaded_weight, "shard_id", None)
492 + if ".indexer.wk." in name:
493 + name = name.replace(".indexer.wk.", ".indexer.wk_weights_proj.")
494 + stacked_shard_id = 0
495 + # Skip loading extra bias for GPTQ models.
496 + if name.endswith(".bias") and name not in params_dict:
497 + continue
498 +
499 + # Determine per-weight whether this is fused or split format.
500 + is_fused_expert = _is_fused_expert_weight(name)
501 + expert_params_mapping = (
502 + fused_expert_params_mapping
503 + if is_fused_expert
504 + else split_expert_params_mapping
505 + )
506 +
507 + is_expert_weight = False
508 + loaded_expert_param_names: set[str] = set()
509 + for mapping in expert_params_mapping:
510 + param_name, weight_name, expert_id, shard_id = mapping
511 + if weight_name not in name:
512 + continue
513 +
514 + # This is an expert weight and must not be attempted as any
515 + # other kind of weight later on.
516 + is_expert_weight = True
517 +
518 + # Do not modify `name`: the loop may continue past this point.
519 + name_mapped = name.replace(weight_name, param_name)
520 + if is_pp_missing_parameter(name_mapped, self):
521 + continue
522 + if is_fused_expert:
523 + if "experts.gate_up_proj" in name:
524 + chunks = loaded_weight.chunk(2, dim=-2)
525 + success_w1 = self.load_fused_expert_weights(
526 + name_mapped, params_dict, chunks[0], "w1", num_experts
527 + )
528 + success_w3 = self.load_fused_expert_weights(
529 + name_mapped, params_dict, chunks[1], "w3", num_experts
530 + )
531 + success = success_w1 and success_w3
532 + else:
533 + success = self.load_fused_expert_weights(
534 + name_mapped,
535 + params_dict,
536 + loaded_weight,
537 + shard_id,
538 + num_experts,
539 + )
540 + if success:
541 + name = name_mapped
542 + break
543 + else:
544 + # Split per-expert format (V3 style).
545 + if name_mapped not in params_dict:
546 + if _should_skip_missing_param(name_mapped):
547 + continue
548 + logger.warning_once(
549 + "Skipping unknown checkpoint weight: %s",
550 + name_mapped,
551 + )
552 + continue
553 + param = params_dict[name_mapped]
554 + # Ask the weight loader whether it succeeded, otherwise we
555 + # may skip experts that have other available replicas.
556 + weight_loader: Callable[..., typing.Any] = param.weight_loader
557 + success = weight_loader(
558 + param,
559 + loaded_weight,
560 + name_mapped,
561 + shard_id=shard_id,
562 + expert_id=expert_id,
563 + return_success=True,
564 + )
565 + if success:
566 + if not is_fused_expert:
567 + loaded_expert_param_names.add(name_mapped)
568 + continue
569 + name = name_mapped
570 + break
571 + else:
572 + if loaded_expert_param_names:
573 + loaded_params.update(loaded_expert_param_names)
574 + continue
575 + if "learnable_sink_param" in name:
576 + if is_pp_missing_parameter(name, self):
577 + continue
578 + narrow_weight = loaded_weight[head_rank_start:head_rank_end]
579 + n = narrow_weight.shape[0]
580 + with torch.no_grad():
581 + params_dict[name][:n].copy_(narrow_weight)
582 + else:
583 + if is_expert_weight:
584 + # An expert weight that is not mapped to this rank.
585 + continue
586 + if "gate.e_score_correction_bias" in name:
587 + name = name.replace(
588 + "gate.e_score_correction_bias", "expert_bias"
589 + )
590 + if "router.gate." in name:
591 + name = name.replace("router.", "")
592 + if "hc_fn" in name:
593 + name = name.replace("hc_fn", "hc_fn.weight")
594 + if "hc_head_fn" in name:
595 + name = name.replace("hc_head_fn", "hc_head_fn.weight")
596 +
597 + if is_pp_missing_parameter(name, self):
598 + continue
599 +
600 + if name not in params_dict:
601 + if _should_skip_missing_param(name):
602 + continue
603 + logger.warning_once(
604 + "Skipping unknown checkpoint weight: %s",
605 + name,
606 + )
607 + continue
608 +
609 + param = params_dict[name]
610 + weight_loader = getattr(
611 + param, "weight_loader", default_weight_loader
612 + )
613 + if stacked_shard_id is None:
614 + weight_loader(param, loaded_weight)
615 + else:
616 + weight_loader(param, loaded_weight, stacked_shard_id)
617 + loaded_params.add(name)
618 +
619 + return loaded_params
620 +
621 +
622 +class HYV4ForCausalLM(nn.Module, SupportsPP, SupportsLoRA):
623 + hf_to_vllm_mapper = WeightsMapper(
624 + orig_to_new_stacked={
625 + ".q_a_proj": (".fused_qkv_a_proj", 0),
626 + ".kv_a_proj_with_mqa": (".fused_qkv_a_proj", 1),
627 + ".mlp.gate_proj": (".mlp.gate_up_proj", 0),
628 + ".mlp.up_proj": (".mlp.gate_up_proj", 1),
629 + ".shared_experts.gate_proj": (".shared_experts.gate_up_proj", 0),
630 + ".shared_experts.up_proj": (".shared_experts.gate_up_proj", 1),
631 + ".indexer.weights_proj.": (".indexer.wk_weights_proj.", 1),
632 + }
633 + )
634 + packed_modules_mapping = {
635 + "gate_up_proj": ["gate_proj", "up_proj"],
636 + # MLA runs both latent down-projections as one GEMM.
637 + "fused_qkv_a_proj": ["q_a_proj", "kv_a_proj_with_mqa"],
638 + # The indexer fuses wk and weights_proj into one GEMM.
639 + "wk_weights_proj": ["wk", "weights_proj"],
640 + }
641 +
642 + def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
643 + super().__init__()
644 + config = vllm_config.model_config.hf_config
645 + quant_config = vllm_config.quant_config
646 + self.config = config
647 + self.quant_config = quant_config
648 + if quant_config is not None:
649 + quant_config.packed_modules_mapping = self.packed_modules_mapping
650 +
651 + parallel_config = vllm_config.parallel_config
652 + eplb_config = parallel_config.eplb_config
653 + self.num_redundant_experts = eplb_config.num_redundant_experts
654 +
655 + self.model = HYV4Model(
656 + vllm_config=vllm_config, prefix=maybe_prefix(prefix, "model")
657 + )
658 + if get_pp_group().is_last_rank:
659 + # The head stays in the model dtype; ``enable_lm_head_fp32`` is
660 + # surfaced as ``head_dtype`` on the config (see HYV4Config), which
661 + # makes LogitsProcessor accumulate the projection straight into
662 + # fp32 instead of materializing an fp32 copy of the weight.
663 + self.lm_head = ParallelLMHead(
664 + config.vocab_size,
665 + config.hidden_size,
666 + quant_config=quant_config,
667 + prefix=maybe_prefix(prefix, "lm_head"),
668 + )
669 + # With tie_word_embeddings, embed_tokens is kept on the last rank
670 + # (see HYV4Model.__init__) so the weight can be shared here.
671 + if self.config.tie_word_embeddings:
672 + self.lm_head.weight = self.model.embed_tokens.weight
673 + else:
674 + self.lm_head = PPMissingLayer()
675 + self.logits_processor = LogitsProcessor(config.vocab_size)
676 + self.make_empty_intermediate_tensors = (
677 + self.model.make_empty_intermediate_tensors
678 + )
679 +
680 + def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:
681 + return self.model.embed_input_ids(input_ids)
682 +
683 + def forward(
684 + self,
685 + input_ids: torch.Tensor,
686 + positions: torch.Tensor,
687 + intermediate_tensors: IntermediateTensors | None = None,
688 + inputs_embeds: torch.Tensor | None = None,
689 + ) -> torch.Tensor | IntermediateTensors:
690 + return self.model(input_ids, positions, intermediate_tensors, inputs_embeds)
691 +
692 + def compute_logits(
693 + self,
694 + hidden_states: torch.Tensor,
695 + ) -> torch.Tensor | None:
696 + logits = self.logits_processor(self.lm_head, hidden_states)
697 +
698 + if getattr(self.config, "soft_logits_capping", False):
699 + soft_cap = self.config.soft_logits_capping_logits
700 + logits = soft_cap * torch.nn.functional.tanh(logits / soft_cap)
701 +
702 + return logits
703 +
704 + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
705 + # MTP layers are loaded by the MTP head, not by the target model. They
706 + # appear as model.mtp_layers.<i>.* and as the layer ids past the backbone.
707 + mtp_start = self.config.num_hidden_layers
708 + drop_prefixes: dict[str, str | None] = {"model.mtp_layers.": None}
709 + for i in range(getattr(self.config, "num_nextn_predict_layers", 0)):
710 + drop_prefixes[f"model.layers.{mtp_start + i}."] = None
711 + if self.config.tie_word_embeddings:
712 + drop_prefixes["lm_head."] = None
713 + drop_mapper = WeightsMapper(orig_to_new_prefix=drop_prefixes)
714 + loader = AutoWeightsLoader(self)
715 + return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper | drop_mapper)
716 +
717 + def get_expert_mapping(self) -> list[tuple[str, str, int, str]]:
718 + return self.model.get_expert_mapping()