[TurboQuant] enable FA3/FA4 for prefill paths

tests/evals/gsm8k/configs/Qwen3-4B-TQ-k3v4nc.yaml

@@ -2,4 +2,4 @@ model_name: "Qwen/Qwen3-4B"
 accuracy_threshold: 0.78
 num_questions: 1319
 num_fewshot: 5
-server_args: "--kv-cache-dtype turboquant_k3v4_nc --enforce-eager --max-model-len 4096"
+server_args: "--kv-cache-dtype turboquant_k3v4_nc --max-model-len 4096"

tests/evals/gsm8k/configs/Qwen3-4B-TQ-k8v4.yaml

@@ -2,4 +2,4 @@ model_name: "Qwen/Qwen3-4B"
 accuracy_threshold: 0.80
 num_questions: 1319
 num_fewshot: 5
-server_args: "--kv-cache-dtype turboquant_k8v4 --enforce-eager --max-model-len 4096"
+server_args: "--kv-cache-dtype turboquant_k8v4 --max-model-len 4096"

tests/evals/gsm8k/configs/Qwen3-4B-TQ-t3nc.yaml

@@ -2,4 +2,4 @@ model_name: "Qwen/Qwen3-4B"
 accuracy_threshold: 0.75
 num_questions: 1319
 num_fewshot: 5
-server_args: "--kv-cache-dtype turboquant_3bit_nc --enforce-eager --max-model-len 4096"
+server_args: "--kv-cache-dtype turboquant_3bit_nc --max-model-len 4096"

tests/evals/gsm8k/configs/Qwen3-4B-TQ-t4nc.yaml

@@ -2,4 +2,4 @@ model_name: "Qwen/Qwen3-4B"
 accuracy_threshold: 0.80
 num_questions: 1319
 num_fewshot: 5
-server_args: "--kv-cache-dtype turboquant_4bit_nc --enforce-eager --max-model-len 4096"
+server_args: "--kv-cache-dtype turboquant_4bit_nc --max-model-len 4096"

vllm/v1/attention/backends/flash_attn.py

@@ -255,11 +255,16 @@ class FlashAttentionMetadata:
 def _get_sliding_window_configs(
     vllm_config: VllmConfig,
 ) -> set[tuple[int, int] | None]:
-    """Get the set of all sliding window configs used in the model."""
+    """Get the set of all sliding window configs used in the model.
+
+    Only inspects FlashAttentionImpl layers. Other backends (e.g.
+    TurboQuant, MLA) use their own metadata builders and are skipped.
+    """
     sliding_window_configs: set[tuple[int, int] | None] = set()
     layers = get_layers_from_vllm_config(vllm_config, Attention)
     for layer in layers.values():
-        assert isinstance(layer.impl, FlashAttentionImpl)
+        if not isinstance(layer.impl, FlashAttentionImpl):
+            continue
         sliding_window_configs.add(layer.impl.sliding_window)
     return sliding_window_configs
 

vllm/v1/attention/backends/turboquant_attn.py

@@ -39,6 +39,7 @@
     MultipleOf,
 )
 from vllm.v1.attention.backends.fa_utils import (
+    get_flash_attn_version,
     is_flash_attn_varlen_func_available,
 )
 from vllm.v1.attention.backends.utils import split_decodes_and_prefills
@@ -271,6 +272,9 @@
         self._val_data_bytes = math.ceil(head_size * cfg.effective_value_quant_bits / 8)
         self._n_centroids = cfg.n_centroids if not cfg.key_fp8 else 1
 
+        # Detect flash-attn version (FA2/3/4) for prefill paths.
+        self.fa_version = get_flash_attn_version(head_size=head_size)
+
         # Fixed NUM_KV_SPLITS (grid dims must be constant for cudagraph,
         # and benchmarks show no regression vs dynamic in eager mode).
         vllm_config = get_current_vllm_config()
@@ -513,6 +517,7 @@
                 max_seqlen_k=attn_metadata.max_query_len,
                 softmax_scale=self.scale,
                 causal=True,
+                fa_version=self.fa_version,
             )
 
         # Continuation or no flash_attn: per-request attention.
@@ -562,6 +567,7 @@
                         max_seqlen_k=q_len,
                         softmax_scale=self.scale,
                         causal=True,
+                        fa_version=self.fa_version,
                     )
                 else:
                     q_t = q_seq.transpose(0, 1).contiguous()
@@ -736,6 +742,7 @@
                 max_seqlen_k=seq_len,
                 softmax_scale=self.scale,
                 causal=True,
+                fa_version=self.fa_version,
             )
         else:
             # SDPA fallback: expand KV for GQA, build causal mask