[BugFix] Potential Fix for FA3 full-cudagraph IMA

[LucasWilkinson]

@WoosukKwon reported an IMA with FA3 full-CG that was fixed by doing https://github.com/vllm-project/vllm/compare/woosuk/fa3-ima?expand=1

the theory here is that get_scheduler_metadata was being called with a different max_num_splits than what was being passed to FlashAttentionMetadata

this is an alternative solution that doesn't lose the logic to use max_num_splits=0 (i.e. use the heuristic) for batches larger then max_cudagraph_size

we do not currently have a repo so cannot confirm this resolves @WoosukKwon 's IMA but this should be resolved regardless; we should alway make sure the arguments to get_scheduler_metadata and FlashAttentionMetadata are inline

vllm serve meta-llama/Meta-Llama-3-8B-Instruct -O.cudagraph_mode=FULL


lm_eval --model local-completions --model_args "base_url=http://0.0.0.0:8000/v1/completions,model=meta-llama/Meta-Llama-3-8B-Instruct,num_concurrent=256" --tasks gsm8k
...
local-completions (base_url=http://0.0.0.0:8000/v1/completions,model=meta-llama/Meta-Llama-3-8B-Instruct,num_concurrent=256), gen_kwargs: (None), limit: None, num_fewshot: None, batch_size: 1
|Tasks|Version|     Filter     |n-shot|  Metric   |   |Value |   |Stderr|
|-----|------:|----------------|-----:|-----------|---|-----:|---|-----:|
|gsm8k|      3|flexible-extract|     5|exact_match|↑  |0.7544|±  |0.0119|
|     |       |strict-match    |     5|exact_match|↑  |0.7559|±  |0.0118|

[gemini-code-assist]

Code Review

This pull request aims to fix a potential Invalid Memory Access in FlashAttention 3 with full CUDA graphs by ensuring the max_num_splits parameter is consistent. The change refactors the logic for setting max_num_splits to a common location. However, the current implementation introduces a critical flaw: it can lead to an UnboundLocalError because max_num_splits is not defined in all code paths. My review provides a fix for this issue to ensure the variable is always initialized. Addressing this will also help achieve the PR's goal of making the parameter consistent.

[WoosukKwon]

Thanks for the fix!

[WoosukKwon]

@LucasWilkinson Can you please check the CI again?