Nice work!

ref #1156

hold on, we should also verify on H100

Note: the current fp8 prefill kernels in flashinfer is a mixed-precision kernel (upcast fp8 K/V to fp16 inside kernels and use fp16 tensor cores), which is slow because of the online dequantization overhead, and it's different from the triton kernel logic here (where you downcast fp16 query to fp8 and use fp8 tensor cores).

The native fp8 prefill attention kernel in flashinfer will land together with fa3, where we use fp8 tensor cores, but the accuracy landscape might change (and might align with your triton implementation).

Note: the current fp8 prefill kernels in flashinfer is a mixed-precision kernel (upcast fp8 K/V to fp16 inside kernels and use fp16 tensor cores), which is slow because of the online dequantization overhead, and it's different from the triton kernel logic here (where you downcast fp16 query to fp8 and use fp8 tensor cores).

The native fp8 prefill attention kernel in flashinfer will land together with fa3, where we use fp8 tensor cores, but the accuracy landscape might change (and might align with your triton implementation).

Looking forward to it very much!