From e963325b61c02fb9a3115542e5c728f61eaae5bf Mon Sep 17 00:00:00 2001 From: biondizzle Date: Mon, 11 May 2026 05:19:49 +0000 Subject: [PATCH] WIP: MegaMoE NVFP4 kernel + diagnostics MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - 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 --- patches/deepseek_v4.py | 239 ++++++++++++++++++++++++++++++++++++----- 1 file changed, 210 insertions(+), 29 deletions(-) diff --git a/patches/deepseek_v4.py b/patches/deepseek_v4.py index b4c7126..97ed479 100644 --- a/patches/deepseek_v4.py +++ b/patches/deepseek_v4.py @@ -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