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

Fix n_h reference before assignment in single_shot

Browse Source
This commit is contained in:
biondizzle
2026-05-31 23:14:24 +00:00
parent c5adbbfde6
commit 050b5ee449
1 changed files with 7 additions and 3 deletions
Download Patch File Download Diff File Expand all files Collapse all files

10
single_shot_inference.py
View File

@@ -698,16 +698,20 @@ def main():
if fn_k in all_w: ffn_norms[li] = all_w[fn_k].to(dev, torch.float32)
# Production Nvfp4Linear for attention projections
n_h = cfg["num_attention_heads"]
q_comp_dim = cfg.get('query_compression_dim', 1536)
o_groups = cfg.get('o_groups', 16)
o_lora_rank = cfg.get('o_lora_rank', 1024)
prod_lins = {}
for li in range(n_layers):
dev = f"cuda:{li % NUM_GPUS}"
pfx = f"model.layers.{li}.self_attn"
plin = {}
for proj, in_f, out_f in [
('q_a', H, cfg.get('query_compression_dim', 1536)),
('q_b', cfg.get('query_compression_dim', 1536), n_h * hd),
('q_a', H, q_comp_dim),
('q_b', q_comp_dim, n_h * hd),
('kv', H, hd),
('o_b', cfg.get('o_groups', 16) * cfg.get('o_lora_rank', 1024), H),
('o_b', o_groups * o_lora_rank, H),
]:
wt, ws, ws2, isc = get_nvfp4_weight(all_w, pfx, proj)
if wt is not None and ws is not None: