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 15, 09:45 UTC
[Kernel] [V1] Fix performance regression for triton unified attention (#18161)
Signed-off-by: Thomas Parnell <tpa@zurich.ibm.com> Co-authored-by: Lucas Wilkinson <lwilkinson@neuralmagic.com>
Thomas Parnell committed
May 15, 2025 at 15:39 UTC
01c22335baa03dbbc3ae662c76a089ff1c5b5742
2 files changed
+18
-5
vllm/attention/ops/triton_unified_attention.py
+2
-2
@@ -56,11 +56,11 @@ def kernel_unified_attention_2d(
56
stride_k_cache_0: tl.int64, # int
57
stride_k_cache_1: tl.int64, # int
58
stride_k_cache_2: tl.int64, # int
59
- stride_k_cache_3: tl.int64, # int
59
+ stride_k_cache_3: tl.constexpr, # int
60
stride_v_cache_0: tl.int64, # int
61
stride_v_cache_1: tl.int64, # int
62
stride_v_cache_2: tl.int64, # int
63
- stride_v_cache_3: tl.int64, # int
63
+ stride_v_cache_3: tl.constexpr, # int
64
query_start_len_ptr, # [num_seqs+1]
65
BLOCK_Q: tl.constexpr, # int
66
num_seqs: tl.int32,
vllm/v1/attention/backends/triton_attn.py
+16
-3
@@ -1,6 +1,6 @@
1
# SPDX-License-Identifier: Apache-2.0
2
"""Attention layer with PagedAttention and Triton prefix prefill."""
3
-from typing import Any, Optional
3
+from typing import TYPE_CHECKING, Any, Optional
4
5
import torch
6
@@ -12,10 +12,23 @@ from vllm.logger import init_logger
12
from vllm.platforms import current_platform
13
from vllm.v1.attention.backends.flash_attn import (
14
FlashAttentionMetadata, FlashAttentionMetadataBuilder)
15
+from vllm.v1.kv_cache_interface import AttentionSpec
16
+from vllm.v1.worker.block_table import BlockTable
17
+
18
+if TYPE_CHECKING:
19
+ from vllm.v1.worker.gpu_model_runner import GPUModelRunner
20
21
logger = init_logger(__name__)
22
23
24
+class TritonAttentionMetadataBuilder(FlashAttentionMetadataBuilder):
25
+
26
+ def __init__(self, runner: "GPUModelRunner", kv_cache_spec: AttentionSpec,
27
+ block_table: BlockTable):
28
+ super().__init__(runner, kv_cache_spec, block_table)
29
+ self.aot_schedule = False
30
+
31
+
32
class TritonAttentionBackend(AttentionBackend):
33
34
accept_output_buffer: bool = True
@@ -52,8 +65,8 @@ class TritonAttentionBackend(AttentionBackend):
65
return False
66
67
@staticmethod
55
- def get_builder_cls() -> type["FlashAttentionMetadataBuilder"]:
56
- return FlashAttentionMetadataBuilder
68
+ def get_builder_cls() -> type["TritonAttentionMetadataBuilder"]:
69
+ return TritonAttentionMetadataBuilder
70
71
72
class TritonAttentionImpl(AttentionImpl):