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.
This commit is contained in:
2026-05-18 19:27:49 +00:00
parent 82ac648563
commit 1836e5fdc7

View File

@@ -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]]: