[Feature][Attention][UX]: Incorporate Features into Attention Selection

[robertgshaw2-redhat]

🚀 The feature, motivation and pitch

SUMMARY:

• we have default attention backends by priority and a notion of which backend supports what hw
• however, certain features are not considered in this (e.g. fp8 kv cache, e.g. attention sinks)

Recent example, we had test failures because we updated the logic to load kv cache quantization from the model config. But since CUTLASS_MLA is the default backend on B200, we started seeing test failures (since CUTLASS MLA does not support fp8 kv cache) because we were not automatically falling back to FLASHINFER_MLA (which does)

So the proposal is to:

• make sure all attention backends report what features are supported
• update the attention selector to consider these features in the selection

Alternatives

No response

Additional context

No response

Before submitting a new issue...

• Make sure you already searched for relevant issues, and asked the chatbot living at the bottom right corner of the documentation page, which can answer lots of frequently asked questions.

[yifant-code]

Hi I would like to work on this. Here is my initial thought:

1. Add a capabilities descriptor for each attention backend (e.g. flags for fp8 kv cache support, etc.).
2. Update the attention selection logic to use these capabilities in addition to the existing hardware-based priority.
3. Ensure graceful fallback/error when no backend supports the requested features
4. Add tests to cover cases like: fp8 kv cache requested → skip CUTLASS_MLA, pick FLASHINFER_MLA.

Does this direction sound good? If so, could I get assigned to this issue? Thanks.

[robertgshaw2-redhat]

Sounds good @yifant-code!

Could you please work with @MatthewBonanni and @LucasWilkinson on the design? There is some recent PR from Matt on the attention selection logic that this can build upon.

[mgoin]

We already have entries to manage features like kv cache and attention sink support, so something is just registered wrong there if it doesn't actually support fp8

vllm/vllm/v1/attention/backends/mla/cutlass_mla.py

Lines 39 to 43 in 174e39e

	supported_kv_cache_dtypes: ClassVar[list[CacheDType]] = [
	"auto",
	"fp8",
	"fp8_e4m3",
	]

[robertgshaw2-redhat]

my bad, I was not aware of this

[danielafrimi]

Hi, I opened the reverted PR, which should support fetching the KV cache dtype from the config file.
Didn't fully understand if the cutlass MLA supports FP8 as @mgoin wrote, or it's just a registered mistake.
If it's a mistake, can we simply remove the FP8 CacheDType from the supported list?

seems that if we remove it (and it is actually not supported), this backend on B200 won't be chosen.

[MatthewBonanni]

I implemented something like this in #24794, the logic is here:

vllm/vllm/platforms/cuda.py

Line 259 in 185c22b

def get_valid_backends(

First, get_attn_backend_cls validates that the user-selected backend (if there is one) is valid for the configuration, and if not, get_valid_backends produces a list of all backends supporting the configuration, and then the highest-priority backend for the particular hardware is selected.

KV cache dtype is accounted for; supported_kv_cache_dtypes just needs to be set properly as @mgoin mentioned.

Regardless, CUTLASS_MLA should support FP8 KV cache; I added it in #23289. Is it not working correctly?

[danielafrimi]

Hi @MatthewBonanni

As @robertgshaw2-redhat mentions, my reverted PR casue the CI fail
CI: https://buildkite.com/vllm/ci/builds/43393/steps/canvas?jid=019b176c-8881-47f4-aaf9-5dac38740906

It fails on those tests:

1. tests/quantization/test_blackwell_moe.py::test_deepseek_nvfp4_moe_flashinfer_cutlass
2. tests/quantization/test_blackwell_moe.py::test_deepseek_nvfp4_moe_flashinfer_trtllm

My PR fetched from the model configuration the kv_cache dtype (modelopt format) and set it, instead of auto.

[DarkLight1337]

Related PR that cleans things up: #30212

[danielafrimi]

OK, I think I understand the issue with my PR now:

• My PR explicitly reads the kv_cache dtype from the model configuration (ModelOpt format) and sets it accordingly. Previously, if it wasn’t specified, it defaulted to auto.

• The CI failure happens for the model nvidia/DeepSeek-R1-0528-NVFP4-v2, where the checkpoint is NVFP4 but the configured kv_cache dtype is FP8.

• Before my PR, CI was passing because the kv_cache dtype was not correctly read from hf_quant_config.json, so it fell back to auto. In practice, this meant the model was running with an NVFP4 checkpoint and an FP16/BF16 KV cache, which is consistent with what was described in @MatthewBonanni PR (DeepSeek models were checked with FP8 KV cache, but the checkpoints themselves were not NVFP4).

The remaining question is whether Cutlass MLA actually supports running an NVFP4 checkpoint with an FP8 KV cache. This seems to be the real compatibility issue exposed by my change.

Do you agree with that? @robertgshaw2-redhat @mgoin