A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 10, 17:45 UTC
Triton Attention: Support cross-layers blocks (#30687)
Signed-off-by: Or Ozeri <oro@il.ibm.com>
Or Ozeri committed
Jan 5, 2026 at 21:29 UTC
d8e38d493907882eb247a2ba4ced32610020ac27
2 files changed
+14
-3
tests/v1/kv_offload/test_cpu_offloading.py
+1
-3
@@ -15,12 +15,10 @@ from vllm.distributed.kv_events import BlockStored, KVEventBatch
15
from vllm.platforms import current_platform
16
17
CPU_BLOCK_SIZES = [48]
18
-ATTN_BACKENDS = ["FLASH_ATTN"]
18
+ATTN_BACKENDS = ["FLASH_ATTN", "TRITON_ATTN"]
19
20
if current_platform.is_cuda():
21
ATTN_BACKENDS.append("FLASHINFER")
22
-elif current_platform.is_rocm():
23
- ATTN_BACKENDS = ["TRITON_ATTN"]
22
23
24
class MockSubscriber:
vllm/v1/attention/backends/triton_attn.py
+13
@@ -290,6 +290,19 @@ class TritonAttentionBackend(AttentionBackend):
290
raise ValueError("Block size must be a multiple of 16.")
291
return (num_blocks, 2, block_size, num_kv_heads, head_size)
292
293
+ @staticmethod
294
+ def get_kv_cache_stride_order(
295
+ include_num_layers_dimension: bool = False,
296
+ ) -> tuple[int, ...]:
297
+ # `stride_order` indicates the permutation that gets
298
+ # us from `get_kv_cache_shape` to the actual memory layout we want.
299
+ if include_num_layers_dimension:
300
+ # (num_blocks, num_layers, 2, block_size, num_kv_heads, head_size)
301
+ return (1, 0, 2, 3, 4, 5)
302
+
303
+ # (num_blocks, 2, block_size, num_kv_heads, head_size)
304
+ return (0, 1, 2, 3, 4)
305
+
306
@staticmethod
307
def use_cascade_attention(*args, **kwargs) -> bool:
308
return False