biondizzle/deepseek-v4-quant
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

WIP: MegaMoE NVFP4 kernel + diagnostics

Browse Source 
- 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
This commit is contained in:
biondizzle
2026-05-11 05:19:49 +00:00
parent 7e2f219259
commit e963325b61
1 changed files with 210 additions and 29 deletions
Download Patch File Download Diff File Expand all files Collapse all files

239
patches/deepseek_v4.py
View File

@@ -425,8 +425,24 @@ def make_deepseek_v4_expert_params_mapping(
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,
@@ -451,6 +467,8 @@ class DeepseekV4MegaMoEExperts(nn.Module):
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,
@@ -462,18 +480,34 @@ class DeepseekV4MegaMoEExperts(nn.Module):
)
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 // 32,
dtype=torch.uint8,
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,
@@ -485,18 +519,31 @@ class DeepseekV4MegaMoEExperts(nn.Module):
)
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 // 32,
dtype=torch.uint8,
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
@@ -525,6 +572,11 @@ class DeepseekV4MegaMoEExperts(nn.Module):
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:
@@ -573,36 +625,162 @@ class DeepseekV4MegaMoEExperts(nn.Module):
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(
self._ue8m0_uint8_to_float(self.w13_weight_scale.data).contiguous(),
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(
self._ue8m0_uint8_to_float(self.w2_weight_scale.data).contiguous(),
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(
(self.w13_weight.data.view(torch.int8).contiguous(), w13_scale),
(self.w2_weight.data.view(torch.int8).contiguous(), w2_scale),
(w13_mxfp4_weight.view(torch.int8).contiguous(), w13_scale),
(w2_mxfp4_weight.view(torch.int8).contiguous(), w2_scale),
)
)
# Drop the original loader-side parameters: the MegaMoE kernels only
# consume the transformed views above. transform_weights_for_mega_moe
# allocates a fresh tensor for the L1 weight (see _interleave_l1_weights)
# and fresh SF tensors for L1/L2; the L2 weight is the only tensor that
# aliases the original storage, and _transformed_l2_weights still holds
# it, so the storage stays live after we drop the Parameter.
# 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
@@ -751,9 +929,7 @@ class DeepseekV4MoE(nn.Module):
config = vllm_config.model_config.hf_config
quant_config = vllm_config.quant_config
self.prefix = prefix
self.use_mega_moe = (
vllm_config.kernel_config.moe_backend == "deep_gemm_mega_moe"
)
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. "
@@ -774,12 +950,7 @@ class DeepseekV4MoE(nn.Module):
raise NotImplementedError(
"DeepSeek V4 MegaMoE currently supports sqrtsoftplus routing only."
)
if self.use_mega_moe and getattr(config, "expert_dtype", "fp4") != "fp4":
raise NotImplementedError(
"DeepSeek V4 MegaMoE only supports fp4 experts; got expert_dtype="
f"{config.expert_dtype!r}. Drop --kernel-config moe_backend="
"deep_gemm_mega_moe for this checkpoint."
)
# NVFP4 experts work with mega_moe via NVFP4→MXFP4 conversion in finalize_weights
self.gate = GateLinear(
config.hidden_size,
@@ -1262,9 +1433,7 @@ class DeepseekV4Model(nn.Module):
config = vllm_config.model_config.hf_config
quant_config = vllm_config.quant_config
self.config = config
self.use_mega_moe = (
vllm_config.kernel_config.moe_backend == "deep_gemm_mega_moe"
)
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. "
@@ -1648,7 +1817,7 @@ class DeepseekV4Model(nn.Module):
- 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: Stay in native NVFP4 (ModelOptNvFp4FusedMoE)
- MoE experts: Handled by DeepseekV4MegaMoEExperts (NVFP4→MXFP4)
"""
E2M1_LUT = torch.tensor(
[0, 0.5, 1, 1.5, 2, 3, 4, 6], dtype=torch.bfloat16
@@ -1667,6 +1836,7 @@ class DeepseekV4Model(nn.Module):
fp8_from_bf16 = 0
bf16_converted = 0
compressor_converted = 0
diag_printed = False
for layer_idx, layer in enumerate(self.layers):
attn = layer.attn
@@ -1679,13 +1849,20 @@ class DeepseekV4Model(nn.Module):
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:
# modelopt did NOT quantize o_a_proj — it's bf16 already.
# Convert bf16 -> FP8 directly for fp8_einsum path.
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:
@@ -1694,6 +1871,10 @@ class DeepseekV4Model(nn.Module):
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