[V1][TPU] Integrate the new ragged paged attention kernel with vLLM v1 on TPU #13379

vanbasten23 · 2025-02-17T06:28:30Z

This PR integrates the new ragged paged attention kernel with vLLM v1 on TPU. In particular, this PR

Update torch_xla pin to the latest
Update pallas.py in v1 to use the new ragged paged attention kernel instead of the 3 separate kernels in v0.
Combine prompt and decode steps into one single step in tpu_model_runner.py, similar to what GPU does today.

Test plan:

$ VLLM_USE_V1=1 python vllm/examples/offline_inference/basic.py 2>&1 | tee out.txt
$ VLLM_USE_V1=1 pytest -s -v vllm/tests/entrypoints/llm/test_accuracy.py::test_lm_eval_accuracy_v1_engine 2>&1 | tee out.txt

github-actions · 2025-02-17T06:28:39Z

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

🚀

mergify · 2025-02-18T06:20:29Z

This pull request has merge conflicts that must be resolved before it can be
merged. Please rebase the PR, @vanbasten23.

https://docs.github.com/en/pull-requests/collaborating-with-pull-requests/working-with-forks/syncing-a-fork

…eal attn_metadata in dummy_run and basic.py is still working fine.

…elp much about the dynamo compilation

…rV1.forward

vllm/v1/core/scheduler.py

examples/offline_inference/basic/basic.py

vanbasten23 · 2025-02-21T19:43:00Z

@alexm-redhat @WoosukKwon , there are 2 issues currently. One is running vllm/tests/entrypoints/llm/test_accuracy.py::test_lm_eval_accuracy_v1_engine is very slow (2h, probably due to excessive compiling), which I'm investigating. Another issue is the test vllm/tests/entrypoints/llm/test_accuracy.py::test_lm_eval_accuracy_v1_engine actually fails. Any suggestion on how to find a smaller repro in order to debug?

vanbasten23 · 2025-02-21T20:02:05Z

vllm/v1/worker/tpu_model_runner.py

        self.model = torch.compile(model,
                                   backend="openxla",
                                   fullgraph=True,
                                   dynamic=False)


@WoosukKwon , do you remember the reason why we set dynamic=False?

vanbasten23 · 2025-02-27T19:48:27Z

cc @bvrockwell

vanbasten23 · 2025-02-27T19:50:02Z

@vllm-v-team@google.com

vanbasten23 · 2025-02-27T20:09:31Z

hey @mgoin , it's my first PR in vLLM repo. I see "pre-commit / pre-commit (pull_request)" in the CI is red, it seems that it complains the format and mypy. For formatting, is there a linter I can use in vLLM?

requirements-tpu.txt

mgoin · 2025-02-27T20:27:52Z

Hey @vanbasten23 please install precommit using these directions https://docs.vllm.ai/en/latest/contributing/overview.html#testing

pip install -r requirements-dev.txt
pre-commit install --hook-type pre-commit --hook-type commit-msg

Then on your next commit it will apply

vanbasten23 · 2025-02-27T21:16:02Z

Hey @vanbasten23 please install precommit using these directions https://docs.vllm.ai/en/latest/contributing/overview.html#testing
pip install -r requirements-dev.txt
pre-commit install --hook-type pre-commit --hook-type commit-msg
Then on your next commit it will apply

Thanks. I followed it. Somehow, running pre-commit run --all-files removed https://github.com/vanbasten23/vllm/blob/58d1b2aa772deb166355423997fbf5c1b6b186a1/vllm/v1/attention/backends/pallas.py#L7 which is important even though it is not directly used. I've added it manually.

Signed-off-by: Xiongfei Wei <isaacwxf23@gmail.com>

Signed-off-by: mgoin <mgoin64@gmail.com>

mgoin

I am able to get good evaluations with Qwen/Qwen2-1.5B-Instruct on gsm8k with a small number of samples and with a small batch size. The attention kernel clearly has issues at the moment with compilation time/memory usage, but this interface is where we should go. So I think we should land this as-is and quickly iterate with new kernels to improve the usability and fix bugs, rather than stay in the emulated V0 style we have on main currently.

Thank you for your great work!

vanbasten23 · 2025-02-28T05:29:44Z

Thanks @mgoin for reviewing the PR.

vanbasten23 · 2025-02-28T15:07:18Z

Hi @mgoin , could you help merge the PR? I don't see a merge button on my side.

mgoin · 2025-02-28T15:09:59Z

Thanks for the ping, yes you need committer status to merge, which I'll handle. Let me quickly chat with @alexm-redhat before merging

alexm-redhat

@vanbasten23 thanks for the hard work on this, the PR looks good to land! Just some small comments, but you can address them also after landing.

vllm/v1/attention/backends/pallas.py

…1 on TPU (vllm-project#13379) Signed-off-by: Xiongfei Wei <isaacwxf23@gmail.com> Signed-off-by: mgoin <mgoin64@gmail.com> Co-authored-by: mgoin <mgoin64@gmail.com>

…1 on TPU (vllm-project#13379) Signed-off-by: Xiongfei Wei <isaacwxf23@gmail.com> Signed-off-by: mgoin <mgoin64@gmail.com> Co-authored-by: mgoin <mgoin64@gmail.com> Signed-off-by: Louis Ulmer <ulmerlouis@gmail.com>

…1 on TPU (vllm-project#13379) Signed-off-by: Xiongfei Wei <isaacwxf23@gmail.com> Signed-off-by: mgoin <mgoin64@gmail.com> Co-authored-by: mgoin <mgoin64@gmail.com>

merge prompt and decode

8e5bbb3

mergify bot added ci/build v1 labels Feb 17, 2025

add more comments

d0eac0f

mergify bot added the needs-rebase label Feb 18, 2025

vanbasten23 added 2 commits February 19, 2025 07:04

cleaned up a bit

2830ed4

disable print, enable torch.compile

2316f14

vanbasten23 force-pushed the xiowei/tpu_v1_kernel_integration_take2 branch from f5f21f2 to 2316f14 Compare February 19, 2025 17:55

vanbasten23 added 7 commits February 19, 2025 20:54

pad block_table 2nd dim to a multiple of 128 to accomodate the kernel.

f5d5429

Updated the torch_xla pin again: the smem oom is gone. Also use the r…

08cda8f

…eal attn_metadata in dummy_run and basic.py is still working fine.

remove total_num_scheduled_tokens from attn_metadata. But it didn't h…

89ea8f1

…elp much about the dynamo compilation

pull total_num_scheduled_tokens and logits_indices out of ModelWrappe…

6520319

…rV1.forward

change position_ids to use int32 instead of int64

e272741

fix merge conflict

8f28a14

removed my comments

ea03585

vanbasten23 commented Feb 21, 2025

View reviewed changes

vllm/v1/core/scheduler.py Outdated Show resolved Hide resolved

mergify bot removed the needs-rebase label Feb 21, 2025

vanbasten23 commented Feb 21, 2025

View reviewed changes

examples/offline_inference/basic/basic.py Outdated Show resolved Hide resolved

vanbasten23 marked this pull request as ready for review February 21, 2025 17:18

vanbasten23 requested review from WoosukKwon, alexm-redhat, comaniac, njhill, robertgshaw2-redhat and ywang96 as code owners February 21, 2025 17:18

vanbasten23 commented Feb 21, 2025

View reviewed changes

remove some comments. Fix mypy annotation.

84ec082

mgoin reviewed Feb 27, 2025

View reviewed changes

requirements-tpu.txt Outdated Show resolved Hide resolved

run linter

1a942d5

vanbasten23 force-pushed the xiowei/tpu_v1_kernel_integration_take2 branch from 3690070 to 1a942d5 Compare February 27, 2025 21:10

vanbasten23 and others added 4 commits February 27, 2025 21:51

Bump torch_xla version again

0592890

Signed-off-by: Xiongfei Wei <isaacwxf23@gmail.com>

Fix lint issues

b9bac9a

Signed-off-by: mgoin <mgoin64@gmail.com>

Revert basic

6bf9e68

Signed-off-by: mgoin <mgoin64@gmail.com>

Keep import torch_xla.experimental.custom_kernel

175224a

Signed-off-by: mgoin <mgoin64@gmail.com>

mgoin approved these changes Feb 28, 2025

View reviewed changes

mgoin added ready ONLY add when PR is ready to merge/full CI is needed tpu Related to Google TPUs labels Feb 28, 2025

mgoin changed the title ~~Integrate the new ragged paged attention kernel with vLLM v1 on TPU~~ [V1][TPU] Integrate the new ragged paged attention kernel with vLLM v1 on TPU Feb 28, 2025

alexm-redhat approved these changes Feb 28, 2025

View reviewed changes

vllm/v1/attention/backends/pallas.py Show resolved Hide resolved

vllm/v1/attention/backends/pallas.py Show resolved Hide resolved

mgoin merged commit c3b6559 into vllm-project:main Feb 28, 2025
47 checks passed

vanbasten23 mentioned this pull request Mar 3, 2025

[v1] Add comments to the new ragged paged attention Pallas kernel #14155

Merged

mgoin mentioned this pull request Mar 3, 2025

[V1][TPU] TPU multimodal model support for ragged attention #14158

Merged

miladm assigned vanbasten23 Mar 13, 2025

ckhordiasma mentioned this pull request Apr 17, 2025

[do not merge] pr test for nm changes into 2.20 red-hat-data-services/vllm#107

Closed

Uh oh!

[V1][TPU] Integrate the new ragged paged attention kernel with vLLM v1 on TPU #13379

[V1][TPU] Integrate the new ragged paged attention kernel with vLLM v1 on TPU #13379

Uh oh!

Conversation

vanbasten23 commented Feb 17, 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 Feb 17, 2025

Uh oh!

mergify bot commented Feb 18, 2025

Uh oh!

Uh oh!

Uh oh!

vanbasten23 commented Feb 21, 2025

Uh oh!

vanbasten23 Feb 21, 2025

Choose a reason for hiding this comment

Uh oh!

vanbasten23 commented Feb 27, 2025

Uh oh!

vanbasten23 commented Feb 27, 2025

Uh oh!

vanbasten23 commented Feb 27, 2025

Uh oh!

Uh oh!

mgoin commented Feb 27, 2025

Uh oh!

vanbasten23 commented Feb 27, 2025

Uh oh!

mgoin left a comment

Choose a reason for hiding this comment

Uh oh!

vanbasten23 commented Feb 28, 2025

Uh oh!

vanbasten23 commented Feb 28, 2025

Uh oh!

mgoin commented Feb 28, 2025

Uh oh!

alexm-redhat left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

vanbasten23 commented Feb 17, 2025 •

edited by github-actions bot

Loading