single_shot: use reference dequant for attn projections, focus on MoE+FMHA

Nvfp4Linear causing CUDA context corruption (likely CuTeDSL JIT
triggered by _ensure_initialized). Disable for now to validate
the critical paths first:
- Production FMHA with sink bias
- Production MoE (Nvfp4MoE + Nvfp4SharedExpert)
- Production Router (dense/hash)
- Production mHC

Attention projections use reference dequant+matmul for now.
Will re-enable Nvfp4Linear after validating MoE path.

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(