sdpa: fix SM100 frost zero-KV cluster deadlock + enable zero-length seq_len_kv sweeps - #575
Conversation
|
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info⚙️ Run configurationConfiguration used: Path: .coderabbit.yaml Review profile: CHILL Plan: Enterprise Run ID: 📒 Files selected for processing (1)
🚧 Files skipped from review as they are similar to previous changes (1)
📝 WalkthroughWalkthroughThe SDPA SM100 prefill kernel adds Q/O alias slab-free synchronization. The random configuration generator now includes zero-length KV sequences with a 10% probability. ChangesSDPA Q/O alias synchronization
Estimated code review effort: 3 (Moderate) | ~20 minutes Merge Risk: ⚪ Minimal · up to The PR adds a localized SM100 zero-KV deadlock fix and regression coverage; no actionable merge-blocking risk remains at the current head beyond normal checks and review. Sequence Diagram(s)sequenceDiagram
participant QReload
participant QOSlabFree
participant OStore
QReload->>QOSlabFree: signal slab availability
OStore->>QOSlabFree: wait for slab-free return
OStore->>QOSlabFree: advance return phase
Possibly related PRs
Suggested labels: Suggested reviewers: 🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Comment |
|
@cudnn-ci-bot run frost |
|
🚀 Running mirror pipeline Branch: cudnn-gh/pr-575-fd0f300 |
fd0f300 to
a1a6d7f
Compare
|
Note GitHub couldn't provide a complete incremental comparison for this pull request, so CodeRabbit is performing a full review instead. This review may take a little longer. |
…gate 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 <noreply@anthropic.com>
…eeps 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 <noreply@anthropic.com>
a1a6d7f to
096a5bb
Compare
|
@cudnn-ci-bot run frost |
|
🚀 Running mirror pipeline Branch: cudnn-gh/pr-575-096a5bb |
Before submitting
pre-commit runand committed any formatting changes.cat-*, one or moremod-*, and oneorig-*(see label list).Affected area
Python FE OSS kernels (frost SM100 d192 SDPA) + test infrastructure
Summary
Two commits (fix first so the branch stays bisect-green):
1.
frost(sdpa): fix an SM100 cluster deadlock on zero-KV varlen tiles. Zero-lengthseq_len_kvsequences produce EMPTY tiles (kv-loop bounds collapse). On empty tiles the correction→TMA-STG chain that arrives the d192 Q∪O alias gate has no dependency on the TMA-LDG warp, so a delayed LDG loses the race and its 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 and the whole cluster wedges (MMA starves onmb_q_full, softmax onmb_bmm1_done). Manifests as a warm-GPU, test-order-dependent hang; root-caused via cuda-gdb shared-memory dumps of a hung CTA (alias gate at phase 1 with LDG never past its first wait; a per-warp round trace splitting exactly at the two starved warps).Fix: make the alias gate a proper FULL/EMPTY pair — new
mb_qo_slab_freereturn edge (LDG arrives after consuming each alias phase; STG waits it before each alias arrive), bounding either side's lead to one phase by construction. d192 is the onlyQO_ALIASkernel; this mirrors the throttle the d256 flavor already has (mb_tmastg_go). Diff is minimal:_common_sm100.py+ the d192 kernel, +43/−1.2.
tests: re-enable the ~10% zero-lengthseq_len_kvvariant intest/python/sdpa/random_config.py(commented out since the test_mhas_v2 refactor; its original form also had a latentrandintempty-range crash ons_q > s_kvconfigs, fixed by decoupling the lower bound). This is the regression coverage for the fix.CLC-scheduler hardening found during the same investigation was split out to keep this PR minimal — see #610 (implementation preserved there for a head start).
Why
Zero-length KV sequences are valid inputs (cuDNN semantics: O := 0, LSE := -inf for dead rows) that frameworks do produce. The sweeps never exercised them, and the frost SM100 d192 kernel deadlocked on them under real scheduling pressure.
Related issues
Related to #512. Hardening follow-up: #610. The same deadlock signature was hit by the PR #606 Blackwell pipeline (envelope-grid empty tiles); that branch carries the equivalent fix until one of the two merges.
API and compatibility impact
None user-visible. One extra smem mbarrier array + one wait/arrive pair per tile in the aliased (d192) kernel.
Testing
On SM100 (B200), cuDNN 9.26 nightly + 9.24 wheel, frost engines enabled: the previously-hanging ragged pair passes repeatedly (was ~90% hang warm); fwd+bwd ragged suites, lean_attn, and the frost SM100 fwd L0+L1 suite pass; full
test_mhas_v2.pyL0 completes with only pre-existing 9.26-nightly backend rejections of mixed seq-len forms (reproducible with frost disabled). The combined content also passedfrost_tests:sdpa [Blackwell](2830/2830) on GitLab pipeline 62846289 via PR #606's branch.Summary by CodeRabbit
Bug Fixes
Tests