430 lines
17 KiB
Python
430 lines
17 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
|
import typing
|
|
from collections.abc import Callable, Iterable
|
|
|
|
import torch
|
|
import torch.nn as nn
|
|
|
|
from vllm.config import VllmConfig
|
|
from vllm.model_executor.layers.fused_moe import (
|
|
fused_moe_make_expert_params_mapping,
|
|
)
|
|
from vllm.model_executor.layers.layernorm import RMSNorm
|
|
from vllm.model_executor.layers.logits_processor import LogitsProcessor
|
|
from vllm.model_executor.layers.vocab_parallel_embedding import (
|
|
VocabParallelEmbedding,
|
|
)
|
|
from vllm.model_executor.model_loader.weight_utils import (
|
|
default_weight_loader,
|
|
maybe_remap_kv_scale_name,
|
|
)
|
|
from vllm.model_executor.models.deepseek_mtp import SharedHead
|
|
from vllm.model_executor.models.deepseek_v2 import DeepseekV2MixtureOfExperts
|
|
from vllm.model_executor.models.utils import maybe_prefix
|
|
from vllm.platforms import current_platform
|
|
from vllm.sequence import IntermediateTensors
|
|
|
|
from .model import (
|
|
Glm5NextDecoderLayer,
|
|
Glm5NextMLAAttention,
|
|
Glm5NextMoE,
|
|
_try_load_fp8_attn_proj,
|
|
_try_load_fp8_indexer_wk,
|
|
get_spec_layer_idx_from_weight_name,
|
|
)
|
|
from .ops.fused_eh_norm import fused_eh_norm
|
|
|
|
|
|
class Glm5NextMultiTokenPredictorLayer(nn.Module):
|
|
def __init__(self, vllm_config: VllmConfig, prefix: str) -> None:
|
|
super().__init__()
|
|
assert vllm_config.speculative_config is not None
|
|
config = vllm_config.speculative_config.draft_model_config.hf_config
|
|
self.config = config
|
|
quant_config = vllm_config.quant_config
|
|
|
|
self.enorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
|
|
self.hnorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
|
|
self.eh_proj = nn.Linear(config.hidden_size * 2, config.hidden_size, bias=False)
|
|
|
|
# Reserve room for the incomplete pool tail and align the sparse MLA
|
|
# buffer width to BLOCK_N=128.
|
|
topk_tokens = config.index_topk
|
|
assert topk_tokens is not None
|
|
kpool = config.index_kpool
|
|
assert kpool is not None
|
|
buffer_width = topk_tokens + (kpool - 1 if kpool > 1 else 0)
|
|
sparse_topk_block_n = 128
|
|
buffer_width = (
|
|
(buffer_width + sparse_topk_block_n - 1) // sparse_topk_block_n
|
|
) * sparse_topk_block_n
|
|
topk_indices_buffer = torch.empty(
|
|
vllm_config.scheduler_config.max_num_batched_tokens,
|
|
buffer_width,
|
|
dtype=torch.int32,
|
|
device=current_platform.device_type,
|
|
)
|
|
self.shared_head = SharedHead(
|
|
config=config, prefix=prefix, quant_config=quant_config
|
|
)
|
|
# MTP layers sit past the base model's hidden layers; parse the index
|
|
# from the prefix (e.g. "...layers.32") so the decoder builds an MLA
|
|
# (DSA) layer rather than KDA for the MTP path.
|
|
layer_idx = int(prefix.rsplit(".", 1)[-1])
|
|
self.mtp_block = Glm5NextDecoderLayer(
|
|
vllm_config=vllm_config,
|
|
config=config,
|
|
layer_idx=layer_idx,
|
|
prefix=prefix,
|
|
topk_indices_buffer=topk_indices_buffer,
|
|
is_mtp_layer=True,
|
|
)
|
|
|
|
def forward(
|
|
self,
|
|
input_ids: torch.Tensor,
|
|
positions: torch.Tensor,
|
|
previous_hidden_states: torch.Tensor,
|
|
inputs_embeds: torch.Tensor | None = None,
|
|
spec_step_index: int = 0,
|
|
) -> torch.Tensor:
|
|
assert inputs_embeds is not None
|
|
# Fused: zero pos-0 embeds + enorm(embeds) + hnorm(prev) + cat -> [N, 2H].
|
|
eh_input = fused_eh_norm(
|
|
positions,
|
|
inputs_embeds,
|
|
previous_hidden_states,
|
|
self.enorm.weight,
|
|
self.hnorm.weight,
|
|
self.enorm.variance_epsilon,
|
|
)
|
|
hidden_states = self.eh_proj(eh_input)
|
|
# Fuse the residual add and final RMSNorm. Glm5NextMoE already performs
|
|
# its all-reduce, so no collective is needed here. The post-norm result
|
|
# feeds both draft logits and the next recycled hidden state.
|
|
hidden_states, residual, _, _ = self.mtp_block(
|
|
positions=positions, hidden_states=hidden_states, residual=None
|
|
)
|
|
hidden_states, _ = self.shared_head.norm(hidden_states, residual=residual)
|
|
return hidden_states, hidden_states
|
|
|
|
|
|
class Glm5NextMultiTokenPredictor(nn.Module):
|
|
def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
|
|
super().__init__()
|
|
config = vllm_config.model_config.hf_config
|
|
self.mtp_start_layer_idx = config.num_hidden_layers
|
|
self.num_mtp_layers = config.num_nextn_predict_layers
|
|
self.layers = torch.nn.ModuleDict(
|
|
{
|
|
str(idx): Glm5NextMultiTokenPredictorLayer(
|
|
vllm_config, f"{prefix}.layers.{idx}"
|
|
)
|
|
for idx in range(
|
|
self.mtp_start_layer_idx,
|
|
self.mtp_start_layer_idx + self.num_mtp_layers,
|
|
)
|
|
}
|
|
)
|
|
self.embed_tokens = VocabParallelEmbedding(
|
|
config.vocab_size,
|
|
config.hidden_size,
|
|
prefix=maybe_prefix(prefix, "embed_tokens"),
|
|
)
|
|
# Plain list for the per-propose lookup: ModuleDict[str(...)] builds a
|
|
# string and hashes it on every draft step.
|
|
self._mtp_layers = list(self.layers.values())
|
|
self._mtp_mla_attns = []
|
|
for layer in self._mtp_layers:
|
|
self_attn = layer.mtp_block.self_attn
|
|
assert isinstance(self_attn, Glm5NextMLAAttention)
|
|
self._mtp_mla_attns.append(self_attn.mla_attn)
|
|
self.logits_processor = LogitsProcessor(config.vocab_size)
|
|
|
|
def set_skip_topk(self, skip: bool):
|
|
# index_share_for_mtp_iteration: step 0 computes top-k, steps 1+ reuse.
|
|
for mla_attn in self._mtp_mla_attns:
|
|
mla_attn.skip_topk = skip
|
|
|
|
def compact_topk_indices(self, slot_ids: torch.Tensor):
|
|
"""Gather the top-k index rows at ``slot_ids`` to the front of the buffer."""
|
|
num_slots = slot_ids.numel()
|
|
for mla_attn in self._mtp_mla_attns:
|
|
topk_indices_buffer = mla_attn.topk_indices_buffer
|
|
assert topk_indices_buffer is not None
|
|
topk_indices_buffer[:num_slots] = topk_indices_buffer[slot_ids]
|
|
|
|
def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:
|
|
return self.embed_tokens(input_ids)
|
|
|
|
def forward(
|
|
self,
|
|
input_ids: torch.Tensor,
|
|
positions: torch.Tensor,
|
|
previous_hidden_states: torch.Tensor,
|
|
inputs_embeds: torch.Tensor | None = None,
|
|
spec_step_idx: int = 0,
|
|
) -> torch.Tensor:
|
|
if inputs_embeds is None:
|
|
inputs_embeds = self.embed_tokens(input_ids)
|
|
current_step_idx = spec_step_idx % self.num_mtp_layers
|
|
return self._mtp_layers[current_step_idx](
|
|
input_ids,
|
|
positions,
|
|
previous_hidden_states,
|
|
inputs_embeds,
|
|
current_step_idx,
|
|
)
|
|
|
|
def compute_logits(
|
|
self,
|
|
hidden_states: torch.Tensor,
|
|
spec_step_idx: int = 0,
|
|
) -> torch.Tensor:
|
|
current_step_idx = spec_step_idx % self.num_mtp_layers
|
|
mtp_layer = self._mtp_layers[current_step_idx]
|
|
# hidden_states is already post-final-norm (produced in the layer
|
|
# forward and recycled as-is); apply the LM head only, without a
|
|
# second RMSNorm.
|
|
return self.logits_processor(mtp_layer.shared_head.head, hidden_states)
|
|
|
|
def get_top_tokens(
|
|
self,
|
|
hidden_states: torch.Tensor,
|
|
spec_step_idx: int = 0,
|
|
) -> torch.Tensor:
|
|
current_step_idx = spec_step_idx % self.num_mtp_layers
|
|
mtp_layer = self._mtp_layers[current_step_idx]
|
|
# Vocab-parallel argmax for the greedy draft: per-rank head projection
|
|
# + local argmax + a [batch, 2*tp] (value, index) reduce, instead of
|
|
# materializing and all-gathering full [N, vocab] logits per draft
|
|
# step. Tie-breaking matches the full argmax (shards are contiguous
|
|
# and rank-ordered, so the lowest-rank winner is the lowest global
|
|
# index), so greedy draft tokens are unchanged.
|
|
return self.logits_processor.get_top_tokens(
|
|
mtp_layer.shared_head.head, hidden_states
|
|
)
|
|
|
|
|
|
class Glm5NextMTP(nn.Module, DeepseekV2MixtureOfExperts):
|
|
def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
|
|
super().__init__()
|
|
self.config = vllm_config.model_config.hf_config
|
|
self.quant_config = vllm_config.quant_config
|
|
self.model = Glm5NextMultiTokenPredictor(
|
|
vllm_config=vllm_config, prefix=maybe_prefix(prefix, "model")
|
|
)
|
|
self.set_moe_parameters()
|
|
|
|
def set_moe_parameters(self):
|
|
self.num_moe_layers = self.config.num_nextn_predict_layers
|
|
self.num_expert_groups = self.config.n_group
|
|
self.moe_layers = []
|
|
self.moe_mlp_layers = []
|
|
example_moe = None
|
|
for layer in self.model.layers.values():
|
|
mlp = layer.mtp_block.mlp
|
|
if isinstance(mlp, Glm5NextMoE):
|
|
example_moe = mlp
|
|
self.moe_mlp_layers.append(mlp)
|
|
self.moe_layers.append(mlp.experts)
|
|
self.extract_moe_parameters(example_moe)
|
|
|
|
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,
|
|
) -> 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 get_top_tokens(
|
|
self,
|
|
hidden_states: torch.Tensor,
|
|
spec_step_idx: int = 0,
|
|
) -> torch.Tensor:
|
|
# Greedy-draft path used when use_local_argmax_reduction is enabled:
|
|
# vocab-parallel argmax, no full-vocab logits.
|
|
return self.model.get_top_tokens(hidden_states, spec_step_idx)
|
|
|
|
def _rewrite_spec_layer_name(self, spec_layer: int, name: str) -> str:
|
|
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
|
|
|
|
def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
|
|
stacked_params_mapping = [
|
|
("gate_up_proj", "gate_proj", 0),
|
|
("gate_up_proj", "up_proj", 1),
|
|
("fused_qkv_a_proj", "q_a_proj", 0),
|
|
("fused_qkv_a_proj", "kv_a_proj_with_mqa", 1),
|
|
("wk_weights_proj", "wk", 0),
|
|
("wk_weights_proj", "weights_proj", 1),
|
|
]
|
|
expert_params_mapping = fused_moe_make_expert_params_mapping(
|
|
self,
|
|
ckpt_gate_proj_name="gate_proj",
|
|
ckpt_down_proj_name="down_proj",
|
|
ckpt_up_proj_name="up_proj",
|
|
num_experts=self.config.n_routed_experts,
|
|
)
|
|
|
|
params_dict = dict(self.named_parameters())
|
|
loaded_params: set[str] = set()
|
|
_pending_wk_fp8: dict = {}
|
|
# GLM-5.3-Flash NoPE checkpoints omit the RoPE rows from
|
|
# ``kv_a_proj_with_mqa``; the FP8-to-BF16 path pads them for the model.
|
|
kv_a_pad_size = 0
|
|
if self.config.mla_nope and self.config.qk_rope_head_dim > 0:
|
|
kv_a_pad_size = self.config.qk_rope_head_dim
|
|
for name, loaded_weight in weights:
|
|
if "rotary_emb.inv_freq" in name:
|
|
continue
|
|
# Multimodal (Glm5NextForConditionalGeneration) checkpoints prefix
|
|
# the text-tower weights with "model.language_model."; the MTP head
|
|
# is built as a text-only model (model.layers.*), so strip the
|
|
# prefix to match.
|
|
if name.startswith("model.language_model."):
|
|
name = name.replace("model.language_model.", "model.", 1)
|
|
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)
|
|
|
|
if _try_load_fp8_indexer_wk(
|
|
name,
|
|
loaded_weight,
|
|
_pending_wk_fp8,
|
|
params_dict,
|
|
loaded_params,
|
|
):
|
|
continue
|
|
|
|
# FP8 checkpoint: dequantize the BF16-kept MLA projections
|
|
# (q_a_proj / kv_a_proj_with_mqa / o_proj) to BF16, mirroring the
|
|
# target model. The model holds fused_qkv_a_proj / o_proj in BF16,
|
|
# so the checkpoint's block-FP8 weight + weight_scale_inv for these
|
|
# has no param home and would KeyError without this dequant.
|
|
if _try_load_fp8_attn_proj(
|
|
name,
|
|
loaded_weight,
|
|
_pending_wk_fp8,
|
|
params_dict,
|
|
loaded_params,
|
|
kv_a_pad_size,
|
|
):
|
|
continue
|
|
|
|
for param_name, weight_name, shard_id in stacked_params_mapping:
|
|
if weight_name not in name:
|
|
continue
|
|
if ("mlp.experts." in name) and name not in params_dict:
|
|
continue
|
|
name_mapped = name.replace(weight_name, param_name)
|
|
if (
|
|
param_name == "fused_qkv_a_proj"
|
|
) and name_mapped not in params_dict:
|
|
continue
|
|
else:
|
|
name = name_mapped
|
|
if name.endswith(".bias") and name not in params_dict:
|
|
continue
|
|
param = params_dict[name]
|
|
weight_loader = param.weight_loader
|
|
weight_loader(param, loaded_weight, shard_id)
|
|
break
|
|
else:
|
|
is_expert_weight = False
|
|
for mapping in expert_params_mapping:
|
|
param_name, weight_name, expert_id, shard_id = mapping # type: ignore[assignment]
|
|
if weight_name not in name:
|
|
continue
|
|
is_expert_weight = True
|
|
name_mapped = name.replace(weight_name, param_name)
|
|
param = params_dict[name_mapped]
|
|
weight_loader = typing.cast(
|
|
Callable[..., bool], param.weight_loader
|
|
)
|
|
success = weight_loader(
|
|
param,
|
|
loaded_weight,
|
|
name_mapped,
|
|
shard_id=shard_id,
|
|
expert_id=expert_id,
|
|
return_success=True,
|
|
)
|
|
if success:
|
|
name = name_mapped
|
|
break
|
|
else:
|
|
if is_expert_weight:
|
|
continue
|
|
if name.endswith(".bias") and name not in params_dict:
|
|
continue
|
|
name = maybe_remap_kv_scale_name(name, params_dict) # type: ignore[assignment]
|
|
if name is None:
|
|
continue
|
|
if (
|
|
spec_layer != self.model.mtp_start_layer_idx
|
|
and ".layers" not in name
|
|
):
|
|
continue
|
|
param = params_dict[name]
|
|
weight_loader = getattr(
|
|
param, "weight_loader", default_weight_loader
|
|
)
|
|
weight_loader(param, loaded_weight)
|
|
loaded_params.add(name)
|
|
|
|
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."
|
|
)
|
|
return loaded_params
|