[jcaraban]

Summary

Stacked on #7754. bench_sage --kernel mha4_mxfp8 --e2e called the producers and mha_v4_packed directly, unlike every other mha4_* recipe; it now calls mha_v4, so it includes the K-smoothing gate and host path real callers pay. Kernel-only mode is unchanged.

Performance

MI355X, mha4_mxfp8 hd128 --e2e, median of 3: b1 h16 s4096 92.6 -> 122.2 us, b1 h5 s65536 3.788 -> 3.814 ms (the cost now measured, not a regression).

PR Self-Review

Reviewed revision

619ac9ce562c

Behavior and intent

For BF16 inputs, --e2e for mha4_mxfp8 measures mha_v4, including the K-smoothing gate, like the other mha4_* recipes. FP16 retains the previous packed-producer path.

Validation and unverified scope

MI355X on f926eca (before the FP16 fallback; current head has not been GPU-validated): bench_sage.py --kernel mha4_mxfp8 --b 1 --hq 16 --hk 16 --sq 4096 --sk 4096 --d 128 --dv 128 --report-quality, with and without --e2e: cosine 0.99956 both, hd128 MXFP8 object loaded. Black/Ruff pass. Unverified: CI does not run bench_sage.py; MXFP8 is gfx950-only, so no MI300X run.

PR-specific question

bench_sage.py is a benchmark, not a correctness test: before this PR its mha4_mxfp8 --e2e skipped the K-smoothing gate and under-reported end-to-end time against real mha_v4 callers.

Changed op tests

op_tests/op_benchmarks/triton/bench_sage.py: benchmark only; routes the mha4_mxfp8 --e2e runner through mha_v4 like the other recipes.

[Copilot]

🟡 Changes recommended

The raw mha_v4 path rejects FP16 inputs, breaking the previously supported --kernel mha4_mxfp8 --e2e --dtype fp16 benchmark mode.

1 open issue

• Preserve FP16 support for the MHA v4 MXFP8 raw path · New

What changed in this PR

This PR routes the MXFP8 end-to-end attention benchmark through mha_v4, so timing includes the K-smoothing gate and public wrapper overhead.

Changes:

• Replace direct quantization and packed execution in mha4_mxfp8 --e2e with mha_v4.
• Pass the MXFP8 scale modes as named arguments.
• Leave kernel-only benchmarking unchanged.

File	Description
op_tests/​op_benchmarks/​triton/​bench_sage.py	Route MXFP8 end-to-end timing through the public attention wrapper.

_{Review effort: Balanced}

[aiter-bot]

WAITING ON AUTHOR

PR Self-Review

Revision 8a551759ce36 · mode shadow

Waiting for the PR author. The readiness report will appear after the self-review answers are complete for this revision.

Questions to answer in the PR description

1. What behavior does this PR change, and why is the existing behavior insufficient?
2. What did you validate, and what remains unverified? Include commands, hardware/configuration and results.
3. What behavior would each changed op test fail on before this PR, and why is that coverage not already provided by an existing canonical test?
4. For every changed op_tests/ file, what behavior does it cover and why is it not redundant? Write None if no op test changed.

Edit the PR Self-Review section in the PR description. Include the current revision SHA; this comment refreshes automatically.

Missing or invalid answers

• FAIL — Reviewed revision must contain at least the first 7 characters of the current head SHA

_{Shadow mode reports what would block but keeps the status successful. Set AITER_PR_SELF_REVIEW_MODE=enforce only after reviewing real false positives.}