frost(sdpa): optimize and enhance sm120 sdpa forward fp8 kernel - #562
Conversation
Signed-off-by: Haobin Guo <haobing@nvidia.com>
Signed-off-by: Haobin Guo <haobing@nvidia.com>
Aneureka
added
orig-nv-eng
|
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 (3)
💤 Files with no reviewable changes (2)
🚧 Files skipped from review as they are similar to previous changes (1)
📝 WalkthroughWalkthroughSM120 FP8 support now accepts independent QK and value head dimensions from 32 through 256. Heuristics use FP8-specific storage sizes. The kernel updates Scale_S handling, empty-row behavior, sink validation, and Amax_O tracking. Tests cover expanded dimensions and masks. ChangesSM120 FP8 head-dimension contracts
Estimated code review effort: 4 (Complex) | ~45 minutes Mergeability Score: ⚪ Minimal · up to This PR optimizes and expands SM120 FP8 SDPA support while preserving numerical behavior and API compatibility; the reported tests pass, and no actionable merge-blocking risk remains beyond normal checks and review. Sequence Diagram(s)sequenceDiagram
participant Test
participant offers_engine
participant Heuristics
participant prefill_fp8_sm120
Test->>offers_engine: submit QK and V head dimensions
offers_engine->>Heuristics: filter FP8 tile pairs by SMEM capacity
Heuristics-->>offers_engine: return fitting configurations
offers_engine->>prefill_fp8_sm120: launch selected FP8 kernel
prefill_fp8_sm120->>prefill_fp8_sm120: apply log2 Scale_S and normalize rows
prefill_fp8_sm120-->>Test: return output, LSE, and Amax_O
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 (2)
python/cudnn/sdpa/fwd/kernels/prefill_fp8_sm120.py (2)
227-234: 🎯 Functional Correctness | 🔵 Trivial | ⚡ Quick winValidate the head tiles in
__init__. The docstring now requireshead_tile_qkandhead_tile_vto be multiples of 32 between 32 and 256. The constructor does not check this._setup_attributescomputesself.qk_d_frags = self.head_tile_qk // self.MMA_TILER[2]andself.pv_d_frags = self.head_tile_v // self.MMA_TILER[1], so a non-multiple silently truncates the fragment count and drops head columns.get_swizzlecatches only the sub-32-byte case. The graph path is gated bycheck_support, but the direct template-loader path used bytest_fp8_sm120_head_dim_tail_directcallscompile()with caller-suppliedd_qkandd_v.🛡️ Proposed validation
if has_sink: raise ValueError("has_sink is not supported by the fp8 cell (Amax_S semantics)") + for label, tile in (("head_tile_qk", head_tile_qk), ("head_tile_v", head_tile_v)): + if tile % 32 != 0 or not 0 < tile <= 256: + raise ValueError(f"{label} must be a multiple of 32 between 32 and 256; got {tile}") if thd_varlen and (thd_batch < 1 or thd_max_sq < 1):🤖 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_fp8_sm120.py` around lines 227 - 234, Update the constructor validation in the fp8 cell’s __init__ to require head_tile_qk and head_tile_v to be multiples of 32 within the inclusive range 32–256. Reject invalid values before _setup_attributes computes qk_d_frags and pv_d_frags, while preserving the existing dtype, sink, and varlen validations.
1080-1081: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueDocument why
mask_steps = 0stays correct. Withmask_steps = 0, phase 1 never runs, so no iteration receivesis_first_kv_tile=True. Every tile then takes the "not first" path inonline_softmax. That path stays correct only becauserow_maxstarts at-infandrow_sumstarts at0.0(lines 1024-1025):old_scaleevaluates toexp2(-inf) == 0, which zeroes the already-zeroo_regsand the zerorow_sum. The behavior is correct but implicit. A future change to theold_scalepredicate would break it silently.Add a short comment that records this dependency.
🤖 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_fp8_sm120.py` around lines 1080 - 1081, In the mask_steps assignment within the prefill kernel, add a brief comment documenting why zero is valid: phase 1 is skipped, all tiles use the non-first online_softmax path, and its initial row_max/row_sum values make old_scale zero while output and sum registers are already zero. Note that this relies on the current old_scale predicate.
🤖 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/engines.py`:
- Around line 813-822: Correct the SM120 FP8 THD Stats documentation to reflect
the enforced head-major layout: in python/cudnn/sdpa/fwd/engines.py lines
813-822, state that head-major ragged Stats is served and token-major Stats is
f16-only; in test/python/sdpa/frost/test_sdpa_fwd_fp8_sm120.py lines 11-19,
replace token-major with head-major and no head-major with no token-major.
Update documentation only; implementation symbols such as
SdpaFwdDslSm120.check_support already enforce the correct behavior.
In `@python/cudnn/sdpa/fwd/kernels/prefill_fp8_sm120.py`:
- Around line 1061-1062: Update _execute_fp8 in api_dsl.py to validate the
scalar scale_s value after _scalar(scale_s) applies its default and before
launching the kernel; reject every scale_s <= 0 with the existing host-side
validation/error mechanism, while preserving valid positive values and the
default of 1.0.
---
Nitpick comments:
In `@python/cudnn/sdpa/fwd/kernels/prefill_fp8_sm120.py`:
- Around line 227-234: Update the constructor validation in the fp8 cell’s
__init__ to require head_tile_qk and head_tile_v to be multiples of 32 within
the inclusive range 32–256. Reject invalid values before _setup_attributes
computes qk_d_frags and pv_d_frags, while preserving the existing dtype, sink,
and varlen validations.
- Around line 1080-1081: In the mask_steps assignment within the prefill kernel,
add a brief comment documenting why zero is valid: phase 1 is skipped, all tiles
use the non-first online_softmax path, and its initial row_max/row_sum values
make old_scale zero while output and sum registers are already zero. Note that
this relies on the current old_scale predicate.
🪄 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: c5209537-8f9b-4702-a66f-b5d23d1ebd3e
📒 Files selected for processing (6)
python/cudnn/sdpa/fwd/api_dsl.pypython/cudnn/sdpa/fwd/config_sm120.pypython/cudnn/sdpa/fwd/engines.pypython/cudnn/sdpa/fwd/heuristics.pypython/cudnn/sdpa/fwd/kernels/prefill_fp8_sm120.pytest/python/sdpa/frost/test_sdpa_fwd_fp8_sm120.py
Signed-off-by: Haobin Guo <haobing@nvidia.com>
Aneureka
changed the title
|
@cudnn-ci-bot run frost |
|
🚀 Running mirror pipeline Branch: cudnn-gh/pr-562-96800e1 |
Signed-off-by: Haobin Guo <haobing@nvidia.com>
4 participants
Before submitting
pre-commit runand committed any formatting changes.cat-*, one or moremod-*, and oneorig-*(see label list).Affected area
Summary
old_scalelowers to a single predicatedex2via inline PTXScale_Sinto the exp2 bias instead of a per-element multiply before the e4m3 cast;row_sumun-folds once per row in the epiloguemask_steps = 0)Amax_Oreduction: four independent accumulators instead of one serial fmax chain, and the row-validity gate applied once per row half instead of per elementd128to every multiple of 32 up to 256 (different from 16 for f16's engine, both aligning with 32B TMA swizzled chunk and MMA K dimension)D_QKandD_V, e.g.,(192, 128)engines.py/heuristics.py/api_dsl.py(decline-reason ordering, fp8 SMEM sizing in the tile filter, stale docstrings, dead sink-epilogue code)Why
To enhance frost sdpa forward FP8 engine for SM120.
Related issues
Related to #381.
API and compatibility impact
D_QK/D_Vin{32, 64, …, 256}independently (previously exactly 128)Amax_S/Amax_O/ LSE verified against fp32 references atscale_s ∈ {1, 448}Performance
On RTX PRO 6000: ~2% speedup
Testing
Summary by CodeRabbit
Summary by CodeRabbit
New Features
Bug Fixes
Tests