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 13, 05:45 UTC
[BugFix] Fix FA3 IMA with FULL_AND_PIECEWISE and cascade attention (default) (#28702)
Lucas Wilkinson committed
Nov 14, 2025 at 07:19 UTC
db56a59970a84842da2adc3aa64e436f42448b48
2 files changed
+5
-2
tests/kernels/attention/test_cascade_flash_attn.py
+1
@@ -170,6 +170,7 @@ def test_cascade(
170
logits_soft_cap=soft_cap if soft_cap is not None else 0,
171
block_table=block_tables,
172
common_prefix_len=common_prefix_len,
173
+ max_num_splits=0, # no max
174
fa_version=fa_version,
175
)
176
vllm/v1/attention/backends/flash_attn.py
+4
-2
@@ -704,6 +704,7 @@ class FlashAttentionImpl(AttentionImpl):
704
logits_soft_cap=self.logits_soft_cap,
705
block_table=attn_metadata.block_table,
706
common_prefix_len=attn_metadata.common_prefix_len,
707
+ max_num_splits=attn_metadata.max_num_splits,
708
fa_version=self.vllm_flash_attn_version,
709
prefix_scheduler_metadata=attn_metadata.prefix_scheduler_metadata,
710
suffix_scheduler_metadata=attn_metadata.scheduler_metadata,
@@ -950,6 +951,7 @@ def cascade_attention(
951
logits_soft_cap: float,
952
block_table: torch.Tensor,
953
common_prefix_len: int,
954
+ max_num_splits: int,
955
fa_version: int,
956
prefix_scheduler_metadata: torch.Tensor | None = None,
957
suffix_scheduler_metadata: torch.Tensor | None = None,
@@ -994,7 +996,7 @@ def cascade_attention(
996
# s_aux is incorporated into prefix_lse inside the GPU kernel,
997
# enabling its effect during the final attention merge.
998
s_aux=s_aux,
997
- num_splits=1 if vllm_is_batch_invariant() else 0,
999
+ num_splits=1 if vllm_is_batch_invariant() else max_num_splits,
1000
)
1001
1002
descale_shape = (cu_query_lens.shape[0] - 1, key_cache.shape[-2])
@@ -1019,7 +1021,7 @@ def cascade_attention(
1021
q_descale=q_descale.expand(descale_shape) if q_descale is not None else None,
1022
k_descale=k_descale.expand(descale_shape) if k_descale is not None else None,
1023
v_descale=v_descale.expand(descale_shape) if v_descale is not None else None,
1022
- num_splits=1 if vllm_is_batch_invariant() else 0,
1024
+ num_splits=1 if vllm_is_batch_invariant() else max_num_splits,
1025
)
1026
1027
# Merge prefix and suffix outputs, and store the result in output.