sdpa fp8 sm107: port the has_lse specialization; add a static SMEM guard - #579
Merged
Merged
Conversation
Two follow-ups to NVIDIA#576: 1. NVIDIA#574 (has_lse specialization for the FP8/MXFP8 SM100 flavors) merged between NVIDIA#577 and NVIDIA#576, so the shared adapter now passes has_lse to the fp8 compile() while the SM107 sibling still had the pre-NVIDIA#574 signature - a TypeError on any Rubin fp8 compile at develop tip. This ports NVIDIA#574's fp8-kernel hunks onto the sibling verbatim (LSE None-specialization, compile(has_lse), the specialized epilogue), restoring signature parity with the shared call site. 2. Responding to review on NVIDIA#576: the 9-stage ring with BF16/FP16 O (~242 KiB) exceeds the STANDARD sm_10x 227 KiB per-CTA opt-in and is legal on GR100 only through the sm107 oversized-SMEM launch mode (function attribute 16), which the required internal cutlass-dsl toolchain enables for its sm_107a kernels - board-validated e4m3->bf16 across the full suite. Rather than rejecting output dtypes that demonstrably work on the target stack, a static import-time guard now accounts the geometry against the GR100 hardware budget so a future stage/width bump fails with a clear message instead of at launch. Validated: SM100 box 34 passed; SM107 board 34 passed (full fp8 e2e suite through the sibling, has_lse and no-lse populations both exercised via the suite's generate_stats matrix). Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
vedaanta
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 (1)
📝 WalkthroughWalkthroughThe SM107 FP8 prefill kernel adds optional LSE support. Compilation can omit the LSE tensor and stores while retaining amax outputs. Rubin shared-memory usage is validated at import time. ChangesSM107 FP8 prefill kernel
Estimated code review effort: 3 (Moderate) | ~20 minutes Mergeability Score: ⚪ Minimal · up to The PR restores SM107 FP8 signature parity and adds a static shared-memory guard; no actionable merge-blocking risk remains after normal checks and review. Sequence Diagram(s)sequenceDiagram
participant compile
participant host_launcher
participant prefill_kernel
compile->>host_launcher: Pass LSE tensor or None
host_launcher->>prefill_kernel: Launch with optional LSE tensor
prefill_kernel->>prefill_kernel: Store LSE only when present
prefill_kernel->>prefill_kernel: Update amax independently
Possibly related PRs
🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Comment |
Collaborator
Author
|
@cudnn-ci-bot run frost |
|
🚀 Running mirror pipeline Branch: cudnn-gh/pr-579-bfb8cba |
jhjpark
approved these changes
Aug 13, 2026
vedaanta
added a commit
to vedaanta/cudnn-frontend
that referenced
this pull request
Aug 13, 2026
Two Rubin levers, baked into the SM107 sibling (the Blackwell kernel is untouched; the only shared-file change is a guarded is_exclusive parameter on tile_dsl.tmem_alloc that <=512-column callers never trace): 1. Fused LDTM row-max: the unmasked softmax path loads S_acc and reduces the row max in ONE tcgen05.ld.red.f32.max per 64-column chunk (tmem_load_max_reduction_x64, the mechanism the MXFP8 kernel uses on cc10.3) instead of a manual load + software reduction; masked iters keep the software path (the fused max reduces before a mask could apply). 2. Row-sum in MMA: an N=16 ones-BMM (P x ones) accumulates the softmax denominator into 16 Sigma TMEM columns right after each sub-tile's O, where the correction warp's per-iter alpha-rescale keeps it online-softmax-consistent for free (rescale span widens 128 -> 144); the softmax warps stop reading P back through the register file to sum it, and the correction epilogue reads Sigma (one ld x1) instead of the stats sum word. B rides a constant 8-row x 128 B K-major all-ones tile written once at init (any N-split of all-ones is all-ones, so the cga2 halves are symmetric for free). Sigma now sums the fp8-QUANTIZED P, exactly matching O's numerator; empty-kv tiles guard the stale Sigma read via max == -inf (LSE = -inf convention kept). TMEM grows to the full 576-col sm107 EXCLUSIVE grant (dealloc must return the grant - a smaller dealloc leaks columns as cudaErrorTensorMemoryLeak). Same-session A/B on w2u1g (cc10.7, B8 H16 D128 e4m3->bf16 dense, vs the sibling at NVIDIA#579): 16k nomask 4.017 -> 3.471 ms (-13.6%), 16k causal 2.387 -> 2.070 (-13.3%), 65k nomask 60.780 -> 51.842 (-14.7%, 5430 TFLOPS eff), 65k causal 31.777 -> 27.013 (-15.0%). Board suite 34/34; SM100 box 34/34 (Blackwell untouched). Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
vedaanta
added a commit
to vedaanta/cudnn-frontend
that referenced
this pull request
Aug 13, 2026
Two Rubin levers, baked into the SM107 sibling (the Blackwell kernel is untouched; the only shared-file change is a guarded is_exclusive parameter on tile_dsl.tmem_alloc that <=512-column callers never trace): 1. Fused LDTM row-max: the unmasked softmax path loads S_acc and reduces the row max in ONE tcgen05.ld.red.f32.max per 64-column chunk (tmem_load_max_reduction_x64, the mechanism the MXFP8 kernel uses on cc10.3) instead of a manual load + software reduction; masked iters keep the software path (the fused max reduces before a mask could apply). 2. Row-sum in MMA: an N=16 ones-BMM (P x ones) accumulates the softmax denominator into 16 Sigma TMEM columns right after each sub-tile's O, where the correction warp's per-iter alpha-rescale keeps it online-softmax-consistent for free (rescale span widens 128 -> 144); the softmax warps stop reading P back through the register file to sum it, and the correction epilogue reads Sigma (one ld x1) instead of the stats sum word. B rides a constant 8-row x 128 B K-major all-ones tile written once at init (any N-split of all-ones is all-ones, so the cga2 halves are symmetric for free). Sigma now sums the fp8-QUANTIZED P, exactly matching O's numerator; empty-kv tiles guard the stale Sigma read via max == -inf (LSE = -inf convention kept). TMEM grows to the full 576-col sm107 EXCLUSIVE grant (dealloc must return the grant - a smaller dealloc leaks columns as cudaErrorTensorMemoryLeak). Same-session A/B on w2u1g (cc10.7, B8 H16 D128 e4m3->bf16 dense, vs the sibling at NVIDIA#579): 16k nomask 4.017 -> 3.471 ms (-13.6%), 16k causal 2.387 -> 2.070 (-13.3%), 65k nomask 60.780 -> 51.842 (-14.7%, 5430 TFLOPS eff), 65k causal 31.777 -> 27.013 (-15.0%). Board suite 34/34; SM100 box 34/34 (Blackwell untouched). Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
vedaanta
added a commit
to vedaanta/cudnn-frontend
that referenced
this pull request
Aug 13, 2026
Two Rubin levers, baked into the SM107 sibling (the Blackwell kernel is untouched; the only shared-file change is a guarded is_exclusive parameter on tile_dsl.tmem_alloc that <=512-column callers never trace): 1. Fused LDTM row-max: the unmasked softmax path loads S_acc and reduces the row max in ONE tcgen05.ld.red.f32.max per 64-column chunk (tmem_load_max_reduction_x64, the mechanism the MXFP8 kernel uses on cc10.3) instead of a manual load + software reduction; masked iters keep the software path (the fused max reduces before a mask could apply). 2. Row-sum in MMA: an N=16 ones-BMM (P x ones) accumulates the softmax denominator into 16 Sigma TMEM columns right after each sub-tile's O, where the correction warp's per-iter alpha-rescale keeps it online-softmax-consistent for free (rescale span widens 128 -> 144); the softmax warps stop reading P back through the register file to sum it, and the correction epilogue reads Sigma (one ld x1) instead of the stats sum word. B rides a constant 8-row x 128 B K-major all-ones tile written once at init (any N-split of all-ones is all-ones, so the cga2 halves are symmetric for free). Sigma now sums the fp8-QUANTIZED P, exactly matching O's numerator; empty-kv tiles guard the stale Sigma read via max == -inf (LSE = -inf convention kept). TMEM grows to the full 576-col sm107 EXCLUSIVE grant (dealloc must return the grant - a smaller dealloc leaks columns as cudaErrorTensorMemoryLeak). Same-session A/B on w2u1g (cc10.7, B8 H16 D128 e4m3->bf16 dense, vs the sibling at NVIDIA#579): 16k nomask 4.017 -> 3.471 ms (-13.6%), 16k causal 2.387 -> 2.070 (-13.3%), 65k nomask 60.780 -> 51.842 (-14.7%, 5430 TFLOPS eff), 65k causal 31.777 -> 27.013 (-15.0%). Board suite 34/34; SM100 box 34/34 (Blackwell untouched). Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
vedaanta
added a commit
that referenced
this pull request
Aug 13, 2026
Two Rubin levers, baked into the SM107 sibling (the Blackwell kernel is untouched; the only shared-file change is a guarded is_exclusive parameter on tile_dsl.tmem_alloc that <=512-column callers never trace): 1. Fused LDTM row-max: the unmasked softmax path loads S_acc and reduces the row max in ONE tcgen05.ld.red.f32.max per 64-column chunk (tmem_load_max_reduction_x64, the mechanism the MXFP8 kernel uses on cc10.3) instead of a manual load + software reduction; masked iters keep the software path (the fused max reduces before a mask could apply). 2. Row-sum in MMA: an N=16 ones-BMM (P x ones) accumulates the softmax denominator into 16 Sigma TMEM columns right after each sub-tile's O, where the correction warp's per-iter alpha-rescale keeps it online-softmax-consistent for free (rescale span widens 128 -> 144); the softmax warps stop reading P back through the register file to sum it, and the correction epilogue reads Sigma (one ld x1) instead of the stats sum word. B rides a constant 8-row x 128 B K-major all-ones tile written once at init (any N-split of all-ones is all-ones, so the cga2 halves are symmetric for free). Sigma now sums the fp8-QUANTIZED P, exactly matching O's numerator; empty-kv tiles guard the stale Sigma read via max == -inf (LSE = -inf convention kept). TMEM grows to the full 576-col sm107 EXCLUSIVE grant (dealloc must return the grant - a smaller dealloc leaks columns as cudaErrorTensorMemoryLeak). Same-session A/B on w2u1g (cc10.7, B8 H16 D128 e4m3->bf16 dense, vs the sibling at #579): 16k nomask 4.017 -> 3.471 ms (-13.6%), 16k causal 2.387 -> 2.070 (-13.3%), 65k nomask 60.780 -> 51.842 (-14.7%, 5430 TFLOPS eff), 65k causal 31.777 -> 27.013 (-15.0%). Board suite 34/34; SM100 box 34/34 (Blackwell untouched). Co-authored-by: Claude Fable 5 <noreply@anthropic.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.
Summary
Two follow-ups to #576, one urgent:
Signature-parity fix (Rubin fp8 is broken at develop tip without it). frost(sdpa): has_lse specialization for the FP8/MXFP8 SM100 flavors #574 merged between sdpa fp8 sm100: fix THD O-descriptor row stride (latent) #577 and sdpa: enable per-tensor FP8 prefill on SM107 (Rubin) via a sibling kernel #576, giving the SM100 fp8 kernel
compile(has_lse=...)— which the shared adapter call site now passes — while the SM107 sibling (copied pre-frost(sdpa): has_lse specialization for the FP8/MXFP8 SM100 flavors #574) kept the old signature:TypeErroron any Rubin fp8 compile. This ports frost(sdpa): has_lse specialization for the FP8/MXFP8 SM100 flavors #574's fp8-kernel hunks onto the sibling verbatim. Exactly the sibling-refresh coordination flagged in sdpa: enable per-tensor FP8 prefill on SM107 (Rubin) via a sibling kernel #576's description, realized within hours of the merges.SMEM accounting guard, responding to this review thread: the reviewer's arithmetic is right (BF16/FP16-O at 9 stages ≈ 242 KiB > the standard 227 KiB opt-in) but the conclusion isn't — the configuration is legal on GR100 through the sm107 oversized-SMEM launch mode (function attribute 16), which the required internal cutlass-dsl toolchain enables for sm_107a kernels, and it has been board-validated e4m3→bf16 across the full suite. Instead of rejecting working dtypes, an import-time guard now accounts the geometry against the GR100 hardware budget with a clear failure message for future geometry bumps.
Validation
generate_statsmatrix exercises both has_lse populations🤖 Generated with Claude Code
Summary by CodeRabbit
New Features
Bug Fixes