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

Switch lm_head from NVFP4 to BF16 GEMM

Browse Source 
Checkpoint lm_head.weight is BF16 (129280x7168). NVFP4 quantization
of lm_head can flatten/reorder vocabulary logits — wrong experts in
the final projection directly corrupt token selection. BF16 preserves
the full logit distribution for accurate token selection.

Tagged pure-nvfp4 before this change for rollback.
This commit is contained in:
biondizzle
2026-06-03 11:37:41 +00:00
parent 3320abfe24
commit bdd9ab9669
1 changed files with 15 additions and 16 deletions
Download Patch File Download Diff File Expand all files Collapse all files

31
single_shot_inference.py
View File

@@ -1397,21 +1397,15 @@ def main():
torch.cuda.set_device(0)
embed_w = all_w.get("model.embed_tokens.weight")
embed = torch.nn.Embedding.from_pretrained(embed_w.bfloat16().to('cuda:0'))
# lm_head: NVFP4 production GEMM
# lm_head: BF16 GEMM — checkpoint weight is BF16, no quantization needed.
# NVFP4 quantization of lm_head can flatten/reorder vocabulary logits;
# BF16 preserves the full logit distribution for accurate token selection.
lm_w_raw = all_w.get("lm_head.weight", embed_w).bfloat16().to('cuda:0')
from dsv4.layers.linear import Nvfp4Linear
lm_head_lin = Nvfp4Linear(lm_w_raw.shape[1], lm_w_raw.shape[0], max_num_tokens=8192, device='cuda:0')
from dsv4.ops.quantize import quantize_weight_to_nvfp4
lm_fp4, lm_sf, lm_gs = quantize_weight_to_nvfp4(lm_w_raw.T.contiguous())
lm_head_lin.fp4 = [lm_fp4.permute(1, 0).contiguous()]
lm_head_lin.sf = [lm_sf.permute(1, 0).contiguous()]
lm_head_lin.gs = [lm_gs]
lm_head_lin.ws2 = [None]
lm_head_lin._activation_global_scale = 1.0 / (6.0 * 448.0)
lm_head_lin._use_runtime_gsa = True
lm_head_lin.finalize_weights()
# Store as (vocab, hidden) BF16 — direct cuBLAS GEMM
lm_head_weight = lm_w_raw.contiguous() # (129280, 7168) BF16
lm_head_lin = None # flag: use BF16 path
lm_w = None
print(" lm_head: NVFP4 production GEMM")
print(" lm_head: BF16 GEMM (checkpoint weight, no quantization)")
final_norm_w = all_w.get("model.norm.weight")
if final_norm_w is not None: final_norm_w = final_norm_w.to('cuda:0', torch.float32)
@@ -1660,8 +1654,8 @@ def main():
gl._activation_global_scale = fixed_gsa
gl._use_runtime_gsa = False
n_fixed += 1
# lm_head
if hasattr(lm_head_lin, '_gsa_buf') and hasattr(lm_head_lin, '_use_runtime_gsa') and lm_head_lin._use_runtime_gsa:
# lm_head (skip if BF16 path)
if lm_head_lin is not None and hasattr(lm_head_lin, '_gsa_buf') and hasattr(lm_head_lin, '_use_runtime_gsa') and lm_head_lin._use_runtime_gsa:
fixed_gsa = lm_head_lin._gsa_buf.item()
lm_head_lin._activation_global_scale = fixed_gsa
lm_head_lin._use_runtime_gsa = False
@@ -1669,7 +1663,12 @@ def main():
print(f" Warmup gsa: fixed {n_fixed} projection gsa values from step 0 (MoE/SE keep runtime gsa)", flush=True)
x_out = hc_head.forward(X) if hc_head is not None else X[:, 0, :]
if final_norm_w is not None: x_out = rmsnorm(x_out, final_norm_w)
logits = lm_head_lin(x_out)
# lm_head: BF16 GEMM (cuBLAS) or NVFP4 GEMM
if lm_head_lin is not None:
logits = lm_head_lin(x_out)
else:
# BF16 path: direct matmul with (vocab, hidden) weight
logits = torch.nn.functional.linear(x_out, lm_head_weight)
if profile: torch.cuda.synchronize()
t_lm = time.perf_counter()
# Check thinking start token logit on first step