[v1] Add comments to the new ragged paged attention Pallas kernel (vllm-project#14155)

Signed-off-by: Xiongfei Wei <[email protected]>
Co-authored-by: Michael Goin <[email protected]>

Author: 2 people

vllm/v1/attention/backends/pallas.py

@@ -11,6 +11,7 @@
                                               AttentionLayer, AttentionType)
 from vllm.attention.backends.utils import CommonAttentionState
 
+# These are the 2 tunable parameters of the paged attention Pallas kernel.
 NUM_QUERIES_PER_BLOCK = 16
 NUM_KV_PAGES_PER_BLOCK = 128
 
@@ -154,6 +155,9 @@ def forward(
             write_to_kv_cache(key, value, key_cache, value_cache, slot_mapping)
 
         query = query * self.scale
+        # use_kernel switches between using kernel or reference implementation
+        # (non kernel: https://github.com/pytorch/xla/blob/cee0820e78fc9675e2d0511db891fd44342e890d/torch_xla/experimental/custom_kernel.py#L890).
+        use_kernel = False
         output = torch.ops.xla.ragged_paged_attention(
             query,
             key_cache,
@@ -164,7 +168,7 @@ def forward(
             attn_metadata.num_seqs,
             num_kv_pages_per_block=NUM_KV_PAGES_PER_BLOCK,
             num_queries_per_block=NUM_QUERIES_PER_BLOCK,
-            use_kernel=False,
+            use_kernel=use_kernel,
         )
 
         return output.reshape(num_tokens, hidden_size)