diff --git a/single_shot_inference.py b/single_shot_inference.py index 6be8ba92..718ce59f 100644 --- a/single_shot_inference.py +++ b/single_shot_inference.py @@ -17,7 +17,7 @@ This is the ground truth for vLLM / SGLang integration. """ import os, sys, time, json, math, argparse, logging import torch -os.environ['CUDA_LAUNCH_BLOCKING'] = '1' # Catch CUDA errors synchronously +# os.environ['CUDA_LAUNCH_BLOCKING'] = '1' # Disabled — was for debugging import torch.nn.functional as F from pathlib import Path @@ -685,38 +685,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, q_comp_dim), - ('q_b', q_comp_dim, n_h * hd), - ('kv', H, hd), - ('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: - lin = make_nvfp4_linear(in_f, out_f, dev, wt, ws, ws2, isc) - plin[proj] = lin - if plin: - prod_lins[li] = plin - if (li+1) % 10 == 0: - print(f" Built Nvfp4Linear {li+1}/{n_layers} layers") - # Sync to catch errors early - torch.cuda.set_device(li % NUM_GPUS) - torch.cuda.synchronize() + # Nvfp4Linear for attention projections (deferred — use reference for now) + # Production MoE + Router + FMHA are the critical paths. + # Nvfp4Linear for small projections can be enabled once MoE is validated. + prod_lins = {} # Empty = use reference dequant path + print(" Using reference dequant for attention projections") # Routers, MoE, shared experts routers, moe_runners, se_runners = {}, {}, {} for li in range(n_layers): dev = f"cuda:{li % NUM_GPUS}" pfx = f"model.layers.{li}.mlp" + torch.cuda.set_device(li % NUM_GPUS) + # Verify GPU is in good state before MoE loading + torch.cuda.synchronize() is_hash = (li < cfg.get("num_hash_layers", 3)) and (f"{pfx}.gate.tid2eid" in all_w) router = Router(