biondizzle/deepseek-v4-quant
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity
Files
c8564caf9d8551fe6cd8e5af478d6daec6b5ff51
deepseek-v4-quant/patches/deepseek_v4.py
biondizzle e963325b61 WIP: MegaMoE NVFP4 kernel + diagnostics 
- Force use_mega_moe=True for NVFP4 pipeline
- DeepseekV4MegaMoEExperts: load NVFP4 params (float8 block scales,
  float32 global/input scales), convert NVFP4→BF16→MXFP4 in
  finalize_weights for the DeepGEMM mega_moe kernel
- Add _nvfp4_to_bf16 and _bf16_to_mxfp4 conversion methods
- Remove expert_dtype check blocking mega_moe
- Add diagnostics for wo_a and bf16 layer conversion
- Still WIP: attention layer bugs under investigation
2026-05-11 05:19:49 +00:00

2457 lines
99 KiB
Python
Raw Blame History

# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
# ==============================================================================
# DeepSeek V4 NVFP4 Patch — Version Banner (printed at import time)
# ==============================================================================
import datetime as _dt
import os as _os
_git_commit = _os.popen("git -C /root/nvidia-meeting/deepseek-v4-quant rev-parse --short HEAD 2>/dev/null || echo 'unknown'").read().strip()
print(f"""
{'='*70}
DeepSeek V4 NVFP4 Patch
{'='*70}
Commit: {_git_commit}
Loaded: {_dt.datetime.now().strftime('%Y-%m-%d %H:%M:%S UTC')}
Node: {_os.uname().nodename}
Architecture:
wo_a → FP8 + DeepGEMM block scale (BMM einsum)
wq_b/wo_b → BF16 (UnquantizedLinearMethod)
fused_wqa → BF16 (stacked q_a + kv, dequantized from NVFP4)
compressor → BF16 (reconstructed from separate kv_proj+gate_proj)
shared_exp → FP8 (Fp8LinearMethod, DeepGEMM)
MoE experts → NVFP4 (FusedMoE, FLASHINFER_TRTLLM) — NOT converted
Bugs fixed:
#1 DeepGEMM sf.dim() — block scale format (deepgemm_post_process)
#2 fused_skip_regex — q_b/o_a/o_b scales no longer skipped
#3 input_scale — removed from weight dequant (activations only)
#4 compressor indexer — sub_path for .indexer keys
#5 block scale dtype — must be float32, not float8_e4m3fn
#6 block scale values — torch.full(fp8_scale) not torch.ones
#7 UE8M0 block scale — .to(float32) misinterprets E8M0 as E4M3
{'='*70}
""")
# ==============================================================================
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 always use FP8 block quantization for
linear/attention layers. The MoE expert weights vary by checkpoint:
- ``expert_dtype="fp4"`` (e.g. DeepSeek-V4-Flash): MXFP4 experts
with ue8m0 (e8m0fnu) FP8 linear scales.
- ``expert_dtype="fp8"`` (e.g. DeepSeek-V4-Flash-Base): FP8 block
experts with float32 FP8 linear scales.
The dispatch and the linear scale dtype are both keyed off
``expert_dtype`` from the model's hf_config; missing values default
to ``"fp4"`` so existing FP4 checkpoints stay unchanged.
NOTE: ``expert_dtype`` is resolved lazily because this config is
constructed during VllmConfig setup, before ``set_current_vllm_config``
is active. Reading hf_config eagerly in ``__init__`` would always see
the default ``"fp4"`` and silently misroute Flash-Base checkpoints.
"""
def __init__(self, *args, **kwargs):
super().__init__(*args, **kwargs)
self._resolved_expert_dtype: str | None = None
# ``is_scale_e8m0`` is a property that resolves on first read,
# by which time the current vllm_config has been set.
@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; defer the decision until a
# later call lands inside set_current_vllm_config.
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
from vllm.logger import init_logger
init_logger(__name__).info_once(
"DeepSeek V4 expert_dtype resolved to %r", expert_dtype
)
return self._resolved_expert_dtype
@property
def is_scale_e8m0(self) -> bool:
# FP4 checkpoints store FP8 linear scales as e8m0fnu; FP8 expert
# checkpoints (Flash-Base) store them as float32.
return self.expert_dtype == "fp4"
@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"
@triton.jit
def _deepseek_v4_stage_mega_moe_inputs_kernel(
hidden_states,
x_fp8,
x_sf,
topk_ids,
topk_weights,
topk_idx_out,
topk_weights_out,
hidden_stride_m: tl.constexpr,
hidden_stride_k: tl.constexpr,
x_stride_m: tl.constexpr,
x_stride_k: tl.constexpr,
x_sf_stride_m: tl.constexpr,
x_sf_stride_k: tl.constexpr,
topk_ids_stride_m: tl.constexpr,
topk_ids_stride_k: tl.constexpr,
topk_weights_stride_m: tl.constexpr,
topk_weights_stride_k: tl.constexpr,
topk_idx_stride_m: tl.constexpr,
topk_idx_stride_k: tl.constexpr,
topk_weights_out_stride_m: tl.constexpr,
topk_weights_out_stride_k: tl.constexpr,
hidden_size: tl.constexpr,
top_k: tl.constexpr,
BLOCK_K: tl.constexpr,
GROUP_K: tl.constexpr,
BLOCK_TOPK: tl.constexpr,
) -> None:
token_id = tl.program_id(0)
k_block_id = tl.program_id(1)
k_offsets = k_block_id * BLOCK_K + tl.arange(0, BLOCK_K)
k_mask = k_offsets < hidden_size
hidden = tl.load(
hidden_states + token_id * hidden_stride_m + k_offsets * hidden_stride_k,
mask=k_mask,
other=0.0,
).to(tl.float32)
num_groups: tl.constexpr = BLOCK_K // GROUP_K
hidden_groups = tl.reshape(tl.abs(hidden), [num_groups, GROUP_K])
amax = tl.max(hidden_groups, axis=1)
amax = tl.maximum(amax, 1.0e-4)
scale = amax / 448.0
scale_bits = scale.to(tl.uint32, bitcast=True)
scale_exp = ((scale_bits >> 23) & 0xFF) + ((scale_bits & 0x7FFFFF) != 0).to(
tl.uint32
)
scale_exp = tl.minimum(tl.maximum(scale_exp, 1), 254)
rounded_scale = (scale_exp << 23).to(tl.float32, bitcast=True)
hidden_groups = tl.reshape(hidden, [num_groups, GROUP_K])
scaled = hidden_groups * (1.0 / rounded_scale)[:, None]
scaled = tl.reshape(scaled, [BLOCK_K])
fp8 = scaled.to(tl.float8e4nv)
tl.store(
x_fp8 + token_id * x_stride_m + k_offsets * x_stride_k,
fp8,
mask=k_mask,
)
scale_offsets = tl.arange(0, num_groups)
packed_scale = tl.sum(scale_exp << (scale_offsets * 8), axis=0).to(tl.int32)
tl.store(
x_sf + token_id * x_sf_stride_m + k_block_id * x_sf_stride_k,
packed_scale,
)
if k_block_id == 0:
topk_offsets = tl.arange(0, BLOCK_TOPK)
topk_mask = topk_offsets < top_k
ids = tl.load(
topk_ids + token_id * topk_ids_stride_m + topk_offsets * topk_ids_stride_k,
mask=topk_mask,
other=0,
).to(tl.int64)
tl.store(
topk_idx_out
+ token_id * topk_idx_stride_m
+ topk_offsets * topk_idx_stride_k,
ids,
mask=topk_mask,
)
weights = tl.load(
topk_weights
+ token_id * topk_weights_stride_m
+ topk_offsets * topk_weights_stride_k,
mask=topk_mask,
other=0.0,
)
tl.store(
topk_weights_out
+ token_id * topk_weights_out_stride_m
+ topk_offsets * topk_weights_out_stride_k,
weights,
mask=topk_mask,
)
def _stage_deepseek_v4_mega_moe_inputs(
hidden_states: torch.Tensor,
topk_weights: torch.Tensor,
topk_ids: torch.Tensor,
x_fp8: torch.Tensor,
x_sf: torch.Tensor,
topk_idx_out: torch.Tensor,
topk_weights_out: torch.Tensor,
) -> None:
num_tokens, hidden_size = hidden_states.shape
if num_tokens == 0:
return
if hidden_size % 128 != 0:
raise ValueError(
"DeepSeek V4 MegaMoE input staging requires hidden_size to be "
"a multiple of 128."
)
top_k = topk_ids.shape[1]
if topk_weights.shape != topk_ids.shape:
raise ValueError(
"DeepSeek V4 MegaMoE input staging requires topk_weights and "
"topk_ids to have the same shape."
)
block_k = 128
grid = (num_tokens, triton.cdiv(hidden_size, block_k))
block_topk = triton.next_power_of_2(top_k)
_deepseek_v4_stage_mega_moe_inputs_kernel[grid](
hidden_states,
x_fp8,
x_sf,
topk_ids,
topk_weights,
topk_idx_out,
topk_weights_out,
hidden_states.stride(0),
hidden_states.stride(1),
x_fp8.stride(0),
x_fp8.stride(1),
x_sf.stride(0),
x_sf.stride(1),
topk_ids.stride(0),
topk_ids.stride(1),
topk_weights.stride(0),
topk_weights.stride(1),
topk_idx_out.stride(0),
topk_idx_out.stride(1),
topk_weights_out.stride(0),
topk_weights_out.stride(1),
hidden_size,
top_k,
BLOCK_K=block_k,
GROUP_K=32,
BLOCK_TOPK=block_topk,
num_warps=4,
)
def make_deepseek_v4_expert_params_mapping(
num_experts: int,
) -> list[tuple[str, str, int, str]]:
return [
(
"experts.w13_" if shard_id in ("w1", "w3") else "experts.w2_",
f"experts.{expert_id}.{weight_name}.",
expert_id,
shard_id,
)
for expert_id in range(num_experts)
for shard_id, weight_name in [
("w1", "w1"),
("w2", "w2"),
("w3", "w3"),
]
]
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 converts them to MXFP4 format for the
DeepGEMM fp8_fp4_mega_moe kernel at finalize_weights time.
NVFP4 → MXFP4 conversion:
1. Unpack E2M1 FP4 → BF16
2. Dequantize with UE8M0 block_scale * float32 global_scale
3. Re-quantize BF16 → MXFP4 (E2M1 + UE8M0, group_size=32)
4. Feed to deep_gemm transform_weights_for_mega_moe
"""
_symm_buffer_cache: dict[tuple[int, int, int, int, int, int, int], object] = {}
# 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
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.uint8,
),
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
self.w13_weight_scale_2 = nn.Parameter(
torch.zeros(num_local_experts, 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.uint8,
),
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._transformed_l1_weights: tuple[torch.Tensor, torch.Tensor] | None = None
self._transformed_l2_weights: tuple[torch.Tensor, torch.Tensor] | None = 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,
return_success: bool = False,
) -> bool | None:
local_expert_id = self._map_global_expert_id(expert_id)
if local_expert_id == -1:
return False if return_success else None
# Scalar params (weight_scale_2, input_scale): 1D per-expert
if "weight_scale_2" in weight_name or "input_scale" in weight_name:
param.data[local_expert_id].copy_(loaded_weight)
return True if return_success else None
expert_data = param.data[local_expert_id]
if shard_id in ("w1", "w3"):
if "w13_" not in weight_name:
return False if return_success else None
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 if return_success else None
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 if return_success else None
@staticmethod
def _ue8m0_uint8_to_float(sf: torch.Tensor) -> torch.Tensor:
return (sf.to(torch.int32) << 23).view(torch.float32)
def _check_runtime_supported(self) -> None:
if not torch.cuda.is_available():
raise NotImplementedError("DeepSeek V4 MegaMoE requires CUDA.")
device = self.w13_weight.device
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._transformed_l1_weights is not None:
return
self._check_runtime_supported()
import vllm.third_party.deep_gemm as deep_gemm
# === NVFP4 → BF16 → MXFP4 conversion ===
# The DeepGEMM mega_moe kernel expects MXFP4 format:
# - E2M1 packed uint8 (same as NVFP4)
# - UE8M0 uint8 block scales, group_size=32
# NVFP4 has:
# - E2M1 packed uint8 (same)
# - E8M0 float8_e4m3fn block scales, group_size=16
# - float32 global_scale and input_scale
# We dequant NVFP4→BF16 then requant BF16→MXFP4.
w13_bf16 = self._nvfp4_to_bf16(
self.w13_weight.data, self.w13_weight_scale.data,
self.w13_weight_scale_2.data, self.w13_input_scale.data,
)
w2_bf16 = self._nvfp4_to_bf16(
self.w2_weight.data, self.w2_weight_scale.data,
self.w2_weight_scale_2.data, self.w2_input_scale.data,
)
# Re-quantize BF16 → MXFP4 (E2M1 + UE8M0, group_size=32)
MXFP4_GROUP_SIZE = 32
w13_mxfp4_weight, w13_mxfp4_scale = self._bf16_to_mxfp4(
w13_bf16, MXFP4_GROUP_SIZE)
w2_mxfp4_weight, w2_mxfp4_scale = self._bf16_to_mxfp4(
w2_bf16, MXFP4_GROUP_SIZE)
# Transform into DeepGEMM mega_moe layout
w13_scale = deep_gemm.transform_sf_into_required_layout(
w13_mxfp4_scale.contiguous(),
2 * self.intermediate_size,
self.hidden_size,
(1, 32),
self.num_local_experts,
)
w2_scale = deep_gemm.transform_sf_into_required_layout(
w2_mxfp4_scale.contiguous(),
self.hidden_size,
self.intermediate_size,
(1, 32),
self.num_local_experts,
)
self._transformed_l1_weights, self._transformed_l2_weights = (
deep_gemm.transform_weights_for_mega_moe(
(w13_mxfp4_weight.view(torch.int8).contiguous(), w13_scale),
(w2_mxfp4_weight.view(torch.int8).contiguous(), w2_scale),
)
)
# Drop the original loader-side parameters
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
@staticmethod
def _ue8m0_to_float32(sf: torch.Tensor) -> torch.Tensor:
"""Convert E8M0 (float8_e4m3fn bytes, semantically UE8M0) to float32."""
return (sf.view(torch.uint8).to(torch.int32) << 23).view(torch.float32)
def _nvfp4_to_bf16(
self,
w_uint8: torch.Tensor, # [E, M, K//2]
w_scale_f8: torch.Tensor, # [E, M, K//16] (float8_e4m3fn, E8M0 format)
w_scale_2: torch.Tensor, # [E] float32 global scale
w_input_scale: torch.Tensor, # [E] float32 activation scale
) -> torch.Tensor:
"""Dequantize NVFP4 expert weights to BF16.
Formula: weight_bf16 = e2m1_value * block_scale_ue8m0 * global_scale
(input_scale is for activations, not weights)
"""
device = w_uint8.device
E, M, K2 = w_uint8.shape
K = K2 * 2 # unpacked dim
# Unpack E2M1 FP4 → BF16
e2m1_lut = torch.tensor(
[0.0, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0],
dtype=torch.bfloat16, device=device,
)
even_raw = (w_uint8 & 0x0F).int()
odd_raw = ((w_uint8 >> 4) & 0x0F).int()
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)
even_vals = even_sign * e2m1_lut[even_raw & 0x07]
odd_vals = odd_sign * e2m1_lut[odd_raw & 0x07]
w_bf16 = torch.stack([even_vals, odd_vals], dim=-1).reshape(E, M, K)
# Dequantize: e2m1 * block_scale * global_scale
block_scale = self._ue8m0_to_float32(w_scale_f8) # [E, M, K//16]
GROUP_SIZE = 16
# Expand block scale to match weight elements
block_scale_expanded = block_scale.unsqueeze(-1).expand(
-1, -1, -1, GROUP_SIZE
).reshape(E, M, K)
# Global scale: [E] → [E, 1, 1]
global_scale = w_scale_2.view(E, 1, 1)
w_dequant = w_bf16.float() * block_scale_expanded * global_scale
return w_dequant.to(torch.bfloat16)
def _bf16_to_mxfp4(
self,
w_bf16: torch.Tensor, # [E, M, K]
group_size: int = 32,
) -> tuple[torch.Tensor, torch.Tensor]:
"""Re-quantize BF16 → MXFP4 (E2M1 packed uint8 + UE8M0 uint8 scales).
Returns (w_packed_uint8, w_scale_uint8) where:
w_packed_uint8: [E, M, K//2]
w_scale_uint8: [E, M, K//group_size] as uint8 UE8M0 bytes
"""
device = w_bf16.device
E, M, K = w_bf16.shape
# Block quantization
n_groups = K // group_size
w_groups = w_bf16.reshape(E, M, n_groups, group_size)
# Compute block amax
amax = w_groups.abs().amax(dim=-1) # [E, M, n_groups]
# UE8M0 scale: floor to nearest power of 2
# value = 2^(exp-127), so exp = floor(log2(amax)) + 127
amax_clamped = amax.clamp(min=2**-126)
scale_exp = torch.floor(torch.log2(amax_clamped)).to(torch.int32) + 127
scale_exp = scale_exp.clamp(1, 254).to(torch.uint8) # avoid 0 and 255
# Recover float32 scale for quantization
scale_f32 = (scale_exp.to(torch.int32) << 23).view(torch.float32) # [E, M, n_groups]
# Quantize each element: value / scale → nearest E2M1
E2M1_POS = torch.tensor(
[0.0, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0],
dtype=torch.float32, device=device,
)
scaled = w_groups.float() / scale_f32.unsqueeze(-1) # [E, M, n_groups, group_size]
scaled_abs = scaled.abs()
diff = (scaled_abs.unsqueeze(-1) - E2M1_POS).abs()
fp4_idx = diff.argmin(dim=-1) # [E, M, n_groups, group_size]
sign = (scaled < 0).int()
fp4_val = (sign << 3) | fp4_idx.int()
# Pack 2 FP4 values per byte
fp4_flat = fp4_val.reshape(E, M, K)
even = fp4_flat[:, :, 0::2]
odd = fp4_flat[:, :, 1::2]
w_packed = ((odd << 4) | even).to(torch.uint8)
return w_packed, scale_exp
def get_symm_buffer(self):
import vllm.third_party.deep_gemm as deep_gemm
group = get_ep_group().device_group
device = torch.accelerator.current_device_index()
key = (
id(group),
device,
self.num_experts,
self.max_num_tokens,
self.top_k,
self.hidden_size,
self.intermediate_size,
)
symm_buffer = self._symm_buffer_cache.get(key)
if symm_buffer is None:
symm_buffer = deep_gemm.get_symm_buffer_for_mega_moe(
group,
self.num_experts,
self.max_num_tokens,
self.top_k,
self.hidden_size,
self.intermediate_size,
)
self._symm_buffer_cache[key] = symm_buffer
return symm_buffer
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 vllm.third_party.deep_gemm as deep_gemm
symm_buffer = self.get_symm_buffer()
num_tokens = hidden_states.shape[0]
_stage_deepseek_v4_mega_moe_inputs(
hidden_states,
topk_weights,
topk_ids,
symm_buffer.x[:num_tokens],
symm_buffer.x_sf[:num_tokens],
symm_buffer.topk_idx[:num_tokens],
symm_buffer.topk_weights[:num_tokens],
)
# 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._transformed_l1_weights is not None
assert self._transformed_l2_weights is not None
deep_gemm.fp8_fp4_mega_moe(
y,
self._transformed_l1_weights,
self._transformed_l2_weights,
symm_buffer,
activation_clamp=activation_clamp,
fast_math=fast_math,
)
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→MXFP4 conversion 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,
)
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 = 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),
]
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:
for param_name, weight_name, shard_id in stacked_params_mapping:
# Skip non-stacked layers and experts (experts handled below).
if ".experts." in name:
continue
if weight_name not in name:
continue
name = name.replace(weight_name, param_name)
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:
print(f'[DEBUG-STACK] 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} err={e}')
raise
loaded_params.add(name)
break
else:
if ".experts." in name:
# E8M0 scales are stored as float8_e8m0fnu in
# checkpoints but the MoE param is uint8. copy_()
# would do a numeric conversion (e.g. 2^-7 → 0),
# destroying the raw exponent bytes.
if (
"weight_scale" in name
and loaded_weight.dtype == torch.float8_e8m0fnu
):
loaded_weight = loaded_weight.view(torch.uint8)
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,
return_success=True,
)
if success:
name = name_mapped
break
loaded_params.add(name_mapped)
continue
elif "attn_sink" in name:
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): UE8M0 format
# scale_fmt=ue8m0: block_scale = 2^(exp-127), stored as
# uint8 exponent byte viewed as float8_e4m3fn
block_scale = amax / (6.0 * weight_scale_2_val)
# Convert to UE8M0: floor to nearest power of 2
# UE8M0 exponent = floor(log2(block_scale)) + 127
block_scale_clamped = block_scale.clamp(min=2**-127)
block_scale_exp = torch.floor(torch.log2(block_scale_clamped)).to(torch.int32) + 127
block_scale_exp = block_scale_exp.clamp(0, 254).to(torch.uint8)
weight_scale = block_scale_exp.view(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,
)
# For each block, dequantize the block scale from UE8M0
block_scale_f32 = (block_scale_exp.to(torch.int32) << 23).view(torch.float32)
# Scale the weight values: normalized = w / (block_scale * weight_scale_2)
# We need to find the nearest FP4 value
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)
# 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
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:
for layer in islice(self.layers, self.start_layer, self.end_layer):
layer.ffn.finalize_mega_moe_weights()
def _convert_nvfp4_post_load(self):
"""Post-load conversion of NVFP4 weights for vLLM compatibility.
Strategy:
- wo_a: Convert to FP8 (attention forward reads weight/weight_scale_inv
directly and passes to deepseek_v4_fp8_einsum, bypassing quant_method)
- fused_wqa_wkv, wq_b, wo_b: Dequant NVFP4->bf16 (called via
.forward() which goes through quant_method; FP8 would dtype-mismatch)
- compressor.fused_wkv_wgate: Dequant NVFP4->bf16 (used via direct
torch.mm in attention parallel stream)
- shared_experts (gate_up_proj, down_proj): Dequant NVFP4->bf16
- MoE experts: Handled by DeepseekV4MegaMoEExperts (NVFP4→MXFP4)
"""
E2M1_LUT = torch.tensor(
[0, 0.5, 1, 1.5, 2, 3, 4, 6], dtype=torch.bfloat16
)
FP8_MAX = torch.finfo(torch.float8_e4m3fn).max
# wo_a: attention forward reads .weight and .weight_scale_inv directly
# for fp8_einsum. Only layer that needs FP8 conversion.
fp8_proj_names = {"wo_a"}
# Attention layers called via .forward() — need bf16
bf16_proj_names = {"fused_wqa_wkv", "wq_b", "wo_b"}
# Shared expert layers called via .forward() — need bf16
bf16_shared_names = {"gate_up_proj", "down_proj"}
fp8_converted = 0
fp8_from_bf16 = 0
bf16_converted = 0
compressor_converted = 0
diag_printed = False
for layer_idx, layer in enumerate(self.layers):
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 == torch.uint8:
# NVFP4 -> dequant to bf16 -> requant to FP8
if not diag_printed and layer_idx == 0:
ws = getattr(mod, 'weight_scale', None)
ws2 = getattr(mod, 'weight_scale_2', None)
print(f"[DIAG-wo_a:0] dtype={mod.weight.dtype} shape={mod.weight.shape} "
f"ws_dtype={ws.dtype if ws is not None else None} ws_shape={ws.shape if ws is not None else None} "
f"ws2_val={ws2.data.item() if ws2 is not None else None}")
self._convert_nvfp4_to_fp8(mod, E2M1_LUT, FP8_MAX)
fp8_converted += 1
elif mod.weight.dtype == torch.bfloat16:
if not diag_printed and layer_idx == 0:
print(f"[DIAG-wo_a:0] dtype=bf16 shape={mod.weight.shape} (direct bf16→fp8)")
self._convert_bf16_to_fp8(mod, FP8_MAX)
fp8_from_bf16 += 1
diag_printed = True
# BF16 conversion: attention layers via .forward()
for proj_name in bf16_proj_names:
if not hasattr(attn, proj_name):
continue
mod = getattr(attn, proj_name)
if not hasattr(mod, "weight") or mod.weight.dtype != torch.uint8:
continue
if not diag_printed and layer_idx == 0:
ws = getattr(mod, 'weight_scale', None)
print(f"[DIAG-bf16:0/{proj_name}] dtype={mod.weight.dtype} shape={mod.weight.shape} "
f"ws_dtype={ws.dtype if ws is not None else None} ws_shape={ws.shape if ws is not None else None}")
self._dequant_nvfp4_to_bf16(mod, E2M1_LUT)
bf16_converted += 1
# Compressor: fused_wkv_wgate used via direct torch.mm
# Compressor weights were SKIPPED during loading (skip patterns)
# because the stacking weight_loader corrupts NVFP4 uint8 data.
# We reconstruct the bf16 weight from the individual sub-weights
# that were loaded separately before stacking.
# Note: compressor.kv_proj.weight and compressor.gate_proj.weight
# are skipped, so fused_wkv_wgate.weight is zeros (empty tensor).
# We need to manually create it.
mla_attn = getattr(attn, "mla_attn", None)
if mla_attn is not None:
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)
# Indexer compressor (C4A layers only)
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")
# Shared experts
ffn = layer.ffn
if hasattr(ffn, "shared_experts") and ffn.shared_experts is not None:
for proj_name in bf16_shared_names:
if not hasattr(ffn.shared_experts, proj_name):
continue
mod = getattr(ffn.shared_experts, proj_name)
if not hasattr(mod, "weight") or mod.weight.dtype != torch.uint8:
continue
self._dequant_nvfp4_to_bf16(mod, E2M1_LUT)
bf16_converted += 1
total_fp8 = fp8_converted + fp8_from_bf16
total_bf16 = bf16_converted + compressor_converted
if total_fp8 > 0 or total_bf16 > 0:
print(f"NVFP4 post-load: {fp8_converted} NVFP4->FP8, "
f"{fp8_from_bf16} BF16->FP8, "
f"{bf16_converted} attn/shared->BF16, "
f"{compressor_converted} compressor->BF16, "
f"MoE experts stay NVFP4")
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"):
# scale_fmt=ue8m0: weight_scale bytes are E8M0 format (power-of-2 only).
# A simple .to(float32) misinterprets them as E4M3. Must reinterpret
# the raw uint8 bits as IEEE 754 exponent field.
block_scale = self._ue8m0_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
# Replace weight with bf16 version
mod.weight = torch.nn.Parameter(w_dequant, requires_grad=False)
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"):
# scale_fmt=ue8m0: reinterpret E8M0 bytes as float32
block_scale = self._ue8m0_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,
)
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()
def _reconstruct_compressor_weight(self, fused_mod, parent_mod, layer_idx, e2m1_lut, sub_path=""):
"""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.
"""
import glob
from safetensors.torch import load_file
# 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}"
# Find which shard contains this layer's compressor weights
wkv_key = f"{layer_prefix}.kv_proj.weight"
wgate_key = f"{layer_prefix}.gate_proj.weight"
wkv_scale_key = f"{layer_prefix}.kv_proj.weight_scale"
wgate_scale_key = f"{layer_prefix}.gate_proj.weight_scale"
wkv_scale2_key = f"{layer_prefix}.kv_proj.weight_scale_2"
wgate_scale2_key = f"{layer_prefix}.gate_proj.weight_scale_2"
wkv_iscale_key = f"{layer_prefix}.kv_proj.input_scale"
wgate_iscale_key = f"{layer_prefix}.gate_proj.input_scale"
# Load from safetensors
wkv_uint8 = None
wgate_uint8 = None
wkv_block_scale = None
wgate_block_scale = None
wkv_global_scale = None
wgate_global_scale = None
wkv_input_scale = None
wgate_input_scale = None
shard_files = sorted(glob.glob(os.path.join(ckpt_dir, "model-*.safetensors")))
for shard_file in shard_files:
try:
shard_data = load_file(shard_file)
except Exception:
continue
if wkv_key in shard_data:
wkv_uint8 = shard_data[wkv_key]
wkv_block_scale = shard_data.get(wkv_scale_key)
wkv_global_scale = shard_data.get(wkv_scale2_key)
wkv_input_scale = shard_data.get(wkv_iscale_key)
if wgate_key in shard_data:
wgate_uint8 = shard_data[wgate_key]
wgate_block_scale = shard_data.get(wgate_scale_key)
wgate_global_scale = shard_data.get(wgate_scale2_key)
wgate_input_scale = shard_data.get(wgate_iscale_key)
if wkv_uint8 is not None and wgate_uint8 is not None:
break
if wkv_uint8 is None or wgate_uint8 is None:
# Layer might not have a compressor (compress_ratio=1 layers)
return 0
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:
# scale_fmt=ue8m0: reinterpret E8M0 bytes as float32
block_scale = self._ue8m0_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)
# DEBUG: log shapes to diagnose compressor weight mismatch
print(f"NVFP4 compressor layer {layer_idx}: wkv={wkv_dequant.shape}, wgate={wgate_dequant.shape}, fused={w_fused.shape}, existing_param={fused_mod.weight.shape}")
# 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
return 0
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
def _ue8m0_to_float32(sf: torch.Tensor) -> torch.Tensor:
"""Convert UE8M0 (E8M0 power-of-2) scale bytes to float32.
NVFP4 checkpoints with scale_fmt=ue8m0 store per-block weight scales as
E8M0 format (8-bit exponent, no mantissa). The value = 2^(raw_byte - 127).
The bytes are loaded as float8_e4m3fn by safetensors, but a simple
.to(float32) misinterprets them as E4M3 (which has mantissa bits).
Correct conversion: place the raw uint8 bits into the exponent field
of an IEEE 754 float32 (bits 23-30), yielding 2^(raw-127) * implicit_1.
"""
raw_uint8 = sf.view(torch.uint8)
return (raw_uint8.to(torch.int32) << 23).view(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)
".shared_experts.gate_proj.": ".shared_experts.w1.",
".shared_experts.up_proj.": ".shared_experts.w3.",
# modelopt uses mlp, vllm uses ffn internally
".mlp.": ".ffn.",
},
)
class DeepseekV4ForCausalLM(nn.Module):
model_cls = DeepseekV4Model
# Default mapper assumes the original FP4-expert checkpoint layout.
# Overridden per-instance in __init__ when expert_dtype != "fp4".
hf_to_vllm_mapper = _make_deepseek_v4_weights_mapper("fp4")
def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
super().__init__()
config = vllm_config.model_config.hf_config
self.config = config
expert_dtype = getattr(config, "expert_dtype", "fp4")
if expert_dtype != "fp4":
self.hf_to_vllm_mapper = _make_deepseek_v4_weights_mapper(expert_dtype)
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]:
loader = AutoWeightsLoader(self, skip_substrs=["mtp."])
loaded_params = loader.load_weights(weights, mapper=self.hf_to_vllm_mapper)
self.model.finalize_mega_moe_weights()
self.model._convert_nvfp4_post_load()
return loaded_params
def get_expert_mapping(self) -> list[tuple[str, str, int, str]]:
return self.model.get_expert_mapping()
Reference in New Issue View Git Blame Copy Permalink