biondizzle/nvfp4-megamoe-kernel
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Fix input_scale BEFORE process_weights_after_loading runs

Browse Source 
Instead of dequantizing to BF16 (which gets overwritten by
process_weights_after_loading), fix the input_scale parameter
on the module before the quant method reads it. The quant method
computes input_global_scale_inv = input_scale.max(), so fixing
input_scale propagates the correct activation scale.

Computes correct input_scale by temporarily dequantizing weight
to BF16, running warmup forward, and computing act_amax.
input_scale = 1/(act_amax * headroom).
This commit is contained in:
biondizzle
2026-05-18 16:43:44 +00:00
parent 2fc81ccac4
commit 2835cb040b
1 changed files with 69 additions and 18 deletions
Download Patch File Download Diff File Expand all files Collapse all files

87
vllm/patches/deepseek_v4.py
View File

@@ -1685,37 +1685,26 @@ class DeepseekV4Model(nn.Module):
def _convert_nvfp4_post_load(self):
"""Post-load conversion of NVFP4 weights for vLLM compatibility.
All attention NVFP4 projections (except wo_a) are dequantized to BF16.
The checkpoint input_scale values cause NaN during activation quantization
in FlashInferCutlassNvFp4LinearKernel. BF16 bypasses this entirely.
Fixes the attention input_scale values BEFORE
process_weights_after_loading runs. The checkpoint input_scale
values are wrong and cause NaN during activation quantization.
We compute correct values by dequantizing to BF16 temporarily
and running a warmup forward.
wo_a is converted to FP8 for fp8_einsum (no input_scale needed).
Compressor weights are reconstructed from checkpoint sub-weights.
"""
bf16_proj_names = {"fused_wqa_wkv", "wq_b", "wo_b"}
fp8_proj_names = {"wo_a"}
bf16_converted = 0
fp8_converted = 0
compressor_converted = 0
input_scale_fixes = 0
_shard_index = self._build_shard_index("/model") if os.path.isdir("/model") else None
from tqdm import tqdm
for layer_idx, layer in tqdm(enumerate(self.layers), total=len(self.layers), desc=" (upcast)NVFP4→BF16 attn projs", unit="layer"):
for layer_idx, layer in tqdm(enumerate(self.layers), total=len(self.layers), desc=" (fix)NVFP4 attn input_scale", unit="layer"):
attn = layer.attn
# BF16 dequantization: attention projections (except wo_a)
for proj_name in bf16_proj_names:
if not hasattr(attn, proj_name):
continue
mod = getattr(attn, proj_name)
if not hasattr(mod, "weight"):
continue
if mod.weight.dtype in (torch.uint8, torch.int8):
E2M1_LUT = torch.tensor([0, 0.5, 1, 1.5, 2, 3, 4, 6], dtype=torch.bfloat16)
self._dequant_nvfp4_to_bf16(mod, E2M1_LUT)
bf16_converted += 1
# FP8 conversion: wo_a (used by fp8_einsum, no input_scale)
FP8_MAX = torch.finfo(torch.float8_e4m3fn).max
for proj_name in fp8_proj_names:
@@ -1729,6 +1718,68 @@ class DeepseekV4Model(nn.Module):
self._convert_nvfp4_to_fp8(mod, E2M1_LUT, FP8_MAX)
fp8_converted += 1
# Fix input_scale for attention NVFP4 projections
# process_weights_after_loading reads input_scale and computes
# input_global_scale_inv = 1/input_scale. By fixing input_scale
# here, the quant method will propagate the correct value.
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, "input_scale"):
continue
if not hasattr(mod, "weight") or mod.weight.dtype not in (torch.uint8, torch.int8):
continue
# Temporarily dequantize weight to BF16 for warmup
E2M1_LUT = torch.tensor([0, 0.5, 1, 1.5, 2, 3, 4, 6], dtype=torch.bfloat16)
w_uint8 = mod.weight.data
w_bf16_unpacked = self._unpack_nvfp4_to_bf16(w_uint8, E2M1_LUT, w_uint8.device)
if hasattr(mod, "weight_scale") and hasattr(mod, "weight_scale_2"):
block_scale = self._block_scale_to_float32(mod.weight_scale.data)
if block_scale.dim() == 2 and w_bf16_unpacked.dim() == 2:
block_size = w_bf16_unpacked.shape[1] // block_scale.shape[1]
block_scale_expanded = block_scale.unsqueeze(-1).expand(-1, -1, block_size).reshape(w_bf16_unpacked.shape)
else:
block_scale_expanded = block_scale
global_scale = mod.weight_scale_2.data.max().item()
w_bf16_dequant = (w_bf16_unpacked.float() * block_scale_expanded * global_scale).to(torch.bfloat16)
else:
w_bf16_dequant = w_bf16_unpacked
# Compute correct input_scale from warmup
with torch.no_grad():
in_features = w_bf16_dequant.shape[-1]
dummy_input = torch.randn(256, in_features, dtype=torch.bfloat16, device=mod.weight.device) * 2.0
ref_output = torch.nn.functional.linear(dummy_input, w_bf16_dequant)
act_amax = ref_output.amax().item()
del w_bf16_unpacked, w_bf16_dequant, ref_output
# input_scale should be 1/(amax * headroom) — this is the
# activation global scale that maps activations to FP4 range.
# process_weights_after_loading computes:
# input_global_scale_inv = input_scale.max()
# input_global_scale = 1 / input_global_scale_inv
headroom = 1.2
new_input_scale = 1.0 / (act_amax * headroom) if act_amax > 0 else mod.input_scale.data
if layer_idx == 0:
old_input_scale = mod.input_scale.data.item() if mod.input_scale.data.numel() == 1 else mod.input_scale.data.max().item()
print(f"[CLAWMINE] Layer 0: {proj_name} input_scale: {old_input_scale:.8f} → {new_input_scale:.8f} (act_amax={act_amax:.4f})")
mod.input_scale = torch.nn.Parameter(
torch.tensor([new_input_scale] * mod.input_scale.data.numel(), dtype=mod.input_scale.data.dtype, device=mod.input_scale.data.device),
requires_grad=False
)
input_scale_fixes += 1
_shard_index = self._build_shard_index("/model") if os.path.isdir("/model") else None
from tqdm import tqdm
for layer_idx, layer in tqdm(enumerate(self.layers), total=len(self.layers), desc=" (upcast)NVFP4→BF16 attn projs", unit="layer"):
attn = layer.attn
# Compressor: still needs BF16 reconstruction
mla_attn = getattr(attn, "mla_attn", None)
if mla_attn is not None: