From e3c24769e2f8471b6434bedeac091134da9ed368 Mon Sep 17 00:00:00 2001
From: biondizzle <biondizzle@gmail.com>
Date: Mon, 18 May 2026 23:41:39 +0000
Subject: [PATCH] Handle wo_a as bfloat16 (unquantized in NVFP4 checkpoint)

o_a_proj is NOT quantized by modelopt in the checkpoint (bfloat16),
but the attention forward pass expects FP8 (weight + weight_scale_inv).

- Create wo_a with quant_config=None to load bfloat16 weights
- Add FP8 quantization of wo_a in finalize_mega_moe_weights:
  per-tensor symmetric quantization to float8_e4m3fn + weight_scale_inv
- This matches what the fused_inv_rope_fp8_quant + einsum expects
---
 vllm/patches/deepseek_v4.py | 34 +++++++++++++++++++++++++++++++++-
 1 file changed, 33 insertions(+), 1 deletion(-)

diff --git a/vllm/patches/deepseek_v4.py b/vllm/patches/deepseek_v4.py
index 30778e60..bee7fc75 100644
--- a/vllm/patches/deepseek_v4.py
+++ b/vllm/patches/deepseek_v4.py
@@ -989,11 +989,15 @@ class DeepseekV4Attention(nn.Module):
         )
 
         self.kv_norm = RMSNorm(self.head_dim, self.eps)
+        # wo_a is NOT quantized in the NVFP4 checkpoint (modelopt left it as bfloat16),
+        # but the attention forward pass expects FP8 (weight + weight_scale_inv).
+        # Pass quant_config=None to load bfloat16, then process_weights_after_loading
+        # will handle the FP8 quantization.
         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,
+            quant_config=None,
             return_bias=False,
             prefix=f"{prefix}.wo_a",
         )
@@ -1624,6 +1628,34 @@ class DeepseekV4Model(nn.Module):
     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()
+            # Quantize wo_a to FP8 (checkpoint has bfloat16, forward expects FP8)
+            attn = layer.self_attn
+            if hasattr(attn, 'wo_a') and attn.wo_a.weight.dtype == torch.bfloat16:
+                self._quantize_wo_a_to_fp8(attn.wo_a)
+
+    @staticmethod
+    def _quantize_wo_a_to_fp8(wo_a: ColumnParallelLinear) -> None:
+        """Quantize wo_a weight from bfloat16 to float8_e4m3fn.
+
+        The attention forward pass (fused_inv_rope_fp8_quant + einsum)
+        expects wo_a.weight as FP8 and wo_a.weight_scale_inv as float32.
+        The NVFP4 checkpoint stores wo_a as bfloat16, so we quantize here.
+        Uses per-tensor symmetric quantization (same as modelopt FP8).
+        """
+        weight_bf16 = wo_a.weight.data
+        # Per-tensor FP8 quantization: scale = amax / fp8_max
+        fp8_max = torch.finfo(torch.float8_e4m3fn).max  # 448.0
+        amax = weight_bf16.abs().max().float()
+        scale = amax / fp8_max
+        # Avoid division by zero
+        if scale == 0:
+            scale = torch.tensor(1.0, device=scale.device)
+        scale_inv = 1.0 / scale
+        weight_fp8 = (weight_bf16.float() * scale).to(torch.float8_e4m3fn)
+        wo_a.weight = torch.nn.Parameter(weight_fp8, requires_grad=False)
+        wo_a.weight_scale_inv = torch.nn.Parameter(
+            scale_inv.clone(), requires_grad=False
+        )
 
 
 @torch.compile(backend=current_platform.simple_compile_backend)