fp8 support

[endurehero]

Functionality

Support FP8 WGMMA based on the async pipeline design of FlashMLA. The TransV part draws on the implementation of SmemTranspose64x64 in Fa3.
Currently, Q/K/V only support symmetric PerTensor quantization. Since the maximum value of P does not exceed 1, the f32tofp8_cast is directly used for quantization.

Performance

cuda driver version: 535.183.06
nvcc version: 12.8
torch version: 2.6

On the H20, MLA typically demonstrate a high degree of arithmetic intensity. Consequently, the Memory Floating - point Utilization (MFU) is employed as a performance metric.

On the H800, MLA typically encounter memory-bound situations. Consequently, the Memory Bandwidth Utilization (MBU) metric is adopted to evaluate the performance of the kernel. There is still a lot of room for optimization on the H800. Look forward to working together.

Reproduction

python3 ./tests/test_flash_mla.py --dtype e4m3

[sijiac]

awesome, did you mind adding a compile flag to save the time when FP8 is not needed? Thanks

[endurehero]

awesome, did you mind adding a compile flag to save the time when FP8 is not needed? Thanks

Of course. Already Done

[beginlner]

Great work! However, I can’t merge this PR at the moment because, based on our tests, per-sequence kvcache scaling significantly reduces accuracy for MLA.

[endurehero]

Great work! However, I can’t merge this PR at the moment because, based on our tests, per-sequence kvcache scaling significantly reduces accuracy for MLA.

What about the granularity of PerPageBlock? I can easily adapt it

[beginlner]

What about the granularity of PerPageBlock? I can easily adapt it

We think PerPageBlock is neither enough. kv_rope (64) needs to be bf16.

[endurehero]

What about the granularity of PerPageBlock? I can easily adapt it

We think PerPageBlock is neither enough. kv_rope (64) needs to be bf16.

Got it!

[moses3017]

What about the granularity of PerPageBlock? I can easily adapt it

We think PerPageBlock is neither enough. kv_rope (64) needs to be bf16.

How about Qnope and Knope using 8-bit quantization, while Qrope and Krope maintain 16-bit data types?

[beginlner]

It's acceptable for Qnope and Knope to use per-(1 token × 128 channel) 8-bit quantization, while Qrope and Krope retain 16-bit precision.

[shinezyy]

It's acceptable for Qnope and Knope to use per-(1 token × 128 channel) 8-bit quantization, while Qrope and Krope retain 16-bit precision.

The outliners in RoPE cache are also discussed in this paper https://arxiv.org/pdf/2502.01563

Can we add a hadamard transform right after RoPE to distribute outliners to multiple head dims? (https://arxiv.org/abs/2404.00456)

[TheTinyTeddy]

It's acceptable for Qnope and Knope to use per-(1 token × 128 channel) 8-bit quantization, while Qrope and Krope retain 16-bit precision.

What precision should S×V be? BF16×BF16 or BF16×FP8 or FP8×FP8 per(1 token × 128 channel)?

[hypdeb]

Great work! However, I can’t merge this PR at the moment because, based on our tests, per-sequence kvcache scaling significantly reduces accuracy for MLA.

@beginlner Hello there, on which workload did you observe these accuracy issues?

[MicroZHY]

Hi @endurehero,

Thank you very much for sharing this impressive FP8 implementation and for the detailed performance numbers!

I have two quick questions that would help me understand the design better:

1. Branch difference
   Could you kindly clarify the relationship between will_fp8_mr and the earlier will_fp8 branch? I noticed both names appear in the commit history, and I’d love to know what motivated the new branch and which improvements or fixes it contains.

2. Necessity of TransV
   I see that the new FP8 path relies on the TransV routine, which borrows the 64×64 shared-memory transpose from FlashAttention-3. Would it be possible to briefly explain why TransV is indispensable for FP8 correctness or performance? I’m curious whether the same result could be achieved with a different layout or if this is a hard requirement for the WGMMA pipeline.

Thanks again for your time and for open-sourcing this work!