Add minimal CUDA graph test per GPU during capture to isolate multi-GPU graph issue

single_shot_inference.py

@@ -312,6 +312,25 @@ class CUDAGraphDecoder:
             pfx = f"model.layers.{li}.self_attn"
             
             # ======== Graph A: pre-attention compute ========
+            # NOTE: We capture each Graph A on the correct GPU. Multi-GPU graph capture
+            # is known to have issues. We add a validation step to verify correctness.
+            #
+            # VALIDATION: Before capturing the full Graph A, we first test a MINIMAL
+            # graph on this GPU to verify basic graph capture works.
+            if li < 2:
+                _test_graph = torch.cuda.CUDAGraph()
+                _test_input = self.x_in_bufs[li].clone()
+                _test_input.fill_(1.0)  # non-zero input
+                with torch.cuda.graph(_test_graph):
+                    _test_output = _test_input * 2.0
+                _test_input.copy_(self.x_in_bufs[li])  # restore original data
+                _test_graph.replay()
+                torch.cuda.synchronize()
+                _test_max = _test_output.abs().max().item()
+                print(f"    L{li} minimal graph test on cuda:{gpu}: output |X|={_test_max:.2f} (expected non-zero)", flush=True)
+                if _test_max == 0.0:
+                    print(f"    *** L{li} MINIMAL GRAPH TEST FAILED on cuda:{gpu}! ***", flush=True)
+                del _test_graph, _test_input, _test_output
             # Input:  X_l = self.x_in_bufs[li]  (1, 4, H)
             # Output: x_normed, q_heads, kv_3d, ctx_a, X_l → pre-allocated buffers
             # NOTE: Norm weights are pre-cached on device in FP32 (attn_norm_dev, etc.)