Skip to content

[Hardware] Initial TPU integration #5292

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom
Merged
merged 193 commits into from
Jun 12, 2024
Merged

[Hardware] Initial TPU integration #5292

merged 193 commits into from
Jun 12, 2024

Conversation

@WoosukKwon
Copy link
Collaborator

@WoosukKwon WoosukKwon commented Jun 5, 2024 •
edited
Loading

This PR implements the initial integration of the Google TPU backend. It uses PyTorch XLA for maximal reuse of the existing code base.

The PR features:

  • Seamless support for popular HF models such as Llama, Mistral, Gemma, etc. The model's head size must be either 128 or 256.
  • Basic functionalities of vLLM, including continuous batching
  • Optimized pallas kernels for FlashAttention and PagedAttention

TODOs (next steps):

  • (Fast) top-p sampling (disabled for now due to performance issues)
  • Distributed (tensor-parallel) inference
  • INT8 quantization
  • MoE
  • Support best_of > 1

WoosukKwon added 30 commits April 1, 2024 02:01
@WoosukKwon
Copy link
Collaborator Author

@alanwaketan Please take a look!

@WoosukKwon WoosukKwon changed the title [WIP][Hardware] Initial TPU integration [Hardware] Initial TPU integration Jun 11, 2024
@WoosukKwon WoosukKwon marked this pull request as ready for review June 11, 2024 17:47
@WoosukKwon WoosukKwon requested a review from JackCaoG June 11, 2024 17:57
JackCaoG
JackCaoG approved these changes Jun 11, 2024
View reviewed changes
Copy link

@JackCaoG JackCaoG left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

LGTM, thanks!

rkooo567
rkooo567 approved these changes Jun 12, 2024
View reviewed changes
Copy link
Collaborator

@rkooo567 rkooo567 left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Mostly asking for comments for clarifying some code!

For testing, are we planning to add relevant CI in the future?

@WoosukKwon
Copy link
Collaborator Author

@rkooo567 Thanks for the quality review!

@WoosukKwon WoosukKwon merged commit 1a8bfd9 into main Jun 12, 2024
@WoosukKwon WoosukKwon deleted the torch-xla branch June 12, 2024 18:53
@WoosukKwon WoosukKwon mentioned this pull request Jun 14, 2024
Closed
6 tasks
robertgshaw2-redhat pushed a commit to neuralmagic/nm-vllm that referenced this pull request Jun 16, 2024
@WoosukKwon
[Hardware] Initial TPU integration (vllm-project#5292)
134aadc
joerunde pushed a commit to joerunde/vllm that referenced this pull request Jun 17, 2024
@WoosukKwon @joerunde
[Hardware] Initial TPU integration (vllm-project#5292)
89b6334
xjpang pushed a commit to xjpang/vllm that referenced this pull request Jun 27, 2024
@WoosukKwon
[Hardware] Initial TPU integration (vllm-project#5292)
5b00d3f
xjpang pushed a commit to xjpang/vllm that referenced this pull request Jul 8, 2024
@WoosukKwon
[Hardware] Initial TPU integration (vllm-project#5292)
3b22eed
xjpang pushed a commit to xjpang/vllm that referenced this pull request Jul 24, 2024
@WoosukKwon
[Hardware] Initial TPU integration (vllm-project#5292)
2e9fb8a
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

@rkooo567 rkooo567 rkooo567 approved these changes

+1 more reviewer

@JackCaoG JackCaoG JackCaoG approved these changes

Reviewers whose approvals may not affect merge requirements

Assignees

No one assigned

Labels

tpu Related to Google TPUs

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

3 participants

@WoosukKwon @JackCaoG @rkooo567