[WIP] [V1] TPU support #11936
[WIP] [V1] TPU support #11936
Conversation
alexm-neuralmagic
requested review from
WoosukKwon,
robertgshaw2-neuralmagic,
njhill,
ywang96 and
comaniac
as code owners
|
👋 Hi! Thank you for contributing to the vLLM project. 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 do one of these:
🚀 |
|
This pull request has merge conflicts that must be resolved before it can be |
alexm-neuralmagic
requested review from
DarkLight1337 and
simon-mo
as code owners
| def __init__( | ||
| self, | ||
| vllm_config: VllmConfig, | ||
| device: torch.device, | ||
| ): |
There was a problem hiding this comment.
the function implementation is almost identical to gpu_model_runner.py, it would be better if we build a ModelRunnerBase class and derive from the base class instead of duplicating.
| return PrefillInputData( | ||
| request_ids=prefill_request_ids, | ||
| prompt_lens=prefill_prompt_lens, | ||
| token_ids=prefill_token_ids, | ||
| position_ids=prefill_position_ids, | ||
| attn_metadata=prefill_attn_metadata, | ||
| ) |
There was a problem hiding this comment.
remove the PrefillInputData data structure, and make it consistent with gpu_model_runner ?
| def _prepare_prefill_inputs( | ||
| self, | ||
| num_scheduled_tokens: List[int], | ||
| ) -> PrefillInputData: |
There was a problem hiding this comment.
do we still need _prepare_prefill_inputs in V1 ?
(i'm assuming this can be handled with _prepare_inputs already.)
There was a problem hiding this comment.
We need to run separate prefill and decode for TPU since we don't have the attention kernel support yet. This is on the way so we hope to remove soon
| attn_metadata=prefill_attn_metadata, | ||
| ) | ||
|
|
||
| def _prepare_decode_inputs(self) -> DecodeInputData: |
There was a problem hiding this comment.
again, do we really need _prepare_decode_inputs in V1 architecture?
(i'm assuming this can be handled with _prepare_inputs already.)
| effective_query_lens=None, | ||
| )) | ||
|
|
||
| def _prepare_inputs(self, scheduler_output: "SchedulerOutput"): |
There was a problem hiding this comment.
this is almost identical with current gpu_model_runner implementation, consider reusing instead of duplicating ?
This PR is a rebase and modification of @robertgshaw2-neuralmagic original PR for TPU support from 1.5 months ago #10241
TODOs: