[TPU] Implement prefix caching for TPUs (#10307)

Signed-off-by: Woosuk Kwon <[email protected]>

requirements-tpu.txt

@@ -16,8 +16,8 @@ ray[default]
 --find-links https://storage.googleapis.com/libtpu-releases/index.html
 --find-links https://storage.googleapis.com/jax-releases/jax_nightly_releases.html
 --find-links https://storage.googleapis.com/jax-releases/jaxlib_nightly_releases.html
-torch==2.6.0.dev20241028+cpu
-torchvision==0.20.0.dev20241028+cpu
-torch_xla[tpu] @ https://storage.googleapis.com/pytorch-xla-releases/wheels/tpuvm/torch_xla-2.6.0.dev20241028-cp310-cp310-linux_x86_64.whl
+torch==2.6.0.dev20241114+cpu
+torchvision==0.20.0.dev20241114+cpu
+torch_xla[tpu] @ https://storage.googleapis.com/pytorch-xla-releases/wheels/tpuvm/torch_xla-2.6.0.dev20241114-cp310-cp310-linux_x86_64.whl
 jaxlib==0.4.32.dev20240829
 jax==0.4.32.dev20240829

vllm/attention/backends/pallas.py

@@ -65,15 +65,14 @@ class PallasMetadata(AttentionMetadata):
     # or all decoding.
     block_tables: Optional[torch.Tensor] = None
     context_lens: Optional[torch.Tensor] = None
+    effective_query_lens: Optional[torch.Tensor] = None
 
     @property
     def prefill_metadata(self) -> Optional["PallasMetadata"]:
         if self.num_prefills == 0:
             return None
 
         assert self.num_decode_tokens == 0
-        assert self.block_tables is None
-        assert self.context_lens is None
         return self
 
     @property
@@ -186,29 +185,50 @@ def forward(
 
         query = query * self.scale
         if attn_metadata.num_prefills > 0:
-            assert seq_len % 16 == 0, (
-                "Pallas FlashAttention kernel requires seq_len to be a "
-                f"multiple of 16 but got {seq_len}")
-
-            # Handle GQA/MQA.
-            if self.num_kv_heads != self.num_heads:
-                key = key.repeat_interleave(self.num_queries_per_kv, dim=-2)
-                key = key.view(batch_size, seq_len, self.num_heads,
-                               self.head_size)
-                value = value.repeat_interleave(self.num_queries_per_kv,
+            if attn_metadata.block_tables is None:
+                # Prefill without paged KV cache.
+                assert seq_len % 16 == 0, (
+                    "Pallas FlashAttention kernel requires seq_len to be a "
+                    f"multiple of 16 but got {seq_len}")
+
+                # Handle GQA/MQA.
+                if self.num_kv_heads != self.num_heads:
+                    key = key.repeat_interleave(self.num_queries_per_kv,
                                                 dim=-2)
-                value = value.view(batch_size, seq_len, self.num_heads,
+                    key = key.view(batch_size, seq_len, self.num_heads,
                                    self.head_size)
-            # FlashAttention requires [batch_size, num_heads, seq_len, d_model]
-            # while the input is [batch_size, seq_len, num_heads, d_model].
-            # Permute the input to match the required format.
-            output = torch.ops.xla.flash_attention(
-                query.permute(0, 2, 1, 3),
-                key.permute(0, 2, 1, 3),
-                value.permute(0, 2, 1, 3),
-                True,
-            )
-            output = output.permute(0, 2, 1, 3)
+                    value = value.repeat_interleave(self.num_queries_per_kv,
+                                                    dim=-2)
+                    value = value.view(batch_size, seq_len, self.num_heads,
+                                       self.head_size)
+                # FlashAttention kernel requires the input shape to be
+                # [batch_size, num_heads, seq_len, d_model]
+                # while the input is [batch_size, seq_len, num_heads, d_model].
+                # Permute the input to match the required format.
+                output = torch.ops.xla.flash_attention(
+                    query.permute(0, 2, 1, 3),
+                    key.permute(0, 2, 1, 3),
+                    value.permute(0, 2, 1, 3),
+                    True,
+                )
+                output = output.permute(0, 2, 1, 3)
+            else:
+                # Prefill with paged KV cache.
+                # TODO(woosuk): Tune the below knobs.
+                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(
+                    query,
+                    key_cache,
+                    value_cache,
+                    attn_metadata.context_lens,
+                    attn_metadata.block_tables,
+                    attn_metadata.effective_query_lens,
+                    num_kv_pages_per_compute_block,
+                    num_queries_per_compute_block,
+                    use_kernel=True,
+                )
         else:
             # Decoding run.
             assert kv_cache[0].numel() > 0