class Qwen4ExpMultiTokenPredictor(nn.Module):
hf_to_vllm_mapper = Qwen3_5Model.hf_to_vllm_mapper | _EXTRA_WEIGHTS_MAPPER
def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None:
super().__init__()
model_config = vllm_config.model_config
config: Qwen4ExpTextConfig = model_config.hf_text_config
self.config = config
self.vocab_size = config.vocab_size
self.mtp_start_layer_idx = config.num_hidden_layers
self.num_mtp_layers = getattr(config, "mtp_num_hidden_layers", 1)
self.hidden_size = config.hidden_size
self.hc_count = config.hc_count
self.embed_tokens = VocabParallelEmbedding(self.vocab_size, self.hidden_size)
draft_vllm_config = _make_draft_vllm_config(
vllm_config,
self.mtp_start_layer_idx,
)
with set_current_vllm_config(draft_vllm_config, prefix=prefix):
# residual_linear_shared fusion: fc_embedding projects the token
# embedding, fc_hidden (shared across HC branches) projects the
# backbone hidden; the embedding is added as a residual to every
# branch (see mtp_residual_linear_shared.md).
self.fc_embedding = ColumnParallelLinear(
self.hidden_size,
self.hidden_size,
gather_output=True,
bias=False,
return_bias=False,
quant_config=draft_vllm_config.quant_config,
prefix=f"{prefix}.fc_embedding",
)
self.fc_hidden = ColumnParallelLinear(
self.hidden_size,
self.hidden_size,
gather_output=True,
bias=False,
return_bias=False,
quant_config=draft_vllm_config.quant_config,
prefix=f"{prefix}.fc_hidden",
)
self.layers = nn.ModuleList(
Qwen4ExpDecoderLayer(
draft_vllm_config,
layer_type="full_attention",
prefix=f"{prefix}.layers.{self.mtp_start_layer_idx + idx}",
)
for idx in range(self.num_mtp_layers)
)
self.is_fused_shared_expert_enabled = is_model_fused_shared_expert_compatible(
self.layers,
Qwen4ExpSparseMoeBlock,
"mlp",
)
self.pre_fc_norm_embedding = GemmaRMSNorm(
self.hidden_size, eps=config.rms_norm_eps
)
self.pre_fc_norm_hidden = GemmaRMSNorm(
self.hidden_size * self.hc_count, eps=config.rms_norm_eps
)
# HC final mixer collapses the multi stream into [T, H] for the LM head.
hc_config = HyperConnectionConfig(
hc_count=config.hc_count,
hidden_size=config.hidden_size,
params_dtype=torch.bfloat16,
hc_lowrank=config.hc_lowrank,
rms_norm_eps=config.rms_norm_eps,
hc_per_branch_norm=True,
)
self.hyper_connection_mixer = GatedResidual(
hc_config,
use_combine=False,
prefix=maybe_prefix(prefix, "hyper_connection_mixer"),
)
self.make_empty_intermediate_tensors = make_empty_intermediate_tensors_factory(
["hidden_states"], self.hidden_size * self.hc_count
)
def _iter_qsa_attentions(self):
"""Yield MTP attention modules that own a QSA indexer."""
for layer in self.layers:
attention = getattr(layer, "self_attn", None)
if (
attention is not None
and getattr(attention, "indexer", None) is not None
):
yield attention
def set_skip_topk(self, skip: bool) -> None:
"""Select on MTP step 0 and reuse its QSA indices on later steps."""
for attention in self._iter_qsa_attentions():
attention.indexer.skip_topk = skip
def compact_topk_indices(self, row_indices: torch.Tensor) -> None:
"""Keep each request's target-aligned step-0 sparse-index row."""
num_rows = row_indices.numel()
for attention in self._iter_qsa_attentions():
buffer = attention.topk_indices_buffer
selected = buffer.index_select(0, row_indices)
buffer[:num_rows].copy_(selected)
def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:
return self.embed_tokens(input_ids)
def forward(
self,
input_ids: torch.Tensor | None,
positions: torch.Tensor,
hidden_states: torch.Tensor | None = None,
intermediate_tensors: IntermediateTensors | None = None,
inputs_embeds: torch.Tensor | None = None,
spec_step_idx: int = 0,
) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor] | IntermediateTensors:
hc_count = self.hc_count
hidden_size = self.hidden_size
prev_block_output: torch.Tensor | None = None
if get_pp_group().is_first_rank:
assert hidden_states is not None
if inputs_embeds is None:
assert input_ids is not None
inputs_embeds = self.embed_input_ids(input_ids)
# Embedding branch: pre-norm -> fc_embedding -> [T, H].
inputs_embeds = self.pre_fc_norm_embedding(inputs_embeds)
inputs_embeds = self.fc_embedding(inputs_embeds)
# Backbone hidden is multi-stream [T, hc_count*H] (scheme A:
# the main model truly emits the pre-final-mixer multi stream
# on the first step; subsequent steps reuse the prior draft
# step's multi stream).
num_tokens = hidden_states.shape[0]
hidden_states = hidden_states.view(num_tokens, hc_count, hidden_size)
hidden_states = self.pre_fc_norm_hidden(hidden_states.flatten(-2)).view(
num_tokens, hc_count, hidden_size
)
hidden_states = self.fc_hidden(hidden_states)
hidden_states = hidden_states.flatten(-2)
prev_block_output = inputs_embeds
else:
assert intermediate_tensors is not None
hidden_states = intermediate_tensors["hidden_states"]
current_step_idx = spec_step_idx % self.num_mtp_layers
layer = self.layers[current_step_idx]
hidden_states, block_output, injection = layer(
hidden_states=hidden_states,
prev_block_output=prev_block_output,
prev_injection=None,
positions=positions,
input_ids=None,
query_start_loc=None,
ngram_context=None,
)
if not get_pp_group().is_last_rank:
# As in the target model, PP carries a materialized tensor rather
# than the delayed hidden/output/injection tuple.
hidden_states = layer.mlp_hyper_connection.combine(
hidden_states, block_output, injection
)
return IntermediateTensors({"hidden_states": hidden_states})
# Last PP rank finalize. Keep both:
# (A) sample_hidden_states [T, H] -> single stream for the LM head
# (B) multi_hidden [T, hc_count*H] -> pre-final-mixer multi stream
# for the next draft step (zero extra compute, just kept).
multi_hidden, sample_hidden_states, _ = (
self.hyper_connection_mixer.combine_and_mix(
hidden_states, block_output, injection
)
)
return sample_hidden_states, multi_hidden
def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
weights = maybe_fuse_shared_experts(
weights,
enabled=self.is_fused_shared_expert_enabled,
n_routed_experts=getattr(self.config, "num_experts", 0) or 0,
n_shared_experts=1,
ckpt_prefix="mlp.shared_expert",
)
mapper = self.hf_to_vllm_mapper | WeightsMapper(
orig_to_new_substr={"hyper_connection_mixer.block_inject_weight": None}
)
loader = AutoWeightsLoader(
self,
ignore_unexpected_suffixes=_QWEN4_EXP_IGNORED_MISSING_SUFFIXES.copy(),
)
return loader.load_weights(weights, mapper=mapper)