biondizzle
22e0370e6e
Get swiglu_limit from vllm_config.model_config.hf_config instead of self (it was only set on the parent DeepseekV4MoE class).
2348 lines
99 KiB
Python
2348 lines
99 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
|
|
|
import typing
|
|
from collections.abc import Callable, Iterable
|
|
from itertools import islice
|
|
|
|
import regex as re
|
|
import os
|
|
import torch
|
|
import torch.nn as nn
|
|
|
|
from vllm.compilation.decorators import support_torch_compile
|
|
from vllm.config import VllmConfig, get_current_vllm_config
|
|
from vllm.distributed import (
|
|
get_ep_group,
|
|
get_tensor_model_parallel_rank,
|
|
get_tensor_model_parallel_world_size,
|
|
)
|
|
from vllm.forward_context import get_forward_context
|
|
from vllm.model_executor.layers.activation import SiluAndMul, SiluAndMulWithClamp
|
|
from vllm.model_executor.layers.deepseek_v4_attention import (
|
|
DeepseekV4Indexer,
|
|
DeepseekV4MLAModules,
|
|
DeepseekV4MultiHeadLatentAttentionWrapper,
|
|
)
|
|
from vllm.model_executor.layers.fused_moe import FusedMoE, GateLinear
|
|
from vllm.model_executor.layers.fused_moe.layer import UnquantizedFusedMoEMethod
|
|
from vllm.model_executor.layers.fused_moe.router.fused_topk_bias_router import (
|
|
fused_topk_bias,
|
|
)
|
|
from vllm.model_executor.layers.layernorm import RMSNorm
|
|
from vllm.model_executor.layers.linear import (
|
|
ColumnParallelLinear,
|
|
MergedColumnParallelLinear,
|
|
RowParallelLinear,
|
|
)
|
|
from vllm.model_executor.layers.logits_processor import LogitsProcessor
|
|
from vllm.model_executor.layers.quantization import (
|
|
QuantizationConfig,
|
|
QuantizationMethods,
|
|
)
|
|
from vllm.model_executor.layers.quantization.fp8 import Fp8Config
|
|
from vllm.model_executor.layers.quantization.mxfp4 import Mxfp4MoEMethod
|
|
from vllm.model_executor.layers.quantization.utils.quant_utils import (
|
|
is_layer_skipped,
|
|
)
|
|
from vllm.model_executor.layers.rotary_embedding import get_rope
|
|
from vllm.model_executor.layers.vocab_parallel_embedding import (
|
|
ParallelLMHead,
|
|
VocabParallelEmbedding,
|
|
)
|
|
from vllm.model_executor.model_loader.weight_utils import default_weight_loader
|
|
from vllm.model_executor.utils import set_weight_attrs
|
|
from vllm.platforms import current_platform
|
|
from vllm.sequence import IntermediateTensors
|
|
from vllm.triton_utils import tl, triton
|
|
from vllm.utils.torch_utils import direct_register_custom_op
|
|
|
|
from .utils import (
|
|
AutoWeightsLoader,
|
|
WeightsMapper,
|
|
extract_layer_index,
|
|
make_layers,
|
|
maybe_prefix,
|
|
)
|
|
|
|
_DEEPSEEK_V4_EXPERT_DTYPES = ("fp4", "fp8")
|
|
|
|
|
|
class DeepseekV4MLP(nn.Module):
|
|
def __init__(
|
|
self,
|
|
hidden_size: int,
|
|
intermediate_size: int,
|
|
hidden_act: str,
|
|
swiglu_limit: float | None = None,
|
|
quant_config: QuantizationConfig | None = None,
|
|
reduce_results: bool = True,
|
|
is_sequence_parallel: bool = False,
|
|
prefix: str = "",
|
|
) -> None:
|
|
super().__init__()
|
|
|
|
# If is_sequence_parallel, the input and output tensors are sharded
|
|
# across the ranks within the tp_group. In this case the weights are
|
|
# replicated and no collective ops are needed.
|
|
# Otherwise we use standard TP with an allreduce at the end.
|
|
self.gate_up_proj = MergedColumnParallelLinear(
|
|
hidden_size,
|
|
[intermediate_size] * 2,
|
|
bias=False,
|
|
quant_config=quant_config,
|
|
disable_tp=is_sequence_parallel,
|
|
prefix=f"{prefix}.gate_up_proj",
|
|
)
|
|
self.down_proj = RowParallelLinear(
|
|
intermediate_size,
|
|
hidden_size,
|
|
bias=False,
|
|
quant_config=quant_config,
|
|
reduce_results=reduce_results,
|
|
disable_tp=is_sequence_parallel,
|
|
prefix=f"{prefix}.down_proj",
|
|
)
|
|
if hidden_act != "silu":
|
|
raise ValueError(
|
|
f"Unsupported activation: {hidden_act}. Only silu is supported for now."
|
|
)
|
|
if swiglu_limit is not None:
|
|
self.act_fn = SiluAndMulWithClamp(swiglu_limit)
|
|
else:
|
|
self.act_fn = SiluAndMul()
|
|
|
|
def forward(self, x):
|
|
gate_up, _ = self.gate_up_proj(x)
|
|
x = self.act_fn(gate_up)
|
|
x, _ = self.down_proj(x)
|
|
return x
|
|
|
|
|
|
class DeepseekV4FP8Config(Fp8Config):
|
|
"""FP8 config for DeepSeek V4 with expert-dtype-aware MoE dispatch.
|
|
|
|
DeepSeek V4 checkpoints use FP8 block quantization for attention
|
|
layers and NVFP4 (E2M1 + float8_e4m3fn block scales) for MoE experts.
|
|
|
|
``expert_dtype`` from hf_config determines the MoE dispatch path.
|
|
For NVFP4 checkpoints (our case), expert_dtype="fp4" which routes
|
|
to DeepseekV4MegaMoEExperts (native NVFP4 CUTLASS kernel).
|
|
"""
|
|
|
|
def __init__(self, *args, **kwargs):
|
|
super().__init__(*args, **kwargs)
|
|
self._resolved_expert_dtype: str | None = None
|
|
|
|
@property
|
|
def expert_dtype(self) -> str:
|
|
if self._resolved_expert_dtype is None:
|
|
try:
|
|
hf_config = get_current_vllm_config().model_config.hf_config
|
|
except Exception:
|
|
# vllm_config not yet set; return safe default but do NOT
|
|
# cache — a later call inside set_current_vllm_config may
|
|
# resolve differently.
|
|
return "fp4"
|
|
expert_dtype = getattr(hf_config, "expert_dtype", "fp4")
|
|
if expert_dtype not in _DEEPSEEK_V4_EXPERT_DTYPES:
|
|
raise ValueError(
|
|
f"Unsupported DeepSeek V4 expert_dtype={expert_dtype!r}; "
|
|
f"expected one of {_DEEPSEEK_V4_EXPERT_DTYPES}."
|
|
)
|
|
self._resolved_expert_dtype = expert_dtype
|
|
return self._resolved_expert_dtype
|
|
|
|
@classmethod
|
|
def get_name(cls) -> QuantizationMethods:
|
|
return "deepseek_v4_fp8"
|
|
|
|
@classmethod
|
|
def override_quantization_method(
|
|
cls, hf_quant_cfg, user_quant, hf_config=None
|
|
) -> QuantizationMethods | None:
|
|
if not (
|
|
isinstance(hf_quant_cfg, dict)
|
|
and hf_quant_cfg.get("quant_method") in ("fp8", "deepseek_v4_fp8")
|
|
):
|
|
return None
|
|
model_type = getattr(hf_config, "model_type", None)
|
|
if model_type == "deepseek_v4" or user_quant == "deepseek_v4_fp8":
|
|
return "deepseek_v4_fp8"
|
|
return None
|
|
|
|
def get_quant_method(self, layer, prefix):
|
|
if isinstance(layer, FusedMoE):
|
|
if is_layer_skipped(
|
|
prefix=prefix,
|
|
ignored_layers=self.ignored_layers,
|
|
fused_mapping=self.packed_modules_mapping,
|
|
):
|
|
return UnquantizedFusedMoEMethod(layer.moe_config)
|
|
if self.expert_dtype == "fp4":
|
|
return Mxfp4MoEMethod(layer.moe_config)
|
|
# expert_dtype == "fp8": fall through to Fp8Config which
|
|
# returns Fp8MoEMethod with block-wise float32 scales.
|
|
return super().get_quant_method(layer, prefix)
|
|
|
|
def is_mxfp4_quant(self, prefix, layer):
|
|
return isinstance(layer, FusedMoE) and self.expert_dtype == "fp4"
|
|
|
|
|
|
|
|
|
|
def make_deepseek_v4_expert_params_mapping(
|
|
num_experts: int,
|
|
) -> list[tuple[str, str, int, str]]:
|
|
# Checkpoint uses gate_proj/up_proj/down_proj, model params use w13_/w2_
|
|
return [
|
|
(
|
|
"experts.w13_" if shard_id in ("w1", "w3") else "experts.w2_",
|
|
f"experts.{expert_id}.{ckpt_name}.",
|
|
expert_id,
|
|
shard_id,
|
|
)
|
|
for expert_id in range(num_experts)
|
|
for shard_id, ckpt_name in [
|
|
("w1", "gate_proj"),
|
|
("w2", "down_proj"),
|
|
("w3", "up_proj"),
|
|
]
|
|
]
|
|
|
|
|
|
class DeepseekV4MegaMoEExperts(nn.Module):
|
|
"""MegaMoE experts for DeepSeek V4 with NVFP4 quantization.
|
|
|
|
Loads NVFP4 expert weights (E2M1 packed uint8 + float8_e4m3fn block scales
|
|
+ float32 global scales) and runs them through the CuTeDSL NVFP4 kernel.
|
|
|
|
The CuTeDSL kernel is a Python-based CUTLASS kernel compiled via MLIR → PTX.
|
|
It handles NVFP4 natively with full Blackwell pipeline overlap (TMA → MMA → Epilogue).
|
|
This replaces the broken C++ CUTLASS kernel (see README.md for the full story).
|
|
"""
|
|
_cutedsl_runner: 'CuTeDSLMoERunner | None' = None
|
|
_weight_load_count: int = 0
|
|
_weight_load_tqdm: 'tqdm | None' = None
|
|
|
|
# NVFP4 E2M1 lookup table (positive values, sign from bit 3)
|
|
E2M1_LUT = [0.0, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0]
|
|
# MXFP4 E2M1 is the same format
|
|
|
|
def __init__(
|
|
self,
|
|
vllm_config: VllmConfig,
|
|
*,
|
|
num_experts: int,
|
|
num_local_experts: int,
|
|
experts_start_idx: int,
|
|
top_k: int,
|
|
hidden_size: int,
|
|
intermediate_size: int,
|
|
prefix: str = "",
|
|
):
|
|
super().__init__()
|
|
self.prefix = prefix
|
|
self.num_experts = num_experts
|
|
self.num_local_experts = num_local_experts
|
|
self.experts_start_idx = experts_start_idx
|
|
self.experts_end_idx = experts_start_idx + num_local_experts
|
|
self.top_k = top_k
|
|
self.hidden_size = hidden_size
|
|
self.intermediate_size = intermediate_size
|
|
self.max_num_tokens = vllm_config.scheduler_config.max_num_batched_tokens
|
|
self.swiglu_limit = vllm_config.model_config.hf_config.swiglu_limit
|
|
|
|
weight_attrs = {"weight_loader": self.weight_loader}
|
|
|
|
# NVFP4 weights: E2M1 packed as uint8, 2 values per byte
|
|
self.w13_weight = nn.Parameter(
|
|
torch.zeros(
|
|
num_local_experts,
|
|
2 * intermediate_size,
|
|
hidden_size // 2,
|
|
dtype=torch.int8,
|
|
),
|
|
requires_grad=False,
|
|
)
|
|
set_weight_attrs(self.w13_weight, weight_attrs)
|
|
|
|
# NVFP4 block scales: float8_e4m3fn, group_size=16
|
|
# Shape: [num_local_experts, 2*intermediate_size, hidden_size // 16]
|
|
self.w13_weight_scale = nn.Parameter(
|
|
torch.zeros(
|
|
num_local_experts,
|
|
2 * intermediate_size,
|
|
hidden_size // 16,
|
|
dtype=torch.float8_e4m3fn,
|
|
),
|
|
requires_grad=False,
|
|
)
|
|
set_weight_attrs(self.w13_weight_scale, weight_attrs)
|
|
self.w13_weight_scale.quant_method = "block"
|
|
|
|
# NVFP4 global scales: float32, per-expert, per-projection (gate, up)
|
|
# shape (num_local_experts, 2) — one scale for gate_proj, one for up_proj
|
|
self.w13_weight_scale_2 = nn.Parameter(
|
|
torch.zeros(num_local_experts, 2, dtype=torch.float32),
|
|
requires_grad=False,
|
|
)
|
|
set_weight_attrs(self.w13_weight_scale_2, weight_attrs)
|
|
|
|
# NVFP4 activation scales: float32, per-expert
|
|
self.w13_input_scale = nn.Parameter(
|
|
torch.zeros(num_local_experts, dtype=torch.float32),
|
|
requires_grad=False,
|
|
)
|
|
set_weight_attrs(self.w13_input_scale, weight_attrs)
|
|
|
|
self.w2_weight = nn.Parameter(
|
|
torch.zeros(
|
|
num_local_experts,
|
|
hidden_size,
|
|
intermediate_size // 2,
|
|
dtype=torch.int8,
|
|
),
|
|
requires_grad=False,
|
|
)
|
|
set_weight_attrs(self.w2_weight, weight_attrs)
|
|
|
|
# NVFP4 block scales for w2
|
|
self.w2_weight_scale = nn.Parameter(
|
|
torch.zeros(
|
|
num_local_experts,
|
|
hidden_size,
|
|
intermediate_size // 16,
|
|
dtype=torch.float8_e4m3fn,
|
|
),
|
|
requires_grad=False,
|
|
)
|
|
set_weight_attrs(self.w2_weight_scale, weight_attrs)
|
|
self.w2_weight_scale.quant_method = "block"
|
|
|
|
self.w2_weight_scale_2 = nn.Parameter(
|
|
torch.zeros(num_local_experts, dtype=torch.float32),
|
|
requires_grad=False,
|
|
)
|
|
set_weight_attrs(self.w2_weight_scale_2, weight_attrs)
|
|
|
|
self.w2_input_scale = nn.Parameter(
|
|
torch.zeros(num_local_experts, dtype=torch.float32),
|
|
requires_grad=False,
|
|
)
|
|
set_weight_attrs(self.w2_input_scale, weight_attrs)
|
|
|
|
self._cutedsl_runner = None
|
|
|
|
# Register in the static forward context so the custom-op wrapper
|
|
# can look up this module by name from within a torch.compile graph.
|
|
compilation_config = vllm_config.compilation_config
|
|
if prefix in compilation_config.static_forward_context:
|
|
raise ValueError(f"Duplicate layer name: {prefix}")
|
|
compilation_config.static_forward_context[prefix] = self
|
|
|
|
def _map_global_expert_id(self, expert_id: int) -> int:
|
|
if expert_id < self.experts_start_idx or expert_id >= self.experts_end_idx:
|
|
return -1
|
|
return expert_id - self.experts_start_idx
|
|
|
|
def weight_loader(
|
|
self,
|
|
param: nn.Parameter,
|
|
loaded_weight: torch.Tensor,
|
|
weight_name: str,
|
|
shard_id: str,
|
|
expert_id: int,
|
|
) -> bool:
|
|
# Progress bar for k8s/docker liveness during GPU upload
|
|
if DeepseekV4MegaMoEExperts._weight_load_count == 0:
|
|
from tqdm import tqdm as _tqdm
|
|
DeepseekV4MegaMoEExperts._weight_load_tqdm = _tqdm(
|
|
total=self.num_local_experts * 20, # ~20 tensors per expert
|
|
desc=" Loading Native NVFP4 Expert Weights",
|
|
unit="tensor",
|
|
)
|
|
DeepseekV4MegaMoEExperts._weight_load_count += 1
|
|
DeepseekV4MegaMoEExperts._weight_load_tqdm.update(1)
|
|
|
|
local_expert_id = self._map_global_expert_id(expert_id)
|
|
if local_expert_id == -1:
|
|
return False
|
|
|
|
# Scalar params (weight_scale_2, input_scale): per-expert
|
|
if "weight_scale_2" in weight_name or "input_scale" in weight_name:
|
|
if "w13_" in weight_name and "weight_scale_2" in weight_name:
|
|
# w13 is fused gate+up — store gate and up scales separately
|
|
# shard_id tells us which projection: w1=gate, w3=up
|
|
proj_idx = 0 if shard_id == "w1" else 1
|
|
param.data[local_expert_id, proj_idx].copy_(loaded_weight)
|
|
else:
|
|
# w2 or input_scale — single scalar per expert
|
|
param.data[local_expert_id].copy_(loaded_weight)
|
|
return True
|
|
|
|
expert_data = param.data[local_expert_id]
|
|
if shard_id in ("w1", "w3"):
|
|
if "w13_" not in weight_name:
|
|
return False
|
|
shard_offset = 0 if shard_id == "w1" else self.intermediate_size
|
|
expert_data = expert_data.narrow(0, shard_offset, self.intermediate_size)
|
|
elif shard_id == "w2":
|
|
if "w2_" not in weight_name:
|
|
return False
|
|
else:
|
|
raise ValueError(f"Unsupported expert shard id: {shard_id}")
|
|
|
|
if expert_data.shape != loaded_weight.shape:
|
|
raise ValueError(
|
|
f"DeepSeek V4 MegaMoE expert weight shape mismatch for "
|
|
f"{weight_name}: parameter shard {tuple(expert_data.shape)} "
|
|
f"vs checkpoint {tuple(loaded_weight.shape)}"
|
|
)
|
|
expert_data.copy_(loaded_weight)
|
|
return True
|
|
|
|
def _check_runtime_supported(self) -> None:
|
|
if not torch.cuda.is_available():
|
|
raise NotImplementedError("DeepSeek V4 MegaMoE requires CUDA.")
|
|
# After finalize_weights, w13_weight is freed — get device from
|
|
# the stacked tensors in the cutedsl runner instead.
|
|
if self.w13_weight is not None:
|
|
device = self.w13_weight.device
|
|
elif self._cutedsl_runner is not None and self._cutedsl_runner._l1_mat_b is not None:
|
|
device = self._cutedsl_runner._l1_mat_b.device
|
|
else:
|
|
device = torch.device("cuda") # fallback
|
|
if device.type != "cuda":
|
|
raise NotImplementedError(
|
|
"DeepSeek V4 MegaMoE expert weights must be loaded on CUDA."
|
|
)
|
|
if torch.cuda.get_device_capability(device)[0] < 10:
|
|
raise NotImplementedError("DeepGEMM MegaMoE requires SM100 GPUs.")
|
|
if self.hidden_size % 128 != 0 or self.intermediate_size % 128 != 0:
|
|
raise ValueError(
|
|
"DeepGEMM MegaMoE requires hidden and intermediate sizes "
|
|
"to be multiples of 128."
|
|
)
|
|
|
|
def finalize_weights(self) -> None:
|
|
if self._cutedsl_runner is not None and (self._cutedsl_runner.l1_fp4 is not None or self._cutedsl_runner._l1_mat_b is not None):
|
|
return # Already finalized
|
|
|
|
self._check_runtime_supported()
|
|
|
|
# ── Direct NVFP4 path (no BF16 round-trip) ──
|
|
# Checkpoint stores:
|
|
# weight: uint8 packed E2M1 (2 FP4 values/byte) → view as float4_e2m1fn_x2
|
|
# weight_scale: float8_e4m3fn block scales → use directly
|
|
# weight_scale_2: float32 global scale → use directly
|
|
# The only conversion is uint8 → float4_e2m1fn_x2 (byte-preserving view cast).
|
|
#
|
|
# L1 complication: gate and up have different global scales, but the
|
|
# kernel takes one global_scale_b per expert. Solution: normalize to
|
|
# max(gate_gs, up_gs) and fold the ratio into block scales via float32
|
|
# (one multiply + float8 round-trip on the *ratio only* — much better
|
|
# than dequantizing the entire weight matrix through BF16).
|
|
|
|
from vllm.model_executor.models.nvfp4_cutedsl import CuTeDSLMoERunner
|
|
|
|
l1_fp4, l1_sf, l1_gs = [], [], []
|
|
l2_fp4, l2_sf, l2_gs = [], [], []
|
|
|
|
for e in range(self.num_local_experts):
|
|
# ── L1: gate + up (fused) ──
|
|
gate_w = self.w13_weight.data[e, :self.intermediate_size] # (intermediate, hidden//2) uint8
|
|
up_w = self.w13_weight.data[e, self.intermediate_size:] # (intermediate, hidden//2) uint8
|
|
gate_sf = self.w13_weight_scale.data[e, :self.intermediate_size] # (intermediate, hidden//16) float8
|
|
up_sf = self.w13_weight_scale.data[e, self.intermediate_size:]
|
|
gate_gs = self.w13_weight_scale_2.data[e, 0].item() # float32 scalar
|
|
up_gs = self.w13_weight_scale_2.data[e, 1].item()
|
|
|
|
# Fuse gate+up along N dim, then transpose to K-major (K_packed, N)
|
|
# Checkpoint is (N, K_packed) → permute to (K_packed, N) = (hidden//2, 2*intermediate)
|
|
fused_w = torch.cat([gate_w, up_w], dim=0) # (2*intermediate, hidden//2)
|
|
fused_w_fp4 = fused_w.view(torch.float4_e2m1fn_x2).permute(1, 0).contiguous()
|
|
# shape: (hidden//2, 2*intermediate) — K=hidden packed, N=2*intermediate
|
|
|
|
# Fuse block scales: checkpoint is (N, K_sf), bridge expects (K_sf, N)
|
|
fused_sf = torch.cat([gate_sf, up_sf], dim=0) # (2*intermediate, hidden//16) = (N, K_sf)
|
|
# Transpose to (K_sf, N) for assemble_scales_3d_side
|
|
fused_sf = fused_sf.permute(1, 0).contiguous()
|
|
|
|
# Handle dual global scales: normalize to max, fold ratio into block scales
|
|
l1_max_gs = max(gate_gs, up_gs)
|
|
if gate_gs != up_gs:
|
|
fused_sf_f32 = fused_sf.float()
|
|
# After transpose to (K_sf, N): gate is first intermediate cols, up is next
|
|
fused_sf_f32[:, :self.intermediate_size] *= (gate_gs / l1_max_gs)
|
|
fused_sf_f32[:, self.intermediate_size:] *= (up_gs / l1_max_gs)
|
|
fused_sf = fused_sf_f32.to(torch.float8_e4m3fn)
|
|
|
|
l1_fp4.append(fused_w_fp4)
|
|
l1_sf.append(fused_sf)
|
|
l1_gs.append(l1_max_gs)
|
|
|
|
# ── L2: down (single projection, straightforward) ──
|
|
down_w = self.w2_weight.data[e] # (hidden, intermediate//2) uint8
|
|
down_sf = self.w2_weight_scale.data[e] # (hidden, intermediate//16) float8
|
|
down_gs = self.w2_weight_scale_2.data[e].item() # float32 scalar
|
|
|
|
# Checkpoint is (N, K_packed) → permute to (K_packed, N)
|
|
# K=intermediate (packed dim), N=hidden
|
|
down_w_fp4 = down_w.view(torch.float4_e2m1fn_x2).permute(1, 0).contiguous()
|
|
# shape: (intermediate//2, hidden) — K=intermediate packed, N=hidden
|
|
|
|
# Block scales: checkpoint is (N, K_sf), bridge expects (K_sf, N)
|
|
down_sf = down_sf.permute(1, 0).contiguous()
|
|
|
|
l2_fp4.append(down_w_fp4)
|
|
l2_sf.append(down_sf)
|
|
l2_gs.append(down_gs)
|
|
|
|
# Create CuTeDSL runner with directly-cast weights
|
|
self._cutedsl_runner = CuTeDSLMoERunner(
|
|
num_experts=self.num_local_experts,
|
|
hidden_size=self.hidden_size,
|
|
intermediate_size=self.intermediate_size,
|
|
max_num_tokens=self.max_num_tokens,
|
|
top_k=self.top_k,
|
|
device=l1_fp4[0].device,
|
|
experts_start_idx=self.experts_start_idx,
|
|
)
|
|
self._cutedsl_runner.l1_fp4 = l1_fp4
|
|
self._cutedsl_runner.l1_sf = l1_sf
|
|
self._cutedsl_runner.l1_gs = l1_gs
|
|
self._cutedsl_runner.l2_fp4 = l2_fp4
|
|
self._cutedsl_runner.l2_sf = l2_sf
|
|
self._cutedsl_runner.l2_gs = l2_gs
|
|
|
|
# Set activation global scales from checkpoint input_scale
|
|
# The input_scale is the pre-computed activation normalization factor.
|
|
# w13_input_scale shape: (num_experts, 2) for gate+up, but may be (num_experts,) after EP split
|
|
# w2_input_scale shape: (num_experts, 1) or (num_experts,)
|
|
w13_igs = self.w13_input_scale.data
|
|
w2_igs = self.w2_input_scale.data
|
|
if w13_igs.dim() == 2:
|
|
l1_igs = w13_igs[:, 0] # gate input_scale
|
|
else:
|
|
l1_igs = w13_igs # already 1D per expert
|
|
if w2_igs.dim() == 2:
|
|
l2_igs = w2_igs[:, 0]
|
|
else:
|
|
l2_igs = w2_igs
|
|
# Use checkpoint input_scale as initial guess, then warmup will override
|
|
self._cutedsl_runner._l1_activation_global_scale = l1_igs.mean().item()
|
|
self._cutedsl_runner._l2_activation_global_scale = l2_igs.mean().item()
|
|
|
|
# Drop the original loader-side parameters
|
|
self._w13_input_scale = self.w13_input_scale.data.clone()
|
|
self._w2_input_scale = self.w2_input_scale.data.clone()
|
|
self.w13_weight = None
|
|
self.w13_weight_scale = None
|
|
self.w13_weight_scale_2 = None
|
|
self.w13_input_scale = None
|
|
self.w2_weight = None
|
|
self.w2_weight_scale = None
|
|
self.w2_weight_scale_2 = None
|
|
self.w2_input_scale = None
|
|
|
|
# Warmup: compute actual activation global scales from sample data.
|
|
# The checkpoint input_scale is a calibration value that doesn't match
|
|
# runtime activation magnitudes. We run a small forward pass to observe
|
|
# the actual amax and compute correct gs values.
|
|
self._warmup_activation_global_scales()
|
|
|
|
# Set swiglu_limit for activation clamping in the runner
|
|
if self.swiglu_limit is not None:
|
|
self._cutedsl_runner.set_swiglu_limit(float(self.swiglu_limit))
|
|
|
|
def _warmup_activation_global_scales(self) -> None:
|
|
"""Run a warmup forward pass to compute correct activation global scales.
|
|
|
|
Called once per layer during finalize_weights, before cudagraph capture.
|
|
Uses quantize_to_nvfp4 (which calls .max()) to get the exact gs
|
|
from real activation magnitudes, then stores them for use by
|
|
quantize_activation_nvfp4 (no .max(), cudagraph-safe).
|
|
"""
|
|
import torch
|
|
runner = self._cutedsl_runner
|
|
device = runner.device
|
|
num_tokens = min(8, runner.max_num_tokens)
|
|
top_k = runner.top_k
|
|
|
|
with torch.no_grad():
|
|
# Sample hidden states: typical BF16 activations have amax ~1-10
|
|
hidden_states = torch.randn(num_tokens, runner.hidden_size,
|
|
dtype=torch.bfloat16, device=device)
|
|
# Assign all tokens to local experts (0..num_local_experts-1)
|
|
# compute_activation_global_scales expects local IDs, not global
|
|
topk_ids = torch.zeros(num_tokens, top_k, dtype=torch.int64, device=device)
|
|
for i in range(num_tokens):
|
|
for j in range(top_k):
|
|
topk_ids[i, j] = j % runner.num_experts
|
|
topk_weights = torch.ones(num_tokens, top_k, dtype=torch.float32, device=device) / top_k
|
|
|
|
runner.compute_activation_global_scales(
|
|
hidden_states, topk_weights, topk_ids
|
|
)
|
|
|
|
# Note: No explicit CuTeDSL warmup here. With FULL_AND_PIECEWISE
|
|
# CUDA graph mode, the kernel compiles during graph capture (startup).
|
|
# In eager mode, the first inference triggers JIT compilation.
|
|
|
|
def forward(
|
|
self,
|
|
hidden_states: torch.Tensor,
|
|
topk_weights: torch.Tensor,
|
|
topk_ids: torch.Tensor,
|
|
*,
|
|
activation_clamp: float | None,
|
|
fast_math: bool = True,
|
|
) -> torch.Tensor:
|
|
if hidden_states.shape[0] > self.max_num_tokens:
|
|
raise ValueError(
|
|
f"DeepSeek V4 MegaMoE got {hidden_states.shape[0]} tokens, "
|
|
f"but the symmetric buffer was sized for {self.max_num_tokens}."
|
|
)
|
|
y = torch.empty_like(hidden_states, dtype=torch.bfloat16)
|
|
torch.ops.vllm.deepseek_v4_mega_moe_experts(
|
|
hidden_states,
|
|
topk_weights,
|
|
topk_ids,
|
|
y,
|
|
self.prefix,
|
|
activation_clamp,
|
|
fast_math,
|
|
)
|
|
return y
|
|
|
|
def _run_mega_moe(
|
|
self,
|
|
hidden_states: torch.Tensor,
|
|
topk_weights: torch.Tensor,
|
|
topk_ids: torch.Tensor,
|
|
y: torch.Tensor,
|
|
activation_clamp: float | None,
|
|
fast_math: bool,
|
|
) -> None:
|
|
import os
|
|
|
|
# This method must have been already called during the weight loading phase.
|
|
# We call it again here to cover the dummy weight loading case.
|
|
self.finalize_weights()
|
|
|
|
assert self._cutedsl_runner is not None
|
|
# After _ensure_stacked, per-expert lists are freed and stacked
|
|
# tensors live in _l1_mat_b / _l2_mat_b instead.
|
|
assert (self._cutedsl_runner.l1_fp4 is not None
|
|
or self._cutedsl_runner._l1_mat_b is not None)
|
|
|
|
# Build expert indices list for this rank
|
|
expert_indices = list(range(self.num_local_experts))
|
|
|
|
try:
|
|
result = self._cutedsl_runner.run(
|
|
hidden_states, topk_weights, topk_ids,
|
|
expert_indices=expert_indices,
|
|
)
|
|
y.copy_(result)
|
|
except Exception as exc:
|
|
import traceback
|
|
traceback.print_exc()
|
|
raise
|
|
|
|
|
|
DeepseekV4MegaMoEExperts.weight_loader.supports_moe_loading = True # type: ignore[attr-defined]
|
|
|
|
|
|
def _deepseek_v4_mega_moe_experts_op(
|
|
hidden_states: torch.Tensor,
|
|
topk_weights: torch.Tensor,
|
|
topk_ids: torch.Tensor,
|
|
out: torch.Tensor,
|
|
layer_name: str,
|
|
activation_clamp: float | None,
|
|
fast_math: bool,
|
|
) -> None:
|
|
self = get_forward_context().no_compile_layers[layer_name]
|
|
self._run_mega_moe(
|
|
hidden_states,
|
|
topk_weights,
|
|
topk_ids,
|
|
out,
|
|
activation_clamp,
|
|
fast_math,
|
|
)
|
|
|
|
|
|
def _deepseek_v4_mega_moe_experts_op_fake(
|
|
hidden_states: torch.Tensor,
|
|
topk_weights: torch.Tensor,
|
|
topk_ids: torch.Tensor,
|
|
out: torch.Tensor,
|
|
layer_name: str,
|
|
activation_clamp: float | None,
|
|
fast_math: bool,
|
|
) -> None:
|
|
return None
|
|
|
|
|
|
direct_register_custom_op(
|
|
op_name="deepseek_v4_mega_moe_experts",
|
|
op_func=_deepseek_v4_mega_moe_experts_op,
|
|
mutates_args=["out"],
|
|
fake_impl=_deepseek_v4_mega_moe_experts_op_fake,
|
|
)
|
|
|
|
|
|
class DeepseekV4MoE(nn.Module):
|
|
def __init__(
|
|
self,
|
|
vllm_config: VllmConfig,
|
|
prefix: str = "",
|
|
):
|
|
super().__init__()
|
|
|
|
self.tp_size = get_tensor_model_parallel_world_size()
|
|
config = vllm_config.model_config.hf_config
|
|
quant_config = vllm_config.quant_config
|
|
self.prefix = prefix
|
|
self.use_mega_moe = True # Force mega_moe for NVFP4 pipeline
|
|
if self.use_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.routed_scaling_factor = getattr(config, "routed_scaling_factor", 1.0)
|
|
self.hidden_size = config.hidden_size
|
|
|
|
self.n_routed_experts = config.n_routed_experts
|
|
self.n_activated_experts = config.num_experts_per_tok
|
|
self.moe_intermediate_size = config.moe_intermediate_size
|
|
self.swiglu_limit = config.swiglu_limit
|
|
self.renormalize = config.norm_topk_prob
|
|
self.scoring_func = getattr(config, "scoring_func", "sqrtsoftplus")
|
|
if self.use_mega_moe and self.scoring_func != "sqrtsoftplus":
|
|
raise NotImplementedError(
|
|
"DeepSeek V4 MegaMoE currently supports sqrtsoftplus routing only."
|
|
)
|
|
# NVFP4 experts work with mega_moe via NVFP4 weight transformation in finalize_weights
|
|
|
|
self.gate = GateLinear(
|
|
config.hidden_size,
|
|
config.n_routed_experts,
|
|
out_dtype=torch.float32,
|
|
bias=False,
|
|
prefix=f"{prefix}.gate",
|
|
)
|
|
self.gate.e_score_correction_bias = None
|
|
self.gate.tid2eid = None
|
|
is_hash_moe = extract_layer_index(prefix) < config.num_hash_layers
|
|
self.hash_indices_dtype = torch.int64 if self.use_mega_moe else torch.int32
|
|
|
|
if is_hash_moe:
|
|
# hash MoE doesn't use e_score_correction_bias
|
|
# Use randint instead of empty to avoid garbage values causing
|
|
# invalid memory access in dummy mode (--load-format="dummy")
|
|
self.gate.tid2eid = nn.Parameter(
|
|
torch.randint(
|
|
0,
|
|
config.n_routed_experts,
|
|
(config.vocab_size, config.num_experts_per_tok),
|
|
dtype=self.hash_indices_dtype,
|
|
),
|
|
requires_grad=False,
|
|
)
|
|
elif getattr(config, "topk_method", None) == "noaux_tc":
|
|
self.gate.e_score_correction_bias = nn.Parameter(
|
|
torch.empty(config.n_routed_experts, dtype=torch.float32),
|
|
requires_grad=False,
|
|
)
|
|
|
|
if config.n_shared_experts is None:
|
|
self.shared_experts = None
|
|
else:
|
|
intermediate_size = config.moe_intermediate_size * config.n_shared_experts
|
|
|
|
self.shared_experts = DeepseekV4MLP(
|
|
hidden_size=config.hidden_size,
|
|
intermediate_size=intermediate_size,
|
|
hidden_act=config.hidden_act,
|
|
swiglu_limit=self.swiglu_limit,
|
|
quant_config=quant_config,
|
|
reduce_results=self.use_mega_moe,
|
|
prefix=f"{prefix}.shared_experts",
|
|
)
|
|
|
|
if self.use_mega_moe:
|
|
self._init_mega_moe_experts(vllm_config, config, prefix)
|
|
else:
|
|
self._init_fused_moe_experts(config, quant_config, prefix)
|
|
|
|
def _init_mega_moe_experts(
|
|
self,
|
|
vllm_config: VllmConfig,
|
|
config,
|
|
prefix: str,
|
|
) -> None:
|
|
self.ep_group = get_ep_group()
|
|
self.ep_size = self.ep_group.world_size
|
|
self.ep_rank = self.ep_group.rank_in_group
|
|
assert config.n_routed_experts % self.ep_size == 0
|
|
|
|
self.n_local_experts = config.n_routed_experts // self.ep_size
|
|
self.experts_start_idx = self.ep_rank * self.n_local_experts
|
|
self.experts_end_idx = self.experts_start_idx + self.n_local_experts
|
|
|
|
self.experts = DeepseekV4MegaMoEExperts(
|
|
vllm_config,
|
|
num_experts=config.n_routed_experts,
|
|
num_local_experts=self.n_local_experts,
|
|
experts_start_idx=self.experts_start_idx,
|
|
top_k=config.num_experts_per_tok,
|
|
hidden_size=config.hidden_size,
|
|
intermediate_size=config.moe_intermediate_size,
|
|
prefix=f"{prefix}.experts",
|
|
)
|
|
|
|
def _init_fused_moe_experts(
|
|
self,
|
|
config,
|
|
quant_config,
|
|
prefix: str,
|
|
) -> None:
|
|
self.tp_rank = get_tensor_model_parallel_rank()
|
|
assert config.n_routed_experts % self.tp_size == 0
|
|
|
|
self.n_local_experts = config.n_routed_experts // self.tp_size
|
|
self.experts_start_idx = self.tp_rank * self.n_local_experts
|
|
self.experts_end_idx = self.experts_start_idx + self.n_local_experts
|
|
|
|
self.experts = FusedMoE(
|
|
shared_experts=self.shared_experts,
|
|
gate=self.gate,
|
|
num_experts=config.n_routed_experts,
|
|
top_k=config.num_experts_per_tok,
|
|
hidden_size=config.hidden_size,
|
|
intermediate_size=config.moe_intermediate_size,
|
|
renormalize=config.norm_topk_prob,
|
|
quant_config=quant_config,
|
|
prefix=f"{prefix}.experts",
|
|
scoring_func=self.scoring_func,
|
|
routed_scaling_factor=self.routed_scaling_factor,
|
|
e_score_correction_bias=self.gate.e_score_correction_bias,
|
|
hash_indices_table=self.gate.tid2eid,
|
|
swiglu_limit=self.swiglu_limit,
|
|
router_logits_dtype=torch.float32,
|
|
)
|
|
|
|
def forward(
|
|
self, hidden_states: torch.Tensor, input_ids: torch.Tensor | None = None
|
|
) -> torch.Tensor:
|
|
if self.gate.tid2eid is not None and input_ids is None:
|
|
raise ValueError("DeepSeek V4 hash MoE routing requires input_ids.")
|
|
|
|
if not self.use_mega_moe:
|
|
return self._forward_fused_moe(hidden_states, input_ids)
|
|
|
|
org_shape = hidden_states.shape
|
|
router_logits, _ = self.gate(hidden_states)
|
|
topk_weights, topk_ids = fused_topk_bias(
|
|
hidden_states=hidden_states,
|
|
gating_output=router_logits,
|
|
scoring_func=self.scoring_func,
|
|
e_score_correction_bias=self.gate.e_score_correction_bias.data
|
|
if self.gate.e_score_correction_bias is not None
|
|
else None,
|
|
topk=self.n_activated_experts,
|
|
renormalize=self.renormalize,
|
|
indices_type=self.hash_indices_dtype,
|
|
input_tokens=input_ids,
|
|
hash_indices_table=self.gate.tid2eid,
|
|
routed_scaling_factor=self.routed_scaling_factor,
|
|
)
|
|
activation_clamp = (
|
|
float(self.swiglu_limit) if self.swiglu_limit is not None else None
|
|
)
|
|
final_hidden_states = self.experts(
|
|
hidden_states,
|
|
topk_weights,
|
|
topk_ids,
|
|
activation_clamp=activation_clamp,
|
|
)
|
|
|
|
# EP all-reduce: each rank only computes its local experts,
|
|
# so we must sum across EP ranks to get the full routed output.
|
|
torch.distributed.all_reduce(
|
|
final_hidden_states, group=self.ep_group.device_group
|
|
)
|
|
|
|
if self.shared_experts is not None:
|
|
shared_output = self.shared_experts(hidden_states)
|
|
final_hidden_states += shared_output
|
|
|
|
return final_hidden_states.view(org_shape)
|
|
|
|
def _forward_fused_moe(
|
|
self, hidden_states: torch.Tensor, input_ids: torch.Tensor | None = None
|
|
) -> torch.Tensor:
|
|
org_shape = hidden_states.shape
|
|
if self.experts.is_internal_router:
|
|
# In this case, the gate/router runs inside the FusedMoE class
|
|
final_hidden_states = self.experts(
|
|
hidden_states=hidden_states,
|
|
router_logits=hidden_states,
|
|
input_ids=input_ids,
|
|
)
|
|
else:
|
|
router_logits, _ = self.gate(hidden_states)
|
|
final_hidden_states = self.experts(
|
|
hidden_states=hidden_states,
|
|
router_logits=router_logits,
|
|
input_ids=input_ids,
|
|
)
|
|
|
|
return final_hidden_states.view(org_shape)
|
|
|
|
def finalize_mega_moe_weights(self) -> None:
|
|
if self.use_mega_moe:
|
|
self.experts.finalize_weights()
|
|
|
|
|
|
class DeepseekV4Attention(nn.Module):
|
|
def __init__(
|
|
self,
|
|
vllm_config: VllmConfig,
|
|
prefix: str,
|
|
topk_indices_buffer: torch.Tensor | None = None,
|
|
aux_stream_list: list[torch.cuda.Stream] | None = None,
|
|
):
|
|
super().__init__()
|
|
config = vllm_config.model_config.hf_config
|
|
quant_config = vllm_config.quant_config
|
|
layer_id = extract_layer_index(prefix)
|
|
|
|
self.layer_id = layer_id
|
|
self.hidden_size = config.hidden_size
|
|
self.n_heads = config.num_attention_heads
|
|
tp_size = get_tensor_model_parallel_world_size()
|
|
assert self.n_heads % tp_size == 0
|
|
|
|
self.n_local_heads = self.n_heads // tp_size
|
|
self.q_lora_rank = config.q_lora_rank
|
|
self.o_lora_rank = config.o_lora_rank
|
|
self.head_dim = config.head_dim
|
|
self.rope_head_dim = config.qk_rope_head_dim
|
|
self.nope_head_dim = self.head_dim - self.rope_head_dim
|
|
self.n_groups = config.o_groups
|
|
self.n_local_groups = self.n_groups // tp_size
|
|
self.window_size = config.sliding_window
|
|
# NOTE(zyongye) Compress ratio can't be 0
|
|
# we do this for because MTP layer is not included
|
|
# in the compress ratio list
|
|
if layer_id < config.num_hidden_layers:
|
|
self.compress_ratio = max(1, config.compress_ratios[layer_id])
|
|
else:
|
|
self.compress_ratio = 1
|
|
self.eps = config.rms_norm_eps
|
|
self.max_position_embeddings = config.max_position_embeddings
|
|
|
|
# Padded to min 64 heads for FlashMLA, initialized to -inf
|
|
# (no sink effect). Weight loading fills the first n_local_heads slots.
|
|
padded_heads = max(self.n_local_heads, 64)
|
|
self.attn_sink = nn.Parameter(
|
|
torch.full((padded_heads,), -float("inf"), dtype=torch.float32),
|
|
requires_grad=False,
|
|
)
|
|
|
|
self.fused_wqa_wkv = MergedColumnParallelLinear(
|
|
self.hidden_size,
|
|
[self.q_lora_rank, self.head_dim],
|
|
bias=False,
|
|
quant_config=quant_config,
|
|
prefix=f"{prefix}.fused_wqa_wkv",
|
|
disable_tp=True, # fused ReplicatedLinear
|
|
)
|
|
self.q_norm = RMSNorm(self.q_lora_rank, self.eps)
|
|
self.wq_b = ColumnParallelLinear(
|
|
self.q_lora_rank,
|
|
self.n_heads * self.head_dim,
|
|
bias=False,
|
|
quant_config=quant_config,
|
|
return_bias=False,
|
|
prefix=f"{prefix}.wq_b",
|
|
)
|
|
|
|
self.kv_norm = RMSNorm(self.head_dim, self.eps)
|
|
self.wo_a = ColumnParallelLinear(
|
|
self.n_heads * self.head_dim // self.n_groups,
|
|
self.n_groups * self.o_lora_rank,
|
|
bias=False,
|
|
quant_config=quant_config,
|
|
return_bias=False,
|
|
prefix=f"{prefix}.wo_a",
|
|
)
|
|
self.wo_a.is_bmm = True
|
|
self.wo_a.bmm_batch_size = self.n_local_groups
|
|
self.wo_b = RowParallelLinear(
|
|
self.n_groups * self.o_lora_rank,
|
|
self.hidden_size,
|
|
bias=False,
|
|
quant_config=quant_config,
|
|
return_bias=False,
|
|
prefix=f"{prefix}.wo_b",
|
|
)
|
|
self.softmax_scale = self.head_dim**-0.5
|
|
self.scale_fmt = config.quantization_config["scale_fmt"]
|
|
|
|
self.rope_parameters = config.rope_scaling
|
|
|
|
# Initialize rotary embedding BEFORE DeepseekV4MLAModules (which needs it)
|
|
rope_parameters = dict(config.rope_parameters)
|
|
rope_parameters["rope_theta"] = (
|
|
config.compress_rope_theta if self.compress_ratio > 1 else config.rope_theta
|
|
)
|
|
if config.rope_parameters["rope_type"] != "default":
|
|
config.rope_parameters["rope_type"] = (
|
|
"deepseek_yarn"
|
|
if config.rope_parameters.get("apply_yarn_scaling", True)
|
|
else "deepseek_llama_scaling"
|
|
)
|
|
rope_parameters["mscale"] = 0 # Disable mscale
|
|
rope_parameters["mscale_all_dim"] = 0 # Disable mscale
|
|
rope_parameters["is_deepseek_v4"] = True
|
|
rope_parameters["rope_dim"] = self.rope_head_dim
|
|
self.rotary_emb = get_rope(
|
|
self.head_dim,
|
|
max_position=self.max_position_embeddings,
|
|
rope_parameters=rope_parameters,
|
|
is_neox_style=False,
|
|
)
|
|
|
|
self.indexer = None
|
|
if self.compress_ratio == 4:
|
|
# Only C4A uses sparse attention and hence has indexer.
|
|
self.indexer = DeepseekV4Indexer(
|
|
vllm_config,
|
|
config=config,
|
|
hidden_size=self.hidden_size,
|
|
q_lora_rank=self.q_lora_rank,
|
|
quant_config=quant_config,
|
|
cache_config=vllm_config.cache_config,
|
|
topk_indices_buffer=topk_indices_buffer,
|
|
compress_ratio=self.compress_ratio,
|
|
prefix=f"{prefix}.indexer",
|
|
)
|
|
|
|
mla_modules = DeepseekV4MLAModules(
|
|
vllm_config=vllm_config,
|
|
fused_wqa_wkv=self.fused_wqa_wkv,
|
|
q_norm=self.q_norm,
|
|
wq_b=self.wq_b,
|
|
kv_norm=self.kv_norm,
|
|
wo_a=self.wo_a,
|
|
wo_b=self.wo_b,
|
|
attn_sink=self.attn_sink,
|
|
rotary_emb=self.rotary_emb,
|
|
indexer=self.indexer,
|
|
indexer_rotary_emb=self.rotary_emb,
|
|
topk_indices_buffer=topk_indices_buffer,
|
|
aux_stream_list=aux_stream_list,
|
|
)
|
|
self.mla_attn = DeepseekV4MultiHeadLatentAttentionWrapper(
|
|
hidden_size=self.hidden_size,
|
|
num_heads=self.n_local_heads,
|
|
head_dim=self.head_dim,
|
|
scale=self.softmax_scale,
|
|
qk_nope_head_dim=self.nope_head_dim,
|
|
qk_rope_head_dim=self.rope_head_dim,
|
|
v_head_dim=self.head_dim,
|
|
q_lora_rank=self.q_lora_rank,
|
|
kv_lora_rank=self.head_dim,
|
|
o_lora_rank=self.o_lora_rank,
|
|
mla_modules=mla_modules,
|
|
window_size=self.window_size,
|
|
compress_ratio=self.compress_ratio,
|
|
cache_config=vllm_config.cache_config,
|
|
quant_config=quant_config,
|
|
prefix=prefix,
|
|
)
|
|
|
|
def forward(
|
|
self,
|
|
positions: torch.Tensor,
|
|
hidden_states: torch.Tensor,
|
|
llama_4_scaling: torch.Tensor | None,
|
|
):
|
|
return self.mla_attn(positions, hidden_states, llama_4_scaling)
|
|
|
|
|
|
class DeepseekV4DecoderLayer(nn.Module):
|
|
def __init__(
|
|
self,
|
|
vllm_config,
|
|
prefix,
|
|
topk_indices_buffer: torch.Tensor | None = None,
|
|
aux_stream_list: list[torch.cuda.Stream] | None = None,
|
|
):
|
|
super().__init__()
|
|
|
|
# Lazy import to avoid top-level tilelang dependency.
|
|
# Registers both torch.ops.vllm.mhc_pre and mhc_post
|
|
import vllm.model_executor.layers.mhc # noqa: F401
|
|
|
|
config = vllm_config.model_config.hf_config
|
|
self.hidden_size = config.hidden_size
|
|
|
|
self.rms_norm_eps = config.rms_norm_eps
|
|
self.attn = DeepseekV4Attention(
|
|
vllm_config,
|
|
prefix=f"{prefix}.attn",
|
|
topk_indices_buffer=topk_indices_buffer,
|
|
aux_stream_list=aux_stream_list,
|
|
)
|
|
self.ffn = DeepseekV4MoE(vllm_config, prefix=f"{prefix}.ffn")
|
|
|
|
self.attn_norm = RMSNorm(self.hidden_size, self.rms_norm_eps)
|
|
self.ffn_norm = RMSNorm(self.hidden_size, self.rms_norm_eps)
|
|
self.hc_mult = config.hc_mult
|
|
self.hc_sinkhorn_iters = config.hc_sinkhorn_iters
|
|
self.hc_eps = config.hc_eps
|
|
self.hc_post_alpha = 2.0
|
|
mix_hc = (2 + self.hc_mult) * self.hc_mult
|
|
hc_dim = self.hc_mult * self.hidden_size
|
|
self.hc_attn_fn = nn.Parameter(
|
|
torch.empty(
|
|
(mix_hc, hc_dim),
|
|
dtype=torch.float32,
|
|
),
|
|
requires_grad=False,
|
|
)
|
|
self.hc_ffn_fn = nn.Parameter(
|
|
torch.empty(
|
|
(mix_hc, hc_dim),
|
|
dtype=torch.float32,
|
|
),
|
|
requires_grad=False,
|
|
)
|
|
self.hc_attn_base = nn.Parameter(
|
|
torch.empty(
|
|
mix_hc,
|
|
dtype=torch.float32,
|
|
),
|
|
requires_grad=False,
|
|
)
|
|
self.hc_ffn_base = nn.Parameter(
|
|
torch.empty(
|
|
mix_hc,
|
|
dtype=torch.float32,
|
|
),
|
|
requires_grad=False,
|
|
)
|
|
self.hc_attn_scale = nn.Parameter(
|
|
torch.empty(
|
|
3,
|
|
dtype=torch.float32,
|
|
),
|
|
requires_grad=False,
|
|
)
|
|
self.hc_ffn_scale = nn.Parameter(
|
|
torch.empty(
|
|
3,
|
|
dtype=torch.float32,
|
|
),
|
|
requires_grad=False,
|
|
)
|
|
|
|
def hc_pre(
|
|
self,
|
|
x: torch.Tensor,
|
|
hc_fn: torch.Tensor,
|
|
hc_scale: torch.Tensor,
|
|
hc_base: torch.Tensor,
|
|
):
|
|
post_mix, res_mix, layer_input = torch.ops.vllm.mhc_pre(
|
|
residual=x,
|
|
fn=hc_fn,
|
|
hc_scale=hc_scale,
|
|
hc_base=hc_base,
|
|
rms_eps=self.rms_norm_eps,
|
|
hc_pre_eps=self.hc_eps,
|
|
hc_sinkhorn_eps=self.hc_eps,
|
|
hc_post_mult_value=self.hc_post_alpha,
|
|
sinkhorn_repeat=self.hc_sinkhorn_iters,
|
|
)
|
|
return layer_input, post_mix, res_mix
|
|
|
|
def hc_post(
|
|
self,
|
|
x: torch.Tensor,
|
|
residual: torch.Tensor,
|
|
post: torch.Tensor,
|
|
comb: torch.Tensor,
|
|
):
|
|
return torch.ops.vllm.mhc_post(x, residual, post, comb)
|
|
|
|
def forward(
|
|
self,
|
|
x: torch.Tensor,
|
|
positions: torch.Tensor,
|
|
input_ids: torch.Tensor | None,
|
|
) -> torch.Tensor:
|
|
residual = x
|
|
x, post, comb = self.hc_pre(
|
|
x, self.hc_attn_fn, self.hc_attn_scale, self.hc_attn_base
|
|
)
|
|
x = self.attn_norm(x)
|
|
x = self.attn(positions, x, None)
|
|
x = self.hc_post(x, residual, post, comb)
|
|
|
|
residual = x
|
|
x, post, comb = self.hc_pre(
|
|
x, self.hc_ffn_fn, self.hc_ffn_scale, self.hc_ffn_base
|
|
)
|
|
x = self.ffn_norm(x)
|
|
x = self.ffn(x, input_ids)
|
|
x = self.hc_post(x, residual, post, comb)
|
|
return x
|
|
|
|
|
|
@support_torch_compile
|
|
class DeepseekV4Model(nn.Module):
|
|
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.use_mega_moe = True # Force mega_moe for NVFP4 pipeline
|
|
if self.use_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
|
|
# DeepseekV4MultiHeadLatentAttentionWrapper.attn_gemm_parallel_execute
|
|
# (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)]
|
|
|
|
self.device = current_platform.device_type
|
|
# 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,
|
|
device=self.device,
|
|
)
|
|
|
|
self.embed_tokens = VocabParallelEmbedding(
|
|
config.vocab_size,
|
|
config.hidden_size,
|
|
quant_config=quant_config,
|
|
prefix=f"{prefix}.embed_tokens",
|
|
)
|
|
|
|
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.norm = RMSNorm(config.hidden_size, self.rms_norm_eps)
|
|
|
|
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,
|
|
)
|
|
|
|
# Pre-hc_head residual stream buffer for the MTP draft. Stable
|
|
# address (outside the cudagraph pool) so the copy_ in forward()
|
|
# refreshes it correctly across captured shapes.
|
|
self._mtp_hidden_buffer = torch.empty(
|
|
vllm_config.scheduler_config.max_num_batched_tokens,
|
|
self.hc_dim,
|
|
dtype=vllm_config.model_config.dtype,
|
|
device=self.device,
|
|
)
|
|
|
|
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,
|
|
intermediate_tensors: IntermediateTensors | None,
|
|
inputs_embeds: torch.Tensor | None = None,
|
|
) -> torch.Tensor | IntermediateTensors:
|
|
hidden_states = self.embed_input_ids(input_ids)
|
|
hidden_states = hidden_states.unsqueeze(-2).repeat(1, self.hc_mult, 1)
|
|
if self.use_mega_moe:
|
|
input_ids = input_ids.to(torch.int64)
|
|
for layer in islice(self.layers, self.start_layer, self.end_layer):
|
|
hidden_states = layer(
|
|
hidden_states,
|
|
positions,
|
|
input_ids,
|
|
)
|
|
|
|
# Stash pre-hc_head residual for the MTP draft (captured copy_).
|
|
num_tokens = hidden_states.shape[0]
|
|
self._mtp_hidden_buffer[:num_tokens].copy_(hidden_states.flatten(1))
|
|
|
|
hidden_states = hc_head(
|
|
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)
|
|
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),
|
|
]
|
|
|
|
# Checkpoint key → model param name substitutions.
|
|
# Applied to each (name, weight) pair before matching against
|
|
# params_dict. Order matters: longer/more-specific patterns first.
|
|
CKPT_KEY_SUBST = {
|
|
# self_attn projection names → vLLM attn attribute names
|
|
".self_attn.q_a_proj.": ".attn.wq_a.",
|
|
".self_attn.q_b_proj.": ".attn.wq_b.",
|
|
".self_attn.q_a_norm.": ".attn.q_norm.",
|
|
".self_attn.o_a_proj.": ".attn.wo_a.",
|
|
".self_attn.o_b_proj.": ".attn.wo_b.",
|
|
".self_attn.sinks": ".attn.attn_sink",
|
|
".self_attn.kv_proj.": ".attn.wkv.",
|
|
".self_attn.kv_norm.": ".attn.kv_norm.",
|
|
# Indexer: self_attn.compressor.indexer → attn.indexer
|
|
# MUST come before the generic .self_attn.compressor. rule
|
|
".self_attn.compressor.indexer.q_b_proj.": ".attn.indexer.wq_b.",
|
|
".self_attn.compressor.indexer.kv_norm.": ".attn.indexer.k_norm.",
|
|
".self_attn.compressor.indexer.position_bias": ".attn.indexer.compressor.ape",
|
|
".self_attn.compressor.indexer.gate_proj.": ".attn.indexer.compressor.wgate.",
|
|
".self_attn.compressor.indexer.kv_proj.": ".attn.indexer.compressor.wkv.",
|
|
".self_attn.compressor.indexer.": ".attn.indexer.",
|
|
# Compressor: self_attn.compressor → attn.mla_attn.compressor
|
|
# Compressor projections for stacking (fused_wkv_wgate)
|
|
".self_attn.compressor.kv_proj.": ".attn.mla_attn.compressor.wkv.",
|
|
".self_attn.compressor.gate_proj.": ".attn.mla_attn.compressor.gate.",
|
|
".self_attn.compressor.kv_norm.": ".attn.kv_norm.",
|
|
".self_attn.compressor.position_bias": ".attn.mla_attn.compressor.ape",
|
|
".self_attn.compressor.": ".attn.mla_attn.compressor.",
|
|
# Shared expert projections (stacking into gate_up_proj)
|
|
# Must include .mlp. prefix since break prevents .mlp.→.ffn. from
|
|
# firing on the same key after these patterns match.
|
|
".mlp.shared_experts.gate_proj.": ".ffn.shared_experts.w1.",
|
|
".mlp.shared_experts.up_proj.": ".ffn.shared_experts.w3.",
|
|
".mlp.shared_experts.down_proj.": ".ffn.shared_experts.down_proj.",
|
|
# Hadamard coding params: checkpoint has .attn_hc.base/fn/scale
|
|
# and .ffn_hc.base/fn/scale; model has hc_attn_base/fn/scale
|
|
# and hc_ffn_base/fn/scale (underscore not dot before base/fn/scale)
|
|
".attn_hc.base": ".hc_attn_base",
|
|
".attn_hc.fn": ".hc_attn_fn",
|
|
".attn_hc.scale": ".hc_attn_scale",
|
|
".ffn_hc.base": ".hc_ffn_base",
|
|
".ffn_hc.fn": ".hc_ffn_fn",
|
|
".ffn_hc.scale": ".hc_ffn_scale",
|
|
"hc_head.hc_base": "hc_head_base",
|
|
"hc_head.hc_fn": "hc_head_fn",
|
|
"hc_head.hc_scale": "hc_head_scale",
|
|
# compressor.position_bias → compressor.ape
|
|
".compressor.position_bias": ".compressor.ape",
|
|
# modelopt uses mlp, vllm uses ffn internally
|
|
".mlp.": ".ffn.",
|
|
}
|
|
|
|
params_dict = dict(self.named_parameters())
|
|
loaded_params: set[str] = set()
|
|
|
|
# TP for attention
|
|
tp_size = get_tensor_model_parallel_world_size()
|
|
tp_rank = get_tensor_model_parallel_rank()
|
|
n_head = self.config.num_attention_heads
|
|
n_local_head = n_head // tp_size
|
|
head_rank_start = n_local_head * tp_rank
|
|
head_rank_end = n_local_head * (tp_rank + 1)
|
|
|
|
# Pre-compute expert mapping ONCE.
|
|
expert_mapping = self.get_expert_mapping()
|
|
|
|
for name, loaded_weight in weights:
|
|
# Strip 'model.' prefix from checkpoint keys.
|
|
# vLLM's weight iteration yields keys like 'model.layers.0...'
|
|
# but named_parameters() on DeepseekV4Model returns 'layers.0...'
|
|
if name.startswith("model."):
|
|
name = name[len("model."):]
|
|
|
|
# Apply checkpoint → model name substitutions
|
|
for ckpt_pat, model_pat in CKPT_KEY_SUBST.items():
|
|
if ckpt_pat in name:
|
|
name = name.replace(ckpt_pat, model_pat)
|
|
break # first match wins (order matters)
|
|
|
|
for param_name, weight_name, shard_id in stacked_params_mapping:
|
|
# Skip MoE routed experts (handled separately below).
|
|
# Use .ffn.experts. (not .experts.) to avoid skipping
|
|
# shared_experts which also contains ".experts.".
|
|
if ".ffn.experts." in name:
|
|
continue
|
|
if weight_name not in name:
|
|
continue
|
|
name_mapped = name.replace(weight_name, param_name)
|
|
if name_mapped not in params_dict:
|
|
continue
|
|
name = name_mapped
|
|
|
|
param = params_dict[name]
|
|
weight_loader = param.weight_loader
|
|
|
|
# ModelOpt NVFP4 packed weight fix for MergedColumnParallelLinear.
|
|
#
|
|
# modelopt exports NVFP4 packed weights as uint8 (2 values/byte
|
|
# along the column dim). But MergedColumnParallelLinear creates
|
|
# the weight param as bfloat16 (ModelWeightParameter), because
|
|
# ModelOptNvFp4Config only patches Linear, not
|
|
# MergedColumnParallelLinear.
|
|
#
|
|
# When loading uint8 packed weights into a bf16 param, we need to
|
|
# unpack them. Each uint8 byte contains 2 E2M1 FP4 values.
|
|
# We unpack using the LUT and return bf16.
|
|
#
|
|
# The weight_scale is loaded separately and process_weights_after_loading
|
|
# will handle the actual NVFP4 quantization.
|
|
if (loaded_weight.dtype == torch.uint8
|
|
and param.data.dtype != torch.uint8
|
|
and loaded_weight.shape[-1] * 2 == param.data.shape[-1]):
|
|
# Unpack NVFP4 (E2M1) → BF16
|
|
# E2M1 LUT: 0→0, 1→0.5, 2→1, 3→1.5, 4→2, 5→3, 6→4, 7→6
|
|
# Sign bit in bit 3 (indices 8-15 are negatives)
|
|
FP4_LUT = torch.tensor([
|
|
0.0, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0,
|
|
-0.0, -0.5, -1.0, -1.5, -2.0, -3.0, -4.0, -6.0,
|
|
], dtype=torch.float32, device=loaded_weight.device)
|
|
lower = FP4_LUT[(loaded_weight & 0x0F).long()] # (..., in_packed, )
|
|
upper = FP4_LUT[((loaded_weight >> 4) & 0x0F).long()]
|
|
# Interleave: [lower_0, upper_0, lower_1, upper_1, ...]
|
|
out = torch.empty(
|
|
*loaded_weight.shape[:-1], loaded_weight.shape[-1] * 2,
|
|
dtype=torch.float32, device=loaded_weight.device,
|
|
)
|
|
out[..., 0::2] = lower
|
|
out[..., 1::2] = upper
|
|
loaded_weight = out.to(torch.bfloat16)
|
|
|
|
try:
|
|
weight_loader(param, loaded_weight, shard_id)
|
|
except (AssertionError, ValueError, RuntimeError) as e:
|
|
raise RuntimeError(
|
|
f'Weight load failed: name={name} shard_id={shard_id} '
|
|
f'param.shape={param.shape} param.dtype={param.data.dtype} '
|
|
f'loaded.shape={loaded_weight.shape} loaded.dtype={loaded_weight.dtype}'
|
|
) from e
|
|
loaded_params.add(name)
|
|
break
|
|
else:
|
|
if ".ffn.experts." in name:
|
|
# NVFP4 checkpoint stores float8_e4m3fn scales, not E8M0.
|
|
# E8M0 would indicate an MXFP4 checkpoint — wrong format.
|
|
if (
|
|
"weight_scale" in name
|
|
and loaded_weight.dtype == torch.float8_e8m0fnu
|
|
):
|
|
raise ValueError(
|
|
f"E8M0 weight_scale in NVFP4 checkpoint ({name}) — "
|
|
f"checkpoint format mismatch"
|
|
)
|
|
for mapping in expert_mapping:
|
|
param_name, weight_name, expert_id, shard_id = mapping
|
|
if weight_name not in name:
|
|
continue
|
|
name_mapped = name.replace(weight_name, param_name)
|
|
if name_mapped not in params_dict:
|
|
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=shard_id,
|
|
expert_id=expert_id,
|
|
)
|
|
if success:
|
|
name = name_mapped
|
|
loaded_params.add(name_mapped)
|
|
break
|
|
else:
|
|
continue
|
|
continue
|
|
elif "attn_sink" in name:
|
|
if name not in params_dict:
|
|
continue
|
|
narrow_weight = loaded_weight[head_rank_start:head_rank_end]
|
|
n = narrow_weight.shape[0]
|
|
params_dict[name][:n].copy_(narrow_weight)
|
|
loaded_params.add(name)
|
|
continue
|
|
else:
|
|
if name not in params_dict:
|
|
# ModelOpt NVFP4 export includes params not in the
|
|
# vllm model (e.g., compressor.position_bias).
|
|
# Skip them silently.
|
|
continue
|
|
param = params_dict[name]
|
|
|
|
# Handle bf16 → uint8 mismatch for o_a_proj:
|
|
# modelopt didn't quantize o_a_proj (bf16, no scales),
|
|
# but ModelOptNvFp4Config creates wo_a with NVFP4 quant
|
|
# (uint8 weight + scales). We quantize the bf16 weight
|
|
# to NVFP4 at load time so the layer runs in NVFP4 path.
|
|
if (name.endswith(".weight")
|
|
and loaded_weight.dtype != torch.uint8
|
|
and param.data.dtype == torch.uint8):
|
|
# Quantize bf16 → NVFP4 (E2M1 packed uint8 + scales)
|
|
w_bf16 = loaded_weight
|
|
out_dim, in_dim = w_bf16.shape
|
|
block_size = 16
|
|
assert in_dim % block_size == 0
|
|
n_blocks = in_dim // block_size
|
|
|
|
# Reshape into blocks
|
|
w_blocks = w_bf16.reshape(out_dim, n_blocks, block_size)
|
|
|
|
# Compute per-block amax
|
|
amax = w_blocks.abs().amax(dim=-1) # [out, n_blocks]
|
|
|
|
# Global scale (weight_scale_2): max amax / (6.0 * 448.0)
|
|
global_amax = amax.max()
|
|
# Use 448.0 as the max e4m3 value for scale computation
|
|
weight_scale_2_val = global_amax / (6.0 * 448.0)
|
|
weight_scale_2 = weight_scale_2_val.to(torch.float32)
|
|
|
|
# Per-block scale (weight_scale): float8_e4m3fn
|
|
# block_scale = amax / (6.0 * weight_scale_2)
|
|
block_scale = amax / (6.0 * weight_scale_2_val)
|
|
weight_scale = block_scale.clamp(0.0, 448.0).to(torch.float8_e4m3fn)
|
|
|
|
# Quantize to FP4 (E2M1)
|
|
# E2M1 LUT: 0, 0.5, 1, 1.5, 2, 3, 4, 6 (positive)
|
|
FP4_POS = torch.tensor(
|
|
[0.0, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0],
|
|
dtype=torch.float32, device=w_bf16.device,
|
|
)
|
|
# Scale the weight values: normalized = w / (block_scale * weight_scale_2)
|
|
block_scale_f32 = block_scale.clamp(0.0, 448.0)
|
|
scaled = w_blocks / (block_scale_f32.unsqueeze(-1) * weight_scale_2_val)
|
|
# Find nearest FP4 index (0-7 for magnitude)
|
|
# Use absolute value for matching, then apply sign
|
|
scaled_abs = scaled.abs()
|
|
# Find closest FP4 value
|
|
diff = (scaled_abs.unsqueeze(-1) - FP4_POS).abs()
|
|
fp4_idx = diff.argmin(dim=-1) # [out, n_blocks, block_size]
|
|
# Apply sign: negative values get bit 3 set
|
|
sign = (scaled < 0).int()
|
|
fp4_val = (sign << 3) | fp4_idx.int()
|
|
# Pack: 2 FP4 values per uint8 byte
|
|
# Even positions → lower nibble, Odd → upper nibble
|
|
fp4_flat = fp4_val.reshape(out_dim, -1) # [out, in_dim]
|
|
assert fp4_flat.shape[1] % 2 == 0
|
|
even = fp4_flat[:, 0::2] # lower nibble
|
|
odd = fp4_flat[:, 1::2] # upper nibble
|
|
packed = (odd << 4) | even
|
|
weight_packed = packed.to(torch.uint8).view(torch.int8)
|
|
|
|
# Reshape weight_scale to [out, n_blocks]
|
|
weight_scale_2d = weight_scale.reshape(out_dim, n_blocks)
|
|
|
|
# Load the quantized weight into the uint8 param
|
|
weight_loader = param.weight_loader
|
|
weight_loader(param, weight_packed)
|
|
loaded_params.add(name)
|
|
|
|
# Load scales into sibling params
|
|
base = name.rsplit(".", 1)[0]
|
|
# weight_scale
|
|
ws_name = f"{base}.weight_scale"
|
|
if ws_name in params_dict:
|
|
ws_param = params_dict[ws_name]
|
|
ws_loader = getattr(ws_param, "weight_loader", default_weight_loader)
|
|
ws_loader(ws_param, weight_scale_2d)
|
|
loaded_params.add(ws_name)
|
|
# weight_scale_2
|
|
ws2_name = f"{base}.weight_scale_2"
|
|
if ws2_name in params_dict:
|
|
ws2_param = params_dict[ws2_name]
|
|
ws2_loader = getattr(ws2_param, "weight_loader", default_weight_loader)
|
|
ws2_loader(ws2_param, weight_scale_2.reshape(1))
|
|
loaded_params.add(ws2_name)
|
|
# input_scale: use 1.0 default (dynamic quant)
|
|
is_name = f"{base}.input_scale"
|
|
if is_name in params_dict:
|
|
is_param = params_dict[is_name]
|
|
is_loader = getattr(is_param, "weight_loader", default_weight_loader)
|
|
is_loader(is_param, torch.tensor(1.0, dtype=torch.float32))
|
|
loaded_params.add(is_name)
|
|
continue
|
|
|
|
# Handle uint8 NVFP4 packed → bf16 unpack for non-stacked
|
|
# params (e.g. indexer.weights_proj). Checkpoint stores
|
|
# NVFP4 as uint8 (2 values/byte), but model param is bf16.
|
|
if (loaded_weight.dtype == torch.uint8
|
|
and param.data.dtype != torch.uint8
|
|
and loaded_weight.shape[-1] * 2 == param.data.shape[-1]):
|
|
FP4_LUT = torch.tensor([
|
|
0.0, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0,
|
|
-0.0, -0.5, -1.0, -1.5, -2.0, -3.0, -4.0, -6.0,
|
|
], dtype=torch.float32, device=loaded_weight.device)
|
|
lower = FP4_LUT[(loaded_weight & 0x0F).long()]
|
|
upper = FP4_LUT[((loaded_weight >> 4) & 0x0F).long()]
|
|
out = torch.empty(
|
|
*loaded_weight.shape[:-1],
|
|
loaded_weight.shape[-1] * 2,
|
|
dtype=torch.float32, device=loaded_weight.device,
|
|
)
|
|
out[..., 0::2] = lower
|
|
out[..., 1::2] = upper
|
|
loaded_weight = out.to(torch.bfloat16)
|
|
|
|
weight_loader = getattr(
|
|
param, "weight_loader", default_weight_loader
|
|
)
|
|
weight_loader(param, loaded_weight)
|
|
loaded_params.add(name)
|
|
continue
|
|
|
|
return loaded_params
|
|
|
|
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_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 FusedMoE.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_mega_moe_weights(self) -> None:
|
|
from tqdm import tqdm
|
|
layers = list(islice(self.layers, self.start_layer, self.end_layer))
|
|
for layer in tqdm(layers, desc=" (JIT compile)NVFP4 MoE layers", unit="layer"):
|
|
layer.ffn.finalize_mega_moe_weights()
|
|
|
|
def _convert_nvfp4_post_load(self):
|
|
"""Post-load conversion of NVFP4 weights for vLLM compatibility.
|
|
|
|
Only wo_a needs FP8 conversion (attention forward uses fp8_einsum
|
|
which requires FP8 inputs). All other NVFP4 weights stay native —
|
|
vLLM's FlashInferCutlassNvFp4LinearKernel handles them directly.
|
|
|
|
Compressor weights are reconstructed from checkpoint sub-weights
|
|
because the stacking weight_loader corrupts NVFP4 uint8 data.
|
|
"""
|
|
FP8_MAX = torch.finfo(torch.float8_e4m3fn).max
|
|
|
|
# Only wo_a needs conversion — fp8_einsum requires FP8 weight + scale
|
|
fp8_proj_names = {"wo_a"}
|
|
fp8_converted = 0
|
|
compressor_converted = 0
|
|
|
|
_shard_index = self._build_shard_index("/model") if os.path.isdir("/model") else None
|
|
|
|
from tqdm import tqdm
|
|
for layer_idx, layer in tqdm(enumerate(self.layers), total=len(self.layers), desc=" (upcast)NVFP4→FP8 wo_a only", unit="layer"):
|
|
attn = layer.attn
|
|
|
|
# FP8 conversion: only wo_a
|
|
for proj_name in fp8_proj_names:
|
|
if not hasattr(attn, proj_name):
|
|
continue
|
|
mod = getattr(attn, proj_name)
|
|
if not hasattr(mod, "weight"):
|
|
continue
|
|
if mod.weight.dtype in (torch.uint8, torch.int8):
|
|
E2M1_LUT = torch.tensor([0, 0.5, 1, 1.5, 2, 3, 4, 6], dtype=torch.bfloat16)
|
|
self._convert_nvfp4_to_fp8(mod, E2M1_LUT, FP8_MAX)
|
|
fp8_converted += 1
|
|
|
|
# Compressor: still needs BF16 reconstruction
|
|
mla_attn = getattr(attn, "mla_attn", None)
|
|
if mla_attn is not None:
|
|
E2M1_LUT = torch.tensor([0, 0.5, 1, 1.5, 2, 3, 4, 6], dtype=torch.bfloat16)
|
|
compressor = getattr(mla_attn, "compressor", None)
|
|
if compressor is not None and hasattr(compressor, "fused_wkv_wgate"):
|
|
compressor_converted += self._reconstruct_compressor_weight(
|
|
compressor.fused_wkv_wgate, attn, layer_idx, E2M1_LUT, _shard_index=_shard_index)
|
|
indexer = getattr(mla_attn, "indexer", None)
|
|
if indexer is not None:
|
|
idx_compressor = getattr(indexer, "compressor", None)
|
|
if idx_compressor is not None and hasattr(idx_compressor, "fused_wkv_wgate"):
|
|
compressor_converted += self._reconstruct_compressor_weight(
|
|
idx_compressor.fused_wkv_wgate, indexer, layer_idx, E2M1_LUT, sub_path=".indexer", _shard_index=_shard_index)
|
|
|
|
|
|
|
|
|
|
def _dequant_nvfp4_to_bf16(self, mod, e2m1_lut):
|
|
"""Dequantize NVFP4 weight to bf16 for normal .forward() path."""
|
|
w_uint8 = mod.weight.data
|
|
device = w_uint8.device
|
|
w_bf16 = self._unpack_nvfp4_to_bf16(w_uint8, e2m1_lut, device)
|
|
|
|
# Dequantize with scales
|
|
if hasattr(mod, "weight_scale") and hasattr(mod, "weight_scale_2"):
|
|
block_scale = self._block_scale_to_float32(mod.weight_scale.data)
|
|
if block_scale.dim() == 2 and w_bf16.dim() == 2:
|
|
block_size = w_bf16.shape[1] // block_scale.shape[1]
|
|
block_scale_expanded = block_scale.unsqueeze(-1).expand(
|
|
-1, -1, block_size
|
|
).reshape(w_bf16.shape)
|
|
else:
|
|
block_scale_expanded = block_scale
|
|
global_scale = mod.weight_scale_2.data.max().item()
|
|
input_scale = (
|
|
mod.input_scale.data.max().item()
|
|
if hasattr(mod, "input_scale")
|
|
else 1.0
|
|
)
|
|
# NOTE: input_scale is for ACTIVATIONS, not weights.
|
|
# Weight dequant = e2m1 * block_scale * global_scale (NO input_scale)
|
|
w_dequant = w_bf16.float() * block_scale_expanded * global_scale
|
|
w_dequant = w_dequant.to(torch.bfloat16)
|
|
else:
|
|
w_dequant = w_bf16
|
|
|
|
# Free source tensors eagerly to avoid holding uint8+bf16+fp32 simultaneously
|
|
del w_uint8, w_bf16
|
|
mod.weight = torch.nn.Parameter(w_dequant, requires_grad=False)
|
|
del w_dequant
|
|
from vllm.model_executor.layers.linear import UnquantizedLinearMethod
|
|
mod.quant_method = UnquantizedLinearMethod()
|
|
for attr in ("weight_scale", "weight_scale_2", "input_scale",
|
|
"weight_scale_inv"):
|
|
if hasattr(mod, attr):
|
|
delattr(mod, attr)
|
|
|
|
def _convert_nvfp4_to_fp8(self, mod, e2m1_lut, fp8_max):
|
|
"""Convert NVFP4 weight to FP8 for fp8_einsum path (wo_a only).
|
|
|
|
Uses DeepGEMM's deepgemm_post_process_fp8_weight_block to ensure
|
|
correct weight and scale format for fp8_einsum with BMM.
|
|
"""
|
|
w_uint8 = mod.weight.data
|
|
device = w_uint8.device
|
|
w_bf16 = self._unpack_nvfp4_to_bf16(w_uint8, e2m1_lut, device)
|
|
|
|
# Dequantize with scales
|
|
if hasattr(mod, "weight_scale") and hasattr(mod, "weight_scale_2"):
|
|
|
|
block_scale = self._block_scale_to_float32(mod.weight_scale.data)
|
|
if block_scale.dim() == 2 and w_bf16.dim() == 2:
|
|
block_size = w_bf16.shape[1] // block_scale.shape[1]
|
|
block_scale_expanded = block_scale.unsqueeze(-1).expand(
|
|
-1, -1, block_size
|
|
).reshape(w_bf16.shape)
|
|
else:
|
|
block_scale_expanded = block_scale
|
|
global_scale = mod.weight_scale_2.data.max().item()
|
|
input_scale = (
|
|
mod.input_scale.data.max().item()
|
|
if hasattr(mod, "input_scale")
|
|
else 1.0
|
|
)
|
|
# NOTE: input_scale is for ACTIVATIONS, not weights.
|
|
# Weight dequant = e2m1 * block_scale * global_scale (NO input_scale)
|
|
w_dequant = w_bf16.float() * block_scale_expanded * global_scale
|
|
w_dequant = w_dequant.to(torch.bfloat16)
|
|
else:
|
|
w_dequant = w_bf16
|
|
|
|
# Re-quantize bf16 -> FP8 e4m3 with block quantization
|
|
# DeepGEMM expects block-scale format: weight_scale (FP8 e4m3 block scale)
|
|
# and weight_scale_inv (per-tensor scale).
|
|
# We do per-tensor quantization, so block_scale is all-ones.
|
|
w_amax = w_dequant.abs().amax()
|
|
if w_amax == 0:
|
|
w_amax = torch.tensor(1.0, device=device)
|
|
fp8_scale = w_amax / fp8_max
|
|
w_fp8 = (w_dequant / fp8_scale).to(torch.float8_e4m3fn)
|
|
|
|
# Create block scale filled with the per-tensor fp8_scale value.
|
|
# DeepGEMM divides by the block scale, so each block gets fp8_scale.
|
|
BLOCK_SIZE = 128
|
|
is_bmm = getattr(mod, "is_bmm", False)
|
|
bmm_batch_size = getattr(mod, "bmm_batch_size", 0)
|
|
|
|
# Weight is 2D (output_size, input_size) before BMM reshape
|
|
# Block scale shape: (output_size / BLOCK_SIZE, input_size / BLOCK_SIZE)
|
|
rows = w_fp8.size(0)
|
|
cols = w_fp8.size(1)
|
|
block_rows = rows // BLOCK_SIZE
|
|
block_cols = cols // BLOCK_SIZE
|
|
|
|
# Fill block scale with the per-tensor fp8_scale (NOT all-ones!)
|
|
# This is correct because we requantized with a single per-tensor scale,
|
|
# so every 128x128 block has the same scale = fp8_scale.
|
|
ws = torch.full((block_rows, block_cols), fp8_scale.item(), dtype=torch.float32, device=device)
|
|
|
|
# Use DeepGEMM's post-processing for proper layout transformation
|
|
from vllm.model_executor.layers.quantization.utils.fp8_utils import (
|
|
deepgemm_post_process_fp8_weight_block,
|
|
)
|
|
w_fp8, ws = deepgemm_post_process_fp8_weight_block(
|
|
wq=w_fp8,
|
|
ws=ws,
|
|
quant_block_shape=(BLOCK_SIZE, BLOCK_SIZE),
|
|
use_e8m0=True, # scale_fmt=ue8m0
|
|
is_bmm=is_bmm,
|
|
bmm_batch_size=bmm_batch_size,
|
|
)
|
|
|
|
# Free source tensors eagerly
|
|
del w_uint8, w_bf16, w_dequant
|
|
mod.weight = torch.nn.Parameter(w_fp8, requires_grad=False)
|
|
del w_fp8
|
|
# weight_scale_inv is what the attention runtime reads as b_scale
|
|
# for deepseek_v4_fp8_einsum -> DeepGEMM fp8_einsum.
|
|
# It must be the DeepGEMM-formatted block scale (dg_ws), NOT the
|
|
# per-tensor scalar. See: deepseek_v4_attention.py line 319.
|
|
mod.weight_scale_inv = torch.nn.Parameter(ws, requires_grad=False)
|
|
del ws
|
|
from vllm.model_executor.layers.linear import UnquantizedLinearMethod
|
|
mod.quant_method = UnquantizedLinearMethod()
|
|
for attr in ("weight_scale", "weight_scale_2", "input_scale"):
|
|
if hasattr(mod, attr):
|
|
delattr(mod, attr)
|
|
|
|
@staticmethod
|
|
def _build_shard_index(ckpt_dir: str) -> dict[str, str]:
|
|
"""Build key→shard_path index from safetensors metadata (no tensor I/O)."""
|
|
import glob
|
|
from safetensors import safe_open
|
|
index = {}
|
|
for shard_file in sorted(glob.glob(os.path.join(ckpt_dir, "model-*.safetensors"))):
|
|
try:
|
|
with safe_open(shard_file, framework="pt") as f:
|
|
for key in f.keys():
|
|
index[key] = shard_file
|
|
except Exception:
|
|
continue
|
|
return index
|
|
|
|
def _reconstruct_compressor_weight(self, fused_mod, parent_mod, layer_idx, e2m1_lut, sub_path="", _shard_index=None):
|
|
"""Reconstruct compressor fused_wkv_wgate from checkpoint.
|
|
|
|
Compressor weights are SKIPPED during loading because NVFP4 uint8 data
|
|
can't be loaded into bf16 MergedColumnParallelLinear params (shape mismatch).
|
|
We read the original uint8 data from the safetensors checkpoint, unpack
|
|
E2M1, dequantize, and stack into the fused weight param.
|
|
"""
|
|
from safetensors import safe_open
|
|
|
|
# Find the checkpoint directory
|
|
# The model weights are mounted at /model in Docker
|
|
ckpt_dir = "/model"
|
|
if not os.path.isdir(ckpt_dir):
|
|
print(f"WARNING: layer {layer_idx} compressor: checkpoint dir {ckpt_dir} not found")
|
|
return 0
|
|
|
|
# Determine the layer's compressor key prefix in the checkpoint
|
|
# Before mapper: model.layers.N.self_attn.compressor.{kv_proj,gate_proj}
|
|
# After mapper: model.layers.N.attn.mla_attn.compressor.{wkv,wgate}
|
|
# We read from checkpoint (before mapper), so use original names
|
|
layer_prefix = f"model.layers.{layer_idx}.self_attn.compressor{sub_path}"
|
|
|
|
# All keys we need from the checkpoint
|
|
keys = {
|
|
'wkv_uint8': f"{layer_prefix}.kv_proj.weight",
|
|
'wgate_uint8': f"{layer_prefix}.gate_proj.weight",
|
|
'wkv_block_scale': f"{layer_prefix}.kv_proj.weight_scale",
|
|
'wgate_block_scale': f"{layer_prefix}.gate_proj.weight_scale",
|
|
'wkv_global_scale': f"{layer_prefix}.kv_proj.weight_scale_2",
|
|
'wgate_global_scale': f"{layer_prefix}.gate_proj.weight_scale_2",
|
|
'wkv_input_scale': f"{layer_prefix}.kv_proj.input_scale",
|
|
'wgate_input_scale': f"{layer_prefix}.gate_proj.input_scale",
|
|
}
|
|
|
|
# Read tensors using shard index for targeted access (no full-shard loads)
|
|
tensors = {}
|
|
for name, key in keys.items():
|
|
shard_path = (_shard_index or {}).get(key)
|
|
if shard_path is None:
|
|
continue
|
|
try:
|
|
with safe_open(shard_path, framework="pt") as f:
|
|
if key in f.keys():
|
|
tensors[name] = f.get_tensor(key)
|
|
except Exception:
|
|
continue
|
|
|
|
wkv_uint8 = tensors.get('wkv_uint8')
|
|
wgate_uint8 = tensors.get('wgate_uint8')
|
|
|
|
if wkv_uint8 is None or wgate_uint8 is None:
|
|
# Layer might not have a compressor (compress_ratio=1 layers)
|
|
return 0
|
|
|
|
wkv_block_scale = tensors.get('wkv_block_scale')
|
|
wgate_block_scale = tensors.get('wgate_block_scale')
|
|
wkv_global_scale = tensors.get('wkv_global_scale')
|
|
wgate_global_scale = tensors.get('wgate_global_scale')
|
|
wkv_input_scale = tensors.get('wkv_input_scale')
|
|
wgate_input_scale = tensors.get('wgate_input_scale')
|
|
|
|
device = fused_mod.weight.device
|
|
wkv_uint8 = wkv_uint8.to(device)
|
|
wgate_uint8 = wgate_uint8.to(device)
|
|
|
|
# Unpack E2M1 FP4→bf16
|
|
wkv_bf16 = self._unpack_nvfp4_to_bf16(wkv_uint8, e2m1_lut, device)
|
|
wgate_bf16 = self._unpack_nvfp4_to_bf16(wgate_uint8, e2m1_lut, device)
|
|
|
|
# Dequantize with scales
|
|
def _dequant(w_bf16, block_scale, global_scale, input_scale):
|
|
if block_scale is not None and global_scale is not None:
|
|
|
|
block_scale = self._block_scale_to_float32(block_scale.to(device))
|
|
if block_scale.dim() == 2 and w_bf16.dim() == 2:
|
|
block_size = w_bf16.shape[1] // block_scale.shape[1]
|
|
block_scale_exp = block_scale.unsqueeze(-1).expand(
|
|
-1, -1, block_size
|
|
).reshape(w_bf16.shape)
|
|
else:
|
|
block_scale_exp = block_scale
|
|
gs = global_scale.to(device).max().item()
|
|
# NOTE: input_scale is for activations, not weights.
|
|
# Weight dequant = e2m1 * block_scale * global_scale (NO input_scale)
|
|
w = w_bf16.float() * block_scale_exp * gs
|
|
return w.to(torch.bfloat16)
|
|
return w_bf16
|
|
|
|
wkv_dequant = _dequant(wkv_bf16, wkv_block_scale, wkv_global_scale, wkv_input_scale)
|
|
wgate_dequant = _dequant(wgate_bf16, wgate_block_scale, wgate_global_scale, wgate_input_scale)
|
|
|
|
# Stack: concatenate along output dim (dim 0)
|
|
# fused_wkv_wgate.weight = cat([wkv, wgate], dim=0) → (2*head_dim, hidden_size)
|
|
w_fused = torch.cat([wkv_dequant, wgate_dequant], dim=0)
|
|
|
|
|
|
# Replace the weight
|
|
fused_mod.weight = torch.nn.Parameter(w_fused, requires_grad=False)
|
|
from vllm.model_executor.layers.linear import UnquantizedLinearMethod
|
|
fused_mod.quant_method = UnquantizedLinearMethod()
|
|
for attr in ("weight_scale", "weight_scale_2", "input_scale", "weight_scale_inv"):
|
|
if hasattr(fused_mod, attr):
|
|
delattr(fused_mod, attr)
|
|
return 1
|
|
|
|
def _convert_bf16_to_fp8(self, mod, fp8_max):
|
|
"""Convert BF16 weight to FP8 for fp8_einsum path.
|
|
|
|
Used for wo_a which modelopt did NOT quantize (bf16 in checkpoint)
|
|
but which the attention forward reads as FP8 for deepseek_v4_fp8_einsum.
|
|
Uses DeepGEMM's post-processing for proper BMM + scale format.
|
|
"""
|
|
w_bf16 = mod.weight.data
|
|
device = w_bf16.device
|
|
|
|
# Re-quantize bf16 -> FP8 e4m3 with block quantization
|
|
w_amax = w_bf16.abs().amax()
|
|
if w_amax == 0:
|
|
w_amax = torch.tensor(1.0, device=device)
|
|
fp8_scale = w_amax / fp8_max
|
|
w_fp8 = (w_bf16 / fp8_scale).to(torch.float8_e4m3fn)
|
|
|
|
BLOCK_SIZE = 128
|
|
is_bmm = getattr(mod, "is_bmm", False)
|
|
bmm_batch_size = getattr(mod, "bmm_batch_size", 0)
|
|
|
|
rows = w_fp8.size(0)
|
|
cols = w_fp8.size(1)
|
|
block_rows = rows // BLOCK_SIZE
|
|
block_cols = cols // BLOCK_SIZE
|
|
# Fill block scale with per-tensor fp8_scale (NOT all-ones!)
|
|
ws = torch.full((block_rows, block_cols), fp8_scale.item(), dtype=torch.float32, device=device)
|
|
|
|
from vllm.model_executor.layers.quantization.utils.fp8_utils import (
|
|
deepgemm_post_process_fp8_weight_block,
|
|
)
|
|
w_fp8, ws = deepgemm_post_process_fp8_weight_block(
|
|
wq=w_fp8,
|
|
ws=ws,
|
|
quant_block_shape=(BLOCK_SIZE, BLOCK_SIZE),
|
|
use_e8m0=True, # scale_fmt=ue8m0
|
|
is_bmm=is_bmm,
|
|
bmm_batch_size=bmm_batch_size,
|
|
)
|
|
|
|
mod.weight = torch.nn.Parameter(w_fp8, requires_grad=False)
|
|
# weight_scale_inv is what the attention runtime reads as b_scale
|
|
# for deepseek_v4_fp8_einsum -> DeepGEMM fp8_einsum.
|
|
# It must be the DeepGEMM-formatted block scale (dg_ws), NOT the
|
|
# per-tensor scalar. See: deepseek_v4_attention.py line 319.
|
|
mod.weight_scale_inv = torch.nn.Parameter(ws, requires_grad=False)
|
|
# weight_scale is not used at runtime for BMM layers; remove it
|
|
# to avoid confusing other code paths.
|
|
for attr in ("weight_scale", "weight_scale_2", "input_scale"):
|
|
if hasattr(mod, attr):
|
|
delattr(mod, attr)
|
|
from vllm.model_executor.layers.linear import UnquantizedLinearMethod
|
|
mod.quant_method = UnquantizedLinearMethod()
|
|
|
|
@staticmethod
|
|
@staticmethod
|
|
def _block_scale_to_float32(sf: torch.Tensor) -> torch.Tensor:
|
|
"""Convert NVFP4 block scales (float8_e4m3fn) to float32."""
|
|
return sf.to(torch.float32)
|
|
|
|
def _unpack_nvfp4_to_bf16(self, w_uint8, e2m1_lut, device):
|
|
"""Unpack NVFP4 uint8 packed weights to bf16 using E2M1 format."""
|
|
# Extract 4-bit FP4 values (0-15, bit 3 = sign)
|
|
even_raw = (w_uint8 & 0x0F).int()
|
|
odd_raw = ((w_uint8 >> 4) & 0x0F).int()
|
|
# Sign: 0-7 = positive, 8-15 = negative
|
|
even_sign = torch.where(even_raw >= 8, -1.0, 1.0).to(torch.bfloat16)
|
|
odd_sign = torch.where(odd_raw >= 8, -1.0, 1.0).to(torch.bfloat16)
|
|
# Magnitude index: lower 3 bits (0-7)
|
|
even_vals = even_sign * e2m1_lut.to(device)[even_raw & 0x07]
|
|
odd_vals = odd_sign * e2m1_lut.to(device)[odd_raw & 0x07]
|
|
# Interleave and flatten
|
|
w_bf16 = torch.stack([even_vals, odd_vals], dim=-1)
|
|
w_bf16 = w_bf16.reshape(w_uint8.shape[0], -1).to(torch.bfloat16)
|
|
return w_bf16
|
|
@torch.compile(backend=current_platform.simple_compile_backend)
|
|
def hc_head(
|
|
hidden_states: torch.Tensor,
|
|
hc_fn: torch.Tensor,
|
|
hc_scale: torch.Tensor,
|
|
hc_base: torch.Tensor,
|
|
rms_norm_eps: float,
|
|
hc_eps: float,
|
|
) -> torch.Tensor:
|
|
hc_mult, hidden_size = hidden_states.shape[-2:]
|
|
outer_shape = hidden_states.shape[:-2]
|
|
hs_flat = hidden_states.view(-1, hc_mult, hidden_size)
|
|
num_tokens = hs_flat.shape[0]
|
|
out = torch.empty(
|
|
num_tokens, hidden_size, dtype=torch.bfloat16, device=hidden_states.device
|
|
)
|
|
torch.ops.vllm.hc_head_fused_kernel(
|
|
hs_flat,
|
|
hc_fn,
|
|
hc_scale,
|
|
hc_base,
|
|
out,
|
|
hidden_size,
|
|
rms_norm_eps,
|
|
hc_eps,
|
|
hc_mult,
|
|
)
|
|
return out.view(*outer_shape, hidden_size)
|
|
|
|
|
|
def _make_deepseek_v4_weights_mapper(expert_dtype: str) -> WeightsMapper:
|
|
if expert_dtype == "fp4":
|
|
# MXFP4 experts use Mxfp4MoEMethod, which registers scales as
|
|
# ``w{1,2,3}_weight_scale`` (no _inv suffix). FP8 linear and
|
|
# shared experts use Fp8LinearMethod's block scales, which
|
|
# register as ``weight_scale_inv``.
|
|
scale_regex = {
|
|
re.compile(r"(\.experts\.\d+\.w[123])\.scale$"): r"\1.weight_scale",
|
|
re.compile(r"\.scale$"): ".weight_scale_inv",
|
|
}
|
|
else:
|
|
# FP8 experts use Fp8MoEMethod (block_quant=True), which registers
|
|
# scales as ``w{13,2}_weight_scale_inv``. Map all ``.scale`` keys
|
|
# there.
|
|
scale_regex = {
|
|
re.compile(r"\.scale$"): ".weight_scale_inv",
|
|
}
|
|
|
|
# ── ModelOpt NVFP4 export patches ────────────────────────────────
|
|
# modelopt exports with different naming than the original HF ckpt:
|
|
# - Expert projections: gate_proj/up_proj/down_proj → w1/w3/w2
|
|
# - Shared expert projections: gate_proj/up_proj → w1/w3 (stacking)
|
|
# - Compressor: kv_proj → wkv, gate_proj → wgate (stacking)
|
|
# - Attention: self_attn prefix, kv_proj → wkv (stacking)
|
|
# - modelopt uses mlp, vllm uses ffn
|
|
# Order matters for regex: skip patterns MUST come before renames.
|
|
|
|
# Skip NVFP4 scales for compressor+attention fused params.
|
|
# After substr renaming, these map to stacked params (fused_wkv_wgate,
|
|
# fused_wqa_wkv, gate_up_proj) which don't register NVFP4 scale params
|
|
# because ModelOptNvFp4Config only handles Linear, not
|
|
# MergedColumnParallelLinear. We unpack weights as bf16 and let
|
|
# process_weights_after_loading re-quantize them.
|
|
# Must match ORIGINAL checkpoint key names (before substr renaming).
|
|
fused_skip_regex = {
|
|
# Compressor: SKIP ALL tensors. The compressor uses quant_config=None,
|
|
# so MergedColumnParallelLinear creates bf16 weight params. NVFP4 uint8
|
|
# checkpoint data can't be loaded into these params (shape mismatch:
|
|
# uint8 (head_dim, hidden_size//2) vs bf16 (head_dim, hidden_size)).
|
|
# The stacking weight_loader silently skips the sub-weights, leaving
|
|
# random bf16 initialization. We reconstruct the compressor weights
|
|
# manually in post-load conversion by reading from the checkpoint.
|
|
re.compile(r"\.compressor\.kv_proj\.weight$"): None,
|
|
re.compile(r"\.compressor\.gate_proj\.weight$"): None,
|
|
re.compile(r"\.compressor\.kv_proj\.weight_scale$"): None,
|
|
re.compile(r"\.compressor\.gate_proj\.weight_scale$"): None,
|
|
re.compile(r"\.compressor\.kv_proj\.weight_scale_2$"): None,
|
|
re.compile(r"\.compressor\.gate_proj\.weight_scale_2$"): None,
|
|
re.compile(r"\.compressor\.kv_proj\.input_scale$"): None,
|
|
re.compile(r"\.compressor\.gate_proj\.input_scale$"): None,
|
|
# Note: attention and shared expert scale tensors are NO LONGER
|
|
# skipped. After fixing substr mappings, they correctly map to the
|
|
# model's NVFP4 scale parameters (fused_wqa_wkv, wq_b, wo_a,
|
|
# wo_b, gate_up_proj). They load via the stacking logic.
|
|
}
|
|
# Routed expert projections: gate_proj→w1, up_proj→w3, down_proj→w2
|
|
# Regex (not substr) to match ONLY .experts.N. — not .shared_experts.
|
|
expert_rename_regex = {
|
|
re.compile(r"(\.experts\.\d+\.)gate_proj\."): r"\1w1.",
|
|
re.compile(r"(\.experts\.\d+\.)up_proj\."): r"\1w3.",
|
|
re.compile(r"(\.experts\.\d+\.)down_proj\."): r"\1w2.",
|
|
}
|
|
# Merge: skip patterns first, then renames, then original scale_regex
|
|
merged_regex = {}
|
|
merged_regex.update(fused_skip_regex)
|
|
merged_regex.update(expert_rename_regex)
|
|
merged_regex.update(scale_regex)
|
|
|
|
return WeightsMapper(
|
|
orig_to_new_prefix={
|
|
"layers.": "model.layers.",
|
|
"embed.": "model.embed.",
|
|
"norm.": "model.norm.",
|
|
"hc_head": "model.hc_head",
|
|
"mtp.": "model.mtp.",
|
|
},
|
|
orig_to_new_regex=merged_regex,
|
|
orig_to_new_suffix={
|
|
"embed.weight": "embed_tokens.weight",
|
|
".ffn.gate.bias": ".ffn.gate.e_score_correction_bias",
|
|
},
|
|
orig_to_new_substr={
|
|
".attn.compressor.": ".attn.mla_attn.compressor.",
|
|
".shared_experts.w2": ".shared_experts.down_proj",
|
|
# ── ModelOpt NVFP4 substr patches ──
|
|
# Attention: self_attn → attn (projections at attn level, not mla_attn)
|
|
".self_attn.q_a_proj.": ".attn.wq_a.",
|
|
".self_attn.q_b_proj.": ".attn.wq_b.",
|
|
".self_attn.q_a_norm.": ".attn.q_norm.",
|
|
".self_attn.o_a_proj.": ".attn.wo_a.",
|
|
".self_attn.o_b_proj.": ".attn.wo_b.",
|
|
".self_attn.sinks": ".attn.attn_sink",
|
|
# kv_proj → wkv (for stacking into fused_wqa_wkv)
|
|
".self_attn.kv_proj.": ".attn.wkv.",
|
|
".self_attn.kv_norm.": ".attn.kv_norm.",
|
|
# kv_norm is at attention level, not compressor/mla_attn level in vllm
|
|
# Must come before the general compressor mapping
|
|
".self_attn.compressor.kv_norm.": ".attn.kv_norm.",
|
|
# Compressor: self_attn.compressor → attn.mla_attn.compressor
|
|
".self_attn.compressor.": ".attn.mla_attn.compressor.",
|
|
# Compressor projections for stacking (fused_wkv_wgate)
|
|
".compressor.kv_proj.": ".compressor.wkv.",
|
|
".compressor.gate_proj.": ".compressor.wgate.",
|
|
# Shared expert projections (stacking into gate_up_proj)
|
|
# Checkpoint has .shared_experts. but model has .ffn.shared_experts.
|
|
".shared_experts.gate_proj.": ".ffn.shared_experts.w1.",
|
|
".shared_experts.up_proj.": ".ffn.shared_experts.w3.",
|
|
# modelopt uses mlp, vllm uses ffn internally
|
|
".mlp.": ".ffn.",
|
|
},
|
|
)
|
|
|
|
|
|
class DeepseekV4ForCausalLM(nn.Module):
|
|
model_cls = DeepseekV4Model
|
|
|
|
# NOTE: We do NOT set hf_to_vllm_mapper here because our custom
|
|
# load_weights handles all checkpoint→model name remapping inline.
|
|
# If hf_to_vllm_mapper is set, vLLM's AutoWeightsLoader may be invoked
|
|
# INSTEAD of our load_weights, silently dropping NVFP4 weight loading.
|
|
|
|
def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
|
|
super().__init__()
|
|
|
|
config = vllm_config.model_config.hf_config
|
|
self.config = config
|
|
|
|
self.model = self.model_cls(
|
|
vllm_config=vllm_config, prefix=maybe_prefix(prefix, "model")
|
|
)
|
|
self.lm_head = ParallelLMHead(
|
|
config.vocab_size,
|
|
config.hidden_size,
|
|
prefix=maybe_prefix(prefix, "lm_head"),
|
|
)
|
|
self.logits_processor = LogitsProcessor(config.vocab_size)
|
|
|
|
def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:
|
|
return self.model.embed_input_ids(input_ids)
|
|
|
|
def compute_logits(
|
|
self,
|
|
hidden_states: torch.Tensor,
|
|
) -> torch.Tensor | None:
|
|
logits = self.logits_processor(self.lm_head, hidden_states)
|
|
return logits
|
|
|
|
def forward(
|
|
self,
|
|
input_ids: torch.Tensor,
|
|
positions: torch.Tensor,
|
|
intermediate_tensors: IntermediateTensors | None = None,
|
|
inputs_embeds: torch.Tensor | None = None,
|
|
) -> torch.Tensor | IntermediateTensors:
|
|
hidden_states = self.model(
|
|
input_ids, positions, intermediate_tensors, inputs_embeds
|
|
)
|
|
return hidden_states
|
|
|
|
def get_mtp_target_hidden_states(self) -> torch.Tensor | None:
|
|
"""Pre-hc_head residual stream buffer (max_num_batched_tokens,
|
|
hc_mult * hidden_size) for the MTP draft model. Populated by
|
|
forward(); valid after each target step."""
|
|
return getattr(self.model, "_mtp_hidden_buffer", None)
|
|
|
|
def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
|
|
# lm_head lives on this outer model, not on the inner DeepseekV4Model.
|
|
# The inner load_weights silently drops lm_head.weight via
|
|
# "if name not in params_dict: continue". Extract it here.
|
|
rest = []
|
|
for name, loaded_weight in weights:
|
|
if name == "lm_head.weight" or name.endswith(".lm_head.weight"):
|
|
param = self.lm_head.weight
|
|
weight_loader = getattr(param, "weight_loader",
|
|
default_weight_loader)
|
|
weight_loader(param, loaded_weight)
|
|
else:
|
|
rest.append((name, loaded_weight))
|
|
# Use the model-level loader which handles NVFP4 expert mapping,
|
|
# uint8→bf16 unpacking for MergedColumnParallelLinear, and
|
|
# bf16→NVFP4 quantization for unquantized layers.
|
|
# AutoWeightsLoader bypasses this logic and would break NVFP4 loading.
|
|
loaded_params = self.model.load_weights(rest)
|
|
loaded_params.add("lm_head.weight")
|
|
print(" Checkpoint loaded. Preparing NVFP4...", flush=True)
|
|
self.model.finalize_mega_moe_weights()
|
|
self.model._convert_nvfp4_post_load()
|
|
print(" Warming up tilelang kernels...", flush=True)
|
|
self._warmup_tilelang()
|
|
print(" NVFP4 model ready ✓", flush=True)
|
|
|
|
return loaded_params
|
|
|
|
def _warmup_tilelang(self) -> None:
|
|
"""Force-compile all tilelang JIT kernels with dummy data.
|
|
|
|
tilelang's @jit decorator compiles lazily on first call. In eager mode
|
|
(no cudagraphs), the HTTP server comes up before the first inference
|
|
triggers compilation — and any request hitting the model during
|
|
compilation crashes vLLM. This warmup ensures all kernels are compiled
|
|
before the server accepts traffic.
|
|
|
|
We call the custom ops directly with 1-token dummy tensors to populate
|
|
the tilelang kernel cache.
|
|
"""
|
|
import torch
|
|
config = self.model.config
|
|
hc_mult = config.hc_mult
|
|
hidden_size = config.hidden_size
|
|
device = next(self.model.parameters()).device
|
|
hc_mult3 = hc_mult * (2 + hc_mult)
|
|
|
|
# Warmup mhc_pre
|
|
residual = torch.randn(1, hc_mult, hidden_size, dtype=torch.bfloat16,
|
|
device=device)
|
|
fn = torch.randn(hc_mult3, hc_mult * hidden_size, dtype=torch.float32,
|
|
device=device)
|
|
hc_scale = torch.randn(3, dtype=torch.float32, device=device)
|
|
hc_base = torch.randn(hc_mult3, dtype=torch.float32, device=device)
|
|
|
|
try:
|
|
torch.ops.vllm.mhc_pre(
|
|
residual=residual,
|
|
fn=fn,
|
|
hc_scale=hc_scale,
|
|
hc_base=hc_base,
|
|
rms_eps=config.rms_norm_eps,
|
|
hc_pre_eps=config.hc_eps,
|
|
hc_sinkhorn_eps=config.hc_eps,
|
|
hc_post_mult_value=2.0,
|
|
sinkhorn_repeat=config.hc_sinkhorn_iters,
|
|
)
|
|
print(" mhc_pre ✓", flush=True)
|
|
except Exception as e:
|
|
print(f" mhc_pre warmup failed (non-fatal): {e}", flush=True)
|
|
|
|
# Warmup mhc_post
|
|
x = torch.randn(1, hidden_size, dtype=torch.bfloat16, device=device)
|
|
post_mix = torch.randn(1, hc_mult, 1, dtype=torch.float32, device=device)
|
|
comb_mix = torch.randn(1, hc_mult, hc_mult, dtype=torch.float32,
|
|
device=device)
|
|
|
|
try:
|
|
torch.ops.vllm.mhc_post(x, residual, post_mix, comb_mix)
|
|
print(" mhc_post ✓", flush=True)
|
|
except Exception as e:
|
|
print(f" mhc_post warmup failed (non-fatal): {e}", flush=True)
|
|
|
|
# Free dummy tensors
|
|
del residual, fn, hc_scale, hc_base, x, post_mix, comb_mix
|
|
torch.cuda.empty_cache()
|
|
|
|
def get_expert_mapping(self) -> list[tuple[str, str, int, str]]:
|
|
return self.model.get_expert_mapping()
|
|
|