Skip to content

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

Merged
vedaanta merged 2 commits into
NVIDIA:developfrom
vedaanta:mhas-zero-kv-v2
Aug 15, 2026
Merged

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

Conversation

@vedaanta

@vedaanta vedaanta commented Aug 13, 2026

Copy link
Copy Markdown
Collaborator

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.

@vedaanta 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
vedaanta requested a review from Aneureka August 13, 2026 07:17
@coderabbitai

coderabbitai Bot commented Aug 13, 2026

Copy link
Copy Markdown
Contributor

Review Change Stack

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
Loading

Possibly related PRs

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

Copy link
Copy Markdown
Collaborator Author

@cudnn-ci-bot run frost

@cudnn-ci-bot

Copy link
Copy Markdown

🚀 Running mirror pipeline

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

@coderabbitai

coderabbitai Bot commented Aug 15, 2026

Copy link
Copy Markdown
Contributor

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.

@vedaanta 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 and others added 2 commits August 15, 2026 11:46
…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

Copy link
Copy Markdown
Collaborator Author

@cudnn-ci-bot run frost

@cudnn-ci-bot

Copy link
Copy Markdown

🚀 Running mirror pipeline

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

@vedaanta
vedaanta merged commit a5ce256 into NVIDIA:develop Aug 15, 2026
1 check passed
@vedaanta
vedaanta deleted the mhas-zero-kv-v2 branch August 18, 2026 17:14
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

cat-bugfix 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.

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants