frost(sdpa): SM80 backward native strided-LSE reads + THD max_s_kv grid hint - #766
Draft
egilliam-nv wants to merge 3 commits into
Draft
frost(sdpa): SM80 backward native strided-LSE reads + THD max_s_kv grid hint#766egilliam-nv wants to merge 3 commits into
egilliam-nv wants to merge 3 commits into
Conversation
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
|
@cudnn-ci-bot run frost,oss,python_tests |
|
🏁 Pipeline finished SHA: |
SM80 now follows the SM120 backward lowering shape exactly (design doc S5 / F-3, the backward analogue of NVIDIA#682): one adapter class (SdpaBwdDslSm80 in api_dsl.py) implementing the SdpaBwdDsl contract, lowered through the shared lower_dsl_bwd; bwd/api.py is deleted. - SdpaBwdDslSm80: descriptor-level check_support (flavor pick, mask resolution, dense_flex + strided-stats acceptance), no-op compile (the kernels self-cache until the TemplateParams conversion), execute with the full issue NVIDIA#514 carving (pad/gather staging, strided-stats gather, kernel workspace tail) and the d64 fast-path routing. SM80-only operands (bias -> dBias, RoPE) are extra optional keywords, as the contract permits. - lower_dsl_bwd is parameterized by api_type (mirrors lower_dsl_prefill) and now drives both backward cells; SM80-only constructor facts and execute operands forward via signature introspection, so the SM120 adapter is untouched. lower_sm80_bwd (the plan-time APIBase half-way house) is gone. - sdpa_bwd_wrapper_sm80 keeps its public signature (dense via the adapter, packed THD via the kernel varlen path); SdpabwdSm80 is replaced by SdpaBwdDslSm80 (experimental API, exports/tests updated). - The SM80 suites declare their backward output strides explicitly: the shared lowering honors DECLARED port geometry (IR-inferred output strides are provisional row-major -- the layout invariant), where the old adapter leniently trusted caller tensor metadata. Verified on A100: SM80 suites all levels 118 passed; test_mhas_v2 fwd+bwd L0 = 421/0 at 100% FROST routing (fwd 623, bwd 176) through the ported path. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
…ule (issue NVIDIA#604) The NVIDIA#689 analogue for the backward: bprop_f16_sm80 becomes a TemplateParams template (bwd/config_sm80 params + validator) loaded per-specialization via frost.template_loader, with one module-level compile(...) per shape returning the full kernel chain (do_dot / main / dQ cast / GQA reduces / dSink); the host backward() entry point is gone — launch marshaling lives in the adapter (SdpaBwdDslSm80.compile/execute) and the THD functional wrapper. THD packed token totals compile as cute.sym_int DYNAMICS (issue NVIDIA#604): one artifact per (params, n_seqs) re-binds any token totals — the static PARTIAL_Q/PARTIAL_KV gates fold False under THD_VARLEN (per-sequence bounds ride GATE_Q/GATE_KV), and THD+deterministic is now rejected in the validator (the dQ-relay semaphore has no plan-time size under a dynamic sq). Also: RoPE preconditions the old backward() asserted move to the adapter (rope_max_s coverage, tile alignment) and rope_max_s now reaches the plan (it's part of the compiled table's shape); the d64 fast path keeps its dedicated self-caching module (dense-only) and the d64-vs-generic test runs the generic side through the adapter with the gate forced off; a bwd twin of the THD compile-key regression test guards the plan-time-only key. Dead code deleted: the legacy per-piece _compile_* wrappers (except _compile_do_dot, which d64 imports) and the module dummy cache. Closes NVIDIA#604. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
…id hint The NVIDIA#712 analogue for the backward's LOADS: the main-bprop and dSink kernels read the Stats/LSE input stride-aware, compiled against the DECLARED (B, H, SQ) strides (module compile() gains a plan-time lse_stride key; a contiguous plan keeps the packed compact fake — byte-identical codegen). The adapter drops the strided-stats gather staging entirely: sizing loses the staging term, execute binds the caller's storage through a checked declared-layout view (the reader twin of the forward's _checked_lse_view), and the SM120 bwd adapter's existing _lse_strides convention now has its SM80 counterpart. The d64 fast path keeps legacy packed reads, so a strided Stats declaration routes to the generic module (plan-time gate). THD keeps packed math — the grid batch index is the logical sequence, not the tensor's batch-1 dim. THD wrapper: a max_s_kv grid hint (any upper bound — short kv-tiles early-out) removes the wrapper's only D2H sync; n_seq was already shape metadata. Hint-less calls keep the documented host read of cu_k. The dense path rejects the hint loudly. Tests: strided-stats bitwise-vs-contiguous with equal workspace bytes (no staging), and max_s_kv exact/over-provisioned bitwise checks plus the dense-path rejection. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
egilliam-nv
force-pushed
the
sm80-bwd-strided-lse-reads
branch
from
August 28, 2026 16:28
faa8337 to
cb1abe0
Compare
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
2 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.
Important
Stacked on #765 (#716 has merged; this PR is now three commits, the first two being #765's) — review only the last commit:
cb1abe007. This PR will be rebased down to one commit as its bases merge, then marked ready for review.Two follow-ups to #765 that finish the zero-copy story on the SM80 backward:
Native strided-LSE reads (the #712 analogue for the backward's loads)
compile()gains a plan-timelse_stridekey and the device loads go through the fake's declared(B, H, SQ)strides. A contiguous plan keeps the packed compact fake — byte-identical codegen. THD keeps packed math (the grid batch index is the logical sequence, not the tensor's batch-1 dim)._checked_lse_view, and the SM80 counterpart of the SM120 bwd adapter's existing_lse_stridesconvention).THD
max_s_kvgrid hintcu_kto size the over-provisioned grid.sdpa_bwd_wrapper_sm80(..., max_s_kv=...)(any upper bound — short kv-tiles early-out) now removes it;n_seqwas already shape metadata. Hint-less calls keep the documented host read; the dense path rejects the hint loudly.Tests
test_sm80_bwd_strided_stats_native_reads: interleaved-stride stats vs contiguous — bitwise-equal dQ/dK/dV and identicalscratch_workspace_bytes()(proves no staging).test_sm80_bwd_thd_max_s_kv_hint: exact and over-provisioned hints bitwise-equal to the hint-less path; dense-path rejection.Verification (A100, dev backend 9.27)
🤖 Generated with Claude Code