A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 4, 18:45 UTC
[TurboQuant] enable FA3/FA4 for prefill paths (#40092)
Signed-off-by: 墨楼 <huangzhilin.hzl@antgroup.com> Co-authored-by: 墨楼 <huangzhilin.hzl@antgroup.com> Co-authored-by: Claude Opus 4.6 (1M context) <noreply@anthropic.com> Co-authored-by: Codex <codex@openai.com>
Julian Huang committed
Apr 23, 2026 at 12:35 UTC
fe9c3d6c5f66c873d196800384ed6880687b9e52
6 files changed
+55
-15
tests/evals/gsm8k/configs/Qwen3-4B-TQ-k3v4nc.yaml
+1
-1
@@ -2,4 +2,4 @@ model_name: "Qwen/Qwen3-4B"
2
accuracy_threshold: 0.78
3
num_questions: 1319
4
num_fewshot: 5
5
-server_args: "--kv-cache-dtype turboquant_k3v4_nc --enforce-eager --max-model-len 4096"
5
+server_args: "--kv-cache-dtype turboquant_k3v4_nc --max-model-len 4096"
tests/evals/gsm8k/configs/Qwen3-4B-TQ-k8v4.yaml
+1
-1
@@ -2,4 +2,4 @@ model_name: "Qwen/Qwen3-4B"
2
accuracy_threshold: 0.80
3
num_questions: 1319
4
num_fewshot: 5
5
-server_args: "--kv-cache-dtype turboquant_k8v4 --enforce-eager --max-model-len 4096"
5
+server_args: "--kv-cache-dtype turboquant_k8v4 --max-model-len 4096"
tests/evals/gsm8k/configs/Qwen3-4B-TQ-t3nc.yaml
+1
-1
@@ -2,4 +2,4 @@ model_name: "Qwen/Qwen3-4B"
2
accuracy_threshold: 0.75
3
num_questions: 1319
4
num_fewshot: 5
5
-server_args: "--kv-cache-dtype turboquant_3bit_nc --enforce-eager --max-model-len 4096"
5
+server_args: "--kv-cache-dtype turboquant_3bit_nc --max-model-len 4096"
tests/evals/gsm8k/configs/Qwen3-4B-TQ-t4nc.yaml
+1
-1
@@ -2,4 +2,4 @@ model_name: "Qwen/Qwen3-4B"
2
accuracy_threshold: 0.80
3
num_questions: 1319
4
num_fewshot: 5
5
-server_args: "--kv-cache-dtype turboquant_4bit_nc --enforce-eager --max-model-len 4096"
5
+server_args: "--kv-cache-dtype turboquant_4bit_nc --max-model-len 4096"
vllm/v1/attention/backends/flash_attn.py
+7
-2
@@ -255,11 +255,16 @@ class FlashAttentionMetadata:
255
def _get_sliding_window_configs(
256
vllm_config: VllmConfig,
257
) -> set[tuple[int, int] | None]:
258
- """Get the set of all sliding window configs used in the model."""
258
+ """Get the set of all sliding window configs used in the model.
259
+
260
+ Only inspects FlashAttentionImpl layers. Other backends (e.g.
261
+ TurboQuant, MLA) use their own metadata builders and are skipped.
262
+ """
263
sliding_window_configs: set[tuple[int, int] | None] = set()
264
layers = get_layers_from_vllm_config(vllm_config, Attention)
265
for layer in layers.values():
262
- assert isinstance(layer.impl, FlashAttentionImpl)
266
+ if not isinstance(layer.impl, FlashAttentionImpl):
267
+ continue
268
sliding_window_configs.add(layer.impl.sliding_window)
269
return sliding_window_configs
270
vllm/v1/attention/backends/turboquant_attn.py
+44
-9
@@ -39,6 +39,7 @@ from vllm.v1.attention.backend import (
39
MultipleOf,
40
)
41
from vllm.v1.attention.backends.fa_utils import (
42
+ get_flash_attn_version,
43
is_flash_attn_varlen_func_available,
44
)
45
from vllm.v1.attention.backends.utils import split_decodes_and_prefills
@@ -271,6 +272,9 @@ class TurboQuantAttentionImpl(AttentionImpl["TurboQuantMetadata"]):
272
self._val_data_bytes = math.ceil(head_size * cfg.effective_value_quant_bits / 8)
273
self._n_centroids = cfg.n_centroids if not cfg.key_fp8 else 1
274
275
+ # Detect flash-attn version (FA2/3/4) for prefill paths.
276
+ self.fa_version = get_flash_attn_version(head_size=head_size)
277
+
278
# Fixed NUM_KV_SPLITS (grid dims must be constant for cudagraph,
279
# and benchmarks show no regression vs dynamic in eager mode).
280
vllm_config = get_current_vllm_config()
@@ -278,6 +282,43 @@ class TurboQuantAttentionImpl(AttentionImpl["TurboQuantMetadata"]):
282
vllm_config.attention_config.tq_max_kv_splits_for_cuda_graph
283
)
284
285
+ def _flash_attn_varlen(
286
+ self,
287
+ q: torch.Tensor,
288
+ k: torch.Tensor,
289
+ v: torch.Tensor,
290
+ cu_seqlens_q: torch.Tensor,
291
+ cu_seqlens_k: torch.Tensor,
292
+ max_seqlen_q: int,
293
+ max_seqlen_k: int,
294
+ ) -> torch.Tensor:
295
+ # fa_utils.get_flash_attn_version() returns None on backends that
296
+ # should not pass an explicit fa_version kwarg.
297
+ if self.fa_version is None:
298
+ return flash_attn_varlen_func(
299
+ q=q,
300
+ k=k,
301
+ v=v,
302
+ cu_seqlens_q=cu_seqlens_q,
303
+ cu_seqlens_k=cu_seqlens_k,
304
+ max_seqlen_q=max_seqlen_q,
305
+ max_seqlen_k=max_seqlen_k,
306
+ softmax_scale=self.scale,
307
+ causal=True,
308
+ )
309
+ return flash_attn_varlen_func(
310
+ q=q,
311
+ k=k,
312
+ v=v,
313
+ cu_seqlens_q=cu_seqlens_q,
314
+ cu_seqlens_k=cu_seqlens_k,
315
+ max_seqlen_q=max_seqlen_q,
316
+ max_seqlen_k=max_seqlen_k,
317
+ softmax_scale=self.scale,
318
+ causal=True,
319
+ fa_version=self.fa_version,
320
+ )
321
+
322
def _ensure_on_device(self, layer, device):
323
"""One-time derivation of TQ buffers (rotation matrix, midpoints).
324
@@ -503,7 +544,7 @@ class TurboQuantAttentionImpl(AttentionImpl["TurboQuantMetadata"]):
544
# max_query_len == max_seq_len means no request has prior cached KV.
545
# Both are Python ints — no GPU sync.
546
if _HAS_FLASH_ATTN and attn_metadata.max_query_len == attn_metadata.max_seq_len:
506
- return flash_attn_varlen_func(
547
+ return self._flash_attn_varlen(
548
q=query,
549
k=key,
550
v=value,
@@ -511,8 +552,6 @@ class TurboQuantAttentionImpl(AttentionImpl["TurboQuantMetadata"]):
552
cu_seqlens_k=attn_metadata.query_start_loc,
553
max_seqlen_q=attn_metadata.max_query_len,
554
max_seqlen_k=attn_metadata.max_query_len,
514
- softmax_scale=self.scale,
515
- causal=True,
555
)
556
557
# Continuation or no flash_attn: per-request attention.
@@ -552,7 +591,7 @@ class TurboQuantAttentionImpl(AttentionImpl["TurboQuantMetadata"]):
591
if _HAS_FLASH_ATTN:
592
_cu_2[1] = q_len
593
cu = _cu_2
555
- out = flash_attn_varlen_func(
594
+ out = self._flash_attn_varlen(
595
q=q_seq,
596
k=k_seq,
597
v=v_seq,
@@ -560,8 +599,6 @@ class TurboQuantAttentionImpl(AttentionImpl["TurboQuantMetadata"]):
599
cu_seqlens_k=cu,
600
max_seqlen_q=q_len,
601
max_seqlen_k=q_len,
563
- softmax_scale=self.scale,
564
- causal=True,
602
)
603
else:
604
q_t = q_seq.transpose(0, 1).contiguous()
@@ -726,7 +763,7 @@ class TurboQuantAttentionImpl(AttentionImpl["TurboQuantMetadata"]):
763
if _HAS_FLASH_ATTN:
764
cu_seqlens_q = torch.tensor([0, q_len], device=device, dtype=torch.int32)
765
cu_seqlens_k = torch.tensor([0, seq_len], device=device, dtype=torch.int32)
729
- return flash_attn_varlen_func(
766
+ return self._flash_attn_varlen(
767
q=query,
768
k=k_full,
769
v=v_full,
@@ -734,8 +771,6 @@ class TurboQuantAttentionImpl(AttentionImpl["TurboQuantMetadata"]):
771
cu_seqlens_k=cu_seqlens_k,
772
max_seqlen_q=q_len,
773
max_seqlen_k=seq_len,
737
- softmax_scale=self.scale,
738
- causal=True,
774
)
775
else:
776
# SDPA fallback: expand KV for GQA, build causal mask