[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:
🚀 |
Signed-off-by: Woosuk Kwon <[email protected]>
|
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 🎉
Signed-off-by: Woosuk Kwon <[email protected]>
| 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
No description provided.