Triton Attention: Support cross-layers blocks (#30687)

Signed-off-by: Or Ozeri <oro@il.ibm.com>

vllm/v1/attention/backends/triton_attn.py

@@ -290,6 +290,19 @@ class TritonAttentionBackend(AttentionBackend):
             raise ValueError("Block size must be a multiple of 16.")
         return (num_blocks, 2, block_size, num_kv_heads, head_size)
 
+    @staticmethod
+    def get_kv_cache_stride_order(
+        include_num_layers_dimension: bool = False,
+    ) -> tuple[int, ...]:
+        # `stride_order` indicates the permutation that gets
+        # us from `get_kv_cache_shape` to the actual memory layout we want.
+        if include_num_layers_dimension:
+            # (num_blocks, num_layers, 2, block_size, num_kv_heads, head_size)
+            return (1, 0, 2, 3, 4, 5)
+
+        # (num_blocks, 2, block_size, num_kv_heads, head_size)
+        return (0, 1, 2, 3, 4)
+
     @staticmethod
     def use_cascade_attention(*args, **kwargs) -> bool:
         return False