From 86275851d4033e6d1de5ec69b40951078efbe6b4 Mon Sep 17 00:00:00 2001 From: biondizzle Date: Sat, 6 Jun 2026 08:02:35 +0000 Subject: [PATCH] Add minimal CUDA graph test per GPU during capture to isolate multi-GPU graph issue --- single_shot_inference.py | 19 +++++++++++++++++++ 1 file changed, 19 insertions(+) diff --git a/single_shot_inference.py b/single_shot_inference.py index 4e6e6d8d..86e3f803 100644 --- a/single_shot_inference.py +++ b/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.)