Add native SM100 d192/d128 SDPA fprop kernel for DSv3 MLA - #488
Conversation
Port the validated d192/d128 FROST kernel, routing, configuration, scheduler fixes, and correctness coverage onto GitHub develop.
📝 WalkthroughWalkthroughSM100 SDPA flavor selection now supports ChangesSM100 d192/d128 SDPA support
Estimated code review effort: 5 (Critical) | ~90 minutes Sequence Diagram(s)sequenceDiagram
participant SDPA_Test
participant EngineSelection
participant SDPA_API
participant D192Kernel
participant ReferenceSDPA
SDPA_Test->>EngineSelection: request d_qk=192, d_v=128
EngineSelection->>SDPA_API: select d192/d128 flavor
SDPA_API->>D192Kernel: compile and launch kernel
D192Kernel-->>SDPA_Test: return O and LSE
SDPA_Test->>ReferenceSDPA: compute reference output
ReferenceSDPA-->>SDPA_Test: return expected output
Possibly related PRs
Suggested labels: Suggested reviewers: 🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 2
🧹 Nitpick comments (3)
python/cudnn/sdpa/fwd/kernels/prefill_d192_d128_f16_sm100.py (2)
1539-1539: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueRemove the unused local and mark the unused decode result.
Line 1539 assigns
P_COLS_PER_CHUNKin_softmax_warp_group, but the function never reads it. The value is recomputed inside_softmax_kv_bodyat line 1312. Line 1687 also unpackshead_idx, which_softmax_warp_groupnever uses; Ruff reports RUF059 for it.♻️ Proposed cleanup
- CHUNK = 64 - P_COLS_PER_CHUNK = CHUNK // 2 stats_off = LAYOUT.STATS_OFF + sub_tile_id * LAYOUT.STATS_STRIDEApply this at line 1687 and at line 1553:
- q_super_idx, head_idx, batch_idx = _dispatch_decode_payload( + q_super_idx, _head_idx, batch_idx = _dispatch_decode_payload(Note:
CHUNKis also unused in_softmax_warp_grouponceP_COLS_PER_CHUNKis removed. Verify before deleting both.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@python/cudnn/sdpa/fwd/kernels/prefill_d192_d128_f16_sm100.py` at line 1539, Remove the unused P_COLS_PER_CHUNK local from _softmax_warp_group, then remove CHUNK from that function’s parameters if it is no longer referenced. Update the unpacking around the decode result at line 1687 to mark head_idx as intentionally unused, preserving all other values and behavior.Source: Linters/SAST tools
1389-1409: 🗄️ Data Integrity & Integration | 🔵 Trivial | ⚡ Quick winAdd a validation for the softmax chunk count.
CfgD192inheritsTILE_N = 128andN_BMM2_CHUNKS = 2, so the current configuration covers all columns. The validators do not enforceN_BMM2_CHUNKS * 64 == TILE_N. Add this invariant to the applicable validators or deriveN_BMM2_CHUNKSfromTILE_Nto prevent future configurations from silently omitting columns.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@python/cudnn/sdpa/fwd/kernels/prefill_d192_d128_f16_sm100.py` around lines 1389 - 1409, Add validation in the applicable configuration validators for the softmax coverage invariant that N_BMM2_CHUNKS multiplied by 64 equals TILE_N, including CfgD192. Alternatively, derive N_BMM2_CHUNKS directly from TILE_N, ensuring every configured column is covered and no future configuration can silently omit columns.test/python/sdpa/frost/test_sdpa_fwd_dsl_sm100.py (1)
226-243: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winGate the test on the cuDNN backend version. The module-level
pytestmarkalready skips non-SM100 devices._require_dsl()only checks imports. Skip whencudnn.backend_version() < 91200before callingcudnn.pygraph(...).🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@test/python/sdpa/frost/test_sdpa_fwd_dsl_sm100.py` around lines 226 - 243, Update test_dsl_sm100_d192_d128 to query cudnn.backend_version() and skip when it is below 91200, placing this guard before the _run_dsl_graph path invokes cudnn.pygraph. Keep the existing _require_dsl() and test coverage unchanged for supported cuDNN versions.Source: Coding guidelines
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@python/cudnn/sdpa/fwd/kernels/prefill_d192_d128_f16_sm100.py`:
- Around line 1736-1739: Update the comment immediately above O_CHUNK in the
prefill kernel to accurately describe the configured value of 8 and its
reduction from the prior value, while preserving the explanation about
shortening the alpha-rescale live range and avoiding correction-warp register
spills.
- Around line 2125-2126: Update the alignment error message in the d192 envelope
validation to report both element sizes consistently with the checks: use
CFG.BPE for d_qk and CFG.BPE_O for d_v, rather than labeling and printing only
CFG.BPE.
---
Nitpick comments:
In `@python/cudnn/sdpa/fwd/kernels/prefill_d192_d128_f16_sm100.py`:
- Line 1539: Remove the unused P_COLS_PER_CHUNK local from _softmax_warp_group,
then remove CHUNK from that function’s parameters if it is no longer referenced.
Update the unpacking around the decode result at line 1687 to mark head_idx as
intentionally unused, preserving all other values and behavior.
- Around line 1389-1409: Add validation in the applicable configuration
validators for the softmax coverage invariant that N_BMM2_CHUNKS multiplied by
64 equals TILE_N, including CfgD192. Alternatively, derive N_BMM2_CHUNKS
directly from TILE_N, ensuring every configured column is covered and no future
configuration can silently omit columns.
In `@test/python/sdpa/frost/test_sdpa_fwd_dsl_sm100.py`:
- Around line 226-243: Update test_dsl_sm100_d192_d128 to query
cudnn.backend_version() and skip when it is below 91200, placing this guard
before the _run_dsl_graph path invokes cudnn.pygraph. Keep the existing
_require_dsl() and test coverage unchanged for supported cuDNN versions.
🪄 Autofix
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: b0ebd6d4-1732-4ca7-8371-35af19b148ce
📒 Files selected for processing (7)
python/cudnn/sdpa/fwd/api_dsl.pypython/cudnn/sdpa/fwd/config_sm100.pypython/cudnn/sdpa/fwd/engine.pypython/cudnn/sdpa/fwd/engines.pypython/cudnn/sdpa/fwd/kernels/prefill_d192_d128_f16_sm100.pytest/python/sdpa/frost/test_sdpa_fwd_dsl_sm100.pytest/python/sdpa/frost/test_sdpa_graph_analyzer.py
| # O_CHUNK=16 (halved from 32) shortens the alpha-rescale live range — | ||
| # at 32, DSL regalloc spilled correction-warp regs to the stack. | ||
| O_CHUNK = 8 | ||
| N_CHUNKS_O = CFG.TILE_O // O_CHUNK |
There was a problem hiding this comment.
📐 Maintainability & Code Quality | 🟡 Minor | ⚡ Quick win
Fix the stale O_CHUNK comment.
The comment states O_CHUNK=16 (halved from 32), but line 1738 sets O_CHUNK = 8. O_CHUNK controls the alpha-rescale live range and the register pressure of the correction warps. A wrong value in the comment misleads later tuning.
📝 Proposed fix
- # O_CHUNK=16 (halved from 32) shortens the alpha-rescale live range —
- # at 32, DSL regalloc spilled correction-warp regs to the stack.
+ # O_CHUNK=8 shortens the alpha-rescale live range — at larger chunk
+ # sizes DSL regalloc spilled correction-warp regs to the stack.
O_CHUNK = 8📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| # O_CHUNK=16 (halved from 32) shortens the alpha-rescale live range — | |
| # at 32, DSL regalloc spilled correction-warp regs to the stack. | |
| O_CHUNK = 8 | |
| N_CHUNKS_O = CFG.TILE_O // O_CHUNK | |
| # O_CHUNK=8 shortens the alpha-rescale live range — at larger chunk | |
| # sizes DSL regalloc spilled correction-warp regs to the stack. | |
| O_CHUNK = 8 | |
| N_CHUNKS_O = CFG.TILE_O // O_CHUNK |
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@python/cudnn/sdpa/fwd/kernels/prefill_d192_d128_f16_sm100.py` around lines
1736 - 1739, Update the comment immediately above O_CHUNK in the prefill kernel
to accurately describe the configured value of 8 and its reduction from the
prior value, while preserving the explanation about shortening the alpha-rescale
live range and avoiding correction-warp register spills.
| if (d_qk * CFG.BPE) % 16 != 0 or (d_v * CFG.BPE_O) % 16 != 0: | ||
| raise ValueError(f"d192 envelope: d_qk*BPE and d_v*BPE must be 16-byte multiples (TMA global-stride rule); got ({d_qk}, {d_v}) at BPE={CFG.BPE}") |
There was a problem hiding this comment.
📐 Maintainability & Code Quality | 🟡 Minor | ⚡ Quick win
Report the correct element size in the alignment error.
Line 2125 validates d_qk with CFG.BPE and d_v with CFG.BPE_O. The message on line 2126 names only BPE and prints CFG.BPE. If CFG.BPE_O != CFG.BPE, the reported value does not match the check that failed.
📝 Proposed fix
if (d_qk * CFG.BPE) % 16 != 0 or (d_v * CFG.BPE_O) % 16 != 0:
- raise ValueError(f"d192 envelope: d_qk*BPE and d_v*BPE must be 16-byte multiples (TMA global-stride rule); got ({d_qk}, {d_v}) at BPE={CFG.BPE}")
+ raise ValueError(
+ f"d192 envelope: d_qk*BPE and d_v*BPE_O must be 16-byte multiples (TMA global-stride rule); "
+ f"got ({d_qk}, {d_v}) at BPE={CFG.BPE}, BPE_O={CFG.BPE_O}"
+ )📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| if (d_qk * CFG.BPE) % 16 != 0 or (d_v * CFG.BPE_O) % 16 != 0: | |
| raise ValueError(f"d192 envelope: d_qk*BPE and d_v*BPE must be 16-byte multiples (TMA global-stride rule); got ({d_qk}, {d_v}) at BPE={CFG.BPE}") | |
| if (d_qk * CFG.BPE) % 16 != 0 or (d_v * CFG.BPE_O) % 16 != 0: | |
| raise ValueError( | |
| f"d192 envelope: d_qk*BPE and d_v*BPE_O must be 16-byte multiples (TMA global-stride rule); " | |
| f"got ({d_qk}, {d_v}) at BPE={CFG.BPE}, BPE_O={CFG.BPE_O}" | |
| ) |
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@python/cudnn/sdpa/fwd/kernels/prefill_d192_d128_f16_sm100.py` around lines
2125 - 2126, Update the alignment error message in the d192 envelope validation
to report both element sizes consistently with the checks: use CFG.BPE for d_qk
and CFG.BPE_O for d_v, rather than labeling and printing only CFG.BPE.
|
@cudnn-ci-bot check |
|
Backend pipeline not launched Reason: @adshen is not allowlisted to run this bot. Ask an allowlisted maintainer to comment |
|
@cudnn-ci-bot run |
|
🚀 Running mirror pipeline Branch: cudnn-gh/pr-488-75c0b27 |
3 participants
Before submitting
pre-commit runand committed any formatting changes.Affected area
FE OSS kernels or CuTeDSL
Summary
This PR adds a native SM100/Blackwell FROST DSL SDPA forward kernel for the DSv3 MLA logical shape:
D_QK = 192D_V = 128The implementation extends the existing SM100
d=128classic forward pipeline for the logical192/128shape instead of selecting a paddedd=256flavor.The change:
prefill_d192_d128_f16_sm100.py.CfgD192andmake_cfg_d192.d192/d128engine capability row and API flavor route.Why
DSv3 MLA uses Q/K head dimension 192 and V/O head dimension 128. Selecting the existing d256 envelope performs unnecessary padded work. A native logical-shape kernel reduces duration while preserving high compute utilization.
For top-left causal
S=8192, the native kernel reaches 87.09% Compute SOL and 820.0 useful TFLOPS. Against an exact d256 padded-work proxy, it is approximately 1.50x faster by both duration and useful TFLOPS.Related issues
None.
API and compatibility impact
D_QK=192, D_V=128on SM100.4.7.0a0.Testing
Formatting and syntax:
Result: passed. The new kernel was also checked with
python3 -m py_compile.Targeted test coverage:
Local SM100 validation results:
[192, 128].Post-port verification from a clean editable build of this GitHub branch:
test/python/sdpa/frost test/python/test_mhas_v2.py.test_mhas_v2.pyresult on the local cuDNN 9.21.1 backend: 1862 passed, 907 skipped, 116 failed. All failures were native FP8/MXFP8 backward cases outside this forward-kernel diff. Validation against the current development backend is pending the repository CI bot run.Performance methodology
4.7.0a0.B=2,Hq=128,Hkv=128,D_QK=192,D_V=128.0.842-0.846 GHz.Native d192/d128 SOL by mask and size
Comparison against an exact d256 proxy
Before this change, logical
192/128used a d256 envelope. The controlled reference below is a separate exact256/256run used as a padded-work proxy, not a direct logical192/128fallback execution. For DSv3 useful work,(192 + 128) / (256 + 256) = 0.625of the proxy work is useful.For top-left causal
S=8192:256192The native d192/d128 path is approximately 1.50x faster for this case by both duration and useful TFLOPS. This is a proxy estimate rather than a direct fallback A/B measurement.
Existing exact-shape kernel reference
These rows are not direct apples-to-apples comparisons because the logical head dimensions differ. They provide context for current SM100 FROST forward SOL at
S=8192.Optimization notes
The final 15-launch
8192x8192confirmations were:Deep profiling reported zero local/shared spills for all three mask paths, 128 registers per thread, and approximately 25% achieved occupancy. The gains below are context-dependent and are not additive.
mbarrier_wait_parity(..., TRY)loop instead of the timed generic wait192/88instead of200/72sched_res_busy_xu64region together with--ptxas-options -uumnFenceInterferencerather thanFenceCodefor P chunk-0 publicationFenceCodeon P chunk-1 publication192/88to216/40Final-context leave-one-out checks confirmed the dual BMM2-done fences and the no-mask
216/40register split. A previously promising BMM1-done fence was removed after the BMM2 fences subsumed its gain, and no-mask producer unroll 7 was removed after the216/40split made unroll 6 faster again. The no-mask route retains the x64 zero-spill lowering and uses216/40; causal paths use192/88.