class KimiK3MTP(nn.Module):
def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
super().__init__()
self.config = vllm_config.model_config.hf_text_config
self.quant_config = vllm_config.quant_config
self.model = KimiK3MultiTokenPredictor(
vllm_config=vllm_config, prefix=maybe_prefix(prefix, "model")
)
def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:
return self.model.embed_input_ids(input_ids)
def forward(
self,
input_ids: torch.Tensor | None,
positions: torch.Tensor,
hidden_states: torch.Tensor,
intermediate_tensors: IntermediateTensors | None = None,
inputs_embeds: torch.Tensor | None = None,
spec_step_idx: int = 0,
) -> tuple[torch.Tensor, torch.Tensor]:
return self.model(
input_ids,
positions,
hidden_states,
inputs_embeds,
spec_step_idx,
)
def compute_logits(
self,
hidden_states: torch.Tensor,
spec_step_idx: int = 0,
) -> torch.Tensor | None:
return self.model.compute_logits(hidden_states, spec_step_idx)
def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
# Mirror KimiLinearForCausalLM.load_weights naming: leading-dot shard
# names, q_lora-conditional fused QKV, and w1/w2/w3 expert weights.
kda_config = self.config.linear_attn_config
use_full_rank_gate = bool(
kda_config and kda_config.get("use_full_rank_gate", False)
)
beta_shard_id = 5 if use_full_rank_gate else 3
stacked_params_mapping = [
# (param_name, shard_name, shard_id)
(".in_proj_qkvgfab", ".q_proj", 0),
(".in_proj_qkvgfab", ".k_proj", 1),
(".in_proj_qkvgfab", ".v_proj", 2),
(".in_proj_qkvgfab", ".b_proj", beta_shard_id),
(".in_proj_qkvgfab", ".f_a_proj", 4),
(".conv1d", ".q_conv1d", 0),
(".conv1d", ".k_conv1d", 1),
(".conv1d", ".v_conv1d", 2),
(".gate_up_proj", ".gate_proj", 0),
(".gate_up_proj", ".up_proj", 1),
]
if use_full_rank_gate:
stacked_params_mapping.append((".in_proj_qkvgfab", ".g_proj", 3))
if getattr(self.config, "q_lora_rank", None) is not None:
stacked_params_mapping += [
(".fused_qkv_a_proj", ".q_a_proj", 0),
(".fused_qkv_a_proj", ".kv_a_proj_with_mqa", 1),
]
expert_params_mapping = (
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.num_experts,
)
if self.config.is_moe
else []
)
pp_missing_layer_names = get_pp_missing_layer_names(self)
params_dict = dict(self.named_parameters())
# Under the MXFP4 quant interface the routed experts register unpacked
# params (``w13_weight``), while the compressed-tensors checkpoint names
# them ``.weight_packed``. Rebind so the expert mapping resolves; scales
# already share the ``.weight_scale`` suffix.
experts_unpacked = not any(n.endswith("w13_weight_packed") for n in params_dict)
loaded_params: set[str] = set()
for name, loaded_weight in weights:
if "rotary_emb.inv_freq" in name:
continue
# The multimodal checkpoint prefixes text weights with
# ``language_model.``; strip it so names match this draft model's
# parameter paths (``model.layers.{i}.``). Non-text weights
# (vision_tower, mm_projector, ...) never match a spec layer below.
if name.startswith("language_model."):
name = name[len("language_model.") :]
if experts_unpacked and name.endswith(".weight_packed"):
name = name.replace(".weight_packed", ".weight")
spec_layer = get_spec_layer_idx_from_weight_name(self.config, name)
if spec_layer is None:
continue
name = self._rewrite_spec_layer_name(spec_layer, name)
for param_name, weight_name, shard_id in stacked_params_mapping:
if weight_name not in name:
continue
# Routed experts (``.experts.{i}.w1/w2/w3``) are handled by the
# expert mapping below; skip them here. Shared experts
# (``.shared_experts.``) use gate/up_proj and fall through.
if ".experts." in name:
continue
name_mapped = name.replace(weight_name, param_name)
# Only take this mapping if the fused destination actually
# exists (e.g. QKV fusion is only present when q_lora is used).
if name_mapped not in params_dict:
continue
if name_mapped in pp_missing_layer_names:
continue
name = name_mapped
param = params_dict[name]
weight_loader = param.weight_loader
weight_loader(param, loaded_weight, shard_id)
break
else:
for (
expert_param_name,
expert_weight_name,
expert_id,
expert_shard_id,
) in expert_params_mapping:
if expert_weight_name not in name:
continue
name_mapped = name.replace(expert_weight_name, expert_param_name)
if name_mapped in pp_missing_layer_names:
continue
param = params_dict[name_mapped]
weight_loader = param.weight_loader
weight_loader(
param,
loaded_weight,
name_mapped,
shard_id=expert_shard_id,
expert_id=expert_id,
)
name = name_mapped
break
else:
if name.endswith(".bias") and name not in params_dict:
continue
remapped_name = maybe_remap_kv_scale_name(name, params_dict)
if remapped_name is None:
continue
name = remapped_name
# The embedding is shared across MTP layers; only the first
# spec layer carries the hoisted (non-".layers") copy.
if spec_layer != self.model.mtp_start_layer_idx and (
".layers" not in name
):
continue
if name in pp_missing_layer_names:
continue
# The base model uses an attn-residual scheme whose per-layer
# weights (self_attention_res_*, mlp_res_*) are not used by
# the draft block; such names have no matching parameter and
# are safely skipped.
if name not in params_dict:
continue
param = params_dict[name]
weight_loader = getattr(
param, "weight_loader", default_weight_loader
)
weight_loader(param, loaded_weight)
loaded_params.add(name)
# Validate that weights were loaded for each expected MTP layer.
loaded_layers: set[int] = set()
for param_name in loaded_params:
spec_layer = get_spec_layer_idx_from_weight_name(self.config, param_name)
if spec_layer is not None:
loaded_layers.add(spec_layer)
for layer_idx in range(
self.model.mtp_start_layer_idx,
self.model.mtp_start_layer_idx + self.model.num_mtp_layers,
):
if layer_idx not in loaded_layers:
raise ValueError(
f"MTP speculative decoding layer {layer_idx} weights "
f"missing from checkpoint. The checkpoint may not include "
f"the MTP layer weights. Use a checkpoint that includes "
f"MTP layer weights, or disable speculative decoding."
)
return loaded_params
def _rewrite_spec_layer_name(self, spec_layer: int, name: str) -> str:
"""Rewrite a checkpoint weight name to this module's parameter path.
Top-level MTP submodules (enorm/hnorm/eh_proj/shared_head) stay under
``model.layers.{spec_layer}.*``; the shared ``embed_tokens`` is hoisted
to ``model.*``; everything else is a transformer-block weight and gets
``.mtp_block`` inserted.
"""
spec_layer_weight_names = [
"embed_tokens",
"enorm",
"hnorm",
"eh_proj",
"shared_head",
]
shared_weight_names = ["embed_tokens"]
spec_layer_weight = False
shared_weight = False
for weight_name in spec_layer_weight_names:
if weight_name in name:
spec_layer_weight = True
if weight_name in shared_weight_names:
shared_weight = True
break
if not spec_layer_weight:
name = name.replace(
f"model.layers.{spec_layer}.",
f"model.layers.{spec_layer}.mtp_block.",
)
elif shared_weight:
name = name.replace(f"model.layers.{spec_layer}.", "model.")
return name