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 15, 04:45 UTC
[BugFix] Fix full-cuda-graph illegal memory access in FA3 (#20057)
Signed-off-by: Lucas Wilkinson <lwilkins@redhat.com>
Lucas Wilkinson committed
Jun 25, 2025 at 04:39 UTC
0f9e7354f508af3fe314cfb709babaaa668f1b04
1 file changed
+7
-18
vllm/v1/attention/backends/flash_attn.py
+7
-18
@@ -158,12 +158,13 @@ class FlashAttentionMetadataBuilder(
158
159
self.aot_schedule = (get_flash_attn_version() == 3)
160
self.use_full_cuda_graph = compilation_config.full_cuda_graph
161
- if self.use_full_cuda_graph and not self.aot_schedule:
162
- raise ValueError("Full CUDA graph mode requires AOT scheduling, "
163
- "which requires FlashAttention 3.")
164
- self.scheduler_metadata = torch.zeros(self.runner.max_num_reqs + 1,
165
- dtype=torch.int32,
166
- device=self.runner.device)
161
+ if self.use_full_cuda_graph:
162
+ # NOTE(lucas): AOT scheduling not supported in full cuda graph mode
163
+ # yet. This is because the scheduler and kernel need to always use
164
+ # the same num_splits (which acts as an upper bound with the
165
+ # dynamic split scheduler) which is currently heuristically decided
166
+ # by the kernel launching code.
167
+ self.aot_schedule = False
168
169
# Sliding window size to be used with the AOT scheduler will be
170
# populated on first build() call.
@@ -299,18 +300,6 @@ class FlashAttentionMetadataBuilder(
300
max_seq_len=max_seq_len,
301
causal=True)
302
302
- if self.use_full_cuda_graph:
303
- assert scheduler_metadata is not None
304
- n = scheduler_metadata.shape[0]
305
- self.scheduler_metadata[:n].copy_(scheduler_metadata,
306
- non_blocking=True)
307
- # NOTE(woosuk): We should zero out the rest of the scheduler
308
- # metadata to guarantee the correctness. Otherwise, some thread
309
- # blocks may use the invalid scheduler metadata and overwrite the
310
- # output buffer.
311
- self.scheduler_metadata[n:] = 0
312
- scheduler_metadata = self.scheduler_metadata[:n]
313
-
303
attn_metadata = FlashAttentionMetadata(
304
num_actual_tokens=num_actual_tokens,
305
max_query_len=max_query_len,