Improve performance of SM100 DSA backward kernel by Butterfingrz · Pull Request #684 · NVIDIA/cudnn-frontend

Butterfingrz · 2026-08-20T16:44:38Z

PR description

Before submitting

I agree to license this contribution under the terms of LICENSE.txt.
I ran pre-commit run on every changed file and committed any formatting changes.
Labels requested: cat-enhancements, mod-cutedsl, and orig-external (the external contributor account does not have permission to add labels to the NVIDIA repository).

Affected area

FE OSS kernels / CuTeDSL, DSA backward performance, tests, documentation, and benchmark discoverability.

Summary

Improve performance of SM100 DSA backward kernel

Main changes:

Rebalanced the sparse KV gather register budget from 40 to 56 registers/thread while preserving the CTA-wide allocation of 96 registers/thread (61,440 registers), eliminating most loader spills without changing occupancy.
Parallelized each warp's 16 top-k index loads across lanes 0--15. Each lane retains one scalar index and broadcasts it with a shuffle only when its row is consumed, replacing the lane-0 serial load and the 16-entry per-thread live set.
Moved all K-consuming dQ MMAs ahead of the K-independent dKV tail and released K immediately afterward, allowing the next tile's indexed gather to overlap the current tile's remaining dKV MMAs.
Released each S TMEM generation immediately after softmax/quantization and the TMEM-load fence, decoupling the S lifetime from the subsequent P SMEM publication.
Reworked the main dQ epilogue around Ld16x256b and transposed STSM (StMatrix8x8x16bOp(transpose=True)). The STSM writer layout is derived from the T2R destination ownership with make_tiled_copy_D, replacing the scalar STS.U16 publication path.
Expanded dQ TMA-store staging from one to four slots by reusing 64 KiB of dead sK storage. The four main dQ stores can be issued back-to-back before the final drain; the D576 tail uses dead sP storage to avoid aliasing the outstanding main stages.

Benchmarked on NVIDIA B200 (SM100) with driver 580.126.09 and the SM clock locked at 1830 MHz; full setup is listed below.

seqlen_q=seqlen_kv=4096

Config	topk	old ms	old TFLOPS	new ms	new TFLOPS	new speedup
QK=V=512,H=64	128	0.593	289.57	0.545	315.15	+8.8%
QK=V=512,H=64	512	1.324	518.92	1.224	561.51	+8.2%
QK=V=512,H=64	1024	2.306	595.88	2.149	639.59	+7.3%
QK=V=512,H=64	2048	4.285	641.47	4.080	673.66	+5.0%
QK=V=512,H=128	128	1.115	308.23	1.022	336.32	+9.1%
QK=V=512,H=128	512	2.579	532.99	2.406	571.31	+7.2%
QK=V=512,H=128	1024	4.548	604.33	4.316	636.81	+5.4%
QK=V=512,H=128	2048	8.573	641.29	8.147	674.78	+5.2%
QK=576,V=512,H=128	128	1.282	288.10	1.147	322.00	+11.8%
QK=576,V=512,H=128	512	3.421	431.92	2.993	493.67	+14.3%
QK=576,V=512,H=128	1024	6.288	469.92	5.468	540.37	+15.0%
QK=576,V=512,H=128	2048	12.000	492.49	10.407	567.87	+15.3%
QK=576,V=512,H=64	128	0.682	270.93	0.614	301.00	+11.1%
QK=576,V=512,H=64	512	1.747	422.88	1.533	481.75	+13.9%
QK=576,V=512,H=64	1024	3.179	464.80	2.769	533.55	+14.8%
QK=576,V=512,H=64	2048	6.049	488.48	5.248	563.08	+15.3%

The first 12 rows are exactly the matrix used in #318. QK=576,V=512,H=64 is added because D576 is where this change helps most and it is the shape the L0 tests cover. Mean speedup is +10.5% over all 16 cells and +9.4% over the 12 cells of #318's matrix.

The seqlen_q=seqlen_kv=8192 sweep (H=64, topk 128/512/2048) agrees: +9.7%/+8.3%/+5.4% on D512 and +15.5%/+14.5%/+15.3% on D576.

Relationship to #318 and the baseline regression

The first 12 rows above are #318's matrix. Taking the arithmetic mean of the 12 per-cell ratios develop_ms / published_ms - 1, develop @ 55b2773e is 9.0% slower than the post-change times published in #318.

Two measurements split that gap. #318's squash-merge commit cfee7248, re-run here, is 3.7% slower than its own published table by the same metric (+0.1% ... +6.0% per cell) -- environment, toolchain and harness drift rather than code. On this same setup develop is a further 5.2% slower than cfee7248 over those 12 cells: D512 1.5%, D576 12.7%. The two approximately reconcile the 9.0% gap (1.037 x 1.052 - 1 = 9.1%); means of ratios do not multiply exactly, the per-cell products do.

Config	topk	#318 published ms	`cfee7248` here ms
QK=V=512,H=64	128	0.560	0.594
QK=V=512,H=64	512	1.233	1.302
QK=V=512,H=64	1024	2.155	2.249
QK=V=512,H=64	2048	4.019	4.199
QK=V=512,H=128	128	1.066	1.115
QK=V=512,H=128	512	2.424	2.532
QK=V=512,H=128	1024	4.303	4.459
QK=V=512,H=128	2048	8.268	8.403
QK=576,V=512,H=128	128	1.175	1.230
QK=576,V=512,H=128	512	2.949	3.027
QK=576,V=512,H=128	1024	5.327	5.424
QK=576,V=512,H=128	2048	10.197	10.208

develop and this PR for these cells are the old ms and new ms columns of the matrix above.

Normalised by tile count (topk / 64), the develop minus cfee7248 difference on QK=576,H=128 is 26 / 49 / 54 / 56 us per tile at topk 128 / 512 / 1024 / 2048 -- settling at 49--56 for topk >= 512 instead of falling as 1/tiles the way a fixed host-side cost would. The same normalisation on QK=V=512,H=128 gives 0.1 / 5.9 / 5.6 / 5.3 us per tile. That shape matches #396, which reports its own cost as 8.0% on head_dim 576 and 0.03% on 512, from one added per-iteration handshake on the MMA warp's critical path. It is not the whole story: the D576 gap measured here is larger than #396's 8.0%, D512 moved 1.5% although both #395 and #396 report it unchanged, and eight commits touched this file between the two endpoints. The endpoint comparison bounds that cost; it does not isolate it to one commit.

This PR keeps every one of those fixes -- #395's staged store views and TMEM-dealloc barrier, #396's dKV2/dKV3 WAR barrier -- and recovers most of the D576 cost instead. Against cfee7248 on QK=576,H=128 it is 7.2% and 1.1% faster at topk 128 and 512 and 0.8% and 1.9% slower at topk 1024 and 2048; on D512 it is 2.9% to 9.1% faster.

performance evidence

The headline speedups are uninstrumented paired A/B medians: 6 samples per cell, 2 passes x 3 rounds with the A/B order reversed between passes. Run-to-run spread is <=0.35% on 15 of 16 baseline cells and 11 of 16 PR cells, pass-to-pass drift <=0.26%; the weakest cell still has 2x more signal than spread, most D576 cells over 100x.

The causal isolation comes from IKET on D576 (seqlen 512, topk 512, 8 tiles). A variant carrying only the dQ-MMA reorder -- the K-consuming dQ MMAs moved ahead of the K-independent dKV tail with K released immediately afterward, and nothing else -- is the early-K-release ablation below. Steady-state medians:

Interval	develop	early-K-release ablation	this PR
`mma_wait_K` -- MMA stalled on the next gathered K	512 ns	64 ns	32 ns
`mma_tile` -- per-tile period	6,560 ns	6,080 ns	6,016 ns
loader index/setup (`ld_tile - ld_acq_K - ld_gather`)	2,528 ns	2,656 ns	480 ns

The ablation alone cuts the next-K wait by 448 ns/tile and the per-tile period by 480 ns/tile -- equal within one 32 ns IKET tick, so the period gain is accounted for by the removed stall. The full PR takes a further 64 ns/tile off the period.

The loader interval moves only in the full PR (still 2,656 ns in the ablation), so it is evidence for the combined loader-side cleanup -- lane-parallel top-k indices plus the 40->56 gather register rebalance -- and not an isolated timing for either one.

Static SASS (nvdisasm 12.8.90) and Nsight Compute 2026.2.1 (seqlen 4096, topk 2048, H=64), develop -> this PR:

Evidence	D512	D576
SASS `STS.U16` / `STSM`	256 -> 0 / 8 -> 40	288 -> 32 / 8 -> 40
local-memory requests (spill ld + st)	11,862,016 -> 0	18,681,856 -> 200,704
`smsp__inst_executed.sum`	-11.4%	-11.9%
`sm__pipe_tensor_cycles_active` (% peak)	42.61 -> 46.39	33.16 -> 38.18

The 32 remaining STS.U16 are the still-scalar store_dQ_64 D576 tail. One counter moves the other way: D512 shared-store bank conflicts rise 71,208 -> 176,561, which is 0.9% of the 19M store wavefronts and leaves net store wavefronts down 9.5%; on D576 they fall, 29,740 -> 21,651.

TMA transaction counts are unchanged (UTMASTG / UTMACMDFLUSH 8/4 on D512 and 9/5 on D576, both sides) -- the staging change alters the wait predicate, not the number of stores. With launch shape and occupancy also unchanged (below), what is left is less per-tile waiting and fewer instructions.

Instrumented durations are not performance numbers -- the selector forks the compile key and the timestamp writes sit on the critical path -- so the IKET table is structural evidence only.

Launch-shape invariants

Confirmed identical between baseline and this PR via Nsight Compute: grid (4096,1,1), block (640,1,1), 96 registers/thread, and shared memory per block 216.064 KB (D512) / 232.448 KB (D576). Occupancy is unchanged, so the gains are efficiency only. The constructor signature, both _get_workspace_size_* staticmethods and the compile key are untouched, which is why no interface change is required and the PR touches a single file.

Benchmark configuration

NVIDIA B200 (SM100, CC 10.0, 148 SM), SM clock locked at 1830 MHz; driver 580.126.09
PyTorch 2.13.0+cu130, nvidia-cutlass-dsl 4.6.0, cuDNN backend 9.20.0
Baseline: upstream develop @ 55b2773e
Primary matrix: seqlen_q = seqlen_kv = 4096, BF16, d_v = 512, H in {64, 128}, topk in {128, 512, 1024, 2048}
Attention sink enabled; topk_length present with every row set to the configured full top-k
CUDA-event timing of one flash_attn_bwd_sm100 call; JIT compilation and reference forward excluded; 20 initial warmups, 5 re-warmups per round, 100 measured calls, and 2 passes x 3 rounds with reversed A/B order
Timing loop and FLOPs formula taken from benchmark/dsa/benchmark_dsa_sparse_attention_backward.py; the custom driver was cross-checked against that script and agreed to 0.2%
Secondary sweep: seqlen 8192, H=64, topk in {128, 512, 2048}

Validation

All 16 benchmark configurations completed with no NaN and with every output element written (outputs were seeded with a sentinel before each launch and the post-run sentinel count is 0).
dQ is bit-identical to the baseline implementation (rel_l2 = 0.000e+00) and dSink is bit-identical, on all 6 paired correctness configurations. dKV differs by rel_l2 = 1.9e-05 ... 3.8e-05, which is FP32 global-atomic accumulation order and is far below one ULP of BF16 (eps 7.8e-3).
test_DSA_sparse_attention_backward.py -m "L0 or L1": 15 passed, 1 skipped -- identical to the -m L0 result and to the baseline. The single skip is SM90-only.

Summary by CodeRabbit

Bug Fixes
- Improved backward-pass handling for sparse attention, including empty top-k rows and synchronization during gradient computation.
- Enhanced gradient storage reliability across sub-tiles and shared-memory operations.
- Added validation to prevent insufficient shared-memory capacity during processing.
Performance
- Optimized KV loading, index handling, and gradient computation to improve execution efficiency.

Six independent changes to FlashAttentionDSABackwardSm100, measured on a B200 at a locked 1830 MHz, S=8192, over topk 128/512/2048 x causal 0/1. All twelve scenarios improve: D512 by 4.9-8.9%, D576 by 11.1-13.3%. dq is bit-identical to the previous kernel in every scenario (relative error exactly 0), and the launch shape is unchanged: grid 4096, block 640, 96 regs/thread, 216 KiB SMEM/CTA on D512 and 232 KiB on D576. Gather-index path. The per-tile top-k indices were read by lane 0 one row at a time into an rmem tensor, then broadcast. Lane i now reads its own row in _load_tile_topk_idx and the consumer shuffles it out, so the indices live in a single register and the next tile's can be fetched at the tail of the current iteration instead of at the head of the next one. Register budget. num_regs_load_KV 40 -> 56 clears the spills on the gather address path (all of them on D512, ~99% on D576). The per-warp counts must exactly exhaust the CTA pool, which the 640-thread launch fixes at 96 regs/thread: 128*56 + 128*128 + 256*128 + 128*40 = 61440 = 96 * 640. MMA order. dQ = K @ dS now issues before dKV = Q @ dS, which puts load_mma_K_pipeline.consumer_release ahead of the dKV GEMMs rather than after them, freeing the K buffer earlier in the iteration. S lifetime. The fenced T2R of S has fully consumed TMEM S and nothing below reads it, so its consumer_release moves up to the fence instead of trailing P's publication. This unbinds S's lifetime from P's, nothing more. dQ epilogue. store_dQ moves from Ld32x32bOp -- the last one in the file -- to the Ld16x256bOp/StMatrix form the P and dS paths already use, and the four 128-dim sub-tiles each get their own staging slot in the dead K buffer, so a store no longer waits on the previous store's SMEM read. The 64-column D576 tail stages through the dead P buffer instead of sK: sK does have spare room, but only at an offset past the four dQ sub-tiles, which would need its own swizzle/TMA-box proof. That alias is an exact fit today (cosize 4096 == 4096), so it and the sK alias both gain an assert to catch a future stage bump. test/python/fe_api/dsa/test_DSA_sparse_attention_backward.py: 15 passed, 1 skipped. Signed-off-by: Butterfingrz <13524387014@163.com>

coderabbitai · 2026-08-20T16:45:03Z

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: 8d6b7557-ebbc-4355-a3a4-0dcb4a7cc07a

📥 Commits

Reviewing files that changed from the base of the PR and between 2aab0cd and cde9c1e.

📒 Files selected for processing (1)

python/cudnn/deepseek_sparse_attention/sparse_attention_backward/dsa_bwd_sm100.py

Included review availability: Your plan provides up to 12 included reviews per hour; 11 remain after this review.

📝 Walkthrough

Walkthrough

The SM100 sparse-attention backward kernel updates KV index loading, shared-memory staging, dKV synchronization, and dQ output stores. It adds capacity checks, handles staged dQ sub-tiles, and drains pipeline reads before kernel exit.

Changes

SM100 backward kernel

Layer / File(s)	Summary
KV loading and shared-memory staging `python/cudnn/deepseek_sparse_attention/sparse_attention_backward/dsa_bwd_sm100.py`	KV loading now uses packed top-k indices with lane shuffles. dQ staging uses one slot per sub-tile. Shared-memory capacity assertions cover dQ layouts.
dKV computation and pipeline synchronization `python/cudnn/deepseek_sparse_attention/sparse_attention_backward/dsa_bwd_sm100.py`	dKV Q·dS accumulation runs after dQ processing. TMEM alias protection and S-consumer pipeline release occur at reordered synchronization points.
dQ staging and output stores `python/cudnn/deepseek_sparse_attention/sparse_attention_backward/dsa_bwd_sm100.py`	dQ output preparation iterates over staged sub-tiles. Stores use `Ld16x256b`, explicit coordinate partitioning, quantization, tiled transpose operations, and a CTA drain before exit.

Estimated code review effort: 4 (Complex) | ~60 minutes

Merge Risk: ⚪ Minimal · up to cde9c

The performance-focused kernel changes are merge-ready after normal checks and review; no actionable merge-blocking risk remains.

Suggested reviewers

Suggested reviewers: zkyue

🚥 Pre-merge checks | ✅ 5

✅ Passed checks (5 passed)

Check name	Status	Explanation
Docstring Coverage	✅ Passed	No functions found in the changed files to evaluate docstring coverage. Skipping docstring coverage check.
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.
Title check	✅ Passed	The title clearly and concisely describes the main change: improving SM100 DSA backward-kernel performance.
Description check	✅ Passed	The description covers the changes, rationale, benchmarks, validation, compatibility impact, testing, and the unavailable label permission.

✨ Finishing Touches

🧪 Generate unit tests (beta)

Create PR with unit tests

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

Butterfingrz · 2026-08-20T16:46:23Z

Hi , could you please take a look at this PR when you have a chance, Thanks!

Anerudhan · 2026-08-24T16:56:58Z

@cudnn-ci-bot run oss

cudnn-ci-bot · 2026-08-24T16:57:01Z

🏁 Pipeline finished

SHA: cde9c1e
Targets: oss
Branch: cudnn-gh/pr-684-cde9c1e
Pipeline: 64310041
Last updated: 2026-08-24 19:19 UTC

Anerudhan · 2026-08-24T16:57:56Z

Thanks @Butterfingrz for the contribution.
I will update this in a day.

Butterfingrz · 2026-08-24T19:25:46Z

Thanks for taking a look at this PR! @Anerudhan

jiayus-nvidia · 2026-08-25T03:26:38Z

Thanks a lot! LGTM.

Anerudhan assigned Butterfingrz Aug 24, 2026

Anerudhan added orig-external Reported or requested by an external user, customer, or community contributor. mod-cutedsl CuTeDSL kernels, generated kernels, examples, or related integration work. cat-enhancements labels Aug 24, 2026

Anerudhan modified the milestones: Frontend 1.29.0, Frontend 1.28.0 Aug 24, 2026

Navigation Menu

Uh oh!

Improve performance of SM100 DSA backward kernel - #684

Improve performance of SM100 DSA backward kernel#684
Butterfingrz wants to merge 1 commit into
NVIDIA:developfrom
Butterfingrz:perf/optimize-dsa-bwd-sm100

Butterfingrz commented Aug 20, 2026 •

edited

Loading

Uh oh!

coderabbitai Bot commented Aug 20, 2026 •

edited

Loading

Walkthrough

Changes

Suggested reviewers

Uh oh!

Butterfingrz commented Aug 20, 2026

Uh oh!

Anerudhan commented Aug 24, 2026

Uh oh!

cudnn-ci-bot commented Aug 24, 2026 •

edited

Loading

Uh oh!

Anerudhan commented Aug 24, 2026

Uh oh!

Butterfingrz commented Aug 24, 2026

Uh oh!

jiayus-nvidia commented Aug 25, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

Uh oh!

Conversation

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

Uh oh!

PR description

Before submitting

Affected area

Summary

Relationship to #318 and the baseline regression

performance evidence

Launch-shape invariants

Benchmark configuration

Validation

Summary by CodeRabbit

Uh oh!

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

Uh oh!

Walkthrough

Changes

Suggested reviewers

Uh oh!

Butterfingrz commented Aug 20, 2026

Uh oh!

Anerudhan commented Aug 24, 2026

Uh oh!

cudnn-ci-bot commented Aug 24, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

Anerudhan commented Aug 24, 2026

Uh oh!

Butterfingrz commented Aug 24, 2026

Uh oh!

jiayus-nvidia commented Aug 25, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

Butterfingrz commented Aug 20, 2026 •

edited

Loading

coderabbitai Bot commented Aug 20, 2026 •

edited

Loading

cudnn-ci-bot commented Aug 24, 2026 •

edited

Loading