[Perf] Refactor cudagraph_support to enable full CUDA graphs for spec decoding with FlashInfer

[benchislett]

Purpose

Revised implementation of #26937

This PR makes _cudagraph_support a private member and uses get_cudagraph_support(vllm_config, kv_cache_spec). Also updates _check_and_update_cudagraph_mode to consider support per-backend, per-kv-group.

TRTLLM-gen kernels support full cuda graphs, but are only used with FlashInfer on Blackwell under certain conditions.
It might not be safe to change FlashInfer's cudagraph_support to UNIFORM_BATCH always, but we can still set it when we know TRTLLM-gen backend will be used.

Also update the docs to reflect the FlashInfer cuda graph compatibility, and fill in the missing entry for FlashInferMLA.

FIX #26856

Test Plan

See #26937 for functional correctness testing / benchmarking. Rerunning on this branch gives the same results.

Local test run passes for tests/v1/attention.

[mergify]

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

[gemini-code-assist]

Code Review

This pull request introduces a well-designed refactoring to enable more flexible and dynamic CUDA graph support for attention backends. By making cudagraph_support a private member and introducing a new get_cudagraph_support method, the code now dynamically determines the CUDA graph capability on a per-backend, per-KV-group basis. This change is crucial for enabling full CUDA graph support for speculative decoding with FlashInfer on specific hardware like Blackwell. The updates to _check_and_update_cudagraph_mode and the corresponding documentation changes are clear and correct. Overall, this is a solid performance enhancement with clean implementation.

[LucasWilkinson]

LGTM

Id like to work towards reverting #27427 (and move back to this being an instance property) in the future; but we need broader cudagraph refactors to get there

[vadiklyutiy]

Before we use use_trtllm_attention for checking both prefill and decode.
Right now seem use_trtllm_attention is using for checking prefill only and can_use_trtllm_attention .

May we refactor:

• use proper name like use_trtllm_prefill_attn and use_trtllm_decode_attn
• remove from use_trtllm_attention processing of decode case

Maybe it's worth to do in separate PR

[vadiklyutiy]

Before we use use_trtllm_attention for checking both prefill and decode.
Right now seem use_trtllm_attention is using for checking prefill only and can_use_trtllm_attention .

May we refactor:

• use proper name like use_trtllm_prefill_attn and use_trtllm_decode_attn
• remove from use_trtllm_attention processing of decode case

Maybe it's worth to do in separate PR

[vadiklyutiy]

One more style
Is there some reason to hold [_]cudagraph_support and get_cudagraph_support in *MetaBuilder classes, maybe *Backend(AttentionBackend) is better place?

[benchislett]

@vadiklyutiy

• Currently, we have to force the use of TRTLLM attention for decodes if it is supported, so that we can statically decide whether the backend can support UNIFORM_BATCH cuda graphs or not. Ideally, we would only lock-in to that decision when we know that UNIFORM_BATCH is actually being used, and otherwise we would switch dynamically. So I am hoping to keep use_trtllm_attention untouched so that we can re-enable in the near future.
• I don't know exactly why it's defined on the builder and not the backend, but this seems like a reasonable convention for now. I expect it will be refactored separately in the future.
  These interfaces are often undergoing changes, so I am trying to keep this PR's diff minimal until they become more cemented.