sdpa fp8 sm107: fused LDTM row-max + row-sum-in-MMA

[vedaanta]

Summary

Two performance levers baked into the SM107 sibling kernel (Blackwell 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 does the S load + row-max in one tcgen05.ld.red.f32.max per 64-col chunk (the MXFP8 kernel's cc10.3 mechanism); masked iters keep the software reduction.
2. Row-sum in MMA — an N=16 ones-BMM (P × ones) lands the softmax denominator in 16 Σ TMEM columns beside O, where the correction warp's α-rescale keeps it online-softmax-consistent for free; the softmax warps stop pulling P back through the RF to sum it. Σ sums the fp8-quantized P, exactly matching O's numerator. TMEM uses the full 576-col sm107 exclusive grant (dealloc must return the grant — smaller deallocs leak as cudaErrorTensorMemoryLeak).

Perf (same-session A/B vs #579, w2u1g cc10.7, B8 H16 D128 e4m3→bf16 dense)

shape	#579 baseline	this PR	Δ
16k nomask	4.017 ms	3.471 ms	−13.6%
16k causal	2.387 ms	2.070 ms	−13.3%
65k nomask	60.780 ms	51.842 ms (5430 TFLOPS eff)	−14.7%
65k causal	31.777 ms	27.013 ms	−15.0%

Validation

SM107 board 34/34 (full fp8 e2e: all masks / sink / stats / both fp8 dtypes / all output dtypes vs the fp32 dequant reference — the amax_s/LSE checks cover the Σ path); SM100 box 34/34.

Stacked on #579 (contains its commit; rebases clean once it merges). The f16x2-MUFU softmax lever is deliberately excluded — it ships separately with its own numerics evidence.

🤖 Generated with Claude Code

Summary by CodeRabbit

• New Features
  • Added configurable exclusive tensor-memory allocation.
  • Enhanced FP8 scaled dot-product attention on Rubin SM107 hardware with improved softmax-denominator accumulation.
  • Added optional log-sum-exp output, which can be disabled when unnecessary.
  • Added support for empty key/value tiles and improved attention accumulation.
  • Updated compilation options to select LSE output while retaining amax reductions.

Update (review hardening): the all-ones B tile and its init loop are now parametrized on CTA_MMA (_ONES_ROWS = 16 // CTA_MMA). At today's cga2 this is binary-identical (8 rows); it future-proofs a cga1 config, where the fixed 8-row tile would have had the descriptor walk 1 KiB past the tile into neighboring SMEM — functionally masked only because the epilogue consumes Σ column 0 alone, which is not an invariant worth leaning on. Re-gated: board 34/34, SM100 quick pass.

[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: 83c26c2d-5364-4e3b-a529-2c9b87ae566e

📥 Commits

Reviewing files that changed from the base of the PR and between 420e43f and 31b7fed.

📒 Files selected for processing (1)

• python/cudnn/sdpa/fwd/kernels/prefill_d128_fp8_sm107.py

🚧 Files skipped from review as they are similar to previous changes (1)

• python/cudnn/sdpa/fwd/kernels/prefill_d128_fp8_sm107.py

---

📝 Walkthrough

Walkthrough

The PR adds compile-time exclusive TMEM allocation and updates the SM107 FP8 prefill kernel to accumulate softmax denominators in Sigma columns. It also adds fused row-max handling, shared FP8 ones data, Sigma rescaling, and empty-KV handling.

Changes

SM107 FP8 SDPA prefill

Layer / File(s)	Summary
Exclusive TMEM and kernel layout python/cudnn/frost/tile_dsl/tmem.py, python/cudnn/sdpa/fwd/kernels/prefill_d128_fp8_sm107.py	tmem_alloc supports compile-time exclusive allocation. The kernel uses exclusive 576-column TMEM storage, initializes shared FP8 ones data, and passes it through the MMA interfaces.
Sigma denominator accumulation python/cudnn/sdpa/fwd/kernels/prefill_d128_fp8_sm107.py	The kernel creates ones and Sigma descriptors. BMM2 mainloop and epilogue paths accumulate denominators into Sigma columns.
Softmax reduction and correction python/cudnn/sdpa/fwd/kernels/prefill_d128_fp8_sm107.py	SM107 fused row-max loading replaces the unmasked manual path. RF row-sum updates are removed. Correction rescales O and Sigma and reads the final denominator from Sigma, including empty-KV tiles.

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

Mergeability Score: ⚪ Minimal · up to 31b7f

The PR is merge-ready after normal checks and review; no actionable merge-blocking risk remains.

Sequence Diagram(s)

sequenceDiagram
  participant SM107Kernel
  participant tmem_alloc
  participant nvvm_tcgen05_alloc
  participant Softmax
  SM107Kernel->>tmem_alloc: request exclusive TMEM
  tmem_alloc->>nvvm_tcgen05_alloc: allocate exclusive storage
  SM107Kernel->>Softmax: accumulate Sigma during BMM2
  Softmax->>SM107Kernel: load fused row maxima
  SM107Kernel->>Softmax: rescale O and Sigma
  Softmax-->>SM107Kernel: read final Sigma denominator

Loading

Possibly related PRs

• NVIDIA/cudnn-frontend#576: Introduces the SM107 FP8 prefill kernel modified by this PR.
• NVIDIA/cudnn-frontend#579: Modifies the same SM107 prefill kernel and its compile(..., has_lse=True) interface.

🚥 Pre-merge checks | ✅ 5

✅ Passed checks (5 passed)

Check name	Status	Explanation
Title check	✅ Passed	The title clearly identifies the two main SM107 FP8 kernel optimizations: fused LDTM row-max and row-sum in MMA.
Description check	✅ Passed	The description covers the changes, rationale, compatibility impact, performance results, and validation, but omits several template headings and checklist items.
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-580-31b7fed
Pipeline: 62575532
Targets: frost