feat: FP8 output support for CUTLASS MLA paged attention (#2779)
Closes #2778
## Summary
Wires `output_scale` and separate `ElementOut` type through the CUTLASS
MLA paged attention kernel stack to enable native FP8 output,
eliminating a separate bf16→FP8 quant kernel launch.
- Template `MlaSm100` on `ElementOut` separately from `Element` (input
type)
- Pass `output_scale` through `args_from_options`, `runMla`, FFI
binding, and Python API
- Dispatch on input/output dtype in FFI layer (bf16→bf16 or bf16→fp8)
- Pass `output_scale` to reduction kernel for correct split-KV handling
- Add `o_scale` parameter to `BatchMLAPagedAttentionWrapper.run()`
matching the existing ragged prefill API convention
### Motivation
The CUTLASS MLA kernel epilogue already supports arbitrary `ElementOut`
and `output_scale` at the template level, but the wrapper/FFI/Python
layers hardcoded `ElementOut = Element` and didn't expose
`output_scale`. This PR plumbs it through.
### Changes
| File | Change |
|------|--------|
| `cutlass_mla.cuh` | Template `MlaSm100<T, TOut>`, add `output_scale`
to `args_from_options` and `runMla` |
| `cutlass_mla.cu` | Accept `output_scale`, dispatch `bf16→bf16` vs
`bf16→fp8` |
| `flashinfer_mla_binding.cu` | Updated FFI signature |
| `sm100_mla.hpp` | Pass `output_scale` to reduction kernel args |
| `sm100_fmha_mla_reduction.hpp` | Add `output_scale` field, apply
before cast to `ElementOut` |
| `mla.py` | Add `o_scale` parameter to `run()` |
## Test plan
```
============================= test session starts ==============================
platform linux -- Python 3.12.3, pytest-9.0.3, pluggy-1.6.0 -- /root/flashinfer/.venv/bin/python3
cachedir: .pytest_cache
rootdir: /root/flashinfer
configfile: pytest.ini
collecting ... collected 16 items
tests/attention/test_cutlass_mla_fp8_output.py::test_cutlass_mla_fp8_output[fp8_dtype0-1-128-1] PASSED [ 6%]
tests/attention/test_cutlass_mla_fp8_output.py::test_cutlass_mla_fp8_output[fp8_dtype0-1-128-4] PASSED [ 12%]
tests/attention/test_cutlass_mla_fp8_output.py::test_cutlass_mla_fp8_output[fp8_dtype0-1-1024-1] PASSED [ 18%]
tests/attention/test_cutlass_mla_fp8_output.py::test_cutlass_mla_fp8_output[fp8_dtype0-1-1024-4] PASSED [ 25%]
tests/attention/test_cutlass_mla_fp8_output.py::test_cutlass_mla_fp8_output[fp8_dtype0-16-128-1] PASSED [ 31%]
tests/attention/test_cutlass_mla_fp8_output.py::test_cutlass_mla_fp8_output[fp8_dtype0-16-128-4] PASSED [ 37%]
tests/attention/test_cutlass_mla_fp8_output.py::test_cutlass_mla_fp8_output[fp8_dtype0-16-1024-1] PASSED [ 43%]
tests/attention/test_cutlass_mla_fp8_output.py::test_cutlass_mla_fp8_output[fp8_dtype0-16-1024-4] PASSED [ 50%]
tests/attention/test_cutlass_mla_fp8_output.py::test_cutlass_mla_fp8_output_validation_no_out PASSED [ 56%]
tests/attention/test_cutlass_mla_fp8_output.py::test_cutlass_mla_fp8_output_validation_wrong_dtype PASSED [ 62%]
tests/attention/test_cutlass_mla_fp8_output.py::test_cutlass_mla_fp8_output_validation_invalid_scale[0.0] PASSED [ 68%]
tests/attention/test_cutlass_mla_fp8_output.py::test_cutlass_mla_fp8_output_validation_invalid_scale[-1.0] PASSED [ 75%]
tests/attention/test_cutlass_mla_fp8_output.py::test_cutlass_mla_fp8_output_validation_invalid_scale[nan] PASSED [ 81%]
tests/attention/test_cutlass_mla_fp8_output.py::test_cutlass_mla_fp8_output_validation_invalid_scale[inf] PASSED [ 87%]
tests/attention/test_cutlass_mla_fp8_output.py::test_cutlass_mla_bf16_output_unchanged PASSED [ 93%]
tests/attention/test_cutlass_mla_fp8_output.py::test_cutlass_mla_fp8_non_cutlass_backend_rejected PASSED [100%]
=============================== warnings summary ===============================
flashinfer/gemm/kernels/grouped_gemm_masked_blackwell.py:2060
/root/flashinfer/flashinfer/gemm/kernels/grouped_gemm_masked_blackwell.py:2060: DeprecationWarning: tcgen05.OperandMajorMode is deprecated, use cute.nvgpu.OperandMajorMode instead
a_major_mode: tcgen05.OperandMajorMode,
flashinfer/gemm/kernels/grouped_gemm_masked_blackwell.py:2062
/root/flashinfer/flashinfer/gemm/kernels/grouped_gemm_masked_blackwell.py:2062: DeprecationWarning: tcgen05.OperandMajorMode is deprecated, use cute.nvgpu.OperandMajorMode instead
b_major_mode: tcgen05.OperandMajorMode,
flashinfer/gdn_kernels/blackwell/gated_delta_net_chunked.py:99
flashinfer/gdn_kernels/blackwell/gated_delta_net_chunked.py:99
/root/flashinfer/flashinfer/gdn_kernels/blackwell/gated_delta_net_chunked.py:99: DeprecationWarning: tcgen05.OperandMajorMode is deprecated, use cute.nvgpu.OperandMajorMode instead
from cutlass.cute.nvgpu.tcgen05 import OperandMajorMode
-- Docs: https://docs.pytest.org/en/stable/how-to/capture-warnings.html
======================= 16 passed, 4 warnings in 34.24s ========================
```
<!-- This is an auto-generated comment: release notes by coderabbit.ai
-->
## Summary by CodeRabbit
* **New Features**
* FP8 fused-output support for MLA paged attention: callers may provide
an output scale to produce FP8 outputs on the CUTLASS backend.
* **Bug Fixes / Validation**
* Added validation for output-scale (finite, positive, representable)
and enforced that using a scale requires an FP8 output tensor;
non-CUTLASS backends reject output-scale.
* **Tests**
* Added tests for FP8 quantization correctness, validation error cases,
and backend compatibility.
[](https://app.coderabbit.ai/change-stack/flashinfer-ai/flashinfer/pull/2779)
<!-- end of auto-generated comment: release notes by coderabbit.ai -->
---------
Co-authored-by: Claude Opus 4.6 <noreply@anthropic.com>
Co-authored-by: mingyangw <mingyangw@nvidia.com> C
Carl Y committed
1bb79d4a2039f9c83caad3855fce2f6db65f6200
Parent: 2d0e0ef
Committed by GitHub <noreply@github.com>
on 5/11/2026, 5:03:17 PM