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 14, 05:30 UTC
[BugFix] Potential Fix for FA3 full-cudagraph IMA (#25490)
Signed-off-by: Lucas Wilkinson <lwilkins@redhat.com>
Lucas Wilkinson committed
Sep 24, 2025 at 05:04 UTC
2338daffd3ec5935f6a4b53fdd62d0bde5ab1367
1 file changed
+11
-11
vllm/v1/attention/backends/flash_attn.py
+11
-11
@@ -194,10 +194,9 @@ class FlashAttentionMetadataBuilder(
194
195
self.use_full_cuda_graph = \
196
self.compilation_config.cudagraph_mode.has_full_cudagraphs()
197
+ self.max_cudagraph_size = self.compilation_config.max_capture_size
198
199
if self.use_full_cuda_graph and self.aot_schedule:
199
- self.max_cudagraph_size = self.compilation_config.max_capture_size
200
-
200
if self.max_cudagraph_size > 992:
201
# This condition derives from FA3's internal heuristic.
202
# TODO(woosuk): Support larger cudagraph sizes.
@@ -259,6 +258,15 @@ class FlashAttentionMetadataBuilder(
258
self.aot_schedule = False
259
aot_schedule = False
260
261
+ max_num_splits = 0 # 0 means use FA3's heuristics, not CG compatible
262
+ if self.use_full_cuda_graph and \
263
+ num_actual_tokens <= self.max_cudagraph_size:
264
+ # NOTE(woosuk): Setting num_splits > 1 may increase the memory
265
+ # usage, because the intermediate buffers of size [num_splits,
266
+ # num_heads, num_tokens, head_size] are allocated. Therefore,
267
+ # we only set num_splits when using cuda graphs.
268
+ max_num_splits = self.max_num_splits
269
+
270
def schedule(batch_size, cu_query_lens, max_query_len, seqlens,
271
max_seq_len, causal):
272
cache_dtype = self.cache_config.cache_dtype
@@ -281,7 +289,7 @@ class FlashAttentionMetadataBuilder(
289
page_size=self.block_size,
290
causal=causal,
291
window_size=self.aot_sliding_window,
284
- num_splits=self.max_num_splits,
292
+ num_splits=max_num_splits,
293
)
294
return None
295
@@ -322,7 +330,6 @@ class FlashAttentionMetadataBuilder(
330
max_seq_len=max_seq_len,
331
causal=causal)
332
# For FA3 + full cudagraph
325
- max_num_splits = 0
333
if self.use_full_cuda_graph and scheduler_metadata is not None:
334
n = scheduler_metadata.shape[0]
335
self.scheduler_metadata[:n] = scheduler_metadata
@@ -333,13 +340,6 @@ class FlashAttentionMetadataBuilder(
340
self.scheduler_metadata[n:] = 0
341
scheduler_metadata = self.scheduler_metadata[:n]
342
336
- if num_actual_tokens <= self.max_cudagraph_size:
337
- # NOTE(woosuk): Setting num_splits > 1 may increase the memory
338
- # usage, because the intermediate buffers of size [num_splits,
339
- # num_heads, num_tokens, head_size] are allocated. Therefore,
340
- # we only set num_splits when using cuda graphs.
341
- max_num_splits = self.max_num_splits
342
-
343
attn_metadata = FlashAttentionMetadata(
344
num_actual_tokens=num_actual_tokens,
345
max_query_len=max_query_len,