Add --warmup-gsa flag: fix attention/router gsa after first decode step to eliminate amax kernel launches

single_shot_inference.py

@@ -27,6 +27,7 @@ def parse_args():
     p.add_argument('--seed', type=int, default=42)
     p.add_argument('--verbose', type=int, default=1)
     p.add_argument('--prefill-only', action='store_true')
+    p.add_argument('--warmup-gsa', action='store_true', help='Fix gsa values after first decode step (eliminates amax kernel launches)')
     p.add_argument('--profile', action='store_true', help='Profile per-component GPU time using CUDA events')
     p.add_argument('--num-gpus', type=int, default=8)
     p.add_argument('--checkpoint', type=str, default="/root/nvidia-meeting/DeepSeek-V4-Pro-NVFP4")
@@ -988,6 +989,7 @@ def main():
     print(f"\nDecoding (max {MAX_NEW_TOKENS} tokens)...")
     in_thinking = False
     profile = _args.profile
+    warmup_gsa = _args.warmup_gsa
     prof_embed_layers = 0.0
     prof_lm_head = 0.0
     prof_sample = 0.0
@@ -1013,6 +1015,39 @@ def main():
         X = X.to('cuda:0'); torch.cuda.set_device(0)
         if profile: torch.cuda.synchronize()
         t_layers = time.perf_counter()
+
+        # After first decode step: fix gsa values from runtime amax
+        # This eliminates amax_gsa kernel launches on subsequent steps
+        # Only applies to attention linears and router gate (fixed per-projection gsa)
+        # MoE/SE keep runtime gsa (gsa varies per token)
+        if warmup_gsa and step == 0:
+            torch.cuda.synchronize()
+            n_fixed = 0
+            for li in range(n_layers):
+                pl = prod_lins.get(li)
+                if pl is None: continue
+                for key, lin in pl.items():
+                    if hasattr(lin, '_gsa_buf') and hasattr(lin, '_use_runtime_gsa') and lin._use_runtime_gsa:
+                        fixed_gsa = lin._gsa_buf.item()  # One-time sync
+                        lin._activation_global_scale = fixed_gsa
+                        lin._use_runtime_gsa = False
+                        n_fixed += 1
+                # Router gate
+                router = routers.get(li)
+                if router and hasattr(router, '_gate_lin') and router._gate_lin is not None:
+                    gl = router._gate_lin
+                    if hasattr(gl, '_gsa_buf') and hasattr(gl, '_use_runtime_gsa') and gl._use_runtime_gsa:
+                        fixed_gsa = gl._gsa_buf.item()
+                        gl._activation_global_scale = fixed_gsa
+                        gl._use_runtime_gsa = False
+                        n_fixed += 1
+            # lm_head
+            if hasattr(lm_head_lin, '_gsa_buf') and hasattr(lm_head_lin, '_use_runtime_gsa') and lm_head_lin._use_runtime_gsa:
+                fixed_gsa = lm_head_lin._gsa_buf.item()
+                lm_head_lin._activation_global_scale = fixed_gsa
+                lm_head_lin._use_runtime_gsa = False
+                n_fixed += 1
+            print(f"  Warmup gsa: fixed {n_fixed} projection gsa values from step 0 (MoE/SE keep runtime gsa)", flush=True)
         x_out = hc_head.forward(X) if hc_head is not None else X[:, 0, :]
         if final_norm_w is not None: x_out = rmsnorm(x_out, final_norm_w)
         logits = lm_head_lin(x_out)