[TPU] Implement prefix caching for TPUs #10307
Conversation
|
👋 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:
🚀 |
|
Nice work! |
| output = output.permute(0, 2, 1, 3) | ||
| else: | ||
| # Prefill with paged KV cache. | ||
| # TODO(woosuk): Tune the below knobs. |
There was a problem hiding this comment.
Thanks Woosuk for writing the PR.
I'm benchmarking the kernel so likely I'll have some recommended num_kv_pages_per_compute_block/num_queries_per_compute_block to share.
Also, the revised paged attention kernel is in torch_xla nightly. Could you try again? I pulled your PR and it seems it needs additional work to get the effective_q_lens and plumb it to the kernel.
cc: @WoosukKwon
There was a problem hiding this comment.
@vanbasten23 Is the fixed kernel available in today's nightly?
There was a problem hiding this comment.
@vanbasten23 After the kernel fix, the model generates correct outputs with prefix caching 🎉
examples/offline_inference_tpu.py
Outdated
| outputs = llm.generate(prompts, sampling_params) | ||
| for output, answer in zip(outputs, answers): | ||
| for output in outputs: |
There was a problem hiding this comment.
I wonder if you need a test for the prefix caching.
|
Btw, which command did you use run |
| num_kv_pages_per_compute_block = 16 | ||
| num_queries_per_compute_block = 16 | ||
| assert seq_len % num_queries_per_compute_block == 0 | ||
| output = torch.ops.xla.multi_queries_paged_attention( |
There was a problem hiding this comment.
@vanbasten23 - does this new kernel have the same SMEM requirements as the original paged_attention where the entire block table is stored in SMEM?
E.g. for the decoding run (see below), we split the batch dimension into smaller chunks and run the kernel multiple times
There was a problem hiding this comment.
hey @robertgshaw2-neuralmagic , yes this new kernel have the same SMEM requirements. I am aware of the SMEM OOM issue you mentioned and we plan to address it.
|
This pull request has merge conflicts that must be resolved before it can be |
This is weird. Which version & TPU are you using? |
|
I will double check, update this PR, and merge it tonight. |
I'm using TPU v5e but I'm not sure if it depends on a specific TPU version. |
Signed-off-by: Woosuk Kwon <[email protected]>
WoosukKwon
force-pushed
the
tpu-prefix-caching
branch
from
274af74 to
f525281
Compare
Signed-off-by: Woosuk Kwon <[email protected]>
Signed-off-by: Woosuk Kwon <[email protected]> Signed-off-by: Tyler Michael Smith <[email protected]>
Signed-off-by: Woosuk Kwon <[email protected]> Signed-off-by: Clay <[email protected]>
Signed-off-by: Woosuk Kwon <[email protected]> Signed-off-by: Maxime Fournioux <[email protected]>
Signed-off-by: Woosuk Kwon <[email protected]>
This PR implements the prefix caching support for the TPU backend.