MXFP8 Perf Improvement

[yanzhuo607]

Before submitting

• I agree to license this contribution under the terms of LICENSE.txt.
• I ran pre-commit run and committed any formatting changes.

Affected area

Summary

Why

Related issues

API and compatibility impact

Testing

Summary by CodeRabbit

• New Features

  • Added support for LPT scheduling with causal MXFP8 workloads on SM100.
  • MXFP8 execution now reports output maximum values directly from the kernel.

• Bug Fixes

  • Improved MXFP8 output scaling metadata accuracy by calculating maximum absolute output values during execution.
  • Ensured FP8 and MXFP8 statistics buffers reset on the active launch stream.

• Documentation

  • Updated SM100 scheduling and synchronization documentation to reflect current supported options.

[coderabbitai]

📝 Walkthrough

Walkthrough

SM100 SDPA compilation adds LPT scheduling support for causal MXFP8 paths. Statistic-buffer resets now run on the kernel launch stream. MXFP8 kernels compute output amplitude maxima and update a provided buffer. Obsolete comments are removed.

Changes

SM100 SDPA updates

Layer / File(s)	Summary
Scheduler policy selection python/cudnn/sdpa/fwd/config_sm100.py, python/cudnn/sdpa/fwd/api_dsl.py	SM100 validation accepts SCHED_NATURAL and SCHED_LPT. Causal MXFP8 compilation selects LPT when the natural default policy is configured.
Stream-ordered statistic resets python/cudnn/sdpa/fwd/api_dsl.py	PyTorch operations use the launch stream context. MXFP8 amax_o, and per-tensor FP8 amax_s and amax_o, are reset on that stream. Host-side MXFP8 output scanning is removed.
MXFP8 kernel amax reporting python/cudnn/sdpa/fwd/kernels/prefill_d128_mxfp8_sm100.py	The amax_o buffer flows through the launcher and kernel. The correction warp group computes valid-row maximum absolute outputs and atomically updates the buffer.
SM100 kernel documentation cleanup python/cudnn/sdpa/fwd/kernels/prefill_d128_f16_sm100.py, python/cudnn/sdpa/fwd/kernels/prefill_d128_fp8_sm100.py, python/cudnn/sdpa/fwd/kernels/prefill_d128_mxfp8_sm100.py	Obsolete scheduling, compiler-option, and cfence comments are removed or reworded.

Estimated code review effort: 3 (Moderate) | ~25 minutes

Sequence Diagram(s)

sequenceDiagram
  participant api_dsl
  participant _host
  participant _kernel
  participant _correction_warp_group
  api_dsl->>api_dsl: Clear amax_o on launch stream
  api_dsl->>_host: Pass amax_o buffer
  _host->>_kernel: Forward amax_o_tensor
  _kernel->>_correction_warp_group: Forward amax_o_tensor
  _correction_warp_group->>_correction_warp_group: Compute valid-row maximum absolute output
  _correction_warp_group->>amax_o: Atomically update global amax buffer

Loading

🚥 Pre-merge checks | ✅ 4 | ❌ 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Description check	⚠️ Warning	The description includes the checklist but leaves all required change, rationale, impact, issue, and testing sections incomplete.	Complete the affected area, summary, rationale, related issues, API impact, and testing sections with specific project details.

✅ Passed checks (4 passed)

Check name	Status	Explanation
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.
Title check	✅ Passed	The title clearly identifies the main change as an MXFP8 performance improvement.

✨ Finishing Touches

🧪 Generate unit tests (beta)

• Create PR with unit tests

---

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

[coderabbitai]

Actionable comments posted: 1

🤖 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/api_dsl.py`:
- Around line 978-983: Update the FP8 amax reset in the code constructing
amax_o_buf so amax_o_buf.zero_() is enqueued on current_stream before the MXFP8
kernel launch, matching the stream ordering used by the execute-time handle and
other amax resets.

🪄 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: 861c26f8-da62-40ed-80ab-9a7fdf0c2ba1

📥 Commits

Reviewing files that changed from the base of the PR and between ac9356f and 4e6c7ef.

📒 Files selected for processing (5)

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

[vedaanta]

looks good.

Maybe we dont need amax o here

[vedaanta]

@cudnn-ci-bot run

[cudnn-ci-bot]

🚀 Running mirror pipeline

Branch: cudnn-gh/pr-481-529abf8
Pipeline: 61124241