[v1] Add fp32 support to v1 engine through flex attn

[Isotr0py]

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

Purpose

• After FlexAttention PR landed, we can use it for fp32 inference.
• This PR adds fp32 support for v1 engine through enabling FlexAttention backend.

Test Plan

python examples/offline_inference/basic/generate.py --model facebook/opt-125m --dtype float32

Test Result

INFO 06-07 15:48:36 [cuda.py:237] Using FlexAttenion backend for torch.float32 on V1 engine.
INFO 06-07 15:48:36 [weight_utils.py:292] Using model weights format ['*.bin']
Loading pt checkpoint shards:   0% Completed | 0/1 [00:00<?, ?it/s]
Loading pt checkpoint shards: 100% Completed | 1/1 [00:00<00:00,  1.93it/s]
Loading pt checkpoint shards: 100% Completed | 1/1 [00:00<00:00,  1.93it/s]

INFO 06-07 15:48:37 [default_loader.py:272] Loading weights took 0.53 seconds
INFO 06-07 15:48:37 [gpu_model_runner.py:1618] Model loading took 0.4678 GiB and 0.814874 seconds
INFO 06-07 15:48:40 [backends.py:462] Using cache directory: /root/.cache/vllm/torch_compile_cache/815e1a20d0/rank_0_0 for vLLM's torch.compile
INFO 06-07 15:48:40 [backends.py:472] Dynamo bytecode transform time: 3.35 s
INFO 06-07 15:48:42 [backends.py:135] Directly load the compiled graph(s) for shape None from the cache, took 1.047 s
INFO 06-07 15:48:42 [monitor.py:34] torch.compile takes 3.35 s in total
INFO 06-07 15:48:43 [kv_cache_utils.py:715] GPU KV cache size: 176,464 tokens
INFO 06-07 15:48:43 [kv_cache_utils.py:719] Maximum concurrency for 2,048 tokens per request: 86.16x
INFO 06-07 15:49:02 [gpu_model_runner.py:2012] Graph capturing finished in 19 secs, took 0.22 GiB
INFO 06-07 15:49:02 [core.py:171] init engine (profile, create kv cache, warmup model) took 24.91 seconds
Adding requests: 100%|████████████████████████████████████████████████████████████████████████████████████████████████████████| 4/4 [00:00<00:00, 862.85it/s]
Processed prompts: 100%|██████████████████████████████████████████████████| 4/4 [00:05<00:00,  1.39s/it, est. speed input: 4.66 toks/s, output: 11.47 toks/s]
--------------------------------------------------
Prompt: 'Hello, my name is'
Generated text: ' Joel, I am a BYU lecturer and I would love to talk about your class'
--------------------------------------------------
Prompt: 'The president of the United States is'
Generated text: " stuck painting the country like it's 70 years old again, with behind-the"
--------------------------------------------------
Prompt: 'The capital of France is'
Generated text: " directly behind Germany's winning bid for elections as French business and French defence leaders fired"
--------------------------------------------------
Prompt: 'The future of AI is'
Generated text: ' predicted according to the 2020 report by a economists\nAI stand to be in the'
--------------------------------------------------

[github-actions]

👋 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]

Hello @Isotr0py, 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!

Summary of Changes

Hello! Gemini here, providing a summary of this pull request. The primary goal of this PR, as stated in the description and title, is to add support for float32 inference to the v1 engine. This is achieved by leveraging the recently integrated FlexAttention backend, which is capable of handling dtypes beyond the float16 and bfloat16 previously supported by the v1 engine's attention mechanisms (like Triton). The changes involve removing the explicit dtype check in the argument parsing logic and updating the attention backend selection logic to route float32 (and potentially other non-fp16/bf16 dtypes) to the FlexAttention implementation.

Highlights

• Enable float32 for V1 Engine: This pull request enables the use of torch.float32 for inference when using the v1 engine.
• Utilize FlexAttention for float32: The implementation routes float32 computations within the v1 engine's attention mechanism to the FlexAttention backend.
• Remove Dtype Restriction: The previous check limiting the v1 engine to only float16 and bfloat16 dtypes has been removed from the argument parsing.

Changelog

• vllm/engine/arg_utils.py
  • Removed the V1_SUPPORTED_DTYPES check that restricted the v1 engine to only torch.bfloat16 and torch.float16 (lines 1340-1345).
• vllm/platforms/cuda.py
  • Added a condition in get_attn_backend_cls to select FlexAttentionBackend for the v1 engine when the dtype is not torch.float16 or torch.bfloat16 (lines 236-239).

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.

Footnotes

1. 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 effectively adds fp32 support to the V1 engine by integrating FlexAttention. The changes are concise, well-targeted, and the PR description clearly outlines the purpose, testing, and results. The approach of removing the explicit dtype restriction in arg_utils.py and then selecting FlexAttention by default for non-fp16/bf16 dtypes in cuda.py is logical and aligns with making the V1 engine more flexible while ensuring appropriate backend selection. Overall, the changes look good and enhance the V1 engine's capabilities.

Summary of Findings

• Successful FP32 Enablement for V1 Engine: The PR successfully enables fp32 support for the V1 engine. This is achieved by removing a previous dtype restriction that limited V1 to fp16/bf16 and by updating the attention backend selection logic to default to FlexAttention for other dtypes like fp32. The changes are clear, minimal, and directly address the PR's goal.
• Minor Typo in Log Message: A minor typo ('FlexAttenion' instead of 'FlexAttention') was found in a log message in vllm/platforms/cuda.py (line 238). This was not commented on directly due to review settings for low-severity issues. Suggestion: f"Using FlexAttention backend for {dtype} on V1 engine."

Merge Readiness

The code changes appear to be of high quality and effectively implement the intended fp32 support for the V1 engine. Based on this review, the PR seems ready for merging. However, as an AI assistant, I am not authorized to approve pull requests; please ensure it undergoes further review and approval by the repository maintainers.

[houseroad]

Looks legit, shall we add or turn on some unittest?

[houseroad]

The failed test is related. "NotImplementedError: FlexAttention does not support kv sharing yet." Maybe it's still a bit early to turn on FlexAttention for v1 by default? cc: @drisspg

[Isotr0py]

Given the model_config using fp16 in that test, seems there is something going wrong making Flex Attn is used for fp16 by default:

vllm/tests/v1/worker/test_gpu_model_runner.py

Lines 71 to 79 in b9a1791

	model_config = ModelConfig(
	model="facebook/opt-125m",
	task="generate",
	tokenizer="facebook/opt-125m",
	tokenizer_mode="auto",
	trust_remote_code=True,
	dtype="float16",
	seed=42,
	)

[drisspg]

@houseroad
Yeah I agree it might be a little early. I am currently working on a. High level tracker with the current state of the integration and what we need to do/fix before we consider it robust.

Will link the issue

[houseroad]

Looks good.

[simon-mo]

@Isotr0py this PR broke plugin tests (2 GPU) and LoRA TP test (distributed)

Failure
https://buildkite.com/vllm/ci/builds/21699/steps/canvas?sid=01975508-6b8f-4cde-9a41-5ebf39fb8881

The commit before
https://buildkite.com/vllm/ci/builds/21696/steps/table?sid=019754e5-c0da-42de-8fdb-7eeee3ab29a9

[Isotr0py]

Seems that FlexAttn is used for those tests with incorrect dtype setting. Let me fix it ASAP.

[houseroad]

Yeah, LoRA one is "RuntimeError: Worker failed with error ''FlashAttentionMetadata' object has no attribute 'score_mod'', https://buildkite.com/vllm/ci/builds/21699#01975508-6c54-47d6-a951-82873b896a11.

And another one is "RuntimeError: logits must be contiguous", thrown from flashinfer, is it due to flex attention's output is not contiguous? https://buildkite.com/vllm/ci/builds/21699#01975508-6c52-41ed-ba7a-cc1a464cec91

[Isotr0py]

[2025-06-10T05:51:32Z] INFO 06-09 22:51:32 [core.py:70] Initializing a V1 LLM engine (v0.9.1.dev250+gb8089195b) with config: model='meta-llama/Llama-2-7b-hf', speculative_config=None, tokenizer='meta-llama/Llama-2-7b-hf', skip_tokenizer_init=False, tokenizer_mode=auto, revision=None, override_neuron_config={}, tokenizer_revision=None, trust_remote_code=False, dtype=torch.float16, ...
...
| [2025-06-10T05:51:41Z] (VllmWorker rank=1 pid=10974) INFO 06-09 22:51:41 [gpu_model_runner.py:1597] Loading model from scratch...
  | [2025-06-10T05:51:41Z] (VllmWorker rank=1 pid=10974) INFO 06-09 22:51:41 [cuda.py:237] Using FlexAttenion backend for torch.float32 on V1 engine.
...
  | [2025-06-10T05:51:49Z] INFO 06-09 22:51:49 [kv_cache_utils.py:715] GPU KV cache size: 51,232 tokens
  | [2025-06-10T05:51:49Z] INFO 06-09 22:51:49 [kv_cache_utils.py:719] Maximum concurrency for 4,096 tokens per request: 12.51x
  | [2025-06-10T05:51:49Z] INFO 06-09 22:51:49 [kv_cache_utils.py:715] GPU KV cache size: 52,000 tokens
  | [2025-06-10T05:51:49Z] INFO 06-09 22:51:49 [kv_cache_utils.py:719] Maximum concurrency for 4,096 tokens per request: 12.70x
  | [2025-06-10T05:51:49Z] (VllmWorker rank=0 pid=10973) INFO 06-09 22:51:49 [cuda.py:256] Using Flash Attention backend on V1 engine.
  | [2025-06-10T05:51:49Z] (VllmWorker rank=1 pid=10974) INFO 06-09 22:51:49 [cuda.py:256] Using Flash Attention backend on V1 engine.

The LoRA ones is because Loading model from scratch will use fp32 with flex_attn. However, given the model_config has dtype=torch.float16, perhaps something wrong in tensorizer loader,

[houseroad]

If it's not easy to fix it quickly, we can do revert first, then resubmit to unblock.

[Isotr0py]

If it's not easy to fix it quickly, we can do revert first, then resubmit to unblock.

Agree, unfortunately, I'm outside and unavailable to investigate it this afternoon. Let's revert this first to unblock release. I will further investigate this issue tonight.