sdpa fp8 sm100: fused LDTM row-max on cc10.3 (per-tensor FP8)

[vedaanta]

Summary

cc10.3 has the fused LDTM.STAT row-max, but only the MXFP8 kernel read the fused_ldtm_stat flag — the per-tensor FP8 kernel ran the manual tcgen05_ld + software reduction on every part, leaving the win on the table on SM103. This wires the identical specialization into the per-tensor kernel (unmasked path only; masked iters keep the software reduction since the fused max reduces before a mask could apply). cc10.0 folds to the manual path unchanged.

Three-line summary of the mechanism for reviewers new to it: fused_ldtm_stat is per-specialization state in the frozen TemplateParams (auto-set from the device in the adapter), so one kernel file yields distinct traced variants per device class — the same routing MXFP8 has used since its bring-up, and the SM107 sibling simply bakes.

Evidence

• The added branch is byte-for-byte the one the SM107 sibling bakes, board-proven on cc10.7 (the LDTM lever carried most of a −7.5% step at 65k there).
• The helper (tmem_load_max_reduction_x64) and instruction are the same ones the MXFP8 kernel exercises on cc10.3 in production.
• SM100 (cc10.0): 34 passed, behaviorally identical (flag stays False).
• Direct cc10.3 hardware validation queued on a B300 node; result will be posted in-thread before merge if reviewers want it as a gate.

🤖 Generated with Claude Code

Summary by CodeRabbit

• New Features

  • Improved FP8 attention kernel performance on compatible hardware by enabling fused row-reduction operations for additional FP8 formats.
  • Added hardware-aware dispatch to select optimized reductions on newer compute capabilities while preserving compatibility with older ones.

• Documentation

  • Updated technical documentation to reflect expanded fused-reduction support for FP8 kernels.

[coderabbitai]

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: 67a80486-197b-4daa-8a45-2127deea6d1f

📥 Commits

Reviewing files that changed from the base of the PR and between c1288b6 and 182d3a7.

📒 Files selected for processing (3)

• python/cudnn/sdpa/fwd/api_dsl.py
• python/cudnn/sdpa/fwd/config_sm100.py
• python/cudnn/sdpa/fwd/kernels/prefill_d128_fp8_sm100.py

---

📝 Walkthrough

Walkthrough

SM100 FP8 compilation now enables fused LDTM.STAT on compute capability 10.3 for MXFP8 and per-tensor FP8. The prefill kernel uses fused row-max reduction for unmasked iterations and retains software reduction on compute capability 10.0 and masked paths.

Changes

SM100 FP8 fused reduction path

Layer / File(s)	Summary
FP8 fused-stat capability selection python/cudnn/sdpa/fwd/api_dsl.py, python/cudnn/sdpa/fwd/config_sm100.py	Compilation enables fused LDTM.STAT for MXFP8 and per-tensor FP8 on compute capability 10.3. The parameter documentation reflects both FP8 modes.
Prefill row-max execution python/cudnn/sdpa/fwd/kernels/prefill_d128_fp8_sm100.py	The kernel selects hardware row-max reduction for unmasked 64-column chunks when enabled. Compute capability 10.0 and masked iterations use manual loading and software reduction.

Estimated code review effort: 2 (Simple) | ~10 minutes

Mergeability Score: ⚪ Minimal · up to 182d3

This localized optimization enables the fused row-max path for per-tensor FP8 on supported hardware while preserving the existing SM100 behavior; no actionable merge-blocking risk remains beyond normal checks and review.

Sequence Diagram(s)

sequenceDiagram
  participant compile
  participant TemplateParams
  participant prefill_d128_fp8_sm100
  participant tcgen05_ld_red
  compile->>TemplateParams: enable fused_ldtm_stat for FP8 on cc10.3
  TemplateParams->>prefill_d128_fp8_sm100: provide reduction configuration
  prefill_d128_fp8_sm100->>tcgen05_ld_red: load unmasked S_acc chunk with row-max reduction
  tcgen05_ld_red-->>prefill_d128_fp8_sm100: return data and row maxima

Loading

Possibly related PRs

• NVIDIA/cudnn-frontend#574: Both modify SM100 FP8 SDPA forward compilation and the prefill kernel.
• NVIDIA/cudnn-frontend#576: Both modify SM100 FP8 dispatch and kernel behavior.
• NVIDIA/cudnn-frontend#580: Both use fused row-max reduction for unmasked iterations.

Suggested reviewers: anerudhan

🚥 Pre-merge checks | ✅ 4 | ❌ 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Description check	⚠️ Warning	The description explains the change and evidence, but it omits most required template sections, including affected area, compatibility impact, and exact testing commands.	Complete the required template sections, including affected area, why, related issues, API and compatibility impact, testing commands and results, and submission checklist items.

✅ Passed checks (4 passed)

Check name	Status	Explanation
Title check	✅ Passed	The title clearly identifies the per-tensor FP8 SM100 change and the target cc10.3 fused LDTM row-max optimization.
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.

✨ Finishing Touches

🧪 Generate unit tests (beta)

• Create PR with unit tests

---

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

[vedaanta]

@cudnn-ci-bot run frost

[cudnn-ci-bot]

🚀 Running mirror pipeline

Branch: cudnn-gh/pr-586-182d3a7
Pipeline: 62590986
Targets: frost

[vedaanta]

B300 (sm103, cc 10.3) hardware validation done — the merge gate from the PR description is cleared.

Setup: B300 SXM6 node, NGC pytorch:26.03 container, FE source-built from this branch, cuDNN backend 9.25.0, nvidia-cutlass-dsl 4.7.0.

test_sdpa_fwd_fp8_sm100.py -m "L0 or L1": 32 passed in 97s, cc (10, 3)

The per-tensor fp8 e2e suite runs the fused tcgen05.ld.red row-max branch on sm103 with this PR applied; all cases match reference.