[Bugfix] Fix ASM split-K semaphore deadlock under CUDA graph capture - #4916
[Bugfix] Fix ASM split-K semaphore deadlock under CUDA graph capture#4916JohnQinAMD wants to merge 1 commit into
Conversation
🏷️ CI GuideRuns automatically on every PR:
Extended tests (opt-in via labels):
PR title tags: |
|
please provide UT for repro |
8b61f50 to
0e8c750
Compare
Hi @valarLip , UT for repro has been added. Please help review. Thanks |
There was a problem hiding this comment.
Pull request overview
Fixes a CUDA/HIP graph replay hang for ASM a16w16 split-K GEMM by ensuring the split-K semaphore counter is re-zeroed on every graph replay and that concurrently replayed graphs never share the same counter slot.
Changes:
- Add a per-device preallocated semaphore “ring” for graph capture and record
sema.zero_()inside the captured graph so replay always starts from a clean counter. - Keep eager execution using the existing per-(device, stream) cached semaphore, while priming the capture pool outside capture.
- Add unit tests covering zero-fill recording on replay, unique slot assignment under capture, and pool exhaustion behavior.
Reviewed changes
Copilot reviewed 2 out of 2 changed files in this pull request and generated 1 comment.
| File | Description |
|---|---|
| aiter/ops/gemm_op_a16w16.py | Introduces capture-specific semaphore pooling + in-graph zero-fill to prevent split-K deadlocks on graph replay. |
| op_tests/test_gemm_a16w16_graph.py | Adds regression tests validating capture/replay semaphore semantics without launching kernels. |
💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.
|
please pay some attention on your test, current one looks like pure agent garbage |
0e8c750 to
b9e95e9
Compare
b9e95e9 to
5feace3
Compare
5feace3 to
6009c1e
Compare
25e3968 to
9ec13cb
Compare
There was a problem hiding this comment.
Pull request overview
Copilot reviewed 2 out of 2 changed files in this pull request and generated 1 comment.
Suppressed comments (1)
op_tests/test_gemm_a16w16.py:455
- check_graph() assumes these shapes actually select a split-K ASM kernel, but it never validates that the dispatcher chose split>1. Because get_semaphore_workspace() now always exists even when split==1 (the C++ launcher passes ptr_semaphore=nullptr unless split>1), this check could silently become a no-op if heuristics change. Consider explicitly verifying that split>1 was chosen for this (m,n,k) by seeding the eager semaphore to a sentinel value before the warmup GEMM and asserting the kernel reset it back to 0 after synchronize; otherwise fail loudly / skip with a message to adjust shapes.
# Warm up outside capture: the first call loads the module and allocates the
# semaphore workspace, neither of which may happen inside a capture region.
aiter.gemm_a16w16_asm(x, wshuffle, out, bpreshuffle=wshuffle.is_shuffled)
torch.cuda.synchronize()
eager = out.clone()
|
@valarLip Rewrote UT as $ python3 op_tests/test_gemm_a16w16.py --graph -mnk 64,256,5120 32,512,8192 -d bf16 -o fp32
graph dim: (64, 256, 5120) replay vs eager: [checkAllclose atol=0.01 rtol=0.01 passed~]
graph dim: (32, 512, 8192) replay vs eager: [checkAllclose atol=0.01 rtol=0.01 passed~]
all graph capture/replay checks passedTo see it fail, on the main branch run the same command — it hangs in the replay and holds the GPU until killed, so not on a shared runner: git checkout origin/main -- aiter/ops/gemm_op_a16w16.py |
9ec13cb to
8318515
Compare
8318515 to
eeaecd1
Compare
eeaecd1 to
df336d8
Compare
There was a problem hiding this comment.
Pull request overview
Copilot reviewed 2 out of 2 changed files in this pull request and generated 1 comment.
Suppressed comments (1)
aiter/ops/gemm_op_a16w16.py:118
get_semaphore_workspace()usestorch.deviceobjects as dict keys for_capture_rings/_capture_ring_next. If a caller passes an un-indexed CUDA device (torch.device('cuda'), i.e.device.index is None) while switchingtorch.cuda.current_device()between GPUs, the same dict key can accidentally alias different physical devices, returning/zeroing a ring allocated on the wrong device.
Consider canonicalizing the key to an indexed CUDA device (e.g. device = torch.device('cuda', torch.cuda.current_device()) when device.type=='cuda' and device.index is None) before interacting with _capture_rings/_capture_ring_next/LRU cache, similar to how aiter/ops/flydsl/utils.py:get_shared_memory_per_block() normalizes device.index is None to the current device index.
# is_current_stream_capturing() answers about the current device; only pay
# the device switch (~1us, and this is a per-GEMM path) when it differs.
if device.index is None or device.index == torch.cuda.current_device():
capturing = torch.cuda.is_current_stream_capturing()
else:
with torch.cuda.device(device):
capturing = torch.cuda.is_current_stream_capturing()
if capturing:
return _get_captured_semaphore_workspace(device)
# Allocate here, never under capture: graph-pool memory cannot be freed.
_prime_capture_pool(device)
df336d8 to
4c02f02
Compare
There was a problem hiding this comment.
Pull request overview
Copilot reviewed 2 out of 2 changed files in this pull request and generated 1 comment.
Suppressed comments (1)
aiter/ops/gemm_op_a16w16.py:105
- The docstring says the pool is allocated by the first eager split-K call, but
gemm_a16w16_asm()callsget_semaphore_workspace()wheneversplitK is None or splitK > 1(even if the heuristic later chooses splitK==1 and the semaphore is unused). Consider rewording to reflect the actual trigger (first eager call that requests a semaphore, i.e.splitK is None or > 1), so users understand why a process that never captures can still pay the one-time pool allocation.
Under capture this returns a slot from a per-device pool instead, with the
zero-fill recorded as a graph node so replay restores counter == 0. That
pool has to exist before capture starts, so the first eager splitK call on
a device allocates it: a fixed CAPTURE_SEMAPHORE_POOL_SLOTS * 4 KiB, paid
once per device even by a process that never captures a graph.
The split-K counter must be zero when a launch starts, and aiter zeroes it only where it allocates it. A CUDA graph records launches, not that allocation, so replay starts from whatever the counter held, the reduction never fires and the waves spin. Kimi-K3 TP8 with a graphed draft hangs during warmup on 8x MI355X. Give each recorded launch a slot from a per-device pool allocated on the eager path, and record the zero-fill inside the graph. Exhausting the pool raises rather than reusing a slot, since two graphs sharing a counter deadlock when replayed concurrently. Eager launches are unchanged. Supersedes ROCm#4494, which allocated under capture and therefore had to retain every workspace in a list that never shrank. Signed-off-by: Yanyuan Qin <yanyuan.qin@amd.com> Co-authored-by: Claude Opus 5 (1M context) <noreply@anthropic.com>
4c02f02 to
13aa3d8
Compare
Summary
A split-K a16w16 GEMM recorded into a CUDA graph hangs forever on replay — the GPU wedges, no wrong number, no recovery. Each recorded launch now gets its own counter, and the zero-fill is recorded inside the graph.
Supersedes #4494, reverted by #4709 for unbounded VRAM growth.
Motivation
The kernel needs its split-K counter to be 0 when a launch starts. aiter zeroes it once, where it allocates it. A graph records launches, not that allocation, so replay never re-zeroes it — and a counter that is not 0 means the reduction never fires and the waves spin.
rocgdbon a hung endpoint: 132 waves ofbf16gemm_fp32bf16_tn_64x64_splitk_cleanwaiting on a counter holding2.Technical Details
#4494 allocated the workspace while the graph was capturing. That memory belongs to the graph's private pool, and aiter never learns when a graph dies, so it could not free it — every workspace went into a list that grew with each capture. That list is what #4709 objected to.
This allocates the pool once, on the eager path, from the normal allocator. Nothing to keep alive, so no list.
Running out of slots raises instead of reusing one, since two graphs sharing a counter deadlock when replayed together. So does capturing on a device that never ran eagerly, where the only alternative is to allocate from the graph's pool. Both are loud at startup rather than wedged in production later — the same choice
_check_split_k_semaphore_capacitymakes in the FlyDSL split-K path.Removing the counter from the kernel would make all of this unnecessary. That is #4920.
Test Plan
check_graph()and a--graphflag inop_tests/test_gemm_a16w16.py, followingtest_inverse_rope_group_quant.py. Two assertions per shape:get_semaphore_workspace()alone, write a sentinel into what it returns, replay, require it back at 0. No kernel, so a build that stops recording the zero-fill fails in milliseconds instead of spinning.--graphis off by default, so CI (python3 <file>) runs neither.Test Result
MI355X (gfx950):
--graphsema.zero_()AssertionError: capture recorded no zero-fillThe middle row is the one for #4709: delete the single line this PR turns on and the check goes red, so the property is guarded, not just claimed.
End to end, Kimi-K3 TP8 with a graphed speculative draft on 8x MI355X. Without the fix the draft capture wedges 3.6 minutes in and the server never starts — 25 minutes later, 8 GPUs still pinned at 100% with no worker progress. With it, capture completes on all 8 ranks, the endpoint is up in 330 s, and GSM8K-100 scores 1.000 with 0 invalid at 120.18 output tok/s. #4494 measured 120.1 tok/s on the same stack, so throughput is unchanged and the unbounded memory is gone.
Submission Checklist
Tool assistance
Claude Opus 5 (1M context) assisted with the fix, the tests and this description.