458 lines
17 KiB
Python
458 lines
17 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||
"""MiniCPM4 with PagedAttention + fp32 RoPE/RMSNorm for VoxCPM2.
|
||
|
||
Uses vllm Attention for KV cache, keeps fp32 precision ops from
|
||
minicpm4_hf_compat.py to match native VoxCPM2 numerics.
|
||
"""
|
||
|
||
from __future__ import annotations
|
||
|
||
import math
|
||
from collections.abc import Iterable
|
||
from typing import Any
|
||
|
||
import torch
|
||
import torch.nn as nn
|
||
from vllm.config import CacheConfig, VllmConfig
|
||
from vllm.logger import init_logger
|
||
from vllm.model_executor.layers.attention import Attention
|
||
from vllm.model_executor.layers.layernorm import RMSNorm
|
||
from vllm.model_executor.model_loader.weight_utils import default_weight_loader
|
||
from vllm.model_executor.models.utils import make_empty_intermediate_tensors_factory
|
||
from vllm.sequence import IntermediateTensors
|
||
|
||
from .minicpm4_hf_compat import (
|
||
_apply_rotary_pos_emb,
|
||
_MiniCPMLongRoPE,
|
||
_MiniCPMMLP,
|
||
)
|
||
|
||
logger = init_logger(__name__)
|
||
|
||
|
||
def _resolve_lm_cfg(config: Any) -> Any:
|
||
"""Extract lm_config from VoxCPM2Config, converting dict to namespace if needed."""
|
||
lm_cfg = getattr(config, "lm_config", config)
|
||
if isinstance(lm_cfg, dict):
|
||
|
||
class _Cfg:
|
||
pass
|
||
|
||
c = _Cfg()
|
||
for k, v in lm_cfg.items():
|
||
setattr(c, k, v)
|
||
return c
|
||
return lm_cfg
|
||
|
||
|
||
# ===================================================================
|
||
# Attention with vllm PagedAttention backend
|
||
# ===================================================================
|
||
|
||
|
||
class _PagedMiniCPM4Attention(nn.Module):
|
||
"""PagedAttention + fp32 RoPE with separate q/k/v projections."""
|
||
|
||
def __init__(
|
||
self,
|
||
hidden_size: int,
|
||
num_attention_heads: int,
|
||
num_key_value_heads: int,
|
||
kv_channels: int | None,
|
||
layer_idx: int,
|
||
cache_config: CacheConfig | None = None,
|
||
prefix: str = "",
|
||
) -> None:
|
||
super().__init__()
|
||
self.layer_idx = layer_idx
|
||
self.hidden_size = hidden_size
|
||
self.num_heads = num_attention_heads
|
||
self.head_dim = kv_channels if kv_channels else hidden_size // num_attention_heads
|
||
self.num_kv_heads = num_key_value_heads
|
||
self.q_size = self.num_heads * self.head_dim
|
||
self.kv_size = self.num_kv_heads * self.head_dim
|
||
|
||
self.q_proj = nn.Linear(hidden_size, self.q_size, bias=False)
|
||
self.k_proj = nn.Linear(hidden_size, self.kv_size, bias=False)
|
||
self.v_proj = nn.Linear(hidden_size, self.kv_size, bias=False)
|
||
self.o_proj = nn.Linear(self.q_size, hidden_size, bias=False)
|
||
self._fused_qkv_weight: torch.Tensor | None = None
|
||
|
||
self.attn = Attention(
|
||
self.num_heads,
|
||
self.head_dim,
|
||
scale=self.head_dim**-0.5,
|
||
num_kv_heads=self.num_kv_heads,
|
||
cache_config=cache_config,
|
||
prefix=f"{prefix}.attn",
|
||
)
|
||
|
||
def forward(
|
||
self,
|
||
positions: torch.Tensor,
|
||
hidden_states: torch.Tensor,
|
||
rope_emb: _MiniCPMLongRoPE | None = None,
|
||
) -> torch.Tensor:
|
||
"""Forward: fused QKV → fp32 RoPE → PagedAttention → o_proj."""
|
||
if self._fused_qkv_weight is None:
|
||
self._fused_qkv_weight = torch.cat(
|
||
[
|
||
self.q_proj.weight,
|
||
self.k_proj.weight,
|
||
self.v_proj.weight,
|
||
],
|
||
dim=0,
|
||
).detach()
|
||
qkv = nn.functional.linear(hidden_states, self._fused_qkv_weight)
|
||
q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1)
|
||
|
||
if rope_emb is not None:
|
||
cos, sin = rope_emb(positions)
|
||
bsz = q.shape[0]
|
||
q_r = q.view(bsz, self.num_heads, self.head_dim)
|
||
k_r = k.view(bsz, self.num_kv_heads, self.head_dim)
|
||
q_r = q_r.unsqueeze(0).transpose(1, 2) # [1, heads, n_tokens, dim]
|
||
k_r = k_r.unsqueeze(0).transpose(1, 2) # [1, kv_heads, n_tokens, dim]
|
||
q_r, k_r = _apply_rotary_pos_emb(q_r, k_r, cos, sin)
|
||
q = q_r.transpose(1, 2).squeeze(0).reshape(bsz, -1) # [n_tokens, q_size]
|
||
k = k_r.transpose(1, 2).squeeze(0).reshape(bsz, -1) # [n_tokens, kv_size]
|
||
|
||
attn_output = self.attn(q, k, v)
|
||
|
||
output = self.o_proj(attn_output)
|
||
return output
|
||
|
||
|
||
# ===================================================================
|
||
# Decoder Layer
|
||
# ===================================================================
|
||
|
||
|
||
class _PagedMiniCPM4DecoderLayer(nn.Module):
|
||
"""Decoder layer: PagedAttention + fp32 RMSNorm + muP scale_depth."""
|
||
|
||
def __init__(
|
||
self,
|
||
hidden_size: int,
|
||
intermediate_size: int,
|
||
num_attention_heads: int,
|
||
num_key_value_heads: int,
|
||
kv_channels: int | None,
|
||
rms_norm_eps: float,
|
||
layer_idx: int,
|
||
num_hidden_layers: int,
|
||
use_mup: bool,
|
||
scale_depth: float,
|
||
cache_config: CacheConfig | None = None,
|
||
prefix: str = "",
|
||
) -> None:
|
||
super().__init__()
|
||
self.self_attn = _PagedMiniCPM4Attention(
|
||
hidden_size=hidden_size,
|
||
num_attention_heads=num_attention_heads,
|
||
num_key_value_heads=num_key_value_heads,
|
||
kv_channels=kv_channels,
|
||
layer_idx=layer_idx,
|
||
cache_config=cache_config,
|
||
prefix=f"{prefix}.self_attn",
|
||
)
|
||
self.mlp = _MiniCPMMLP(hidden_size, intermediate_size)
|
||
self.input_layernorm = RMSNorm(hidden_size, eps=rms_norm_eps)
|
||
self.post_attention_layernorm = RMSNorm(hidden_size, eps=rms_norm_eps)
|
||
|
||
self.use_mup = use_mup
|
||
self.scale_depth = scale_depth
|
||
self.num_hidden_layers = num_hidden_layers
|
||
|
||
def _residual_scale(self) -> float:
|
||
if self.use_mup:
|
||
return self.scale_depth / math.sqrt(self.num_hidden_layers)
|
||
return 1.0
|
||
|
||
def forward(
|
||
self,
|
||
positions: torch.Tensor,
|
||
hidden_states: torch.Tensor,
|
||
residual: torch.Tensor | None,
|
||
rope_emb: _MiniCPMLongRoPE | None = None,
|
||
) -> tuple[torch.Tensor, torch.Tensor | None]:
|
||
# Pre-norm + attention
|
||
residual = hidden_states
|
||
hidden_states = self.input_layernorm(hidden_states)
|
||
hidden_states = self.self_attn(positions, hidden_states, rope_emb)
|
||
|
||
scale = self._residual_scale()
|
||
if scale != 1.0:
|
||
hidden_states = residual + hidden_states * scale
|
||
else:
|
||
hidden_states = residual + hidden_states
|
||
|
||
# Pre-norm + FFN
|
||
residual = hidden_states
|
||
hidden_states = self.post_attention_layernorm(hidden_states)
|
||
hidden_states = self.mlp(hidden_states)
|
||
|
||
if scale != 1.0:
|
||
hidden_states = residual + hidden_states * scale
|
||
else:
|
||
hidden_states = residual + hidden_states
|
||
|
||
return hidden_states, None
|
||
|
||
|
||
# ===================================================================
|
||
# Full Model
|
||
# ===================================================================
|
||
|
||
|
||
class MiniCPM4PagedForVoxCPM2(nn.Module):
|
||
"""PagedAttention base_lm (28 layers) for VoxCPM2 scaffold."""
|
||
|
||
def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None:
|
||
super().__init__()
|
||
config = vllm_config.model_config.hf_config
|
||
cache_config = vllm_config.cache_config
|
||
self.config = config
|
||
|
||
lm_cfg = _resolve_lm_cfg(config)
|
||
|
||
hidden_size = lm_cfg.hidden_size
|
||
num_hidden_layers = lm_cfg.num_hidden_layers
|
||
kv_channels = getattr(lm_cfg, "kv_channels", None)
|
||
|
||
self.vocab_size = lm_cfg.vocab_size
|
||
self.embed_tokens = nn.Embedding(self.vocab_size, hidden_size)
|
||
|
||
rope_scaling = getattr(lm_cfg, "rope_scaling", None)
|
||
if isinstance(rope_scaling, dict):
|
||
rope_scaling_dict = rope_scaling
|
||
elif hasattr(rope_scaling, "__dict__"):
|
||
rope_scaling_dict = {
|
||
"short_factor": rope_scaling.short_factor,
|
||
"long_factor": rope_scaling.long_factor,
|
||
"original_max_position_embeddings": rope_scaling.original_max_position_embeddings,
|
||
}
|
||
else:
|
||
rope_scaling_dict = {}
|
||
|
||
no_rope = getattr(lm_cfg, "no_rope", False)
|
||
if not no_rope:
|
||
self.rope_emb = _MiniCPMLongRoPE(
|
||
hidden_size=hidden_size,
|
||
num_attention_heads=lm_cfg.num_attention_heads,
|
||
kv_channels=kv_channels,
|
||
rope_theta=getattr(lm_cfg, "rope_theta", 10000.0),
|
||
max_position_embeddings=getattr(lm_cfg, "max_position_embeddings", 32768),
|
||
rope_scaling=rope_scaling_dict,
|
||
)
|
||
else:
|
||
self.rope_emb = None
|
||
|
||
self.layers = nn.ModuleList(
|
||
[
|
||
_PagedMiniCPM4DecoderLayer(
|
||
hidden_size=hidden_size,
|
||
intermediate_size=lm_cfg.intermediate_size,
|
||
num_attention_heads=lm_cfg.num_attention_heads,
|
||
num_key_value_heads=lm_cfg.num_key_value_heads,
|
||
kv_channels=kv_channels,
|
||
rms_norm_eps=lm_cfg.rms_norm_eps,
|
||
layer_idx=i,
|
||
num_hidden_layers=num_hidden_layers,
|
||
use_mup=getattr(lm_cfg, "use_mup", False),
|
||
scale_depth=getattr(lm_cfg, "scale_depth", 1.0),
|
||
cache_config=cache_config,
|
||
prefix=f"{prefix}.layers.{i}",
|
||
)
|
||
for i in range(num_hidden_layers)
|
||
]
|
||
)
|
||
|
||
self.norm = RMSNorm(hidden_size, eps=lm_cfg.rms_norm_eps)
|
||
|
||
self.make_empty_intermediate_tensors = make_empty_intermediate_tensors_factory(
|
||
["hidden_states", "residual"], hidden_size
|
||
)
|
||
|
||
use_mup = getattr(lm_cfg, "use_mup", False)
|
||
self._scale_emb = getattr(lm_cfg, "scale_emb", 1.0) if use_mup else 1.0
|
||
self._compiled_layers: set[int] = set()
|
||
|
||
def embed_input_ids(self, input_ids: torch.Tensor, **_: Any) -> torch.Tensor:
|
||
return self.embed_tokens(input_ids) * self._scale_emb
|
||
|
||
def forward(
|
||
self,
|
||
input_ids: torch.Tensor | None,
|
||
positions: torch.Tensor,
|
||
intermediate_tensors: IntermediateTensors | None = None,
|
||
inputs_embeds: torch.Tensor | None = None,
|
||
**kwargs: Any,
|
||
) -> torch.Tensor | IntermediateTensors:
|
||
if inputs_embeds is not None:
|
||
hidden_states = inputs_embeds
|
||
else:
|
||
hidden_states = self.embed_input_ids(input_ids)
|
||
|
||
residual = None
|
||
for layer in self.layers:
|
||
hidden_states, residual = layer(
|
||
positions,
|
||
hidden_states,
|
||
residual,
|
||
self.rope_emb,
|
||
)
|
||
|
||
hidden_states = self.norm(hidden_states)
|
||
return hidden_states
|
||
|
||
def precompute_fused_qkv(self) -> None:
|
||
"""Materialize fused QKV weights before CUDA Graph capture."""
|
||
for layer in self.layers:
|
||
attn = layer.self_attn
|
||
if attn._fused_qkv_weight is None:
|
||
attn._fused_qkv_weight = torch.cat(
|
||
[attn.q_proj.weight, attn.k_proj.weight, attn.v_proj.weight],
|
||
dim=0,
|
||
).detach()
|
||
|
||
def compile_selective(self) -> list[str]:
|
||
"""Compile the full model forward as one graph.
|
||
|
||
Earlier versions compiled ``layer.mlp`` + ``layer.self_attn.o_proj``
|
||
(PR #2690) and then the whole ``layer`` (perf/voxcpm2-streaming-vae).
|
||
Both still paid one Dynamo dispatch per layer per decode step.
|
||
V3 profiling showed 1,332 per-layer dispatches (~28 layers × ~47
|
||
decode steps) costing ~726 ms of CPU self-time for a long prompt.
|
||
|
||
Compiling ``forward`` at the model level lets Dynamo unroll the
|
||
28-layer Python loop inside the graph. Graph breaks at
|
||
PagedAttention produce sub-graphs but Dynamo memoises the whole
|
||
trace once, so the per-step dispatch drops from 28 to just a few.
|
||
"""
|
||
if self._compiled_layers:
|
||
return []
|
||
# Null the fused-qkv caches so the compile sees the real weight layout.
|
||
for layer in self.layers:
|
||
layer.self_attn._fused_qkv_weight = None
|
||
self.forward = torch.compile(self.forward, mode="default", fullgraph=False)
|
||
# Mark every layer as compiled so idempotent callers don't double-wrap.
|
||
self._compiled_layers.update(range(len(self.layers)))
|
||
return ["forward (whole model)"]
|
||
|
||
def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
|
||
"""Load weights from native checkpoint (base_lm. prefix pre-stripped)."""
|
||
params_dict = dict(self.named_parameters(remove_duplicate=False))
|
||
loaded: set[str] = set()
|
||
|
||
for name, loaded_weight in weights:
|
||
if "rotary_emb.inv_freq" in name:
|
||
continue
|
||
param = params_dict.get(name)
|
||
if param is None:
|
||
continue
|
||
weight_loader = getattr(param, "weight_loader", default_weight_loader)
|
||
weight_loader(param, loaded_weight)
|
||
loaded.add(name)
|
||
|
||
return loaded
|
||
|
||
|
||
# ===================================================================
|
||
# Residual LM with PagedAttention (no RoPE, 8 layers)
|
||
# ===================================================================
|
||
|
||
|
||
class MiniCPM4PagedResidualLM(nn.Module):
|
||
"""PagedAttention residual LM (8 layers, no RoPE) for VoxCPM2."""
|
||
|
||
def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None:
|
||
super().__init__()
|
||
config = vllm_config.model_config.hf_config
|
||
cache_config = vllm_config.cache_config
|
||
self.config = config
|
||
|
||
lm_cfg = _resolve_lm_cfg(config)
|
||
|
||
hidden_size = lm_cfg.hidden_size
|
||
num_hidden_layers = getattr(config, "residual_lm_num_layers", 8)
|
||
kv_channels = getattr(lm_cfg, "kv_channels", None)
|
||
|
||
self.rope_emb = None
|
||
|
||
self.layers = nn.ModuleList(
|
||
[
|
||
_PagedMiniCPM4DecoderLayer(
|
||
hidden_size=hidden_size,
|
||
intermediate_size=lm_cfg.intermediate_size,
|
||
num_attention_heads=lm_cfg.num_attention_heads,
|
||
num_key_value_heads=lm_cfg.num_key_value_heads,
|
||
kv_channels=kv_channels,
|
||
rms_norm_eps=lm_cfg.rms_norm_eps,
|
||
layer_idx=i,
|
||
num_hidden_layers=num_hidden_layers,
|
||
use_mup=getattr(lm_cfg, "use_mup", False),
|
||
scale_depth=getattr(lm_cfg, "scale_depth", 1.0),
|
||
cache_config=cache_config,
|
||
prefix=f"{prefix}.layers.{i}",
|
||
)
|
||
for i in range(num_hidden_layers)
|
||
]
|
||
)
|
||
|
||
self.norm = RMSNorm(hidden_size, eps=lm_cfg.rms_norm_eps)
|
||
self._compiled_layers: set[int] = set()
|
||
|
||
def forward(
|
||
self,
|
||
positions: torch.Tensor,
|
||
inputs_embeds: torch.Tensor,
|
||
) -> torch.Tensor:
|
||
hidden_states = inputs_embeds
|
||
residual = None
|
||
for layer in self.layers:
|
||
hidden_states, residual = layer(
|
||
positions,
|
||
hidden_states,
|
||
residual,
|
||
self.rope_emb,
|
||
)
|
||
hidden_states = self.norm(hidden_states)
|
||
return hidden_states
|
||
|
||
def precompute_fused_qkv(self) -> None:
|
||
"""Materialize fused QKV weights before CUDA Graph capture."""
|
||
for layer in self.layers:
|
||
attn = layer.self_attn
|
||
if attn._fused_qkv_weight is None:
|
||
attn._fused_qkv_weight = torch.cat(
|
||
[attn.q_proj.weight, attn.k_proj.weight, attn.v_proj.weight],
|
||
dim=0,
|
||
).detach()
|
||
|
||
def compile_selective(self) -> list[str]:
|
||
"""Compile the full residual model forward as one graph (same strategy as base_lm)."""
|
||
if self._compiled_layers:
|
||
return []
|
||
for layer in self.layers:
|
||
layer.self_attn._fused_qkv_weight = None
|
||
self.forward = torch.compile(self.forward, mode="default", fullgraph=False)
|
||
self._compiled_layers.update(range(len(self.layers)))
|
||
return ["forward (whole residual)"]
|
||
|
||
def load_weights_from_native(self, native_residual_lm: nn.Module) -> int:
|
||
"""Load weights from native residual_lm. Returns param count."""
|
||
params_dict = dict(self.named_parameters(remove_duplicate=False))
|
||
loaded = 0
|
||
for name, param in native_residual_lm.named_parameters():
|
||
if "rotary_emb" in name:
|
||
continue
|
||
target = params_dict.get(name)
|
||
if target is None:
|
||
continue
|
||
weight_loader = getattr(target, "weight_loader", default_weight_loader)
|
||
weight_loader(target, param.data)
|
||
loaded += 1
|
||
return loaded
|