[Bugfix] Refactor Flashinfer TRTLLM attention kernel selection logic

[elvischenv]

Purpose

#23647 always quantizes query if kv cache type is set to FP8, and will use TRTLLM attention kernel. However, there are lots of cases that do not support TRTLLM attention.

• For the fix, we still try quantizing query first and see if the TRTLLM attention is finally used or not. If it is not used, then we set the query dtype back to the model dtype.
  • This fix still need [Bugfix] Enable FP8 KV cache for FlashInfer and Triton backend on non-sm100 GPUs #24577 to support FP8 KV cache for flashinfer generally
• Also, clean up the logic for the TRTLLM attn kernel selection. We should separate the logic for "support" and "use". The env VLLM_USE_TRTLLM_ATTENTION should determine the "use" path instead of the "support" path.
• Fixed the hanging issue([Bug]: Llama 3.3 70B hangs with full cuda graph for decode-only #24689) when cudagraph is enabled and --max-model-len is not specified.

Test Plan && Test Result

SM100, kv_cache_dtype=fp8

Using TRTLLM attention (query is quantized).
vllm ({'pretrained': 'nvidia/Llama-3.3-70B-Instruct-FP8', 'kv_cache_dtype': 'fp8', 'tensor_parallel_size': 1, 'max_model_len': 2048, 'trust_remote_code': True}), gen_kwargs: (temperature=0.0), limit: 500.0, num_fewshot: 5, batch_size: 200
|Tasks|Version|     Filter     |n-shot|  Metric   |   |Value|   |Stderr|
|-----|------:|----------------|-----:|-----------|---|----:|---|-----:|
|gsm8k|      3|flexible-extract|     5|exact_match|↑  |0.938|±  |0.0108|
|     |       |strict-match    |     5|exact_match|↑  |0.906|±  |0.0131|

SM100, kv_cache_dtype=fp8, VLLM_USE_TRTLLM_ATTENTION=0

VLLM_USE_TRTLLM_ATTENTION is set to 0
vllm ({'pretrained': 'nvidia/Llama-3.3-70B-Instruct-FP8', 'kv_cache_dtype': 'fp8', 'tensor_parallel_size': 1, 'max_model_len': 2048, 'trust_remote_code': True}), gen_kwargs: (temperature=0.0), limit: 500.0, num_fewshot: 5, batch_size: 200
|Tasks|Version|     Filter     |n-shot|  Metric   |   |Value|   |Stderr|
|-----|------:|----------------|-----:|-----------|---|----:|---|-----:|
|gsm8k|      3|flexible-extract|     5|exact_match|↑  |0.930|±  |0.0114|
|     |       |strict-match    |     5|exact_match|↑  |0.904|±  |0.0132|

SM120, kv_cache_dtype=fp8, VLLM_ATTENTION_BACKEND=FLASHINFER

(EngineCore_DP0 pid=1252) INFO 09-10 18:44:05 [cuda.py:285] Using FlashInfer backend on V1 engine.
vllm ({'pretrained': 'nvidia/Llama-3.3-70B-Instruct-FP8', 'kv_cache_dtype': 'fp8', 'tensor_parallel_size': 1, 'max_model_len': 2048, 'trust_remote_code': True}), gen_kwargs: (temperature=0.0), limit: 500.0, num_fewshot: 5, batch_size: 200
|Tasks|Version|     Filter     |n-shot|  Metric   |   |Value|   |Stderr|
|-----|------:|----------------|-----:|-----------|---|----:|---|-----:|
|gsm8k|      3|flexible-extract|     5|exact_match|↑  |0.948|±  |0.0099|
|     |       |strict-match    |     5|exact_match|↑  |0.904|±  |0.0132|

---

Essential Elements of an Effective PR Description Checklist

• The purpose of the PR, such as "Fix some issue (link existing issues this PR will resolve)".
• The test plan, such as providing test command.
• The test results, such as pasting the results comparison before and after, or e2e results
• (Optional) The necessary documentation update, such as updating supported_models.md and examples for a new model.
• (Optional) Release notes update. If your change is user facing, please update the release notes draft in the Google Doc.

[mergify]

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

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

[elvischenv]

cc @ProExpertProg @mgoin @pavanimajety for review, thanks!

[ProExpertProg]

Looks good overall, just one nit

[ProExpertProg]

Sorry, two more notes. Does this mean that when has_sinks is enabled, attention+quant fusion won't work?

[ProExpertProg]

You shouldn't read envs in cached functions, please separate into force_use_trtllm_attention (uncached) and _force_use_trtllm_attention (cached, takes env var as input)

[elvischenv]

Addressed the comments above.

Does this mean that when has_sinks is enabled, attention+quant fusion won't work?

• If kv=auto, all the things work fine, will use TRTLLM BF16-qkv BF16-out kernel
• If kv=fp8, by default it will always quantize query and use TRTLLM FP8-qkv
  • In this case, we found some accuracy issues in the FP8-qkv BF16-out sinks kernel. That has been fixed in the TRTLLM upstream and need to be propagated to Flahsinfer and vLLM. For now just raise an error and suggest user to use kv=auto.
  • There is a WAR for using BF16-q FP8-kv BF16-out kernel, by setting VLLM_FLASHINFER_DISABLE_Q_QUANTIZATION, introducing by [flashinfer] [kernel] support for fp8 kv cache for trtllm prefill attention #24197.

Back to attention+quant fusion, AFAIR, only gpt-oss need attn sinks, and we haven't started quantizing the attn output for it. If we want to enable attention+quant fusion for gpt-oss, we have to at least

1. Quantize the gpt-oss model.
2. Ensure the FP8-qkv FP8/NVFP4-out attn sinks kernels work and have good accuracy on gpt-oss.

So for now, we need to fix FP8-qkv BF16-out sinks kernel first, and then verify the FP8-qkv FP8/NVFP4-out kernel if we need them.