From b57692587a3515b950cdc8f36485f7a4d927f953 Mon Sep 17 00:00:00 2001 From: Vedaanta Agarwalla Date: Sat, 15 Aug 2026 00:58:07 -0700 Subject: [PATCH 1/2] =?UTF-8?q?frost(sdpa):=20fix=20SM100=20zero-KV=20clus?= =?UTF-8?q?ter=20deadlock=20=E2=80=94=20pair=20the=20Q=E2=88=AAO=20alias?= =?UTF-8?q?=20gate?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Zero-length seq_len_kv sequences (and fully-masked BR tiles) produce EMPTY tiles whose kv-loop bounds collapse. On empty tiles the correction→TMA-STG chain that arrives the d192 Q∪O alias gate (mb_q_o_alias) has no dependency on the TMA-LDG warp, so a delayed LDG loses the race and the gate's parity bootstrap credit is consumed by a real arrive: mbarrier parity waits cannot recover once a producer leads by two phases, so LDG parks forever at the gate. Downstream, MMA starves on mb_q_full and softmax on mb_bmm1_done — whole clusters wedge (reproduced as a test-order-dependent hang in the random ragged sweeps once zero-KV generation is enabled; root-caused via cuda-gdb shared-memory dumps of the hung CTA: alias gate at phase 1 with LDG never past its first wait). Fix: make the alias gate a proper FULL/EMPTY pair. New mb_qo_slab_free return edge — TMA-LDG arrives it after consuming each alias phase and TMA-STG waits it before each alias arrive — bounds either side's lead to one phase by construction. d192 is the only QO_ALIAS kernel; this mirrors the throttle the d256 flavor already has (mb_tmastg_go). Validated on SM100 (cuDNN 9.26 + 9.24): the previously-hanging ragged pair passes repeatedly; fwd+bwd ragged suites, lean_attn, and the frost SM100 fwd suite pass; full mhas L0 completes (only pre-existing 9.26-nightly backend rejections of mixed seq-len forms remain, reproducible with frost disabled). Co-Authored-By: Claude Fable 5 --- .../cudnn/sdpa/fwd/kernels/_common_sm100.py | 16 +++++++++++ .../kernels/prefill_d192_d128_f16_sm100.py | 28 ++++++++++++++++++- 2 files changed, 43 insertions(+), 1 deletion(-) diff --git a/python/cudnn/sdpa/fwd/kernels/_common_sm100.py b/python/cudnn/sdpa/fwd/kernels/_common_sm100.py index 2f2f7242c..c04f2bb6e 100644 --- a/python/cudnn/sdpa/fwd/kernels/_common_sm100.py +++ b/python/cudnn/sdpa/fwd/kernels/_common_sm100.py @@ -45,6 +45,10 @@ class Bars(NamedTuple): mb_empty_mainloop: object mb_q_o_alias: object + # Return edge of the Q∪O alias gate (see the soundness note in + # make_classic_bars): TMA-LDG arrives after consuming each alias-gate + # phase; TMA-STG waits it before its next alias arrive. + mb_qo_slab_free: object class D256Bars(NamedTuple): @@ -140,7 +144,19 @@ def _alloc(n): mb_o_empty=MBarrier(_alloc(CFG.TILES_Q), stages=CFG.TILES_Q, init_count=CFG.ONE_WARP, producer=Producer.THREAD), mb_tmem_dealloc=MBarrier(_alloc(1), stages=1, init_count=CORR_LANES_TOTAL, producer=Producer.THREAD), mb_empty_mainloop=MBarrier(_alloc(1), stages=1, init_count=CORR_LANES_TOTAL, producer=Producer.LEADER, scope=Scope.LEADER), + # Q∪O alias gate FULL/EMPTY pair. mb_q_o_alias alone is UNSOUND: + # mbarrier parity waits deadlock once a producer runs >= 2 phases + # ahead, and on EMPTY tiles (zero-KV varlen sequences) the + # corr -> STG -> alias-arrive chain has NO dependency on TMA-LDG, so + # a delayed LDG warp loses the race and its bootstrap parity credit + # is consumed by a real arrive (observed: LDG parked forever at the + # tile-1 alias wait with the barrier already in phase 1, deadlocking + # the whole cluster). mb_qo_slab_free is the return edge: LDG + # arrives it right after consuming each alias phase and STG waits it + # before each alias arrive, bounding either side's lead to one phase + # by construction. mb_q_o_alias=MBarrier(_alloc(CFG.TILES_Q), stages=CFG.TILES_Q, init_count=CFG.ONE_WARP, producer=Producer.THREAD), + mb_qo_slab_free=MBarrier(_alloc(CFG.TILES_Q), stages=CFG.TILES_Q, init_count=CFG.ONE_WARP, producer=Producer.THREAD), ) diff --git a/python/cudnn/sdpa/fwd/kernels/prefill_d192_d128_f16_sm100.py b/python/cudnn/sdpa/fwd/kernels/prefill_d192_d128_f16_sm100.py index 81693eab8..c5e5cb659 100644 --- a/python/cudnn/sdpa/fwd/kernels/prefill_d192_d128_f16_sm100.py +++ b/python/cudnn/sdpa/fwd/kernels/prefill_d192_d128_f16_sm100.py @@ -428,6 +428,7 @@ def _kernel( bars.mb_o_empty[qs].init() if cutlass.const_expr(IS_QO_ALIAS): bars.mb_q_o_alias[qs].init() + bars.mb_qo_slab_free[qs].init() for c in cutlass.range_constexpr(CFG.N_BMM2_CHUNKS): bars.mb_bmm2_ready[qs * CFG.N_BMM2_CHUNKS + c].init() for ks in cutlass.range_constexpr(CFG.STAGES_KV): @@ -700,8 +701,15 @@ def _tmaldg_warp_group( # TMA-STG advances the Q/O alias gate for every tile, including # empty ones. Consume that transaction even though no Q reload # is needed, or the next nonempty tile observes stale parity. + # mb_qo_slab_free is the RETURN edge: without it, runs of + # empty tiles let STG (whose empty-tile O store depends only + # on correction, never on this warp) race >= 2 alias phases + # ahead of a delayed LDG, and mbarrier parity waits deadlock + # at lead 2 — the observed zero-KV cluster hang. _wait_mbarrier(mb_q_reload[0], q_empty_phase) + bars.mb_qo_slab_free[0].arrive() _wait_mbarrier(mb_q_reload[1], q_empty_phase) + bars.mb_qo_slab_free[1].arrive() q_empty_phase = q_empty_phase ^ 1 else: # Prologue interleave Q[0] -> K[first] -> Q[1] -> V[first] -> mainloop — @@ -709,6 +717,8 @@ def _tmaldg_warp_group( kv_row_base = kv_left * CFG.TILE_N _wait_mbarrier(mb_q_reload[0], q_empty_phase) + if cutlass.const_expr(IS_QO_ALIAS): + bars.mb_qo_slab_free[0].arrive() if cutlass.const_expr(CFG.CTA_MMA == 2): bars.mb_q_full[0].arrive(n_bytes=qTmaTransactionBytes, pred=is_leader & nvvm.elect_sync()) else: @@ -744,6 +754,8 @@ def _tmaldg_warp_group( ) _wait_mbarrier(mb_q_reload[1], q_empty_phase) + if cutlass.const_expr(IS_QO_ALIAS): + bars.mb_qo_slab_free[1].arrive() if cutlass.const_expr(CFG.CTA_MMA == 2): bars.mb_q_full[1].arrive(n_bytes=qTmaTransactionBytes, pred=is_leader & nvvm.elect_sync()) else: @@ -862,6 +874,9 @@ def _tmastg_warp_group( via scheduler warp's clusterlaunchcontrol.try_cancel.async. """ o_full_phase = cutlass.Int32(0) # consumer waits — first-arrive flips 0 → 1 + # Alias-gate return edge (see the arrive site below): starts 0 — the + # first wait consumes LDG's first real slab_free arrive. + slab_free_phase = cutlass.Int32(0) tma_o = GmemTileTma(tma_o_desc) @@ -906,11 +921,22 @@ def _tmastg_warp_group( bars.mb_o_empty[qs].arrive() # QO_ALIAS: O[qs] has drained to GMEM → the shared Q∪O slab is free - # for TMA-LDG to clobber with the next tile's Q[qs]. + # for TMA-LDG to clobber with the next tile's Q[qs]. The + # mb_qo_slab_free wait (return edge) throttles this arrive to at + # most ONE phase ahead of LDG's alias-gate consumption — without + # it, empty-tile runs (zero-KV varlen) let this warp lap a + # delayed LDG by 2+ phases and mbarrier parity waits deadlock. + # Consumer-side bootstrap: slab_free_phase starts 0, so the first + # wait consumes LDG's first REAL arrive (LDG cannot be preceded — + # its arrive follows its own alias wait, which this warp's tile-1 + # arrive has not yet advanced). if cutlass.const_expr(IS_QO_ALIAS): + _wait_mbarrier(bars.mb_qo_slab_free[qs], slab_free_phase) bars.mb_q_o_alias[qs].arrive() o_full_phase = o_full_phase ^ 1 + if cutlass.const_expr(IS_QO_ALIAS): + slab_free_phase = slab_free_phase ^ 1 _wait_ptr(sched.mb_scheduler.subview(sched_state.idx), sched_state.phase) nxt_q = (sched.tile_id_smem.subview(sched_state.idx * cutlass.Int32(8) + cutlass.Int32(0))).load() From 096a5bb80d6ab49ff559b9c9e82d0ac5ab640886 Mon Sep 17 00:00:00 2001 From: Vedaanta Agarwalla Date: Fri, 7 Aug 2026 15:38:19 -0700 Subject: [PATCH 2/2] tests: enable zero-length seq_len_kv generation in the random SDPA sweeps MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The ~10% zero-KV variant for per-batch seq_len_kv has been commented out since the test_mhas_v2 refactor, so no backend zero-KV behavior (fully-dead query rows: O := 0, LSE := -inf) was pinned by the random sweeps — only seq_len_q could go to zero. The commented variant had a latent crash: it drew randint(seq_len_q[i], s_kv) (a kv >= q lower bound), and the sweeps deliberately include s_q > s_kv configs, so any batch with seq_len_q[i] > s_kv raised "empty range for randrange()" in the generator (9/256 fwd L0 configs) — presumably why it was disabled. Re-enable the zeros with the bound decoupled from seq_len_q, matching the seq_len_q line above. Validated on SM100 with cuDNN 9.24: fwd L0 245 passed, ragged fwd+bwd L0 286 passed, dense bwd + bias fwd/bwd + fp8 fwd/bwd (dense+ragged) L0 631 passed — 0 failures, and skip counts identical to a baseline run without the zeros (per-batch lengths are runtime data, invisible to plan-time check_support gating). Co-Authored-By: Claude Fable 5 --- test/python/sdpa/random_config.py | 6 +----- 1 file changed, 1 insertion(+), 5 deletions(-) diff --git a/test/python/sdpa/random_config.py b/test/python/sdpa/random_config.py index a8843cbdd..681fb8060 100644 --- a/test/python/sdpa/random_config.py +++ b/test/python/sdpa/random_config.py @@ -355,11 +355,7 @@ def __call__(self, rng, rng_data_seed, rng_geom_seed=None): # ~10% chance of 0-length sequence for each batch randoms_.seq_len_q = [0 if rng.random() < 0.1 else rng.randint(1, randoms_.s_q) for _ in range(randoms_.batches)] # ~10% chance of 0-length sequence for each batch (independent of seq_len_q) - randoms_.seq_len_kv = [ - # 0 if rng.random() < 0.1 else rng.randint(randoms_.seq_len_q[i], randoms_.s_kv) for i in range(randoms_.batches) - rng.randint(1, randoms_.s_kv) - for i in range(randoms_.batches) - ] + randoms_.seq_len_kv = [0 if rng.random() < 0.1 else rng.randint(1, randoms_.s_kv) for i in range(randoms_.batches)] # Decide the left and right bounds for the sliding window mask (None = no bound) randoms_.left_bound = None