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 9, 00:00 UTC
[Attention][FA3] Update FA3 to include new swizzle optimization (#23465)
Signed-off-by: Lucas Wilkinson <lwilkins@redhat.com>
Lucas Wilkinson committed
Feb 3, 2026 at 09:08 UTC
2267cb1cfd838a192f47ff677d91164fb5cb2862
3 files changed
+13
-3
cmake/external_projects/vllm_flash_attn.cmake
+1
-1
@@ -38,7 +38,7 @@ else()
38
FetchContent_Declare(
39
vllm-flash-attn
40
GIT_REPOSITORY https://github.com/vllm-project/flash-attention.git
41
- GIT_TAG 188be16520ceefdc625fdf71365585d2ee348fe2
41
+ GIT_TAG 2adfc8c2177c5b0e8ddeedfd5a8990d80eb496ff
42
GIT_PROGRESS TRUE
43
# Don't share the vllm-flash-attn build between build types
44
BINARY_DIR ${CMAKE_BINARY_DIR}/vllm-flash-attn
vllm/v1/attention/backends/flash_attn.py
+6
-1
@@ -308,10 +308,15 @@ class FlashAttentionMetadataBuilder(AttentionMetadataBuilder[FlashAttentionMetad
308
self.compilation_config.cudagraph_mode.has_full_cudagraphs()
309
)
310
self.max_cudagraph_size = self.compilation_config.max_cudagraph_capture_size
311
+ max_num_seqs = vllm_config.scheduler_config.max_num_seqs
312
313
if self.use_full_cuda_graph and self.aot_schedule:
314
+ # Times 4 due to:
315
+ # https://github.com/vllm-project/flash-attention/blob/3223650ccabe622a0fcae65eec706a50186a89f7/hopper/flash_api.cpp#L650-L653
316
+ # For some tests max_cudagraph_size > max_num_seqs,
317
+ # so we need to use the larger one.
318
self.scheduler_metadata = torch.zeros(
314
- vllm_config.scheduler_config.max_num_seqs + 1,
319
+ max(self.max_cudagraph_size or 0, max_num_seqs) * 4 + 1,
320
dtype=torch.int32,
321
device=self.device,
322
)
vllm/v1/attention/backends/mla/flashattn_mla.py
+6
-1
@@ -127,10 +127,15 @@ class FlashAttnMLAMetadataBuilder(MLACommonMetadataBuilder[FlashAttnMLAMetadata]
127
self.compilation_config.cudagraph_mode.has_full_cudagraphs()
128
)
129
self.max_cudagraph_size = self.compilation_config.max_cudagraph_capture_size
130
+ max_num_seqs = vllm_config.scheduler_config.max_num_seqs
131
132
if self.use_full_cuda_graph and self.fa_aot_schedule:
133
+ # Times 4 due to:
134
+ # https://github.com/vllm-project/flash-attention/blob/3223650ccabe622a0fcae65eec706a50186a89f7/hopper/flash_api.cpp#L650-L653
135
+ # For some tests max_cudagraph_size > max_num_seqs,
136
+ # so we need to use the larger one.
137
self.scheduler_metadata = torch.zeros(
133
- vllm_config.scheduler_config.max_num_seqs + 1,
138
+ max(self.max_cudagraph_size or 0, max_num_seqs) * 4 + 1,
139
dtype=torch.int32,
140
device=self.device,
141
)