Conversation
jiayus-nvidia
added
cat-bug
Contributor
|
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 (6)
💤 Files with no reviewable changes (3)
Included review availability: Your plan provides up to 12 included reviews per hour; 11 remain after this review. 📝 WalkthroughWalkthroughChangesThe change hardens SM90 and SM100 sparse-attention backward kernels against invalid sparse rows and edge-case sink values. It also expands SM100/SM103 FP8 block-sparse support to positive batch and head counts beyond the former restrictions. Sparse backward validity handling
FP8 block-sparse shape support
Estimated code review effort: 4 (Complex) | ~60 minutes Merge Risk: ⚪ Minimal · up to This change hardens sparse backward handling for invalid and empty rows and expands FP8 block-sparse shape support, with regression coverage for the stated cases. No current merge-blocking risk is identified. Suggested reviewers: 🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
Full details: Docstring CoverageExplanation Docstring coverage is 51.28% which is insufficient. The required threshold is 80.00%. Docstring coverage is scoped to functions touched by this diff. Analyzed 39 functions across 6 files. (2 skipped: 2 unsupported.)
✨ Finishing Touches 💡 1🛠️ Fix failing CI checks 💡
🧪 Generate unit tests (beta)
Comment |
jiayus-nvidia
changed the title
jiayus-nvidia
added
the
cat-feature
Contributor
Author
|
@cudnn-ci-bot run oss |
|
🏁 Pipeline finished SHA: |
Anerudhan
approved these changes
Sep 7, 2026
Merged
4 tasks
thynics
pushed a commit
to thynics/cudnn-frontend
that referenced
this pull request
Sep 14, 2026
The sink-folded LSE of the two-CTA D512 kernel only special-cased lse == +inf, so an infinite sink evaluated inf - inf, a finite sink whose log2(e) rescale overflows (for example 2.4e38) did the same, and both produced NaN dq and dkv. An empty row with the sink disabled (lse and sink both -inf) produced NaN d_sink through the dSink reduction, which also had no saturating-sink handling. Test the rescaled LSE and sink values as the H16/H32 kernels do since NVIDIA#926, and give the dSink reduction the same +inf / -inf sink handling as the generic kernel. Index validity was already aligned with the generic kernel. Related to NVIDIA#676.
jiayus-nvidia
pushed a commit
that referenced
this pull request
Sep 15, 2026
* feat(dsa): add H128/D576 two-CTA SM100 sparse attention backward
Add a (2,1,1)-cluster CuTe DSL backward kernel for BF16 DeepSeek Sparse
Attention with H=128, D=576 and D_v=512 on SM100-class devices, and route
the public backward API to it from tensor metadata.
- Score, dP, dQ, dV and dK run on cta_group::2 tensor cores with the two
CTAs of a cluster splitting the 128 heads; the D64 QK/dQ tails are
separate CG2 rounds and the dK tail is a per-CTA warp MMA with FP32
atomics. P and dS are exchanged between the CTAs with shared::cluster
bulk copies. Slot validity is applied before every score and gradient
operation, including the K rows retained for dQ, so ignored top-k slots
cannot contaminate gradients.
- Short query batches spread each query over several clusters (one dQ
owner plus dV/dK workers); long sequences use one cluster per query.
- Dispatch requires an SM100-class device (compute capability 10.0 or
10.3, validated on B200 and B300), contiguous BF16 inputs and topk_max in
{128, 512, 1024, 1152, 2048}; FP16, other head counts and top-k sizes,
noncontiguous inputs, SM107 and SM90 keep their existing paths.
- The plan compiles from tensor metadata without a priming launch and
executes without allocation: the wrapper allocates dq/dkv/d_sink and a
reusable uint8 workspace whose size is exposed through
SparseAttentionBackward.scratch_workspace_bytes(). Zero-initialization
of the FP32 dKV/dSink accumulators happens in-kernel on the launch
stream, and the wrapper cache key now includes device and strides.
- Compatible with nvidia-cutlass-dsl 4.5 (the declared floor): warp-MMA
fragments are created from the partitioned tensors and loop temporaries
are declared ahead of dynamic branches.
Document the dispatch rule, the workspace contract and the benchmark
invocation.
* fix(dsa): guard the H128/D512 two-CTA sink fold and dSink reduction
The sink-folded LSE of the two-CTA D512 kernel only special-cased
lse == +inf, so an infinite sink evaluated inf - inf, a finite sink whose
log2(e) rescale overflows (for example 2.4e38) did the same, and both
produced NaN dq and dkv. An empty row with the sink disabled (lse and
sink both -inf) produced NaN d_sink through the dSink reduction, which
also had no saturating-sink handling.
Test the rescaled LSE and sink values as the H16/H32 kernels do since
#926, and give the dSink reduction the same +inf / -inf sink handling as
the generic kernel. Index validity was already aligned with the generic
kernel. Related to #676.
* test(dsa): cover the H128/D576 two-CTA backward route
- Dispatch: H128/D576/BF16 with each supported topk_max selects the
two-CTA backend on compute capability 10.0 and 10.3; FP16, D_v=576,
unsupported top-k sizes, SM107, SM90 and noncontiguous inputs fall back
to the generic M64 path.
- Numerics against the PyTorch reference for topk 128..2048 with and
without topk_length, covering both the split-cluster regime used by
short query batches and the single-cluster regime that long sequences
take (s_q=129 exceeds the two-SM cluster count of SM100 parts).
- Ignored slots pointing at NaN KV rows leave dq/dkv/d_sink finite, the
reused workspace is untouched outside the declared size, and
execute() neither compiles nor allocates on the hot path.
- Sink normalization, zero and out-of-range topk_length boundaries, and
non-default-stream ordering are extended to the H128 configuration.
* test(dsa): assert the H128/D576 two-CTA route is executed
The dispatch tests only exercised the selector. Spy on the two-CTA launch
entry so the wrapper-path numerics test and the plan-path ignored-slot
test fail if the generic kernel served the request instead. Requested in
review of #1048.
* test(dsa): verify the D576 route in the out-of-range index test
The D576 parametrization asserted only the selector result, which the
dispatch tests already cover; spy on the two-CTA launch entry instead so
a generic fallback cannot pass it. Also state in the docs that both H128
two-CTA backends clamp topk_length, since the D512 kernel does as well.
* fix(dsa): keep padded slots finite in the H128/D512 two-CTA probability path
The two-CTA D512 kernel zero-fills K and V for padded, negative,
out-of-range, or past-length slots, so their score is exactly zero and
the exp2 argument equals the folded negative LSE. Once that exceeds 128
(every valid logit and the sink far below zero) exp2 overflows, the
infinite dS reaches the dQ UMMA, and inf * 0 turns the whole dQ row into
NaN, while the generic, H16 and D576 kernels stay finite.
Clamp the exp2 argument at 64. A valid slot's argument is the log2 of a
probability and never approaches that bound, so valid values are
unchanged; invalid slots keep P and dS finite, and the zero-filled K and
V rows then contribute exact zeros to dQ and are never scattered to dKV.
Masking P per slot instead cost 5-8% on B200 because the math warps sit
on the critical path. Adds a regression test for both H128 two-CTA
routes. Related to #676.
* test(dsa): scope the H128 sink-normalization case to the two-CTA route
On devices that do not select the H128/D576 two-CTA backend (SM107) the
H128 parametrization exercised the generic M64 kernel, which this change
does not touch and which upstream does not otherwise run for H128/D576.
Skip it there, like the other H128/D576 cases; H32 keeps the existing
coverage of the generic path.
* fix(dsa): propagate NaN through the D512 clamp and assert the two-CTA routes in tests
The exp2-argument clamp used the DSL default fmin (PTX min.f32), which
returns the non-NaN operand and would have turned a NaN LSE or sink into
finite gradients on the D512 two-CTA route while the generic and D576
kernels propagate NaN; use the NaN-preserving form. The sink-limit,
padded-slot and sink-normalization tests now assert that the two-CTA
route actually executed rather than only checking the selector. The
benchmark README comment for the D512 route now lists SM100/SM103 like
its dispatch.
* refactor(dsa): key the backward wrapper cache on contiguity, not strides
Routing is the only stride-dependent decision: check_support() selects
the two-CTA route from all(desc.is_contiguous()), the two-CTA plan is
compiled for the contiguous layout, the fallback consumes the runtime
tensors, and the workspace size does not depend on strides. A single
all-inputs-contiguous flag therefore separates the plans that must not
be shared, without recording every stride tuple in the key. Requested in
review of #1048.
---------
Co-authored-by: Long Cheng <longcheng@nvidia.com>
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.
Before submitting
pre-commit runand committed any formatting changes. (pre-commitis unavailable in the test environment; Black,py_compile, andgit diff --checkresults are listed below.)AGENTS.mdfor each directory this PR touches and the changes comply.cat-*, one or morearea:*/op:*, and oneorig-*.Affected area
FE OSS kernels or CuTeDSL
Summary
topk_lengthto the physical top-k width.cp.asyncin the SM90 loader to suppress invalid global reads while zero-filling their shared-memory rows.B=1,H in {4, 8}support boundary to any positive batch and head counts; retain D=128, MHA, and 64-token sequence alignment requirements.Why
#877 fixed invalid sparse-row handling in the generic SM100 kernel, but the dedicated H16 and H32 implementations still trusted active indices and retained the unsafe sink-limit behavior. #785 fixed the original SM90 empty-row, padded-tail, and sink failures, but active negative indices in compact mode, positive out-of-range indices in either mode, and
topk_lengthvalues wider than the physical index row could still reach KV gathers, probability recomputation, or dKV scatter.The updated kernels use one rule throughout the pipeline: a row participates only when it is inside the bounded active prefix and its KV index is in
[0, S_kv). The validity decision is packed once by the loader and reused by the probability path; the scatter path independently rechecks the destination before issuing atomics.The Sage FP8 BSA quantizer and SM100 launcher already derive grids, scales, and scheduler coordinates from runtime batch/head dimensions. The old B/H checks were conservative API gates rather than kernel requirements; B200 validation confirmed non-4/8 head counts and multi-batch execution in both ordinary and split-KV paths.
Related issues
Related to #676.
Builds on #785 and #877.
API and compatibility impact
There are no public function-signature changes. The accepted input domain of the public Sage FP8 BSA wrapper is expanded.
Behavior changes:
topk_lengthis clamped to[0, topk_idxs.shape[-1]]in-kernel.No newer CuTeDSL API is introduced and the supported architecture set is unchanged.
H100 full-wrapper performance (
BF16,S_q = S_kv = 4096; negative change is faster):topk_lengthEach entry is the median of nine CUDA-event samples after 30 warm-up calls; each sample contains 40-200 complete wrapper calls depending on top-k. The compact
topk=2048regression was reproduced in reverse run order. SM100 DSA performance was not measured in this H100 run; the BSA correctness validation below was run separately on B200.Testing
Environment:
The new SM90 positive logical-OOB regression was first run against the parent revision and failed RED with an all-NaN
dQ.Complete DSA backward L0 file after rebasing onto current
develop:cd test/python PYTHONPATH=/code/github/cudnn-frontend/python:/code/dsl_env/lib/python3.12/site-packages \ python -m pytest -q -m L0 fe_api/dsa/test_DSA_sparse_attention_backward.pyResult:
32 passed, 30 skipped, 26 deselectedin 113.28 s. SM100-only DSA cases skipped on the H100 host and are left to SM100 CI.Sage FP8 BSA validation environment:
The new arbitrary-B/H public API cases were first run against the old guards:
H=3andB=2, H=3failed RED. Targeted quantizer and internal split-workspace cases also failed at their respective old guards.Result:
26 passedin 6.90 s on B200/SM100.Static and formatting checks:
Result: all passed; Black reported all five files unchanged.
Summary by CodeRabbit
New Features
Bug Fixes
Documentation