[Attention][V1] Toggle for v1 attention backend #18275

gshtras · 2025-05-16T18:37:24Z

Expanding on #18093
Adding a toggle to force fallback to the 2 stage attention kernel in V1

Including a small fix for the FP8 kv cache on ROCm in the 2 stage kernel approach

Signed-off-by: Gregory Shtrasberg <Gregory.Shtrasberg@amd.com>

github-actions · 2025-05-16T18:37:32Z

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

🚀

tjtanaa · 2025-05-18T07:27:13Z

vllm/v1/attention/backends/triton_attn.py

        num_queries_per_kv = query.shape[1] // key.shape[1]
-        use_prefill_decode_attn = (num_queries_per_kv &
-                                   (num_queries_per_kv - 1)) != 0
+        use_prefill_decode_attn = envs.VLLM_V1_USE_PREFILL_DECODE_ATTENTION or (


Please avoid using envs.ENVIRON to in forward path that is going to be called in runtime. RFC #17067 has revealed that the overhead is very high.

in the def __init__():, we can pre-evaluate the env.ENVIRON self.use_prefill_decode_attn = envs.VLLM_V1_USE_PREFILL_DECODE_ATTENTION

Should we try to automatically set this environment variable to True if it is on ROCm platform, until the performance gap has been resolved?

Model: Qwen/Qwen3-235B-A22B-FP8
Input: Output= 1000:1000

Use: PREFILL_DECODE_ATTENTION

============ Serving Benchmark Result ============ Successful requests: 500 Benchmark duration (s): 129.30 Total input tokens: 494223 Total generated tokens: 364067 Request throughput (req/s): 3.87 Request goodput (req/s): 0.03 Output token throughput (tok/s): 2815.74 Total Token throughput (tok/s): 6638.12 ---------------Time to First Token---------------- Mean TTFT (ms): 17056.34 Median TTFT (ms): 16655.24 P99 TTFT (ms): 31088.47 -----Time per Output Token (excl. 1st token)------ Mean TPOT (ms): 224.22 Median TPOT (ms): 83.60 P99 TPOT (ms): 2039.01 ---------------Inter-token Latency---------------- Mean ITL (ms): 76.99 Median ITL (ms): 61.88 P99 ITL (ms): 270.91 ==================================================

Without PREFILL_DECODE_ATTENTION (with unified attention)

============ Serving Benchmark Result ============ Successful requests: 500 Benchmark duration (s): 153.32 Total input tokens: 494223 Total generated tokens: 354124 Request throughput (req/s): 3.26 Request goodput (req/s): 0.00 Output token throughput (tok/s): 2309.75 Total Token throughput (tok/s): 5533.28 ---------------Time to First Token---------------- Mean TTFT (ms): 18589.72 Median TTFT (ms): 19132.91 P99 TTFT (ms): 35497.60 -----Time per Output Token (excl. 1st token)------ Mean TPOT (ms): 238.41 Median TPOT (ms): 98.48 P99 TPOT (ms): 2331.83 ---------------Inter-token Latency---------------- Mean ITL (ms): 90.84 Median ITL (ms): 69.57 P99 ITL (ms): 470.61 ==================================================

It is not universally worse, it shows different results on different concurrency settings, so going forward we want to keep it as the default with the toggle left for user discretion.

Signed-off-by: Gregory Shtrasberg <Gregory.Shtrasberg@amd.com>

…ream

ProExpertProg

A few minor comments, looks good otherwise

ProExpertProg · 2025-05-16T21:13:04Z

vllm/v1/attention/backends/triton_attn.py

-        use_prefill_decode_attn = (num_queries_per_kv &
-                                   (num_queries_per_kv - 1)) != 0
+        use_prefill_decode_attn = envs.VLLM_V1_USE_PREFILL_DECODE_ATTENTION or (
+            (num_queries_per_kv & (num_queries_per_kv - 1)) != 0)


Could you extract this into a bool (like is_num_q_pow2)

ProExpertProg · 2025-05-21T19:14:40Z

vllm/envs.py


+    # Use separate prefill and decode kernels for V1 attention instead of
+    # the unified triton kernel.
+    "VLLM_V1_USE_PREFILL_DECODE_ATTENTION":


Nit: VLLM_V1_TRITON_ATTN_FORCE_PREFILL_DECODE sounds slightly more accurate to me, but feel free to use the name that works best in your eyes

ProExpertProg · 2025-05-21T19:14:56Z

vllm/envs.py

    "VLLM_DEBUG_LOG_API_SERVER_RESPONSE":
-    lambda: os.environ.get("VLLM_DEBUG_LOG_API_SERVER_RESPONSE", "False").
-    lower() == "true",
+    lambda: os.environ.get("VLLM_DEBUG_LOG_API_SERVER_RESPONSE", "False"


Nit: accidental change?

ruff or yapf change, reformatting the whole file now results in this

ProExpertProg · 2025-05-21T19:15:54Z

vllm/v1/attention/backends/triton_attn.py

                                      "TritonAttentionImpl")

        self.fp8_dtype = current_platform.fp8_dtype()
+        self.use_prefill_decode_attn = envs.VLLM_V1_USE_PREFILL_DECODE_ATTENTION


Suggested change

self.use_prefill_decode_attn = envs.VLLM_V1_USE_PREFILL_DECODE_ATTENTION

self.force_prefill_decode_attn = envs. VLLM_V1_TRITON_ATTN_FORCE_PREFILL_DECODE

Signed-off-by: Gregory Shtrasberg <Gregory.Shtrasberg@amd.com>

mgoin · 2025-05-27T16:02:10Z

Could you instead reuse the VLLM_ATTENTION_BACKEND variable? Is this override just used for testing at the moment?

gshtras · 2025-05-27T16:07:06Z

Could you instead reuse the VLLM_ATTENTION_BACKEND variable? Is this override just used for testing at the moment?

There's currently no separate backend for the split prefill decode attention path, it has been switched to use the unified attention.
We want the ability to force the fallback to the old one for the cases of poor performance until those are resolved.

…ream

Signed-off-by: Gregory Shtrasberg <Gregory.Shtrasberg@amd.com> Signed-off-by: amit <amit.man@gmail.com>

Toggle for v1 attention

ac2b381

Signed-off-by: Gregory Shtrasberg <Gregory.Shtrasberg@amd.com>

gshtras requested review from WoosukKwon, alexm-redhat, comaniac, njhill, robertgshaw2-redhat and ywang96 as code owners May 16, 2025 18:37

mergify bot added the v1 label May 16, 2025

tjtanaa reviewed May 18, 2025

View reviewed changes

gshtras added 2 commits May 19, 2025 15:25

Caching the env variable in the __init__

9540f72

Signed-off-by: Gregory Shtrasberg <Gregory.Shtrasberg@amd.com>

Merge remote-tracking branch 'origin/main' into attention_toggle_upst…

dc0f0d1

…ream

DarkLight1337 mentioned this pull request May 21, 2025

[Bug][Failing Test]: LoRA 2 - lora/test_lora_functions.py::test_lora_functions_sync #18498

Closed

1 task

ProExpertProg reviewed May 21, 2025

View reviewed changes

Better naming and logic extracted to a variable

f82da97

Signed-off-by: Gregory Shtrasberg <Gregory.Shtrasberg@amd.com>

gshtras requested a review from ProExpertProg May 22, 2025 18:30

tjtanaa mentioned this pull request May 26, 2025

[Doc]: Find out about how to optimize the parameters of vLLM V1 on ROCm EmbeddedLLM/vllm#64

Closed

1 task

ProExpertProg approved these changes May 27, 2025

View reviewed changes

mgoin approved these changes May 27, 2025

View reviewed changes

mgoin added the ready ONLY add when PR is ready to merge/full CI is needed label May 27, 2025

Merge remote-tracking branch 'origin/main' into attention_toggle_upst…

3b264d5

…ream

gshtras merged commit da4b69d into vllm-project:main May 29, 2025
64 checks passed

gshtras deleted the attention_toggle_upstream branch May 29, 2025 14:48

amitm02 pushed a commit to amitm02/vllm that referenced this pull request Jun 1, 2025

[Attention][V1] Toggle for v1 attention backend (vllm-project#18275)

680147b

Signed-off-by: Gregory Shtrasberg <Gregory.Shtrasberg@amd.com> Signed-off-by: amit <amit.man@gmail.com>

amitm02 pushed a commit to amitm02/vllm that referenced this pull request Jun 1, 2025

[Attention][V1] Toggle for v1 attention backend (vllm-project#18275)

7eca5f2

Signed-off-by: Gregory Shtrasberg <Gregory.Shtrasberg@amd.com> Signed-off-by: amit <amit.man@gmail.com>

	self.use_prefill_decode_attn = envs.VLLM_V1_USE_PREFILL_DECODE_ATTENTION
	self.force_prefill_decode_attn = envs. VLLM_V1_TRITON_ATTN_FORCE_PREFILL_DECODE

Uh oh!

[Attention][V1] Toggle for v1 attention backend #18275

[Attention][V1] Toggle for v1 attention backend #18275

Uh oh!

Conversation

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

Uh oh!

Uh oh!

github-actions bot commented May 16, 2025

Uh oh!

tjtanaa May 18, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

tjtanaa May 18, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

gshtras May 19, 2025

Choose a reason for hiding this comment

Uh oh!

ProExpertProg left a comment

Choose a reason for hiding this comment

Uh oh!

ProExpertProg May 16, 2025

Choose a reason for hiding this comment

Uh oh!

ProExpertProg May 21, 2025

Choose a reason for hiding this comment

Uh oh!

ProExpertProg May 21, 2025

Choose a reason for hiding this comment

Uh oh!

gshtras May 22, 2025

Choose a reason for hiding this comment

Uh oh!

ProExpertProg May 21, 2025

Choose a reason for hiding this comment

Uh oh!

mgoin commented May 27, 2025

Uh oh!

gshtras commented May 27, 2025

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

gshtras commented May 16, 2025 •

edited by github-actions bot

Loading

tjtanaa May 18, 2025 •

edited

Loading

tjtanaa May 18, 2025 •

edited

Loading