frost(sdpa_bwd_sm80): serve THD ports at their own head stride - #1108
Draft
egilliam-nv wants to merge 1 commit into
Draft
egilliam-nv wants to merge 1 commit into
egilliam-nv wants to merge 1 commit into
Conversation
egilliam-nv
added
orig-nv-eng
Contributor
|
Important Draft PR not reviewedDraft PRs are not automatically reviewed by default.
To automatically review draft PRs, update your CodeRabbit configuration: reviews:
auto_review:
drafts: trueComment |
Collaborator
Author
The SM80 backward's packed THD path addressed a row as token * token_stride + head * D: the token stride was plan-time per port (the compiled fake carried it) but the head stride was fixed at D, and the engine rule and the adapter declined anything else. PR NVIDIA#960 turned on head-axis stride gaps in the ragged sweeps (each ragged port a head-interleaved record, head stride D + k*16 bytes), so with FROST engines enabled test_sdpa_random_bwd_ragged_L0 on A100 fell from 383 draws served by this row to 48 -- the rest ran nowhere, the cuDNN backend declining every ragged backward on SM80. The head stride now rides the fake like the token stride: compile() takes thd_head_strides (eight ports, 0 = D) and _fake_rows builds (T*ts, ts, hs, 1); the kernels read X.stride[2] where they multiplied by D -- the Q/dO and K/V tile bases, the direct-bound dK/dV stores, the O.dO dot, the bounded dQ cast and the dK/dV fold. A compact port keeps the compact fake (byte-identical codegen); a staged port (head dim inside the flavor envelope) stays compact at the kernel. Rule (engine mismatch and adapter check_support): element stride 1, head stride >= D, token stride >= H * head_stride, each a multiple of 8 elements so every head base stays 16-byte aligned for the cp.async loads; the head stride wildcards on a size-1 head axis as before. The execute-time view check pins the plan's head stride alongside its token stride. Tests: five head-gapped graph cases (every port, head and token gaps together, GQA, a padded head dim, causal + deterministic) with NaN-filled gap cells asserted never read or written -- the gap-cell helper now covers head gaps as well as token gaps -- plus accept / misaligned / overlapping probe cases. Tracker footnote k and the engine docstring updated. Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
egilliam-nv
force-pushed
the
sm80-thd-head-strides
branch
from
September 16, 2026 19:53
1d331f9 to
d1790fc
Compare
Collaborator
Author
|
@cudnn-ci-bot run python_tests,frost |
|
🏁 Pipeline finished SHA: 31 passed, 2 failed, 5 manual
|
This branch has not been deployed
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
3 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
sdpa_bwd_sm80now serves packed THD ports at their own head stride, not only their own token stride. Independent of draft PR #1058 (which will rebase onto this once it lands).Why
PR #960 turned on head-axis stride gaps in the ragged sweeps: each ragged Q/K/V/O independently becomes a head-interleaved record (head stride
D + k·16 bytes). The SM80 THD backward addressed a row astoken × token_stride + head × Dand declined any other head stride, so withCUDNN_FRONTEND_ENABLE_FROST_ENGINES=1on A100 the backward ragged sweep went from 383 draws served by this row to 48. The cuDNN backend declines every ragged backward on SM80, so the other draws ran nowhere.test_sdpa_random_bwd_ragged_L0, A100, FROST onsdpa_bwd_sm80stats_use_log2, which the SM80 forward does not yet declare; those skip before the backward)What changed
headbyDreadsX.stride[2]instead (Q/dO and K/V tile bases, direct-bound dK/dV stores, O·dO dot, bounded dQ cast, dK/dV fold). A compact port keeps the compact fake, byte-identical codegen.thd_head_strides(eight ports, 0 = D), validated with the token strides;_fake_rows(..., ts, hs).mismatchand adaptercheck_support): element stride 1, head stride>= D, token stride>= H × head_stride, each a multiple of 8 elements (every head base 16-byte aligned for thecp.asyncloads). Size-1 head axis wildcards as before. The execute-time view check pins the plan's head stride too.Tests
test_sdpa_bwd_thd_sm80.py:test_graph_thd_gapped_head_strides×5 (every port head-gapped, head and token gaps together, GQA, padded head dim, causal + deterministic) with NaN gap cells asserted never read or written;test_accept_thd_gapped_head_strides,test_reject_thd_misaligned_head_stride,test_reject_thd_head_stride_below_the_head.Local (A100): SM80 THD / analyzer / integration / stream-respect / api-index suites 198 passed; ragged backward sweep with FROST on 291 passed, 93 skipped (log2 draws), 0 failed.
Docs
Tracker footnote ᵏ,
_sm80_specdocstring, kernel and adapter comments.Labels
cat-enhancements·area:frost·area:global_attention·orig-nv-eng🤖 Generated with Claude Code