fix: correct Nvfp4Linear dimension inference from checkpoint weights

Weight shape (N_packed, K_packed) means:
- out_features = N_packed (GEMM output dim in BF16)
- in_features = K_packed * 2 (BF16 input dim, for activation buffer)
This commit is contained in:
2026-06-01 00:32:36 +00:00
parent e671780008
commit 15f45b57c3

View File

@@ -131,12 +131,12 @@ def make_nvfp4_linear(in_features, out_features, device, all_w, pfx, proj_name):
d = device
weight, ws, ws2, isc = get_nvfp4_weight(all_w, pfx, proj_name)
assert weight is not None, f"{pfx}.{proj_name}.weight not found"
# Infer actual dimensions from checkpoint weight shape
# weight is (O, I//2) uint8 where O = out_features//2
actual_out = weight.shape[0] * 2
actual_in = weight.shape[1] * 2
if actual_out != out_features or actual_in != in_features:
log.warning(f"{pfx}.{proj_name}: declared ({in_features},{out_features}) but weight is ({actual_in},{actual_out})")
# Checkpoint weight is (N_packed, K_packed) uint8
# NVFP4 GEMM output dim = N_packed BF16 elements
# Activation buffer needs K_packed FP4 columns = in_features BF16
# So: in_features = K_packed * 2, out_features = N_packed
actual_out = weight.shape[0] # N_packed = GEMM output dimension
actual_in = weight.shape[1] * 2 # K_packed * 2 = BF16 input dim (for buffer allocation)
lin = Nvfp4Linear(actual_in, actual_out, max_num_tokens=8192, device=d)
lin.fp4 = [weight.to(d)]; lin.sf = [ws.to(d)]
lin.ws2 = [ws2.to(d) if ws2 is not None else None] # weight_scale_2