SIGN IN SIGN UP

FlashInfer: Kernel Library for LLM Serving

0 0 188 Python

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.

[![Review Change
Stack](https://storage.googleapis.com/coderabbit_public_assets/review-stack-in-coderabbit-ui.svg)](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