[BugFix] Fix full-cuda-graph illegal memory access in FA3 #20057

LucasWilkinson · 2025-06-25T04:45:27Z

Essential Elements of an Effective PR Description Checklist

The purpose of the PR, such as "Fix some issue (link existing issues this PR will resolve)".
The test plan, such as providing test command.
The test results, such as pasting the results comparison before and after, or e2e results
(Optional) The necessary documentation update, such as updating supported_models.md and examples for a new model.

Purpose

Fix a sporadic illegal memory access discovered by @WoosukKwon

Test Plan

lm-eval + @WoosukKwon repro attempt

Test Result

after 1.5h @WoosukKwon has yet to repo

lm_eval   --model vllm   --model_args '{"pretrained":"Qwen/Qwen2.5-VL-72B-Instruct","tensor_parallel_size":8,"compilation_config":{"full_cuda_graph":true}}'   --tasks gsm8k   --batch_size auto
...
|Tasks|Version|     Filter     |n-shot|  Metric   |   |Value |   |Stderr|
|-----|------:|----------------|-----:|-----------|---|-----:|---|-----:|
|gsm8k|      3|flexible-extract|     5|exact_match|↑  |0.5277|±  |0.0138|
|     |       |strict-match    |     5|exact_match|↑  |0.5353|±  |0.0137|

lm_eval   --model vllm   --model_args '{"pretrained":"Qwen/Qwen2.5-VL-72B-Instruct","tensor_parallel_size":8,"compilation_config":{"full_cuda_graph":false}}'   --tasks gsm8k   --batch_size auto
...
|Tasks|Version|     Filter     |n-shot|  Metric   |   |Value |   |Stderr|
|-----|------:|----------------|-----:|-----------|---|-----:|---|-----:|
|gsm8k|      3|flexible-extract|     5|exact_match|↑  |0.5292|±  |0.0137|
|     |       |strict-match    |     5|exact_match|↑  |0.5406|±  |0.0137|

(Optional) Documentation Update

Signed-off-by: Lucas Wilkinson <lwilkins@redhat.com>

github-actions · 2025-06-25T04:45:35Z

👋 Hi! Thank you for contributing to the vLLM project.

💬 Join our developer Slack at https://slack.vllm.ai to discuss your PR in #pr-reviews, coordinate on features in #feat- channels, or join special interest groups in #sig- channels.

Just a reminder: PRs would not trigger full CI run by default. Instead, it would only run fastcheck CI which starts running only a small and essential subset of CI tests to quickly catch errors. You can run other CI tests on top of those by going to your fastcheck build on Buildkite UI (linked in the PR checks section) and unblock them. If you do not have permission to unblock, ping simon-mo or khluu to add you in our Buildkite org.

Once the PR is approved and ready to go, your PR reviewer(s) can run CI to test the changes comprehensively before merging.

To run CI, PR reviewers can either: Add ready label to the PR or enable auto-merge.

🚀

gemini-code-assist

Summary of Changes

Hello @LucasWilkinson, I'm Gemini Code Assist¹! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed!

This pull request implements a critical bug fix for a sporadic illegal memory access encountered in FlashAttention 3 when operating with full CUDA graphs. The solution involves a targeted adjustment to the attention backend's initialization logic, specifically by disabling AOT scheduling when full CUDA graph mode is active, thereby resolving the underlying memory issue.

Highlights

Bug Fix: Illegal Memory Access: This pull request addresses and fixes a sporadic illegal memory access issue that occurred when using FlashAttention 3 (FA3) with full CUDA graphs enabled.
Configuration Logic Adjustment: The previous validation logic that raised a ValueError if full CUDA graph mode was enabled without AOT (Ahead-of-Time) scheduling has been removed. Instead, when full CUDA graph is active, AOT scheduling is now explicitly disabled (self.aot_schedule = False) to prevent the reported memory access error.

Using Gemini Code Assist

The full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips.

Invoking Gemini

You can request assistance from Gemini at any point in your pull request via creating an issue comment (i.e. comment on the pull request page) using either /gemini <command> or @gemini-code-assist <command>. Below is a summary of the supported commands.

Feature	Command	Description
Code Review	`/gemini review`	Performs a code review for the current pull request in its current state.
Pull Request Summary	`/gemini summary`	Provides a summary of the current pull request in its current state.
Comment	@gemini-code-assist	Responds in comments when explicitly tagged, both in issue comments and review comments.
Help	`/gemini help`	Displays a list of available commands.

Customization

To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a .gemini/ folder in the base of the repository. Detailed instructions can be found here.

Limitations & Feedback

Gemini Code Assist is currently in preview and may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with 👍 and 👎 on @gemini-code-assist comments to provide feedback.

You can also get AI-powered code generation, chat, as well as code reviews directly in the IDE at no cost with the Gemini Code Assist IDE Extension.

Review the Privacy Notices, Generative AI Prohibited Use Policy, Terms of Service, and learn how to configure Gemini Code Assist in GitHub here. Gemini can make mistakes, so double check it and use code with caution. ↩

gemini-code-assist

Code Review

This pull request addresses a sporadic illegal memory access by modifying the logic for full_cuda_graph and aot_schedule. The change correctly disables AOT scheduling when full CUDA graph is enabled, which resolves the immediate issue. However, this new behavior directly contradicts a previous requirement stated in the removed code, indicating a need to review and update any related documentation or code that might still reflect the old, inconsistent requirement.

gemini-code-assist · 2025-06-25T04:46:31Z

vllm/v1/attention/backends/flash_attn.py

+        if self.use_full_cuda_graph:
+            # NOTE(lucas): AOT scheduling not supported in full cuda graph mode
+            #  yet. This is because the scheduler and kernel need to always use
+            #  the same num_splits (which acts as an upper bound with the
+            #  dynamic split scheduler) which is currently heuristically decided
+            #  by the kernel launching code.
+            self.aot_schedule = False


The previous logic (lines 161-163 removed) stated that "Full CUDA graph mode requires AOT scheduling". However, the new code (lines 161-167 added) explicitly sets self.aot_schedule = False when self.use_full_cuda_graph is enabled. This is a direct contradiction of the previous requirement.

While the new comment clarifies that "AOT scheduling not supported in full cuda graph mode yet", it's important to ensure that any other parts of the codebase or documentation that refer to full_cuda_graph requiring aot_schedule are updated to reflect this change. The previous ValueError message was misleading given the current state of support.

Signed-off-by: Lucas Wilkinson <lwilkins@redhat.com>

WoosukKwon · 2025-06-25T04:52:31Z

Thanks so much!

…ct#20057) Signed-off-by: Lucas Wilkinson <lwilkins@redhat.com>

…ct#20057) Signed-off-by: Lucas Wilkinson <lwilkins@redhat.com> Signed-off-by: Will Eaton <weaton@redhat.com>

…ct#20057) Signed-off-by: Lucas Wilkinson <lwilkins@redhat.com>

…ct#20057) Signed-off-by: Lucas Wilkinson <lwilkins@redhat.com> Signed-off-by: avigny <47987522+avigny@users.noreply.github.com>

…ct#20057) Signed-off-by: Lucas Wilkinson <lwilkins@redhat.com>

fix full-cuda-graph IMA

93dd61b

Signed-off-by: Lucas Wilkinson <lwilkins@redhat.com>

LucasWilkinson requested review from WoosukKwon, robertgshaw2-redhat, njhill, ywang96, comaniac and alexm-redhat as code owners June 25, 2025 04:45

gemini-code-assist bot reviewed Jun 25, 2025

View reviewed changes

mergify bot added the v1 label Jun 25, 2025

gemini-code-assist bot reviewed Jun 25, 2025

View reviewed changes

remove uneeded code

c3d80d7

Signed-off-by: Lucas Wilkinson <lwilkins@redhat.com>

WoosukKwon approved these changes Jun 25, 2025

View reviewed changes

WoosukKwon added the ready ONLY add when PR is ready to merge/full CI is needed label Jun 25, 2025

aarnphm enabled auto-merge (squash) June 25, 2025 06:24

aarnphm merged commit 0f9e735 into vllm-project:main Jun 25, 2025
82 checks passed

gmarinho2 pushed a commit to gmarinho2/vllm that referenced this pull request Jun 26, 2025

[BugFix] Fix full-cuda-graph illegal memory access in FA3 (vllm-proje…

da470de

…ct#20057) Signed-off-by: Lucas Wilkinson <lwilkins@redhat.com>

xjpang pushed a commit to xjpang/vllm that referenced this pull request Jun 30, 2025

[BugFix] Fix full-cuda-graph illegal memory access in FA3 (vllm-proje…

38984d2

…ct#20057) Signed-off-by: Lucas Wilkinson <lwilkins@redhat.com>

wseaton pushed a commit to wseaton/vllm that referenced this pull request Jun 30, 2025

[BugFix] Fix full-cuda-graph illegal memory access in FA3 (vllm-proje…

b618206

…ct#20057) Signed-off-by: Lucas Wilkinson <lwilkins@redhat.com> Signed-off-by: Will Eaton <weaton@redhat.com>

wseaton pushed a commit to wseaton/vllm that referenced this pull request Jun 30, 2025

[BugFix] Fix full-cuda-graph illegal memory access in FA3 (vllm-proje…

7f4bda7

…ct#20057) Signed-off-by: Lucas Wilkinson <lwilkins@redhat.com>

wwl2755-google pushed a commit to wwl2755-google/vllm that referenced this pull request Jul 1, 2025

[BugFix] Fix full-cuda-graph illegal memory access in FA3 (vllm-proje…

36b68fd

…ct#20057) Signed-off-by: Lucas Wilkinson <lwilkins@redhat.com>

googlercolin pushed a commit to googlercolin/vllm that referenced this pull request Aug 29, 2025

[BugFix] Fix full-cuda-graph illegal memory access in FA3 (vllm-proje…

33436b0

…ct#20057) Signed-off-by: Lucas Wilkinson <lwilkins@redhat.com>

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Uh oh!

Uh oh!

[BugFix] Fix full-cuda-graph illegal memory access in FA3 #20057

[BugFix] Fix full-cuda-graph illegal memory access in FA3 #20057

Uh oh!

LucasWilkinson commented Jun 25, 2025 •

edited by github-actions bot

Loading

Uh oh!

github-actions bot commented Jun 25, 2025

Uh oh!

gemini-code-assist bot left a comment

Uh oh!

gemini-code-assist bot left a comment

Uh oh!

gemini-code-assist bot Jun 25, 2025

Uh oh!

WoosukKwon commented Jun 25, 2025

Uh oh!

Uh oh!

Uh oh!

Uh oh!

[BugFix] Fix full-cuda-graph illegal memory access in FA3 #20057

[BugFix] Fix full-cuda-graph illegal memory access in FA3 #20057

Uh oh!

Conversation

LucasWilkinson commented Jun 25, 2025 • edited by github-actions bot Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Essential Elements of an Effective PR Description Checklist

Purpose

Test Plan

Test Result

(Optional) Documentation Update

Uh oh!

github-actions bot commented Jun 25, 2025

Uh oh!

gemini-code-assist bot left a comment

Choose a reason for hiding this comment

Summary of Changes

Highlights

Footnotes

Uh oh!

gemini-code-assist bot left a comment

Choose a reason for hiding this comment

Code Review

Uh oh!

gemini-code-assist bot Jun 25, 2025

Choose a reason for hiding this comment

Uh oh!

WoosukKwon commented Jun 25, 2025

Uh oh!

Uh oh!

Uh oh!

LucasWilkinson commented Jun 25, 2025 •

edited by github-actions bot

Loading