diff --git a/single_shot_inference.py b/single_shot_inference.py
index e13b7f86..f9c4443b 100644
--- a/single_shot_inference.py
+++ b/single_shot_inference.py
@@ -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: