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

fix: unpack uint8 NVFP4→bf16 for non-stacked params (weights_proj)

Browse Source 
indexer.weights_proj is uint8 [64,3584] in checkpoint but bf16 [64,7168]
in model. The uint8→bf16 unpack logic only ran in the stacked_params
loop, so non-stacked NVFP4 params hit a size mismatch assertion.
This commit is contained in:
biondizzle
2026-05-15 00:51:50 +00:00
parent e6ed9facf3
commit af6583eb19
1 changed files with 21 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files

21
vllm/patches/deepseek_v4.py
View File

@@ -1534,6 +1534,27 @@ class DeepseekV4Model(nn.Module):
loaded_params.add(is_name)
continue
# Handle uint8 NVFP4 packed → bf16 unpack for non-stacked
# params (e.g. indexer.weights_proj). Checkpoint stores
# NVFP4 as uint8 (2 values/byte), but model param is bf16.
if (loaded_weight.dtype == torch.uint8
and param.data.dtype != torch.uint8
and loaded_weight.shape[-1] * 2 == param.data.shape[-1]):
FP4_LUT = torch.tensor([
0.0, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0,
-0.0, -0.5, -1.0, -1.5, -2.0, -3.0, -4.0, -6.0,
], dtype=torch.float32, device=loaded_weight.device)
lower = FP4_LUT[(loaded_weight & 0x0F).long()]
upper = FP4_LUT[((loaded_weight >> 4) & 0x0F).long()]
out = torch.empty(
*loaded_weight.shape[:-1],
loaded_weight.shape[-1] * 2,
dtype=torch.float32, device=loaded_weight.device,
)
out[..., 0::2] = lower
out[..., 1::2] = upper
loaded_weight = out.to(torch.bfloat16)
weight_loader = getattr(
param, "weight_loader", default_weight_loader
)