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
[V1][Attention] Split triton_attn in triton-only and rocm specific backends (#24648)
Signed-off-by: Burkhard Ringlein <ngl@zurich.ibm.com>
Burkhard Ringlein committed
Sep 22, 2025 at 17:20 UTC
175811e3b53f8f13eb3e8ac6aae02050f6e1412f
5 files changed
+483
-124
vllm/engine/arg_utils.py
+1
@@ -1494,6 +1494,7 @@ class EngineArgs:
1494
"FLEX_ATTENTION",
1495
"TREE_ATTN",
1496
"XFORMERS_VLLM_V1",
1497
+ "ROCM_ATTN_VLLM_V1",
1498
]
1499
if (envs.is_set("VLLM_ATTENTION_BACKEND")
1500
and envs.VLLM_ATTENTION_BACKEND not in V1_BACKENDS):
vllm/platforms/interface.py
+1
@@ -67,6 +67,7 @@ class _Backend(enum.Enum):
67
FLEX_ATTENTION = enum.auto()
68
TREE_ATTN = enum.auto()
69
XFORMERS_VLLM_V1 = enum.auto()
70
+ ROCM_ATTN_VLLM_V1 = enum.auto()
71
72
73
class PlatformEnum(enum.Enum):
vllm/platforms/rocm.py
+10
@@ -231,7 +231,17 @@ class RocmPlatform(Platform):
231
logger.info("Using Flash Attention backend on V1 engine.")
232
return ("vllm.v1.attention.backends."
233
"rocm_aiter_fa.AiterFlashAttentionBackend")
234
+ elif (envs.VLLM_ROCM_USE_AITER and
235
+ envs.VLLM_USE_AITER_UNIFIED_ATTENTION) or \
236
+ envs.VLLM_V1_USE_PREFILL_DECODE_ATTENTION or \
237
+ selected_backend == _Backend.ROCM_ATTN_VLLM_V1:
238
+ # rocm specific backend, with aiter and/or
239
+ # triton prefix-prefill
240
+ logger.info("Using Rocm/Aiter Attention backend on V1 engine.")
241
+ return ("vllm.v1.attention.backends."
242
+ "rocm_attn.RocmAttentionBackend")
243
else:
244
+ # default case, using triton unified attention
245
logger.info("Using Triton Attention backend on V1 engine.")
246
return ("vllm.v1.attention.backends."
247
"triton_attn.TritonAttentionBackend")
vllm/v1/attention/backends/rocm_attn.py
new
+426
@@ -0,0 +1,426 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""Attention layer with PagedAttention and Triton prefix prefill."""
4
+from dataclasses import dataclass
5
+from functools import cache
6
+from typing import ClassVar, Optional
7
+
8
+import torch
9
+
10
+from vllm import _custom_ops as ops
11
+from vllm import envs
12
+from vllm.attention.backends.abstract import (AttentionBackend, AttentionImpl,
13
+ AttentionMetadata, AttentionType)
14
+from vllm.attention.ops.chunked_prefill_paged_decode import (
15
+ chunked_prefill_paged_decode)
16
+from vllm.attention.ops.paged_attn import PagedAttention
17
+from vllm.config import VllmConfig
18
+from vllm.logger import init_logger
19
+from vllm.model_executor.layers.quantization.utils.quant_utils import (
20
+ QuantKey, kFp8StaticTensorSym)
21
+from vllm.platforms import current_platform
22
+from vllm.v1.attention.backends.flash_attn import FlashAttentionMetadata
23
+from vllm.v1.attention.backends.utils import (AttentionCGSupport,
24
+ AttentionMetadataBuilder,
25
+ CommonAttentionMetadata)
26
+from vllm.v1.kv_cache_interface import AttentionSpec
27
+
28
+logger = init_logger(__name__)
29
+
30
+
31
+@dataclass
32
+class RocmAttentionMetadata:
33
+ # NOTE(sang): Definition of context_len, query_len, and seq_len.
34
+ # |---------- N-1 iteration --------|
35
+ # |---------------- N iteration ---------------------|
36
+ # |- tokenA -|......................|-- newTokens ---|
37
+ # |---------- context_len ----------|
38
+ # |-------------------- seq_len ---------------------|
39
+ # |-- query_len ---|
40
+
41
+ num_actual_tokens: int # Number of tokens excluding padding.
42
+ max_query_len: int
43
+ query_start_loc: torch.Tensor
44
+ max_seq_len: int
45
+ seq_lens: torch.Tensor
46
+ block_table: torch.Tensor
47
+ slot_mapping: torch.Tensor
48
+
49
+ # For cascade attention.
50
+ use_cascade: bool
51
+ common_prefix_len: int
52
+ cu_prefix_query_lens: Optional[torch.Tensor]
53
+ prefix_kv_lens: Optional[torch.Tensor]
54
+ suffix_kv_lens: Optional[torch.Tensor]
55
+
56
+ # Optional aot scheduling
57
+ scheduler_metadata: Optional[torch.Tensor] = None
58
+ prefix_scheduler_metadata: Optional[torch.Tensor] = None
59
+
60
+
61
+class RocmAttentionMetadataBuilder(
62
+ AttentionMetadataBuilder[RocmAttentionMetadata]):
63
+ cudagraph_support: ClassVar[AttentionCGSupport] = AttentionCGSupport.ALWAYS
64
+
65
+ def __init__(self, kv_cache_spec: AttentionSpec, layer_names: list[str],
66
+ vllm_config: VllmConfig, device: torch.device):
67
+ super().__init__(kv_cache_spec, layer_names, vllm_config, device)
68
+
69
+ self.block_size = kv_cache_spec.block_size
70
+
71
+ model_config = vllm_config.model_config
72
+ self.num_heads_q = model_config.get_num_attention_heads(
73
+ vllm_config.parallel_config)
74
+ self.num_heads_kv = model_config.get_num_kv_heads(
75
+ vllm_config.parallel_config)
76
+ self.headdim = model_config.get_head_size()
77
+
78
+ def build_for_cudagraph_capture(
79
+ self, common_attn_metadata: CommonAttentionMetadata
80
+ ) -> RocmAttentionMetadata:
81
+ attn_metadata = self.build(0, common_attn_metadata)
82
+ # When doing full graph capture, setting seq_lens to
83
+ # max_model_len will cause graph capture to be extremely
84
+ # slow, so here we set it to 1.
85
+ attn_metadata.seq_lens.fill_(1)
86
+ return attn_metadata
87
+
88
+ def build(self,
89
+ common_prefix_len: int,
90
+ common_attn_metadata: CommonAttentionMetadata,
91
+ fast_build: bool = False) -> RocmAttentionMetadata:
92
+ num_actual_tokens = common_attn_metadata.num_actual_tokens
93
+ max_query_len = common_attn_metadata.max_query_len
94
+
95
+ max_seq_len = common_attn_metadata.max_seq_len
96
+ query_start_loc = common_attn_metadata.query_start_loc
97
+ seq_lens = common_attn_metadata.seq_lens
98
+ block_table_tensor = common_attn_metadata.block_table_tensor
99
+ slot_mapping = common_attn_metadata.slot_mapping
100
+
101
+ use_cascade = common_prefix_len > 0
102
+
103
+ if use_cascade:
104
+ cu_prefix_query_lens = torch.tensor([0, num_actual_tokens],
105
+ dtype=torch.int32,
106
+ device=self.device)
107
+ prefix_kv_lens = torch.tensor([common_prefix_len],
108
+ dtype=torch.int32,
109
+ device=self.device)
110
+ suffix_kv_lens = (common_attn_metadata.seq_lens_cpu -
111
+ common_prefix_len)
112
+ suffix_kv_lens = suffix_kv_lens.to(self.device)
113
+ else:
114
+ cu_prefix_query_lens = None
115
+ prefix_kv_lens = None
116
+ suffix_kv_lens = None
117
+ prefix_scheduler_metadata = None
118
+
119
+ attn_metadata = RocmAttentionMetadata(
120
+ num_actual_tokens=num_actual_tokens,
121
+ max_query_len=max_query_len,
122
+ query_start_loc=query_start_loc,
123
+ max_seq_len=max_seq_len,
124
+ seq_lens=seq_lens,
125
+ block_table=block_table_tensor,
126
+ slot_mapping=slot_mapping,
127
+ use_cascade=use_cascade,
128
+ common_prefix_len=common_prefix_len,
129
+ cu_prefix_query_lens=cu_prefix_query_lens,
130
+ prefix_kv_lens=prefix_kv_lens,
131
+ suffix_kv_lens=suffix_kv_lens,
132
+ prefix_scheduler_metadata=prefix_scheduler_metadata,
133
+ )
134
+ return attn_metadata
135
+
136
+
137
+class RocmAttentionBackend(AttentionBackend):
138
+
139
+ accept_output_buffer: bool = True
140
+
141
+ @classmethod
142
+ def get_supported_dtypes(cls) -> list[torch.dtype]:
143
+ return [torch.float16, torch.bfloat16]
144
+
145
+ @classmethod
146
+ def get_supported_head_sizes(cls) -> list[int]:
147
+ return [32, 64, 96, 128, 160, 192, 224, 256]
148
+
149
+ @classmethod
150
+ def validate_head_size(cls, head_size: int) -> None:
151
+ supported_head_sizes = cls.get_supported_head_sizes()
152
+ if head_size not in supported_head_sizes:
153
+ attn_type = cls.__name__.removesuffix("Backend")
154
+ raise ValueError(
155
+ f"Head size {head_size} is not supported by {attn_type}. "
156
+ f"Supported head sizes are: {supported_head_sizes}. "
157
+ "Set VLLM_ATTENTION_BACKEND=FLEX_ATTENTION to use "
158
+ "FlexAttention backend which supports all head sizes.")
159
+
160
+ @staticmethod
161
+ def get_name() -> str:
162
+ return "ROCM_ATTN_VLLM_V1"
163
+
164
+ @staticmethod
165
+ def get_impl_cls() -> type["RocmAttentionImpl"]:
166
+ return RocmAttentionImpl
167
+
168
+ @staticmethod
169
+ def get_metadata_cls() -> type["AttentionMetadata"]:
170
+ return RocmAttentionMetadata
171
+
172
+ @staticmethod
173
+ def get_kv_cache_shape(
174
+ num_blocks: int,
175
+ block_size: int,
176
+ num_kv_heads: int,
177
+ head_size: int,
178
+ ) -> tuple[int, ...]:
179
+ if block_size % 16 != 0:
180
+ raise ValueError("Block size must be a multiple of 16.")
181
+ return (2, num_blocks, block_size, num_kv_heads, head_size)
182
+
183
+ @staticmethod
184
+ def use_cascade_attention(*args, **kwargs) -> bool:
185
+ return False
186
+
187
+ @staticmethod
188
+ def get_builder_cls() -> type["RocmAttentionMetadataBuilder"]:
189
+ return RocmAttentionMetadataBuilder
190
+
191
+
192
+@cache
193
+def use_aiter_unified_attention() -> bool:
194
+ """Check if aiter unified attention should be used."""
195
+ # VLLM_ROCM_USE_AITER_MHA needs to set to 0 as well as it is set
196
+ # to 1 as default
197
+ return envs.VLLM_ROCM_USE_AITER \
198
+ and envs.VLLM_USE_AITER_UNIFIED_ATTENTION
199
+
200
+
201
+class RocmAttentionImpl(AttentionImpl):
202
+
203
+ def fused_output_quant_supported(self, quant_key: QuantKey):
204
+ return quant_key == kFp8StaticTensorSym
205
+
206
+ def __init__(
207
+ self,
208
+ num_heads: int,
209
+ head_size: int,
210
+ scale: float,
211
+ num_kv_heads: int,
212
+ alibi_slopes: Optional[list[float]],
213
+ sliding_window: Optional[int],
214
+ kv_cache_dtype: str,
215
+ logits_soft_cap: Optional[float] = None,
216
+ attn_type: AttentionType = AttentionType.DECODER,
217
+ kv_sharing_target_layer_name: Optional[int] = None,
218
+ sinks: Optional[torch.Tensor] = None,
219
+ ) -> None:
220
+ self.num_heads = num_heads
221
+ self.head_size = head_size
222
+ self.scale = float(scale)
223
+ self.num_kv_heads = num_kv_heads
224
+ if alibi_slopes is not None:
225
+ alibi_slopes = torch.tensor(alibi_slopes, dtype=torch.float32)
226
+ self.alibi_slopes = alibi_slopes
227
+ if sliding_window is None:
228
+ self.sliding_window = (-1, -1)
229
+ else:
230
+ self.sliding_window = (sliding_window - 1, 0)
231
+ self.kv_cache_dtype = kv_cache_dtype
232
+ if logits_soft_cap is None:
233
+ # In flash-attn, setting logits_soft_cap as 0 means no soft cap.
234
+ logits_soft_cap = 0
235
+ self.logits_soft_cap = logits_soft_cap
236
+ self.kv_sharing_target_layer_name = kv_sharing_target_layer_name
237
+
238
+ self.num_queries_per_kv = self.num_heads // self.num_kv_heads
239
+
240
+ RocmAttentionBackend.validate_head_size(head_size)
241
+
242
+ if attn_type != AttentionType.DECODER:
243
+ raise NotImplementedError("Encoder self-attention and "
244
+ "encoder/decoder cross-attention "
245
+ "are not implemented for "
246
+ "RocmAttentionImpl")
247
+
248
+ self.fp8_dtype = current_platform.fp8_dtype()
249
+ self.force_prefill_decode_attn = \
250
+ envs.VLLM_V1_USE_PREFILL_DECODE_ATTENTION
251
+
252
+ if not self.force_prefill_decode_attn:
253
+ # If not using prefill decode attention, we use the Triton
254
+ # unified attention implementation.
255
+ if use_aiter_unified_attention():
256
+ logger.info_once(
257
+ "Using aiter unified attention for RocmAttentionImpl")
258
+ from aiter.ops.triton.unified_attention import (
259
+ unified_attention)
260
+ self.unified_attention = unified_attention
261
+ else:
262
+ logger.info_once(
263
+ "Using vllm unified attention for RocmAttentionImpl")
264
+ from vllm.attention.ops.triton_unified_attention import (
265
+ unified_attention)
266
+ self.unified_attention = unified_attention
267
+
268
+ self.sinks = sinks
269
+ if sinks is not None:
270
+ assert sinks.shape[0] == num_heads, (
271
+ "Sinks must have the same number of heads as the number of "
272
+ f"heads in the layer. Sinks shape: {sinks.shape}, "
273
+ f"num_heads: {num_heads}.")
274
+
275
+ def forward(
276
+ self,
277
+ layer: torch.nn.Module,
278
+ query: torch.Tensor,
279
+ key: torch.Tensor,
280
+ value: torch.Tensor,
281
+ kv_cache: torch.Tensor,
282
+ attn_metadata: FlashAttentionMetadata,
283
+ output: Optional[torch.Tensor] = None,
284
+ output_scale: Optional[torch.Tensor] = None,
285
+ output_block_scale: Optional[torch.Tensor] = None,
286
+ ) -> torch.Tensor:
287
+ """Forward pass with FlashAttention.
288
+
289
+ Args:
290
+ query: shape = [num_tokens, num_heads, head_size]
291
+ key: shape = [num_tokens, num_kv_heads, head_size]
292
+ value: shape = [num_tokens, num_kv_heads, head_size]
293
+ kv_cache: shape =
294
+ [2, num_blocks, block_size, num_kv_heads, head_size]
295
+ attn_metadata: Metadata for attention.
296
+ Returns:
297
+ shape = [num_tokens, num_heads * head_size]
298
+ """
299
+ assert output is not None, "Output tensor must be provided."
300
+
301
+ if output_block_scale is not None:
302
+ raise NotImplementedError(
303
+ "fused block_scale output quantization is not yet supported"
304
+ " for RocmAttentionImpl")
305
+
306
+ if attn_metadata is None:
307
+ # Profiling run.
308
+ return output
309
+
310
+ assert attn_metadata.use_cascade is False
311
+
312
+ # IMPORTANT!
313
+ # NOTE(woosuk): With piece-wise CUDA graphs, this method is executed in
314
+ # eager-mode PyTorch. Thus, we need to be careful about any CPU overhead
315
+ # in this method. For example, `view` and `slice` (or `[:n]`) operations
316
+ # are surprisingly slow even in the case they do not invoke any GPU ops.
317
+ # Minimize the PyTorch ops in this method as much as possible.
318
+ # Whenever making a change in this method, please benchmark the
319
+ # performance to make sure it does not introduce any overhead.
320
+
321
+ use_prefill_decode_attn = self.force_prefill_decode_attn
322
+ num_actual_tokens = attn_metadata.num_actual_tokens
323
+
324
+ if use_prefill_decode_attn:
325
+ key_cache, value_cache = PagedAttention.split_kv_cache(
326
+ kv_cache, self.num_kv_heads, self.head_size)
327
+ else:
328
+ key_cache, value_cache = kv_cache.unbind(0)
329
+
330
+ if self.kv_sharing_target_layer_name is None:
331
+ # Reshape the input keys and values and store them in the cache.
332
+ # Skip this if sharing KV cache with an earlier attention layer.
333
+ if use_prefill_decode_attn:
334
+ PagedAttention.write_to_paged_cache(
335
+ key,
336
+ value,
337
+ key_cache,
338
+ value_cache,
339
+ attn_metadata.slot_mapping,
340
+ self.kv_cache_dtype,
341
+ layer._k_scale,
342
+ layer._v_scale,
343
+ )
344
+ else:
345
+ ops.reshape_and_cache_flash(
346
+ key,
347
+ value,
348
+ key_cache,
349
+ value_cache,
350
+ attn_metadata.slot_mapping,
351
+ self.kv_cache_dtype,
352
+ layer._k_scale,
353
+ layer._v_scale,
354
+ )
355
+
356
+ if self.kv_cache_dtype.startswith("fp8"):
357
+ key_cache = key_cache.view(self.fp8_dtype)
358
+ value_cache = value_cache.view(self.fp8_dtype)
359
+ num_tokens, num_heads, head_size = query.shape
360
+ assert layer._q_scale_float == 1.0, \
361
+ "A non 1.0 q_scale is not currently supported."
362
+ if current_platform.is_cuda():
363
+ # Skip Q quantization on ROCm and XPU, enable this on cuda
364
+ # only, since dequantizing back to f32 in the attention kernel
365
+ # is not supported.
366
+ query, _ = ops.scaled_fp8_quant(
367
+ query.reshape(
368
+ (num_tokens, num_heads * head_size)).contiguous(),
369
+ layer._q_scale)
370
+ query = query.reshape((num_tokens, num_heads, head_size))
371
+
372
+ cu_seqlens_q = attn_metadata.query_start_loc
373
+ seqused_k = attn_metadata.seq_lens
374
+ max_seqlen_q = attn_metadata.max_query_len
375
+ max_seqlen_k = attn_metadata.max_seq_len
376
+ block_table = attn_metadata.block_table
377
+
378
+ if use_prefill_decode_attn:
379
+ # Compute attention and update output up to `num_actual_tokens`.
380
+ chunked_prefill_paged_decode(
381
+ query=query[:num_actual_tokens],
382
+ key=key[:num_actual_tokens],
383
+ value=value[:num_actual_tokens],
384
+ output=output[:num_actual_tokens],
385
+ kv_cache_dtype=self.kv_cache_dtype,
386
+ key_cache=key_cache,
387
+ value_cache=value_cache,
388
+ block_table=block_table,
389
+ query_start_loc=cu_seqlens_q,
390
+ seq_lens=seqused_k,
391
+ max_seq_len=max_seqlen_k,
392
+ max_query_len=max_seqlen_q,
393
+ k_scale=layer._k_scale,
394
+ v_scale=layer._v_scale,
395
+ alibi_slopes=self.alibi_slopes,
396
+ sliding_window=self.sliding_window[0],
397
+ sm_scale=self.scale,
398
+ output_scale=output_scale,
399
+ sinks=self.sinks,
400
+ )
401
+
402
+ else:
403
+ descale_shape = (cu_seqlens_q.shape[0] - 1, key.shape[1])
404
+
405
+ self.unified_attention(
406
+ q=query[:num_actual_tokens],
407
+ k=key_cache,
408
+ v=value_cache,
409
+ out=output[:num_actual_tokens],
410
+ cu_seqlens_q=cu_seqlens_q,
411
+ max_seqlen_q=max_seqlen_q,
412
+ seqused_k=seqused_k,
413
+ max_seqlen_k=max_seqlen_k,
414
+ softmax_scale=self.scale,
415
+ causal=True,
416
+ alibi_slopes=self.alibi_slopes,
417
+ window_size=self.sliding_window,
418
+ block_table=block_table,
419
+ softcap=self.logits_soft_cap,
420
+ q_descale=None, # Not supported
421
+ k_descale=layer._k_scale.expand(descale_shape),
422
+ v_descale=layer._v_scale.expand(descale_shape),
423
+ sinks=self.sinks,
424
+ output_scale=output_scale)
425
+
426
+ return output
vllm/v1/attention/backends/triton_attn.py
+45
-124
@@ -1,24 +1,19 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
-"""Attention layer with PagedAttention and Triton prefix prefill."""
3
+"""High-Performance Triton-only Attention layer."""
4
from dataclasses import dataclass
5
-from functools import cache
5
from typing import ClassVar, Optional
6
7
import torch
8
10
-from vllm import envs
9
from vllm.attention.backends.abstract import (AttentionBackend, AttentionImpl,
10
AttentionMetadata, AttentionType)
13
-from vllm.attention.ops.chunked_prefill_paged_decode import (
14
- chunked_prefill_paged_decode)
15
-from vllm.attention.ops.paged_attn import PagedAttention
11
+from vllm.attention.ops.triton_unified_attention import unified_attention
12
from vllm.config import VllmConfig
13
from vllm.logger import init_logger
14
from vllm.model_executor.layers.quantization.utils.quant_utils import (
15
QuantKey, kFp8StaticTensorSym)
16
from vllm.platforms import current_platform
21
-from vllm.v1.attention.backends.flash_attn import FlashAttentionMetadata
17
from vllm.v1.attention.backends.utils import (AttentionCGSupport,
18
AttentionMetadataBuilder,
19
CommonAttentionMetadata)
@@ -144,20 +139,15 @@ class TritonAttentionBackend(AttentionBackend):
139
140
@classmethod
141
def get_supported_dtypes(cls) -> list[torch.dtype]:
147
- return [torch.float16, torch.bfloat16]
148
-
149
- @classmethod
150
- def get_supported_head_sizes(cls) -> list[int]:
151
- return [32, 64, 96, 128, 160, 192, 224, 256]
142
+ return [torch.float16, torch.bfloat16, torch.float32]
143
144
@classmethod
145
def validate_head_size(cls, head_size: int) -> None:
155
- supported_head_sizes = cls.get_supported_head_sizes()
156
- if head_size not in supported_head_sizes:
157
- attn_type = cls.__name__.removesuffix("Backend")
146
+ # Triton Attention supports any head size above 32
147
+ if head_size < 32:
148
raise ValueError(
159
- f"Head size {head_size} is not supported by {attn_type}. "
160
- f"Supported head sizes are: {supported_head_sizes}. "
149
+ f"Head size {head_size} is not supported by TritonAttention."
150
+ f"Head sizes need to be larger or equal 32 for this backend. "
151
"Set VLLM_ATTENTION_BACKEND=FLEX_ATTENTION to use "
152
"FlexAttention backend which supports all head sizes.")
153
@@ -182,7 +172,7 @@ class TritonAttentionBackend(AttentionBackend):
172
) -> tuple[int, ...]:
173
if block_size % 16 != 0:
174
raise ValueError("Block size must be a multiple of 16.")
185
- return (2, num_blocks, block_size, num_kv_heads, head_size)
175
+ return (num_blocks, 2, block_size, num_kv_heads, head_size)
176
177
@staticmethod
178
def use_cascade_attention(*args, **kwargs) -> bool:
@@ -193,15 +183,6 @@ class TritonAttentionBackend(AttentionBackend):
183
return TritonAttentionMetadataBuilder
184
185
196
-@cache
197
-def use_aiter_unified_attention() -> bool:
198
- """Check if aiter unified attention should be used."""
199
- # VLLM_ROCM_USE_AITER_MHA needs to set to 0 as well as it is set
200
- # to 1 as default
201
- return envs.VLLM_ROCM_USE_AITER \
202
- and envs.VLLM_USE_AITER_UNIFIED_ATTENTION
203
-
204
-
186
class TritonAttentionImpl(AttentionImpl):
187
188
def fused_output_quant_supported(self, quant_key: QuantKey):
@@ -250,24 +231,6 @@ class TritonAttentionImpl(AttentionImpl):
231
"TritonAttentionImpl")
232
233
self.fp8_dtype = current_platform.fp8_dtype()
253
- self.force_prefill_decode_attn = \
254
- envs.VLLM_V1_USE_PREFILL_DECODE_ATTENTION
255
-
256
- if not self.force_prefill_decode_attn:
257
- # If not using prefill decode attention, we use the Triton
258
- # unified attention implementation.
259
- if use_aiter_unified_attention():
260
- logger.info_once(
261
- "Using aiter unified attention for TritonAttentionImpl")
262
- from aiter.ops.triton.unified_attention import (
263
- unified_attention)
264
- self.unified_attention = unified_attention
265
- else:
266
- logger.info_once(
267
- "Using vllm unified attention for TritonAttentionImpl")
268
- from vllm.attention.ops.triton_unified_attention import (
269
- unified_attention)
270
- self.unified_attention = unified_attention
234
235
self.sinks = sinks
236
if sinks is not None:
@@ -283,19 +246,19 @@ class TritonAttentionImpl(AttentionImpl):
246
key: torch.Tensor,
247
value: torch.Tensor,
248
kv_cache: torch.Tensor,
286
- attn_metadata: FlashAttentionMetadata,
249
+ attn_metadata: TritonAttentionMetadata,
250
output: Optional[torch.Tensor] = None,
251
output_scale: Optional[torch.Tensor] = None,
252
output_block_scale: Optional[torch.Tensor] = None,
253
) -> torch.Tensor:
291
- """Forward pass with FlashAttention.
254
+ """Forward pass with Paged Attention impl. in Triton.
255
256
Args:
257
query: shape = [num_tokens, num_heads, head_size]
258
key: shape = [num_tokens, num_kv_heads, head_size]
259
value: shape = [num_tokens, num_kv_heads, head_size]
260
kv_cache: shape =
298
- [2, num_blocks, block_size, num_kv_heads, head_size]
261
+ [num_blocks, 2, block_size, num_kv_heads, head_size]
262
attn_metadata: Metadata for attention.
263
Returns:
264
shape = [num_tokens, num_heads * head_size]
@@ -322,40 +285,22 @@ class TritonAttentionImpl(AttentionImpl):
285
# Whenever making a change in this method, please benchmark the
286
# performance to make sure it does not introduce any overhead.
287
325
- use_prefill_decode_attn = self.force_prefill_decode_attn
288
num_actual_tokens = attn_metadata.num_actual_tokens
327
-
328
- if use_prefill_decode_attn:
329
- key_cache, value_cache = PagedAttention.split_kv_cache(
330
- kv_cache, self.num_kv_heads, self.head_size)
331
- else:
332
- key_cache, value_cache = kv_cache.unbind(0)
289
+ key_cache, value_cache = kv_cache.unbind(1)
290
291
if self.kv_sharing_target_layer_name is None:
292
# Reshape the input keys and values and store them in the cache.
293
# Skip this if sharing KV cache with an earlier attention layer.
337
- if use_prefill_decode_attn:
338
- PagedAttention.write_to_paged_cache(
339
- key,
340
- value,
341
- key_cache,
342
- value_cache,
343
- attn_metadata.slot_mapping,
344
- self.kv_cache_dtype,
345
- layer._k_scale,
346
- layer._v_scale,
347
- )
348
- else:
349
- ops.reshape_and_cache_flash(
350
- key,
351
- value,
352
- key_cache,
353
- value_cache,
354
- attn_metadata.slot_mapping,
355
- self.kv_cache_dtype,
356
- layer._k_scale,
357
- layer._v_scale,
358
- )
294
+ ops.reshape_and_cache_flash(
295
+ key,
296
+ value,
297
+ key_cache,
298
+ value_cache,
299
+ attn_metadata.slot_mapping,
300
+ self.kv_cache_dtype,
301
+ layer._k_scale,
302
+ layer._v_scale,
303
+ )
304
305
if self.kv_cache_dtype.startswith("fp8"):
306
key_cache = key_cache.view(self.fp8_dtype)
@@ -379,52 +324,28 @@ class TritonAttentionImpl(AttentionImpl):
324
max_seqlen_k = attn_metadata.max_seq_len
325
block_table = attn_metadata.block_table
326
382
- if use_prefill_decode_attn:
383
- # Compute attention and update output up to `num_actual_tokens`.
384
- chunked_prefill_paged_decode(
385
- query=query[:num_actual_tokens],
386
- key=key[:num_actual_tokens],
387
- value=value[:num_actual_tokens],
388
- output=output[:num_actual_tokens],
389
- kv_cache_dtype=self.kv_cache_dtype,
390
- key_cache=key_cache,
391
- value_cache=value_cache,
392
- block_table=block_table,
393
- query_start_loc=cu_seqlens_q,
394
- seq_lens=seqused_k,
395
- max_seq_len=max_seqlen_k,
396
- max_query_len=max_seqlen_q,
397
- k_scale=layer._k_scale,
398
- v_scale=layer._v_scale,
399
- alibi_slopes=self.alibi_slopes,
400
- sliding_window=self.sliding_window[0],
401
- sm_scale=self.scale,
402
- output_scale=output_scale,
403
- sinks=self.sinks,
404
- )
405
-
406
- else:
407
- descale_shape = (cu_seqlens_q.shape[0] - 1, key.shape[1])
408
-
409
- self.unified_attention(
410
- q=query[:num_actual_tokens],
411
- k=key_cache,
412
- v=value_cache,
413
- out=output[:num_actual_tokens],
414
- cu_seqlens_q=cu_seqlens_q,
415
- max_seqlen_q=max_seqlen_q,
416
- seqused_k=seqused_k,
417
- max_seqlen_k=max_seqlen_k,
418
- softmax_scale=self.scale,
419
- causal=True,
420
- alibi_slopes=self.alibi_slopes,
421
- window_size=self.sliding_window,
422
- block_table=block_table,
423
- softcap=self.logits_soft_cap,
424
- q_descale=None, # Not supported
425
- k_descale=layer._k_scale.expand(descale_shape),
426
- v_descale=layer._v_scale.expand(descale_shape),
427
- sinks=self.sinks,
428
- output_scale=output_scale)
327
+ descale_shape = (cu_seqlens_q.shape[0] - 1, key.shape[1])
328
+
329
+ unified_attention(
330
+ q=query[:num_actual_tokens],
331
+ k=key_cache,
332
+ v=value_cache,
333
+ out=output[:num_actual_tokens],
334
+ cu_seqlens_q=cu_seqlens_q,
335
+ max_seqlen_q=max_seqlen_q,
336
+ seqused_k=seqused_k,
337
+ max_seqlen_k=max_seqlen_k,
338
+ softmax_scale=self.scale,
339
+ causal=True,
340
+ alibi_slopes=self.alibi_slopes,
341
+ window_size=self.sliding_window,
342
+ block_table=block_table,
343
+ softcap=self.logits_soft_cap,
344
+ q_descale=None, # Not supported
345
+ k_descale=layer._k_scale.expand(descale_shape),
346
+ v_descale=layer._v_scale.expand(descale_shape),
347
+ sinks=self.sinks,
348
+ output_scale=output_scale,
349
+ )
350
351
return output