[FlyDSL] Keep split-K preshuffle buffers out of the CUDA graph pool

[PerryZhang01]

_get_preshuffle_split_buffers (added in #5007) is keyed on (device, stream).
A CUDA graph capture always runs on a fresh stream, so a capture is a
guaranteed cache miss and the split-K workspace/semaphore are allocated
inside the capture region. That is the defect.

A buffer allocated during capture and then kept alive by the lru_cache
takes over a block that a transient tensor released earlier in that same
capture, or in an earlier capture into the same pool, was using — and the
kernels writing to that address are already recorded in the graph. The
allocator knows the transient was freed; it does not know a graph still writes
there. So every replay of that graph writes its own intermediate values
straight into aiter's live workspace and semaphore.

The semaphore is what turns this into a wrong answer. torch.zeros ran inside
the first capture, so its zeroing kernel is baked into that graph, which
therefore self-heals on every replay. Every later-captured GEMM graph is an
lru_cache hit — same buffers, no zeroing. When a dirty slot reaches one of
those graphs, splitk_reduce_epilogue never sees
arrival == split_k - 1, the reduction for that tile never fires, and the
output tile is silently left holding whatever was in it.

Note the direction: this is not "a later graph steals the live block".
PyTorch's live accounting is correct and a later capture cannot take a block
the cache still holds (arm C below). The cached buffer is the one that
inherited someone else's address.

The fix keys on device alone and refuses to allocate while a graph is
capturing, with preallocate_preshuffle_split_buffers(device) exported so
integrators can warm the cache during their own warmup.

Why dropping the stream is required, not incidental

A capture stream does not exist before the capture starts, so a stream-keyed
cache can never be warmed ahead of a capture — no amount of warmup makes
the first capturing call a hit. Keying on device is what makes preallocation
possible at all.

Why a capture-time allocation is unsafe — minimal standalone proof

Four arms, pure PyTorch, one GPU, 1 MiB, seconds. No aiter, no vLLM, no model.
Each arm allocates a long-lived buffer, fills it with a canary from eager, and
replays.

arm	setup	address reuse	canary after replay
A allocate during capture (status quo)	in-graph transient T freed, then torch.empty for WS	T.ptr == WS.ptr True	262144/262144 clobbered, value 6.0 = T's result
B allocate before capture (this PR)	WS allocated outside the capture	not on the reuse chain False	5 replays, 0/262144, intact
C can a later graph steal a live block?	graph 1 allocates and holds WS; graph 2 captures after	graph 2 cannot get WS's address False	graph 2's replay leaves WS alone; graph 1's replay writes it
D cross-capture (closest to vLLM)	graph 1 frees T during capture; graph 2 allocates WS during its capture	T.ptr == WS.ptr True	replay graph 1 -> 6.0; replay graph 2 -> 15.0; both write it

Arm C is the one worth reading twice: the allocator is not at fault and the
"stolen live block" story does not hold. Arm A shows the second writer can be
the very same graph, with no other graph involved. Arm D shows that with a
shared graph_pool — which is what vLLM uses — any previously captured graph
qualifies.

This is exactly the situation a long-lived lru_cache buffer creates, and it
is why the rule "do not let capture-time allocations outlive the capture"
exists. Arm B is the fix.

repro script (tested on torch 2.10.0+rocm7.2.4, gfx950; plain CUDA works too)

import torch

N, CANARY = 1 << 18, 1234.5
dev = "cuda"


def arm_a():
    pool, g, s = torch.cuda.graph_pool_handle(), torch.cuda.CUDAGraph(), torch.cuda.Stream()
    src = torch.full((N,), 3.0, device=dev)
    with torch.cuda.graph(g, pool=pool, stream=s):
        T = src * 2.0
        t_ptr = T.data_ptr()
        sink = T.sum()
        del T                                  # block returns to the pool free list
        WS = torch.empty(N, device=dev)        # the "lru_cache workspace"
    print("A alias:", t_ptr == WS.data_ptr())
    WS.fill_(CANARY); torch.cuda.synchronize()
    g.replay(); torch.cuda.synchronize()
    print("A clobbered:", int((WS != CANARY).sum()), "/", N, "value:", WS[0].item())


def arm_b():
    pool, g, s = torch.cuda.graph_pool_handle(), torch.cuda.CUDAGraph(), torch.cuda.Stream()
    src = torch.full((N,), 3.0, device=dev)
    WS = torch.empty(N, device=dev)            # allocated BEFORE capture
    with torch.cuda.graph(g, pool=pool, stream=s):
        T = src * 2.0
        sink = T.sum()
        del T
    WS.fill_(CANARY); torch.cuda.synchronize()
    for _ in range(5):
        g.replay()
    torch.cuda.synchronize()
    print("B clobbered:", int((WS != CANARY).sum()), "/", N)


def arm_c():
    pool, s = torch.cuda.graph_pool_handle(), torch.cuda.Stream()
    src = torch.full((N,), 3.0, device=dev)
    g1 = torch.cuda.CUDAGraph()
    with torch.cuda.graph(g1, pool=pool, stream=s):
        T = src * 2.0
        s1 = T.sum()
        del T
        WS = torch.empty(N, device=dev)        # held alive from here on
    g2 = torch.cuda.CUDAGraph()
    with torch.cuda.graph(g2, pool=pool, stream=s):
        U = src * 7.0
        u_ptr = U.data_ptr()
        s2 = U.sum()
        del U
    print("C graph2 got WS's address:", u_ptr == WS.data_ptr())
    WS.fill_(CANARY); torch.cuda.synchronize()
    g2.replay(); torch.cuda.synchronize()
    print("C after graph2 replay, clobbered:", int((WS != CANARY).sum()), "/", N)


def arm_d():
    pool, s = torch.cuda.graph_pool_handle(), torch.cuda.Stream()
    src = torch.full((N,), 3.0, device=dev)
    g1 = torch.cuda.CUDAGraph()
    with torch.cuda.graph(g1, pool=pool, stream=s):
        T = src * 2.0
        t_ptr = T.data_ptr()
        s1 = T.sum()
        del T
    g2 = torch.cuda.CUDAGraph()
    with torch.cuda.graph(g2, pool=pool, stream=s):
        V = src * 5.0
        s2 = V.sum()
        del V
        WS = torch.empty(N, device=dev)
    print("D cross-capture alias:", t_ptr == WS.data_ptr())
    WS.fill_(CANARY); torch.cuda.synchronize()
    g1.replay(); torch.cuda.synchronize()
    print("D after replaying graph 1:", WS[0].item(), "(6.0 = graph 1's T)")
    WS.fill_(CANARY); torch.cuda.synchronize()
    g2.replay(); torch.cuda.synchronize()
    print("D after replaying graph 2:", WS[0].item(), "(15.0 = graph 2's V)")


for f in (arm_a, arm_b, arm_c, arm_d):
    f()

Expected output:

A alias: True
A clobbered: 262144 / 262144 value: 6.0
B clobbered: 0 / 262144
C graph2 got WS's address: False
C after graph2 replay, clobbered: 0 / 262144
D cross-capture alias: True
D after replaying graph 1: 6.0 (6.0 = graph 1's T)
D after replaying graph 2: 15.0 (15.0 = graph 2's V)

The same thing, in aiter, producing a wrong GEMM result

The arms above are pure PyTorch. This one drives the real
flydsl_preshuffle_gemm_a8 — still no model and no vLLM, one GPU, M=1 N=2048 K=2048 with split_k=4 (a real _ks4 row from a tuned table). All graphs are
captured onto one side stream into one shared pool, as
vllm.distributed.parallel_state.graph_capture() does for every batch size.

A decoy graph allocates a transient and frees it mid-capture. Then GEMM graph
g1 is captured (lru MISS — allocates the buffers), then GEMM graph g2
(lru HIT — same buffers, no zeroing kernel of its own).

	workspace aliases decoy's freed block	sem slots dirtied by one decoy replay	wrong outputs	max abs err
status quo, g1 captured first	True	129/256	0/2048	0.0078 (bf16 ulp)
status quo, g2 captured later	True	129/256	64/2048	778.5
fixed, g1	False	0/256	0/2048	0.0078
fixed, g2	False	0/256	0/2048	0.0078

Deterministic over 3 runs. The arithmetic closes exactly: this shape uses 32
semaphore slots (N/tile_n = 2048/64), of which exactly one is dirty, and
one tile is tile_n = 64 columns — so 64 wrong outputs.

The collision is not engineered. The dirty slot holds 1270874112, which is
0x4BC00000, which is float32 25165824.0, which is 6.0 * 4194304 — the
decoy graph's T.sum() scalar. A 4-byte reduction result landed on
aiter's semaphore slot 0.

Note what the workspace alone does not do: it is clobbered in full
(4194304/4194304 elements) and g1 still returns the right answer, because a
launch writes every plane it reduces. The wrong answer comes from the
semaphore. Depending on the garbage value the epilogue either never fires (what
is measured above — a stale output tile) or fires early, in which case the
reduction reads workspace planes the launch has not written yet; the latter
would produce the out-of-range partials seen on K3 below, but this repro
demonstrates the former.

repro script

Run it against an aiter build that still has the stream-keyed cache; the
"fixed" arm monkeypatches this PR's behaviour (device-keyed, preallocated
before any capture) so both arms run in one process.

"""aiter-level repro of ROCm/aiter#5393: capture-time allocation of the flydsl
split-K buffers makes a later-captured GEMM graph return a WRONG result.

Chain:
  1. `_get_preshuffle_split_buffers` is @lru_cache'd on (device, stream). A CUDA
     graph capture runs on a fresh stream, so the first capturing GEMM MISSES and
     allocates its 16 MiB workspace + 1 KiB semaphore inside the capture region.
  2. Those buffers come from the graph's private mempool and land on addresses a
     transient in an EARLIER graph's capture had freed. Replaying that earlier
     graph therefore writes into aiter's live buffers.
  3. `semaphore = torch.zeros(...)` ran inside the first capture, so its zeroing
     kernel is baked into THAT graph -- which self-heals on every replay.
     Every LATER captured GEMM graph is an lru HIT: same buffers, no zeroing.
  4. A later graph replaying with a dirty semaphore never sees
     `arrival == split_k - 1`, so `splitk_reduce_epilogue` never reduces and the
     output tensor is silently left with whatever it held.

Faithful to vLLM: all graphs are captured onto ONE side stream into ONE shared
pool, which is what vllm.distributed.parallel_state.graph_capture() does for
every batch size in a capture context.

No model, no vLLM, one GPU. First run JIT-compiles the kernel (~1 min).
"""
import torch
from aiter.ops.shuffle import shuffle_weight
from aiter.ops.flydsl import gemm_kernels as GK

DEV = torch.device("cuda", 0)
M, N, K = 1, 2048, 2048                       # a real _ks4 shape from a K3 tuned table
TM, TN, TK, ACP, WPE, XCD, LDS = 16, 64, 256, 1, 4, 0, 2
SPLIT_K = 4
WS_ELEMS = GK.PRESHUFFLE_SPLIT_K_WORKSPACE_ELEMS    # 4*256*32*128 fp32 = 16 MiB
SEM_ELEMS = GK.PRESHUFFLE_SPLIT_K_MAX_TILES         # 256 int32 = 1 KiB
POISON = -777.0


def build():
    torch.manual_seed(0)
    a = (torch.randn(M, K, device=DEV) / 8).to(torch.float8_e4m3fn)
    b = (torch.randn(N, K, device=DEV) / 8).to(torch.float8_e4m3fn)
    xs = torch.rand(M, 1, device=DEV, dtype=torch.float32) + 0.5
    wsc = torch.rand(N, 1, device=DEV, dtype=torch.float32) + 0.5
    ref = ((a.to(torch.float32) @ b.to(torch.float32).T) * xs * wsc.T).to(torch.bfloat16)
    return a, shuffle_weight(b, layout=(16, 16)), xs, wsc, ref


def arm(fixed):
    print(f"\n=== {'FIXED: device-keyed + preallocated (PR #5393)' if fixed else 'STATUS QUO: stream-keyed lru_cache'} ===")
    GK._get_preshuffle_split_buffers.cache_clear()
    a, bsh, xs, wsc, ref = build()

    def run(out):
        GK.flydsl_preshuffle_gemm_a8(a, bsh, xs, wsc, out, TM, TN, TK, ACP, WPE,
                                     XCD, lds_stage=LDS, enable_scheduler=True,
                                     split_k=SPLIT_K)

    # Warm up on the compute stream: compiles the kernel and fills the cache for
    # the COMPUTE stream. vLLM does this too -- and it is why the capture below
    # still misses.
    run(torch.empty(M, N, dtype=torch.bfloat16, device=DEV))
    torch.cuda.synchronize()

    if fixed:
        import functools
        _ws = torch.empty(WS_ELEMS, dtype=torch.float32, device=DEV)
        _sem = torch.zeros(SEM_ELEMS, dtype=torch.int32, device=DEV)

        @functools.lru_cache(maxsize=128)
        def devkey(device):
            return _ws, _sem

        GK._get_preshuffle_split_buffers = lambda device, stream: devkey(device)
        devkey(DEV)                                   # preallocate BEFORE capture

    pool, side = torch.cuda.graph_pool_handle(), torch.cuda.Stream(device=DEV)
    src = torch.full((WS_ELEMS,), 3.0, device=DEV)

    # A decoy graph standing in for any other graph vLLM captures into this pool:
    # it allocates a transient and frees it while still capturing.
    g_decoy = torch.cuda.CUDAGraph()
    with torch.cuda.graph(g_decoy, pool=pool, stream=side):
        T = src * 2.0
        t_ptr = T.data_ptr()
        total = T.sum()                                # a 4-byte scalar, also pooled
        del T, total

    out1 = torch.empty(M, N, dtype=torch.bfloat16, device=DEV)
    g1 = torch.cuda.CUDAGraph()
    with torch.cuda.graph(g1, pool=pool, stream=side):    # lru MISS -> allocates
        run(out1)

    out2 = torch.empty(M, N, dtype=torch.bfloat16, device=DEV)
    g2 = torch.cuda.CUDAGraph()
    with torch.cuda.graph(g2, pool=pool, stream=side):    # lru HIT -> no zeroing
        run(out2)

    ws, sem = GK._get_preshuffle_split_buffers(DEV, side)
    print(f"workspace aliases the decoy's freed transient: {ws.data_ptr() == t_ptr}")

    def trial(name, g, out):
        sem.fill_(0)
        torch.cuda.synchronize()
        g_decoy.replay()
        torch.cuda.synchronize()
        dirty = int((sem != 0).sum())
        out.fill_(POISON)
        torch.cuda.synchronize()
        g.replay()
        torch.cuda.synchronize()
        d = (out.float() - ref.float()).abs()
        bad = int((d > 0.05).sum())
        print(f"  {name:34s} sem slots dirtied by decoy: {dirty:3d}/{SEM_ELEMS}"
              f" | wrong outputs: {bad:5d}/{M*N} | max_err={d.max().item():8.4g}"
              f"  {'<-- WRONG' if bad else 'ok'}")
        return bad

    b1 = trial("g1 (captured first, self-zeroes)", g1, out1)
    b2 = trial("g2 (captured later, lru HIT)", g2, out2)
    return b1, b2


if __name__ == "__main__":
    sq = arm(fixed=False)
    fx = arm(fixed=True)
    print(f"\nstatus quo: g1 wrong={sq[0]}, g2 wrong={sq[1]}")
    print(f"fixed     : g1 wrong={fx[0]}, g2 wrong={fx[1]}")

Expected output:

=== STATUS QUO: stream-keyed lru_cache ===
workspace aliases the decoy's freed transient: True
  g1 (captured first, self-zeroes)   sem slots dirtied by decoy: 129/256 | wrong outputs:     0/2048 | max_err=0.007812  ok
  g2 (captured later, lru HIT)       sem slots dirtied by decoy: 129/256 | wrong outputs:    64/2048 | max_err=   778.5  <-- WRONG

=== FIXED: device-keyed + preallocated (PR #5393) ===
workspace aliases the decoy's freed transient: False
  g1 (captured first, self-zeroes)   sem slots dirtied by decoy:   0/256 | wrong outputs:     0/2048 | max_err=0.007812  ok
  g2 (captured later, lru HIT)       sem slots dirtied by decoy:   0/256 | wrong outputs:     0/2048 | max_err=0.007812  ok

Measurement — gfx950, standalone, no vLLM

Three streams plus one real torch.cuda.CUDAGraph() capture:

	buffer sets allocated	cache	inside capture
before	3 distinct	hits=0, misses=3	is_current_stream_capturing() True at the allocation
after	1	hits=3, misses=1	the capturing call is a hit; a cold cache during capture raises

How this was found — Kimi-K3 DSpark, 8xMI355X

NaN draft logits in speculative decoding. The draft's layer-0 gate_up_proj
(N=3584, K=7168) takes a _ks4/_ks2 split-K kernel at small M
(M = num_seqs x 7, so the tail of a dataset run drops into the split-K
buckets while mid-run concurrency 64 does not — which is why it only ever
failed near the end). The reduction returned ~1e30..inf garbage, SiLU
overflowed to inf, and down_proj's per-token quantization turned the whole
row NaN.

A byte histogram of the corrupt GEMM output shows a sign-symmetric
6e29..inf spread, with NaN and inf rows present in the same step
(7f:13535 vs ff:13448 in the high byte). That is foreign data being read
as float32 partial sums — not a missing reduction term, which would be
bounded and one-sided.

Single-variable runs on the same configuration, changing one thing each time:

run	variable	peak abs gate_up	NaN rows	gsm8k
baseline	—	3.39e+38	1917 (first at step 10)	0.9477
A	draft weights -> bf16 (avoids the fp8 kernel)	11.4	0	0.9583
B	split-K disabled only	10.9	0	0.9492
C	only stream dropped from the cache key	11.1	0	0.9515

Run C full acceptance, 5-shot gsm8k, 1319 questions, TP8 + DCP8, 7 speculative
tokens: flexible-extract 0.9515 +-0.0059, strict-match 0.9507,
1319/1319 with zero faults, nan_rows=0 across 3580 steps / 271768 rows.
The score is mid-range for this configuration (repeats span 0.9484..0.9598) —
the fix removes the fault without moving accuracy, because drafts are verified
by the target, so a corrupt draft costs acceptance rate rather than answers.

Tradeoffs — both deliberate, both worth stating

1. Concurrent split-K launches on different streams now share one buffer
   set. The original code already assumed launches are ordered and that the
   reduction hands the semaphore back zeroed; this extends that same
   assumption across streams instead of within one. (That zeroing protocol is
   not what failed here — it only guards against the split-K path's own stale
   state, and the 16 MiB workspace, not the 1 KiB semaphore, is what a graph
   replay overwrites.) Given the alternative is a
   cache that is structurally un-warmable before a capture, this is the right
   side of the trade — but it is a real widening and should be reviewed as one.

2. A first call landing inside a capture now raises instead of silently
   corrupting. For vLLM this is safe (the first call happens during warmup,
   is_current_stream_capturing() is False). For an integrator whose very
   first split-K call is inside a capture, behaviour changes from "produces
   garbage" to "fails at startup with a message naming the fix". I think loud
   is correct here, but it is a behaviour change, not a pure bugfix.

Tests

op_tests/test_flydsl_splitk_buffers.py (new) pins four properties: buffers
are shared across streams; allocation during capture is refused with an
actionable message; preallocation lets a capture reuse them; and a real
CUDAGraph capture allocates nothing (data_ptr() unchanged). All four fail
against the pre-fix module and pass after.

black --check and ruff==0.16.0 check are clean on both files.

cc @XiaobingSuper (author of #5007)

🤖 Generated with Claude Code

[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:gfx1250-ffm-triton	Run the five-shard gfx1250 FFM Triton test suite
ci:triton-300x	Run an additional Triton test job on MI300X in PRs; main branch always runs both MI35X and MI300X
multigpu	Aiter multi-GPU tests on the 8-GPU runner
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 5393 --add-label <label>

PR title tags & labels:
Component tags ([Triton/Gluon], [HIP], [CK], [ASM], ...) are added to the PR title and as PR labels automatically from the changed files and re-synced on every push — change-type tags like [fix]/[Perf], op tags like [MLA], and human labels (ci:*) are left untouched. Add the no-auto-title label to opt this PR out.

[valarLip]

not a correct direction..