sdpa: fix SM100 frost zero-KV cluster deadlock + enable zero-length seq_len_kv sweeps by vedaanta · Pull Request #575 · NVIDIA/cudnn-frontend

vedaanta · 2026-08-13T07:16:59Z

Before submitting

I agree to license this contribution under the terms of LICENSE.txt.
I ran pre-commit run and committed any formatting changes.
I added GitHub labels: one cat-*, one or more mod-*, and one orig-* (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-length seq_len_kv sequences 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 on mb_q_full, softmax on mb_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_free return 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 only QO_ALIAS kernel; 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-length seq_len_kv variant in test/python/sdpa/random_config.py (commented out since the test_mhas_v2 refactor; its original form also had a latent randint empty-range crash on s_q > s_kv configs, 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.py L0 completes with only pre-existing 9.26-nightly backend rejections of mixed seq-len forms (reproducible with frost disabled). The combined content also passed frost_tests:sdpa [Blackwell] (2830/2830) on GitLab pipeline 62846289 via PR #606's branch.

Summary by CodeRabbit

Bug Fixes
- Improved reliability when processing attention workloads with empty or variable-length sequences.
- Prevented synchronization issues that could occur during delayed data transfers or empty-tile processing.
- Improved coordination when reusing temporary memory during attention computation.
Tests
- Expanded randomized testing to include zero-length key/value sequences, improving coverage for edge cases.

coderabbitai · 2026-08-13T07:17:22Z

No actionable comments were generated in the recent review. 🎉

ℹ️ Recent review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Enterprise

Run ID: f5d89903-908b-4ec3-9040-34a8a9b73699

📥 Commits

Reviewing files that changed from the base of the PR and between a1a6d7f and 096a5bb.

📒 Files selected for processing (1)

python/cudnn/sdpa/fwd/kernels/prefill_d192_d128_f16_sm100.py

🚧 Files skipped from review as they are similar to previous changes (1)

python/cudnn/sdpa/fwd/kernels/prefill_d192_d128_f16_sm100.py

📝 Walkthrough

Walkthrough

The 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.

Changes

SDPA Q/O alias synchronization

Layer / File(s)	Summary
Q/O alias barrier contract `python/cudnn/sdpa/fwd/kernels/_common_sm100.py`	`Bars` declares and initializes the staged `mb_qo_slab_free` barrier.
Q/O alias kernel synchronization `python/cudnn/sdpa/fwd/kernels/prefill_d192_d128_f16_sm100.py`, `test/python/sdpa/random_config.py`	The prefill kernel signals slab availability during Q reload, waits for the return during O-store, tracks the barrier phase, and includes zero-length KV sequences in random configurations.

Estimated code review effort: 3 (Moderate) | ~20 minutes

Merge Risk: ⚪ Minimal · up to 096a5

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

NVIDIA/cudnn-frontend#606: Shares the Q/O slab-free barrier changes and addresses zero-KV deadlock behavior.

Suggested labels: mod-frost

Suggested reviewers: aneureka, jhjpark

🚥 Pre-merge checks | ✅ 4 | ❌ 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 42.31% which is insufficient. The required threshold is 80.00%.	Write docstrings for the functions missing them to satisfy the coverage threshold.

✅ Passed checks (4 passed)

Check name	Status	Explanation
Title check	✅ Passed	The title clearly identifies both primary changes: the SM100 zero-KV deadlock fix and zero-length sequence sweep coverage.
Description check	✅ Passed	The description covers all template sections with detailed scope, rationale, compatibility impact, related issues, and testing results.
Linked Issues check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check	✅ Passed	Check skipped because no linked issues were found for this pull request.

✨ Finishing Touches

🧪 Generate unit tests (beta)

Create PR with unit tests

_{Comment @coderabbitai help to get the list of available commands.}

vedaanta · 2026-08-13T07:37:34Z

@cudnn-ci-bot run frost

cudnn-ci-bot · 2026-08-13T07:37:59Z

🚀 Running mirror pipeline

Branch: cudnn-gh/pr-575-fd0f300
Pipeline: 62481977
Targets: frost

coderabbitai · 2026-08-15T08:37:34Z

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>

vedaanta · 2026-08-15T19:07:47Z

@cudnn-ci-bot run frost

cudnn-ci-bot · 2026-08-15T19:08:11Z

🚀 Running mirror pipeline

Branch: cudnn-gh/pr-575-096a5bb
Pipeline: 62887095
Targets: frost

Paged KV caches (cuDNN's own graph contract: graph.sdpa(k=<page pool>, v=<page pool>, use_padding_mask=True, seq_len_q/kv, paged_attention_k_table, paged_attention_v_table, paged_attention_max_seq_len_kv)) are now served by sdpa_fwd_prefill_sm100 / SdpaFwdDslSm100 on the d128 and d256 f16/bf16 kernels' PAGED_KV specialization. FlashInfer's cudnn decode/prefill graphs reach the FROST kernel with no integration work. Kernel (prefill_d128_f16 / prefill_d256_f16, SM100): only the TMA-LDG warp changes — a K/V tile is issued as 128/page_size row boxes (one box inside a taller page), each box's page id read from the block table on device; boxes past a sequence's ceil(len/page_size) pages take page coordinate -1 (TMA-OOB zeros, mbarrier bytes still credited). HND vs NHD pools are only the pool's strides (already in the compile key): the host trace reads the layout off them and swaps the descriptor dim order. Separate K and V tables, in their declared (batch, page) strides. The KV maximum is block_table.shape[1] * page_size, a dynamic extent (num_pages / max_pages never enter the key). Host: analyzer paged facts (page_size, table refs; s_kv = declared max or what the tables address), paged_kv=True on the SM100 f16 row with the d128/d256 / padded / no-sink / page-geometry conditions in mismatch(), the split heuristic and facts x knobs gate exempt paged graphs from the padded no-split rule (paged is padded by construction and the split IS the decode lever), SdpaFwdDslSm100 binds pool-shaped K/V with their strides compiled in and takes block_table / block_table_v at execute. THD (ragged) queries over a paged cache — chunked prefill — ride the same path: Q/O packed, pools dense, the THD setup kernel skips its packed-KV clamp. Pre-existing fix: the d128 cga1 Q/O-alias TMA-LDG skipped the alias-gate wait on EMPTY tiles, so after a persistent CTA drained an empty split its parity was one behind and the next live tile's Q load raced the previous O drain (long batch off by ~1e-2 when B*H*splits > SM count and another batch has empty splits). Ported the d192 mb_qo_slab_free return edge (NVIDIA#575). Tests: test_sdpa_fwd_paged_sm100.py (graph API: both layouts x page sizes, mixed lengths incl. 0/1, bf16, 32k tokens with heuristic splits, B=256 d=64, d256, paged prefill S_q>1, THD ragged queries, batch-innermost tables, declines; kernel-level: page 32/64/256, forced splits with empty ranges at cga1, d256, GQA group 3, adapter CUDA-graph replay under set_sync_debug_mode("error")); analyzer fact/decline/split tests; test_split_kv_cga1_empty_splits_multiwave; test_mhas_v2's paged helper now builds row-major block tables (what every framework hands cuDNN). SUPPORT_MATRIX_TRACKER: paged row. Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>

Paged KV caches (cuDNN's own graph contract: graph.sdpa(k=<page pool>, v=<page pool>, use_padding_mask=True, seq_len_q/kv, paged_attention_k_table, paged_attention_v_table, paged_attention_max_seq_len_kv)) are now served by sdpa_fwd_prefill_sm100 / SdpaFwdDslSm100 on the d128 and d256 f16/bf16 kernels' PAGED_KV specialization. FlashInfer's cudnn decode/prefill graphs reach the FROST kernel with no integration work. Kernel (prefill_d128_f16 / prefill_d256_f16, SM100): only the TMA-LDG warp changes — a K/V tile is issued as 128/page_size row boxes (one box inside a taller page), each box's page id read from the block table on device; boxes past a sequence's ceil(len/page_size) pages take page coordinate -1 (TMA-OOB zeros, mbarrier bytes still credited). HND vs NHD pools are only the pool's strides (already in the compile key): the host trace reads the layout off them and swaps the descriptor dim order. Separate K and V tables, in their declared (batch, page) strides. The KV maximum is block_table.shape[1] * page_size, a dynamic extent (num_pages / max_pages never enter the key). Host: analyzer paged facts (page_size, table refs; s_kv = declared max or what the tables address), paged_kv=True on the SM100 f16 row with the d128/d256 / padded / no-sink / page-geometry conditions in mismatch(), the split heuristic and facts x knobs gate exempt paged graphs from the padded no-split rule (paged is padded by construction and the split IS the decode lever), SdpaFwdDslSm100 binds pool-shaped K/V with their strides compiled in and takes block_table / block_table_v at execute. THD (ragged) queries over a paged cache — chunked prefill — ride the same path: Q/O packed, pools dense, the THD setup kernel skips its packed-KV clamp. Pre-existing fix: the d128 cga1 Q/O-alias TMA-LDG skipped the alias-gate wait on EMPTY tiles, so after a persistent CTA drained an empty split its parity was one behind and the next live tile's Q load raced the previous O drain (long batch off by ~1e-2 when B*H*splits > SM count and another batch has empty splits). Ported the d192 mb_qo_slab_free return edge (#575). Tests: test_sdpa_fwd_paged_sm100.py (graph API: both layouts x page sizes, mixed lengths incl. 0/1, bf16, 32k tokens with heuristic splits, B=256 d=64, d256, paged prefill S_q>1, THD ragged queries, batch-innermost tables, declines; kernel-level: page 32/64/256, forced splits with empty ranges at cga1, d256, GQA group 3, adapter CUDA-graph replay under set_sync_debug_mode("error")); analyzer fact/decline/split tests; test_split_kv_cga1_empty_splits_multiwave; test_mhas_v2's paged helper now builds row-major block tables (what every framework hands cuDNN). SUPPORT_MATRIX_TRACKER: paged row. Co-authored-by: Claude Fable 5.1 <noreply@anthropic.com>

vedaanta added cat-infra Build, packaging, tooling, dependency, release, or repository maintenance work. mod-frontend cuDNN frontend APIs, operation graph construction, plans, and user-facing wrappers. orig-nv-eng Reported or requested by NVIDIA engineering. labels Aug 13, 2026

vedaanta requested a review from Aneureka August 13, 2026 07:17

jhjpark approved these changes Aug 13, 2026

View reviewed changes

vedaanta force-pushed the mhas-zero-kv-v2 branch from fd0f300 to a1a6d7f Compare August 15, 2026 08:37

vedaanta changed the title ~~tests: enable zero-length seq_len_kv generation in the random SDPA sweeps~~ sdpa: fix SM100 frost zero-KV cluster deadlock + enable zero-length seq_len_kv sweeps Aug 15, 2026

vedaanta added the cat-bugfix label Aug 15, 2026

coderabbitai Bot mentioned this pull request Aug 15, 2026

frost(sdpa): THD execute with zero host reads on SM100 — device-built metadata, plan-time envelope grid, CUDA-graph capturable (issue #552) #606

Merged

vedaanta and others added 2 commits August 15, 2026 11:46

vedaanta mentioned this pull request Aug 15, 2026

frost: align the CLC persistent scheduler with the canonical CUTLASS pattern (hardening split out of #575) #610

Open

vedaanta force-pushed the mhas-zero-kv-v2 branch from a1a6d7f to 096a5bb Compare August 15, 2026 19:01

vedaanta merged commit a5ce256 into NVIDIA:develop Aug 15, 2026
1 check passed

vedaanta mentioned this pull request Aug 17, 2026

sdpa fp8: dense head-dim envelope + one engine per arch x dtype family #587

Merged

vedaanta deleted the mhas-zero-kv-v2 branch August 18, 2026 17:14

coderabbitai Bot mentioned this pull request Aug 19, 2026

sdpa fwd: KV split for the SM100/SM120 prefill kernels #658

Merged

3 tasks done

Anerudhan mentioned this pull request Sep 2, 2026

1.28.0 rc #851

Merged

vedaanta mentioned this pull request Sep 9, 2026

sdpa: serve paged KV caches through the FROST SM100 engine (#920) #964

Merged

5 of 7 tasks

Navigation Menu

sdpa: fix SM100 frost zero-KV cluster deadlock + enable zero-length seq_len_kv sweeps - #575

sdpa: fix SM100 frost zero-KV cluster deadlock + enable zero-length seq_len_kv sweeps#575
vedaanta merged 2 commits into
NVIDIA:developfrom
vedaanta:mhas-zero-kv-v2

vedaanta commented Aug 13, 2026 •

edited by coderabbitai Bot

Loading

Uh oh!

coderabbitai Bot commented Aug 13, 2026 •

edited

Loading

Walkthrough

Changes

Sequence Diagram(s)

❌ Failed checks (1 warning)

Uh oh!

vedaanta commented Aug 13, 2026

Uh oh!

cudnn-ci-bot commented Aug 13, 2026

Uh oh!

coderabbitai Bot commented Aug 15, 2026

Uh oh!

vedaanta commented Aug 15, 2026

Uh oh!

cudnn-ci-bot commented Aug 15, 2026

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

Conversation

vedaanta commented Aug 13, 2026 • edited by coderabbitai Bot Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Before submitting

Affected area

Summary

Why

Related issues

API and compatibility impact

Testing

Summary by CodeRabbit

Uh oh!

coderabbitai Bot commented Aug 13, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Walkthrough

Changes

Sequence Diagram(s)

❌ Failed checks (1 warning)

Uh oh!

vedaanta commented Aug 13, 2026

Uh oh!

cudnn-ci-bot commented Aug 13, 2026

Uh oh!

coderabbitai Bot commented Aug 15, 2026

Uh oh!

vedaanta commented Aug 15, 2026

Uh oh!

cudnn-ci-bot commented Aug 15, 2026

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

vedaanta commented Aug 13, 2026 •

edited by coderabbitai Bot

Loading

coderabbitai Bot commented Aug 13, 2026 •

edited

Loading