[Bugfix] Fix KV scales inconsistency in fp8 MLA & FlashInfer kv_cache_dtype "auto" leading to gibberish

[andylolu2]

Purpose

This PR fixes the following issues:

• FlashInfer + kv_cache_dtype "auto" generates giberrish when layer._[qkv]_scale != 1.0
  • Bug is FI applies the layer._[qkv]_scale unconditionally, even when the QKV values are in unscaled bf16.
  • This applies to both normal & MLA attention paths.
• KV cache scales not properly handed when using MLA + fp8.
  • In MLA, the KV latents necessarily must use the same quantization scale for K & V, so only one of layer._k_scale or layer._v_scale should be used, not both. The current implementation sometimes assumes layer._k_scale is used, other times assumes layer._v_scale or layer._k_scale * layer._v_scale is used, which is inconsistent and leads to bad generations.
  • In this PR I choose the only use layer._k_scale. layer._v_scale is completely ignored when using MLA.
• The CUTLASS_MLA backend says it supports fp8 kv cache but there's no logic to handle the quantization scales properly, so disabling its support for fp8 for now until it gets implemented.

To summarize the situation of fp8 MLA scales:

• q_scale -> Meant for quantizing q_mqa, not q_mha. q_mha currently has no corresponding scales and is naively casted to fp8 if use_fp8_prefill (code reference).
• k_scale -> Meant for quantizing kv_latents, not k_mha or v_mha. k_mha and v_mha currently has no corresponding scales and is naively casted to fp8 if use_fp8_prefill (code reference).
• v_scale -> Completely unused.

Test Plan

The current tests are passing because the qkv_scales are mocked to be 1.0, which silently avoids this bug. Updated the tests to remove the assumption that [qkv]_scales are 1.0.

To assert the layer._v_scale is not used in MLA, I set it to NaN in the mla tests to ensure wrong results if they are ever used.

Test Result

Updated tests are passing.

[gemini-code-assist]

Code Review

This pull request addresses a critical bug where FlashInfer attention backends incorrectly applied quantization scales even when the KV cache was not FP8-quantized, resulting in corrupted output. The fix correctly restricts the application of these scales to cases where the KV cache data type is indeed FP8. The accompanying test modifications, which set mock layer scales to non-unity values, are appropriate for verifying the fix. The changes appear correct and effectively resolve the issue.

[mergify]

Documentation preview: https://vllm--37054.org.readthedocs.build/en/37054/

[andylolu2]

@gemini review

[gemini-code-assist]

Code Review

This pull request addresses several correctness issues related to FP8 quantization scales in FlashInfer and MLA backends. The changes correctly handle scales for decode paths and disable FP8 support for the broken CUTLASS MLA backend. However, the fix is incomplete as the MLA prefill path for FlashInfer and Triton backends still lacks proper FP8 scale handling, which is a critical issue. I've added a comment with details on the missing fix.

[mergify]

This pull request has merge conflicts that must be resolved before it can be
merged. Please rebase the PR, @andylolu2.

https://docs.github.com/en/pull-requests/collaborating-with-pull-requests/working-with-forks/syncing-a-fork

[MatthewBonanni]

Thanks for the fix! Just some small comments

[MatthewBonanni]

LGTM, thanks for the fix and thanks for improving the test coverage!