Fix ASM split-K semaphore deadlock under CUDA graph capture

[JohnQinAMD]

Purpose

What breaks

An ASM a16w16 split-K GEMM recorded into a CUDA graph can spin forever on replay. The device wedges; it does not return a wrong number and it does not recover. Kimi-K3 TP8 with DSpark on 8x MI355X hangs during warmup as soon as the draft is graphed, so the server never becomes ready. The workaround was to bypass split-K or keep the draft eager, at a throughput cost.

What this change touches

File	Change
aiter/ops/gemm_op_a16w16.py	semaphore allocation under capture
op_tests/test_gemm_a16w16_graph.py	new regression test

No kernel code is modified. The ASM binaries are untouched; the change is entirely in who owns the semaphore workspace.

The FlyDSL split-K HGEMM shares the same host-side defect in aiter/ops/flydsl/gemm_kernels.py, and it is deliberately not fixed here. It is not reachable in production today -- the family that hangs, small_m, is not dispatched (iter_small_m_registry_configs is commented out at gemm_kernels.py:24, and aiter/configs/ holds 0 small_m entries against 606 FlyDSL split-K entries), and the family that is dispatched, hgemm, tolerates a dirty counter. Because there is no time pressure there, it is being addressed properly rather than with the same patch: those kernels are generated from the Python DSL, so the counter can be removed outright in favour of an fp32 workspace plus a separate reduce kernel, the way csrc/opus_gemm already does it. That work is tracked separately and does not gate this PR.

The defect, in code

Split-K spreads one GEMM's K reduction across blocks that coordinate through an atomic counter. The kernel requires the counter to be zero when a launch starts. Blocks then count in, the last one to arrive does the reduction and writes the counter back to zero for the next launch.

The counter is zeroed exactly once, where it is allocated (aiter/ops/gemm_op_a16w16.py:36):

@functools.lru_cache(maxsize=64)                                   # allocated once
def _get_semaphore_workspace_keyed(device, stream_id) -> Tensor:
    return torch.zeros(_SEMA_SHAPE, dtype=torch.uint32, device=device)   # the only memset

def get_semaphore_workspace(device: torch.device) -> Tensor:
    """..."""                                                      # docstring elided
    stream = torch.cuda.current_stream(device)                     # the cache key
    return _get_semaphore_workspace_keyed(device, stream.cuda_stream)

Under graph capture both properties this rests on are lost:

1. The zero-fill is not in the graph. The cache entry was allocated the first time the shape ran, which is during warmup, before capture began. The recorded graph therefore contains the kernel launch and no memset. Nothing restores the entry state on replay. (If the very first call ever happens inside capture the memset is recorded -- but the second graph captured on that stream gets the cached tensor and no memset, so the exposure returns.)
2. The stream key stops separating anything. It keys on the capture stream, while the graph replays on whatever stream the caller uses, and PyTorch reuses one side stream across captures. Every graph captured on that stream is handed the same counter. Launches that can be in flight together now share one, which is exactly what the key existed to prevent.

So a counter left non-zero stays non-zero forever: clearing it was never part of the graph.

Point 1 is what the observed hang is. Point 2 is a lost guarantee rather than an observed failure -- 200 concurrent replays of two graphs captured on one stream produced no corruption in testing, though that run cannot prove the replays actually overlapped on the device. It is recorded because the cache was keyed by stream precisely to prevent it, and capture removes that protection.

How it was found

rocgdb on the hung endpoint: all 132 waves of bf16gemm_fp32bf16_tn_64x64_splitk_clean spinning on one counter, which held 2 while the kernel waited for 0. That is what named the invariant.

Why CI did not catch it

Three things had to line up, and they explain the shape of the test below.

1. No test crossed split-K with capture. Eight tests use torch.cuda.graph; none of them is a split-K GEMM, and no split-K test captured. The intersection was empty.
2. A naive crossing would have passed anyway. A clean capture plus replays passes on the unpatched build, because the kernel's own reset keeps the counter at zero. The hang needs the counter to be already dirty at entry, so a test must seed the residue deliberately. This is why the test seeds rather than just captures.
3. The production path needs the whole stack. The residue accumulates across repeated capture and replay under vLLM full cudagraphs with a speculative draft, which is why rocgdb on a live server found it and CI could not.

The new test sits at op_tests/ depth 1, so .github/scripts/split_tests.sh (find op_tests -maxdepth 1 -name 'test_*.py') collects it.

The fix

While capture is active, hand out a fresh zero-initialized workspace per launch and keep a reference to it. Eager execution keeps the existing cache:

if torch.cuda.is_current_stream_capturing():
    w = torch.zeros(_SEMA_SHAPE, dtype=torch.uint32, device=device)
    _captured_semaphore_keepalive.append(w)
    return w
stream = torch.cuda.current_stream(device)
return _get_semaphore_workspace_keyed(device, stream.cuda_stream)

The allocation happens inside the capture region, so its zero-fill is recorded as a graph node and the entry state is restored on every replay. Each recorded launch also gets a private counter, restoring the separation the stream key used to provide. GEMM math, kernel selection and the eager path are unchanged.

Why not the smaller fix. Recording a zero_() on the cached workspace is shorter and retains no memory at all. It was built and measured on the FlyDSL twin of this defect, and it does clear the hang. It was not taken because it moves the lru_cache lookup inside the capture region, and the capture stream is always new, so that lookup always misses (measured: lru_cache misses go 1 -> 2 during capture, versus unchanged for this PR's shape). A miss under capture means torch.zeros runs inside the capture region, so the cache ends up holding memory from that graph's private pool, keyed on the capture stream. Any later eager launch on that stream would then use graph-pool memory outside the graph. Returning before the cache is consulted avoids that.

Cost. One memset node and one 4096 B workspace (_SEMA_SHAPE = (16, 64) uint32) per recorded split-K launch, summed over every graph the process captures. Only launches that actually take the split-K path count -- splitK <= 1 returns an empty tensor and never reaches this code.

Measured rather than left to the reader, with temporary counting instrumentation that is not part of this change. Kimi-K3 TP8 with DSpark on 8x MI355X, FULL_AND_PIECEWISE, eight captured graph sizes: 29 workspaces, 0.11 MiB per rank, identical on all eight ranks, against a --gpu-memory-utilization 0.88 budget on 288 GB cards. The count is small because few shapes select splitK > 1.

The list grows only during startup capture. It grows only while capture is active, and in the hosts that reach this kernel capture is a bounded startup phase:

• vLLM raises on any capture outside its startup window. Both graph wrappers call validate_cudagraph_capturing_enabled() before capturing (vllm/compilation/breakable_cudagraph.py, vllm/compilation/cuda_graph.py), which throws RuntimeError: CUDA graph capturing detected at an inappropriate time unless the window is open. The window is opened and closed twice, both inside gpu_model_runner.py during startup: once around the encoder memory estimate, once around capture_model(). A runtime capture is a hard error, not a silent accumulation. vLLM does reach this kernel -- model_executor/layers/utils.py -> aiter.tuned_gemm.tgemm -> gemm_a16w16_asm -- so the path is reachable, just not repeatable.
• SGLang captures once over a fixed batch-size list in its full-cudagraph backend.
• AITER's own harness does not capture on this path at all. run_perftest defaults to testGraph=False and csrc/gemm_a16w16/gemm_a16w16_tune.py never sets it, so the tuner sweep records no graphs regardless of how many candidates it evaluates.

A host that captured repeatedly for the lifetime of the run would grow the list without bound. The mechanism that would do that in the wider ecosystem is inductor's cudagraph_trees re-recording under torch.compile(mode="reduce-overhead"), which requires calling AITER's tuned GEMM directly rather than through vLLM or SGLang. No such caller is known today, and neither serving host can enter that state.

Tensors allocated during capture come from the graph's private memory pool, so they cannot simply be freed while any graph that recorded them may still replay. AITER has no hook for graph destruction, so it cannot know when that stops being true, and holds them for the process lifetime instead. Under vLLM this costs the literal 4096 B and nothing beyond it: every capture shares one process-global pool (_global_graph_pool in vllm/platforms/interface.py) that is never released, so a retained reference pins nothing that was not already permanent for the life of the process.

How other aiter operators handle capture

Three approaches now exist in tree, and the difference is what each one needs from the graph:

Operator	Approach	Per-replay precondition
opus split-K (csrc/opus_gemm, aiter/tuned_gemm.py)	designed away: partials go to an fp32 workspace and a separate reduce kernel sums them -- "no self-clear, no semaphore"	none
custom all-reduce (aiter/dist/device_communicators/custom_all_reduce.py)	defer the side effect: record buffer addresses during capture, register_graph_buffers() after it ends	none
ASM split-K (this PR)	in-kernel handshake on a counter that must be zero on entry	yes

The opus case is the one to compare against, because _opus_prewarm_capture_workspace deliberately does the opposite of this PR: it sizes its workspace on the exact stream torch.cuda.graph will capture on before capture starts, and no-ops if already capturing. That works because opus only needs its workspace to exist and be registered. Pre-warming cannot fix a per-replay precondition: a pre-warmed workspace is zeroed once, the graph still records no memset, and replay N still cannot recover a dirty counter. Only something inside the graph can, which is why this fix allocates under capture rather than ahead of it.

The opus shape -- each split writing its fp32 partial to a disjoint slice of a [split_k, M, N] workspace, with a separate reduce kernel summing along the split axis and ordering coming from the stream -- removes the reason the counter exists rather than protecting it. That option is not available here: these kernels ship as precompiled hsa/gfx950/bf16gemm/*.co binaries, so the semaphore fix is the only one available for the ASM path, not an interim step.

Test plan

op_tests/test_gemm_a16w16_graph.py seeds the capture-stream counter with 2 -- the value rocgdb found -- captures two shipping (M, N, K) = (64, 256, 5120), splitK=13 launches, replays four times, and compares both outputs against eager. Seeding rather than racing for the residue makes it deterministic; the resulting entry state, and the hang, are the production ones.

The GPU work runs in a subprocess behind a timeout, because the unfixed control wedges the device rather than failing.

IMAGE='vllm/vllm-openai-rocm:kimi-k3@sha256:5aa7e626ff73672f5ca7aae46754570488c23d33ca1ac90756a1d2d1a3fe099b'
docker run --rm --ipc=host --shm-size=16g \
  --device=/dev/kfd --device=/dev/dri --group-add=video --group-add=render \
  -e HIP_VISIBLE_DEVICES=0 -e AITER_JIT_DIR=/tmp/aiter-jit-splitk-graph \
  --entrypoint bash "$IMAGE" -lc '
git clone -q https://github.com/ROCm/aiter.git /tmp/aiter && cd /tmp/aiter
git fetch -q origin pull/4494/head && git checkout -q --detach FETCH_HEAD
test "$(git rev-parse HEAD)" = fdce41d4bdfc20e4e37463626441726c9a81b6dc
AITER_USE_SYSTEM_TRITON=1 GPU_ARCHS=gfx950 python3 setup.py develop
python -m pytest -q -s op_tests/test_gemm_a16w16_graph.py
'

Test results

MI355X (gfx950), one idle GPU per run.

Arm	Result
Unpatched parent 3a79aec3b71393a36351e2d0512414d9a1ba7813	did not complete within 180 s
This PR	1 passed in 24.39s, including a fresh JIT build
Ruff check / format, git diff --check	pass

End to end, Kimi-K3 TP8 with DSpark on 8x MI355X:

Configuration	Result
unpatched, draft graphed	hangs during warmup
split-K bypassed (workaround)	serves, eager draft
this patch, split-K enabled	graph capture 8/8 on all ranks, serves

The captured endpoint scored 99/100 on the frozen GSM8K-100 gate with zero invalid responses, reached 120.9 output tok/s against 114.7 with the eager draft workaround, and logged zero memory faults.

Overlap and limits

The change is limited to who owns the semaphore workspace under capture. No kernel is modified, and the eager path is byte-for-byte the same code as before.

The unfixed control intentionally exercises a deadlock and can stall the selected GPU until killed; it is not a routine shared-runner test.

Tool assistance

Claude Opus 5 (1M context) assisted with the initial semaphore fix and drafting this description. OpenAI Codex assisted with the regression test and validation review.

---

Update -- 2026-08-13

A note for anyone arriving here from the revert.

This was merged as c6ce60c on 2026-08-12 at 09:44 UTC and reverted by #4709 (518cfbc) at 16:53 UTC the same day, so main currently carries neither the fix nor op_tests/test_gemm_a16w16_graph.py, and the Kimi-K3 graphed-draft hang is present again.

The concern was VRAM growth from _captured_semaphore_keepalive, which is a fair thing to ask about a list with no eviction. Looking into it more carefully, the growth does seem to be confined to startup in the hosts that reach this kernel; the Cost. section above has been updated with what I found, including vLLM's capture-window check. The original wording there called the repeated-capture case "a real limitation of this approach" without mentioning that no current host can reach that state, which I think reads more alarming than it should. That was my wording to fix rather than anyone's misreading.

If there is appetite for a re-land, a version that sidesteps the question altogether might be easier to review: a small ring of workspaces pre-allocated on the eager path, with zero_() recorded inside the graph -- a fixed 256 KiB per device, no growth, and no graph-pool memory retained. Happy to put that up instead if it would help. Everything above describes the reverted commit as merged, not that version.

[github-actions]

🏷️ CI Guide

Runs automatically on every PR:

• ✅ Pre-checks (submodule verification, code formatting)
• ✅ Aiter op tests (gfx942 + gfx950)
• ✅ Triton tests on MI35X (only when aiter/ops/triton/** or related paths are changed)

Extended tests (opt-in via labels):

Label	Tests
ci:triton-300x	Run an additional Triton test job on MI300X in PRs; main branch always runs both MI35X and MI300X
ci:sglang	SGLang integration tests: DeepSeek-R1-MXFP4 accuracy, Qwen 3.5 accuracy
ci:atom	ATOM benchmark: DeepSeek-R1-0528, GPT-OSS-120B
ci:atom_full	ATOM accuracy suite for PR and main models from ATOM models_accuracy.json
ci:vllm	vLLM benchmark: GPT-OSS-120B, DeepSeek-R1-0528, Kimi-K2.5
ci:all	All standard extended tests (excludes ci:atom_full)

Only add ci:atom_full for FlyDSL or Triton upgrades.
Add labels via the sidebar or gh pr edit 4494 --add-label <label>

[Copilot]

Pull request overview

This PR fixes a GPU deadlock that can occur when capturing and replaying CUDA graphs that include the a16w16 ASM split-K GEMM path, by avoiding reuse of a cached split-K semaphore counter during capture.

Changes:

• Allocate a fresh, zero-initialized semaphore workspace for each launch when CUDA graph capture is active.
• Retain capture-allocated workspaces for process lifetime to prevent reuse hazards with graph-private memory pools.
• Keep the existing per-(device, stream) lru_cache behavior for non-capture (eager) execution.

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.