PR #40092: [TurboQuant] enable FA3/FA4 for prefill paths

• Repository: vllm-project/vllm
• Author: huangzhilin-hzl
• State: open | merged: False
• Link: https://github.com/vllm-project/vllm/pull/40092

Description (problem / solution / changelog)

Purpose

Resolves part of https://github.com/vllm-project/vllm/issues/40069 (Backend Coverage: extend flash_attn_varlen_func support to FA3/4).

Two issues fixed:

1. FA version passthrough: TurboQuant prefill paths call flash_attn_varlen_func without the fa_version kwarg, so on Hopper (SM90) the call defaults to FA2 instead of leveraging FA3, and on Blackwell (SM100) it misses FA4 entirely. The standard FlashAttention backend already detects and passes fa_version at init time; this PR aligns TurboQuant to the same pattern.

2. Mixed-backend assert fix: _get_sliding_window_configs() in flash_attn.py asserts all Attention layers are FlashAttentionImpl. When kv_cache_dtype_skip_layers routes some layers to a different backend (e.g. TurboQuant), this assert fails. Fixed by skipping non-FA layers, since they use their own metadata builders.

Test Plan

# 1. Unit tests
python -m pytest tests/quantization/test_turboquant.py -v

# 2. GSM8K correctness eval (all 4 TQ presets)
python -m pytest -s -v tests/evals/gsm8k/test_gsm8k_correctness.py \
    --config-list-file=tests/evals/gsm8k/configs/models-turboquant.txt

# 3. E2E inference with CUDAGraph (no enforce_eager, validates assert fix)
CUDA_VISIBLE_DEVICES=0 HF_HUB_OFFLINE=1 python -c "
from vllm import LLM, SamplingParams
for dtype in ['turboquant_k8v4', 'turboquant_3bit_nc']:
    llm = LLM(model='Qwen/Qwen3-4B', kv_cache_dtype=dtype,
              max_model_len=2048, gpu_memory_utilization=0.5)
    outputs = llm.generate(['What is 2+2?'], SamplingParams(max_tokens=32))
    print(f'{dtype}: {outputs[0].outputs[0].text[:80]}')
    del llm
"

Test Result

Hardware: NVIDIA H20 (SM90 / Hopper)

FA version detection

FA version for head_size=128: 3   (was: unspecified, defaulting to FA2)
FA version for head_size=256: 3

Unit tests

114 passed, 6 failed (pre-existing rotation matrix atol issues, unrelated)

Confirmed pre-existing: same 6 failures on unmodified code via git stash / re-run.

E2E inference with CUDAGraph (enforce_eager=False)

Validates both the FA3 passthrough and the assert fix (AOT schedule path is entered).

GSM8K correctness eval (Qwen3-4B, 1319 questions, 5-shot)

Note: t3nc failed in batch run due to GPU memory from zombie processes, passed when run alone.

• 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.

Changed files

• tests/evals/gsm8k/configs/Qwen3-4B-TQ-k3v4nc.yaml (modified, +1/-1)
• tests/evals/gsm8k/configs/Qwen3-4B-TQ-k8v4.yaml (modified, +1/-1)
• tests/evals/gsm8k/configs/Qwen3-4B-TQ-t3nc.yaml (modified, +1/-1)
• tests/evals/gsm8k/configs/Qwen3-4B-TQ-t4nc.yaml (modified, +1/-1)
• vllm/v1/attention/backends/flash_attn.py (modified, +7/-2)
• vllm/v1/attention/backends/turboquant_attn.py (modified, +7/-0)