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 14, 05:30 UTC
[CI/Build] tests(v1): feed Triton attention the (num_blocks, 2, …) KV cache layout in backend-correctness tests (#26663)
Signed-off-by: Huamin Li <3ericli@gmail.com> Co-authored-by: Ye (Charlotte) Qi <yeq@meta.com>
Huamin Li committed
Oct 17, 2025 at 21:11 UTC
c312320764193e7d0ffa99d247c61efe5458a635
1 file changed
+3
-2
tests/v1/attention/test_attention_backends.py
+3
-2
@@ -423,13 +423,14 @@ def _test_backend_correctness(
423
for backend_name in backend_to_test:
424
# FlashAttentionm + FlexAttention:
425
# [2, num_blocks, block_size, num_kv_heads, head_size]
426
- # FlashInfer:
426
+ # FlashInfer + Triton:
427
# [num_blocks, 2, block_size, num_kv_heads, head_size]
428
# Select the appropriate KV cache format for each backend
429
kv_cache_for_backend = kv_cache
430
- if backend_name == _Backend.FLASHINFER:
430
+ if backend_name in (_Backend.FLASHINFER, _Backend.TRITON_ATTN):
431
kv_cache_for_backend = kv_cache.transpose(0, 1)
432
433
+ if backend_name == _Backend.FLASHINFER:
434
# For FlashInfer default to HND layout and
435
kv_cache_for_backend = (
436
kv_cache_for_backend.transpose(2, 3).contiguous().transpose(2, 3)