From 1836e5fdc708ed219a403878c75cccf6802202a8 Mon Sep 17 00:00:00 2001
From: biondizzle <biondizzle@gmail.com>
Date: Mon, 18 May 2026 19:27:49 +0000
Subject: [PATCH] Add shared experts to post-quant BF16 dequant fix

Shared experts also use FlashInferCutlassNvFp4LinearKernel with
broken input_scale. They need the same BF16 dequant treatment.
gate_up_proj and down_proj on ffn.shared_experts.
---
 vllm/patches/deepseek_v4.py | 38 ++++++++++++++++++++++++++-----------
 1 file changed, 27 insertions(+), 11 deletions(-)

diff --git a/vllm/patches/deepseek_v4.py b/vllm/patches/deepseek_v4.py
index 41c305e7..82e2a047 100644
--- a/vllm/patches/deepseek_v4.py
+++ b/vllm/patches/deepseek_v4.py
@@ -2409,35 +2409,51 @@ class DeepseekV4ForCausalLM(nn.Module):
 
     def _post_quant_fix(self) -> None:
         """Called by vLLM's process_weights_after_loading AFTER quant methods
-        have set up their attributes. Dequantizes attention NVFP4 weights to BF16
-        because FlashInferCutlassNvFp4LinearKernel uses broken input_global_scale_inv."""
+        have set up their attributes. Dequantizes NVFP4 weights to BF16 for
+        attention projections and shared experts because
+        FlashInferCutlassNvFp4LinearKernel uses broken input_global_scale_inv."""
         from vllm.model_executor.layers.linear import UnquantizedLinearMethod
         
         E2M1_LUT = torch.tensor([0, 0.5, 1, 1.5, 2, 3, 4, 6], dtype=torch.bfloat16)
         fixed = 0
         for layer_idx, layer in enumerate(self.model.layers):
             attn = layer.attn
+            # Attention projections
             for proj_name in ["fused_wqa_wkv", "wq_b", "wo_b"]:
                 if not hasattr(attn, proj_name):
                     continue
                 mod = getattr(attn, proj_name)
                 if not hasattr(mod, "weight") or mod.weight.dtype not in (torch.uint8, torch.int8):
                     continue
-                # Dequantize to BF16
                 self.model._dequant_nvfp4_to_bf16(mod, E2M1_LUT)
-                # Replace quant method
                 mod.quant_method = UnquantizedLinearMethod()
-                # Clean up NVFP4 attributes
                 for attr in ("weight_scale", "weight_scale_2", "input_scale",
                               "input_global_scale", "input_global_scale_inv",
-                              "weight_global_scale", "alpha",
-                              "weight_scale_inv"):
+                              "weight_global_scale", "alpha", "weight_scale_inv"):
                     if hasattr(mod, attr):
-                        try:
-                            delattr(mod, attr)
-                        except (AttributeError, TypeError):
-                            pass
+                        try: delattr(mod, attr)
+                        except: pass
                 fixed += 1
+            
+            # Shared expert projections (also NVFP4 with broken input_scale)
+            ffn = layer.ffn
+            if hasattr(ffn, 'shared_experts'):
+                for proj_name in ["gate_up_proj", "down_proj"]:
+                    se = ffn.shared_experts
+                    if not hasattr(se, proj_name):
+                        continue
+                    mod = getattr(se, proj_name)
+                    if not hasattr(mod, "weight") or mod.weight.dtype not in (torch.uint8, torch.int8):
+                        continue
+                    self.model._dequant_nvfp4_to_bf16(mod, E2M1_LUT)
+                    mod.quant_method = UnquantizedLinearMethod()
+                    for attr in ("weight_scale", "weight_scale_2", "input_scale",
+                                  "input_global_scale", "input_global_scale_inv",
+                                  "weight_global_scale", "alpha", "weight_scale_inv"):
+                        if hasattr(mod, attr):
+                            try: delattr(mod, attr)
+                            except: pass
+                    fixed += 1
         print(f"  [CLAWMINE] Post-quant fix: {fixed} attention projections → BF16 ✓", flush=True)
 
     def get_expert_mapping(self) -> list[tuple[str, str, int, str]]: