vllm - 💡(How to fix) Fix [Bug]: Turboquant attention crashes on A100 when serving BF16 models with FP8 KV cache [1 participants]

File "triton/language/extra/cuda/utils.py", line 104, in convert_custom_float8_sm80 return convert_custom_float8(arg, dst_ty, fp_downcast_rounding, has_minx2=True, ...)

File "triton/language/extra/cuda/utils.py", line 94, in convert_custom_float8 assert arg.type.scalar.is_fp16() or arg.type.scalar.is_fp32() AssertionError

triton.compiler.errors.CompilationError: at 45:12: k_fp8 = k_vals.to(tl.float8e4b15) if FP8_E4B15 else k_vals.to(tl.float8e4nv) ^

TL;DR

Cast the loaded k_vals to FP32 before converting to FP8 in triton_turboquant_store.py to resolve the AssertionError on A100 (SM80) with BF16 models.

Guidance

• The issue arises from the direct casting of KV values to FP8, which fails when the model uses BF16 on A100 (SM80) due to Triton's software FP8 conversion path only accepting FP16 or FP32 inputs.
• To verify the fix, run the provided reproduction command with the modified triton_turboquant_store.py file.
• The suggested fix involves changing the line k_vals = tl.load(Key_ptr + base + d_offs, mask=d_mask, other=0.0) to k_vals = tl.load(Key_ptr + base + d_offs, mask=d_mask, other=0.0).to(tl.float32) to ensure a lossless conversion for all input dtypes.
• This fix should be applied to the triton_turboquant_store.py file at line 188.

Example

The modified code snippet would look like this:

k_vals = tl.load(Key_ptr + base + d_offs, mask=d_mask, other=0.0).to(tl.float32)
k_fp8 = k_vals.to(tl.float8e4b15) if FP8_E4B15 else k_vals.to(tl.float8e4nv)

Notes

This fix is specific to the A100 (SM80) architecture and may not be necessary for other architectures like H100 (SM89+), which have native FP8 hardware support.

Recommendation

Apply the suggested workaround by casting k_vals to FP32 before converting to FP8, as this is a safe and lossless conversion for all input dtypes.