Add logits validation debug before topk sampling

single_shot_inference.py

@@ -988,6 +988,15 @@ def main():
         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)
+        # Validate logits before sampling
+        if step == 0 or torch.isnan(logits.float()).any().item():
+            print(f"  logits: shape={list(logits.shape)} dtype={logits.dtype} "
+                  f"min={logits.float().min().item():.1f} max={logits.float().max().item():.1f} "
+                  f"has_nan={torch.isnan(logits.float()).any().item()} "
+                  f"has_inf={torch.isinf(logits.float()).any().item()}", flush=True)
+            if torch.isnan(logits.float()).any().item() or torch.isinf(logits.float()).any().item():
+                print(f"  NaN/Inf in logits at step {step}, aborting", flush=True)
+                break
 
         # Sampling — fused CUDA kernel (or greedy argmax for temp=0)
         if is_greedy: