class DeepseekV4Model(nn.Module, EagleModelMixin):
supports_aux_hidden_states_over_pp = True
def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
super().__init__()
config = vllm_config.model_config.hf_config
quant_config = vllm_config.quant_config
self.config = config
self.quant_config = quant_config
self.parallel_config = vllm_config.parallel_config
self.use_native_mega_moe = (
vllm_config.kernel_config.moe_backend in NATIVE_MEGA_MOE_BACKENDS
)
self.use_sequence_parallel = _use_sequence_parallel(vllm_config)
if (
self.use_native_mega_moe
and not vllm_config.parallel_config.enable_expert_parallel
):
raise NotImplementedError(
"DeepSeek V4 MegaMoE currently requires expert parallel. "
"Enable it with --enable-expert-parallel, or pick a different "
"moe backend."
)
self.vocab_size = config.vocab_size
self.hc_eps = config.hc_eps
self.hc_mult = config.hc_mult
self.hc_dim = self.hc_mult * config.hidden_size
self.rms_norm_eps = config.rms_norm_eps
# Three aux streams: one per non-default input GEMM in
# DeepseekV4Attention._run_parallel_input_projections
# (compressor kv_score, indexer.weights_proj, indexer.compressor
# kv_score). fused_wqa_wkv stays on the default stream.
aux_stream_list = [torch.cuda.Stream() for _ in range(3)]
# Reserved topk indices buffer for all Indexer layers to reuse.
self.topk_indices_buffer = torch.empty(
vllm_config.scheduler_config.max_num_batched_tokens,
config.index_topk,
dtype=torch.int32,
)
if get_pp_group().is_first_rank or spec_decode_needs_target_embed(vllm_config):
self.embed_tokens = VocabParallelEmbedding(
config.vocab_size,
config.hidden_size,
quant_config=quant_config,
prefix=f"{prefix}.embed_tokens",
)
else:
self.embed_tokens = PPMissingLayer()
self.start_layer, self.end_layer, self.layers = make_layers(
config.num_hidden_layers,
lambda prefix: DeepseekV4DecoderLayer(
vllm_config,
prefix=prefix,
topk_indices_buffer=self.topk_indices_buffer,
aux_stream_list=aux_stream_list,
),
prefix=f"{prefix}.layers",
)
self.has_local_hash_moe = self.start_layer < min(
self.end_layer, config.mlp_layer_types.count("hash_moe")
)
if get_pp_group().is_last_rank:
self.norm = RMSNorm(config.hidden_size, self.rms_norm_eps)
else:
self.norm = PPMissingLayer()
self.hc_head_fn = nn.Parameter(
torch.empty(
self.hc_mult,
self.hc_dim,
dtype=torch.float32,
),
requires_grad=False,
)
self.hc_head_base = nn.Parameter(
torch.empty(
self.hc_mult,
dtype=torch.float32,
),
requires_grad=False,
)
self.hc_head_scale = nn.Parameter(
torch.empty(1, dtype=torch.float32),
requires_grad=False,
)
spec_config = vllm_config.speculative_config
needs_mtp_hidden_states = spec_config is not None and (
spec_config.use_eagle() or spec_config.uses_draft_model()
)
if get_pp_group().is_last_rank and needs_mtp_hidden_states:
self._mtp_hidden_buffer = torch.empty(
vllm_config.scheduler_config.max_num_batched_tokens,
self.hc_dim,
dtype=vllm_config.model_config.dtype,
)
else:
self._mtp_hidden_buffer = None
if vllm_config.kernel_config.enable_jit_warmup and get_pp_group().is_last_rank:
from vllm.model_executor.kernels.mhc.tilelang_kernels import (
_HC_HEAD_FUSED_TILELANG_KERNEL,
)
_HC_HEAD_FUSED_TILELANG_KERNEL.register_warmup(
hidden_size=config.hidden_size,
hc_mult=self.hc_mult,
rms_eps=self.rms_norm_eps,
hc_eps=self.hc_eps,
)
def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:
return self.embed_tokens(input_ids)
def make_empty_intermediate_tensors(
self,
batch_size: int,
dtype: torch.dtype,
device: torch.device,
) -> IntermediateTensors:
# PP intermediate tensors carry the multi-stream hidden_states
# of shape (num_tokens, hc_mult, hidden_size) — V4 expands the
# token embedding to hc_mult streams before the first decoder
# layer and keeps that shape until hc_head() collapses it.
return IntermediateTensors(
{
"hidden_states": torch.zeros(
(batch_size, self.hc_mult, self.config.hidden_size),
dtype=dtype,
device=device,
),
}
)
def forward(
self,
input_ids: torch.Tensor,
positions: torch.Tensor,
intermediate_tensors: IntermediateTensors | None,
inputs_embeds: torch.Tensor | None = None,
) -> torch.Tensor | IntermediateTensors:
if get_pp_group().is_first_rank:
if inputs_embeds is not None:
hidden_states = inputs_embeds
else:
hidden_states = self.embed_input_ids(input_ids)
else:
assert intermediate_tensors is not None
hidden_states = intermediate_tensors["hidden_states"]
if self.use_native_mega_moe:
input_ids = input_ids.to(torch.int64)
full_num_tokens = positions.shape[0]
if self.use_sequence_parallel:
if envs.VLLM_MOE_SKIP_PADDING and is_forward_context_available():
forward_context = get_forward_context()
forward_context.is_padding = sp_padding_mask(
forward_context.is_padding, hidden_states
)
hidden_states = sp_shard(hidden_states)
input_ids = sp_shard(input_ids)
mega_gate_metadata = None
if self.use_native_mega_moe:
mega_gate_metadata = prepare_mega_gate_routing_metadata(
input_ids,
has_hash_routing=self.has_local_hash_moe,
image_sentinel_base_id=IMAGE_SENTINEL_BASE_ID
if getattr(self.config, "vision_n_layers", 0) > 0
else None,
)
residual, post_mix, res_mix = None, None, None
remote_aux = self.collect_remote_aux_hidden_states(intermediate_tensors)
aux_hidden_states: list[torch.Tensor] = []
final_aux_recon: torch.Tensor | None = None # avoid duplicate mhc_post call
for idx, layer in enumerate(
islice(self.layers, self.start_layer, self.end_layer),
start=self.start_layer,
):
hidden_states, residual, post_mix, res_mix = layer(
hidden_states,
positions,
input_ids,
post_mix,
res_mix,
residual,
mega_gate_metadata=mega_gate_metadata,
)
if idx + 1 in self.aux_hidden_state_layers:
# Reconstruct the aux hidden state for draft models
aux_recon = mhc_post_tilelang(
hidden_states, residual, post_mix, res_mix
)
aux_hidden_state = aux_recon.mean(dim=1)
if self.use_sequence_parallel:
aux_hidden_state = sp_all_gather(aux_hidden_state)[:full_num_tokens]
aux_hidden_states.append(aux_hidden_state)
final_aux_recon = aux_recon
if layer is not None:
# Reuse if the last layer was captured as an aux hidden state
if self.end_layer in self.aux_hidden_state_layers:
hidden_states = final_aux_recon
else:
hidden_states = mhc_post_tilelang(
hidden_states, residual, post_mix, res_mix
)
if not get_pp_group().is_last_rank:
return IntermediateTensors(
{
"hidden_states": hidden_states,
**self.pack_local_aux_hidden_states(aux_hidden_states),
}
)
if self.use_sequence_parallel:
hidden_states = sp_all_gather(hidden_states)[:full_num_tokens]
if self._mtp_hidden_buffer is not None:
num_tokens = hidden_states.shape[0]
self._mtp_hidden_buffer[:num_tokens].copy_(hidden_states.flatten(1))
hidden_states = hc_head_fused_kernel_tilelang(
hidden_states,
self.hc_head_fn,
self.hc_head_scale,
self.hc_head_base,
self.rms_norm_eps,
self.hc_eps,
)
hidden_states = self.norm(hidden_states)
aux_hidden_states = remote_aux + aux_hidden_states
if len(aux_hidden_states) > 0:
return hidden_states, aux_hidden_states
return hidden_states
def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
stacked_params_mapping = [
# (param_name, shard_name, shard_id)
("gate_up_proj", "w1", 0),
("gate_up_proj", "w3", 1),
("attn.fused_wqa_wkv", "attn.wq_a", 0),
("attn.fused_wqa_wkv", "attn.wkv", 1),
("compressor.fused_wkv_wgate", "compressor.wkv", 0),
("compressor.fused_wkv_wgate", "compressor.wgate", 1),
]
params_dict = dict(self.named_parameters())
loaded_params: set[str] = set()
# Pre-compute expert mapping ONCE.
expert_mapping = self.get_expert_mapping()
# Block-FP8 shared experts: pad the intermediate up to the TP-uniform
# block count so the standard loaders below slice it evenly (trailing
# ranks land on the zero pad). SP / unquantized ones need no padding.
pad_shared_expert = (
getattr(self.quant_config, "weight_block_size", None) is not None
and not self.use_sequence_parallel
)
for name, loaded_weight in weights:
if ".shared_experts." in name:
if is_pp_missing_parameter(name, self):
continue
experts = self.layers[extract_layer_index(name)].ffn.experts
if getattr(experts, "has_fused_shared_experts", False):
loaded_params.add(
experts.load_shared_expert_weight(name, loaded_weight)
)
continue
if pad_shared_expert and ".shared_experts." in name:
loaded_weight = self._pad_shared_expert_weight(
self.quant_config, name, loaded_weight
)
for param_name, weight_name, shard_id in stacked_params_mapping:
# Skip non-stacked layers and experts (experts handled below).
if ".experts." in name:
continue
if weight_name not in name:
continue
name = name.replace(weight_name, param_name)
if is_pp_missing_parameter(name, self):
break
if name not in params_dict:
head, _, leaf = name.rpartition(".")
suffixed = f"{head}.base_layer.{leaf}"
if suffixed in params_dict:
name = suffixed
param = params_dict[name]
weight_loader = param.weight_loader
weight_loader(param, loaded_weight, shard_id)
loaded_params.add(name)
break
else:
if ".experts." in name:
# E8M0 scales are stored as float8_e8m0fnu in
# checkpoints but the MoE param is uint8. copy_()
# would do a numeric conversion (e.g. 2^-7 → 0),
# destroying the raw exponent bytes.
if (
"weight_scale" in name
and loaded_weight.dtype == torch.float8_e8m0fnu
):
loaded_weight = loaded_weight.view(torch.uint8)
for mapping in expert_mapping:
param_name, weight_name, expert_id, expert_shard_id = mapping
if weight_name not in name:
continue
name_mapped = name.replace(weight_name, param_name)
if is_pp_missing_parameter(name_mapped, self):
continue
param = params_dict[name_mapped]
# We should ask the weight loader to return success or not
# here since otherwise we may skip experts with other
# available replicas.
weight_loader = typing.cast(
Callable[..., bool], param.weight_loader
)
success = weight_loader(
param,
loaded_weight,
name_mapped,
shard_id=expert_shard_id,
expert_id=expert_id,
return_success=True,
)
if success:
name = name_mapped
break
loaded_params.add(name_mapped)
continue
else:
if is_pp_missing_parameter(name, self):
continue
# Non-LoRA params on a LoRA-wrapped module live at
# ``<head>.base_layer.<leaf>``; the checkpoint is plain.
if name not in params_dict:
head, _, leaf = name.rpartition(".")
suffixed = f"{head}.base_layer.{leaf}"
if suffixed in params_dict:
name = suffixed
param = params_dict[name]
weight_loader = getattr(
param, "weight_loader", default_weight_loader
)
weight_loader(param, loaded_weight)
loaded_params.add(name)
continue
return loaded_params
@staticmethod
def _pad_shared_expert_weight(
quant_config: QuantizationConfig | None,
name: str,
loaded_weight: torch.Tensor,
) -> torch.Tensor:
"""Zero-pad a block-FP8 shared-expert weight/scale on its intermediate
axis so the standard TP loaders split it into even, block-aligned shards
(trailing ranks get the zero pad). gate (w1)/up (w3) [I, H] pad dim 0;
down (w2 -> down_proj) [H, I] pads dim 1.
"""
block_size = getattr(quant_config, "weight_block_size", None)
assert block_size is not None
# Round the intermediate axis up to a whole number of TP shards. The axis
# is in elements for weights (step = block) and in blocks for scales.
step = 1 if name.endswith("weight_scale_inv") else block_size[0]
dim = 1 if ".down_proj." in name else 0
mult = get_tensor_model_parallel_world_size() * step
pad = cdiv(loaded_weight.shape[dim], mult) * mult - loaded_weight.shape[dim]
if pad == 0:
return loaded_weight
pad_shape = list(loaded_weight.shape)
pad_shape[dim] = pad
return torch.cat([loaded_weight, loaded_weight.new_zeros(pad_shape)], dim=dim)
def get_expert_mapping(self) -> list[tuple[str, str, int, str]]:
first_layer = next(iter(islice(self.layers, self.start_layer, self.end_layer)))
if first_layer.ffn.use_native_mega_moe:
return make_deepseek_v4_expert_params_mapping(self.config.n_routed_experts)
# Params for weights, fp8 weight scales, fp8 activation scales
# (param_name, weight_name, expert_id, shard_id)
return fused_moe_make_expert_params_mapping(
self,
ckpt_gate_proj_name="w1",
ckpt_down_proj_name="w2",
ckpt_up_proj_name="w3",
num_experts=self.config.n_routed_experts,
)
def finalize_mhc_broadcast_weights(self) -> None:
if not get_pp_group().is_first_rank or self.start_layer >= self.end_layer:
return
layer = self.layers[self.start_layer]
if isinstance(layer, DeepseekV4DecoderLayer):
broadcast = (
layer.hc_attn_fn.detach()
.view(-1, layer.hc_mult, layer.hidden_size)
.sum(dim=1)
)
if layer.hc_attn_fn_broadcast is None:
layer.hc_attn_fn_broadcast = broadcast
else:
layer.hc_attn_fn_broadcast.copy_(broadcast)