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, 06:00 UTC
[Hardware][AMD] integrate aiter chunked prefill into vllm (#18596)
Signed-off-by: fsx950223 <fsx950223@outlook.com> Signed-off-by: charlifu <charlifu@amd.com> Co-authored-by: fsx950223 <fsx950223@outlook.com> Co-authored-by: charlifu <charlifu@amd.com>
Zzz9990 committed
Jun 18, 2025 at 23:46 UTC
8b6e1d639c66d5828d03a7df2c3a500030a5c5cd
3 files changed
+602
-3
vllm/envs.py
+8
@@ -87,6 +87,7 @@ if TYPE_CHECKING:
87
VLLM_ROCM_USE_AITER_MOE: bool = True
88
VLLM_ROCM_USE_AITER_RMSNORM: bool = True
89
VLLM_ROCM_USE_AITER_MLA: bool = True
90
+ VLLM_ROCM_USE_AITER_MHA: bool = True
91
VLLM_ROCM_USE_SKINNY_GEMM: bool = True
92
VLLM_ROCM_FP8_PADDING: bool = True
93
VLLM_ROCM_MOE_PADDING: bool = True
@@ -653,6 +654,13 @@ environment_variables: dict[str, Callable[[], Any]] = {
654
"VLLM_ROCM_USE_AITER_MLA":
655
lambda: (os.getenv("VLLM_ROCM_USE_AITER_MLA", "True").lower() in
656
("true", "1")),
657
+
658
+ # Whether to use aiter mha ops.
659
+ # By default is enabled.
660
+ "VLLM_ROCM_USE_AITER_MHA":
661
+ lambda: (os.getenv("VLLM_ROCM_USE_AITER_MHA", "True").lower() in
662
+ ("true", "1")),
663
+
664
# use rocm skinny gemms
665
"VLLM_ROCM_USE_SKINNY_GEMM":
666
lambda: (os.getenv("VLLM_ROCM_USE_SKINNY_GEMM", "True").lower() in
vllm/platforms/rocm.py
+9
-3
@@ -215,9 +215,15 @@ class RocmPlatform(Platform):
215
selected_backend = _Backend.ROCM_FLASH
216
217
if envs.VLLM_USE_V1:
218
- logger.info("Using Triton Attention backend on V1 engine.")
219
- return ("vllm.v1.attention.backends."
220
- "triton_attn.TritonAttentionBackend")
218
+ if envs.VLLM_ROCM_USE_AITER and envs.VLLM_ROCM_USE_AITER_MHA \
219
+ and on_gfx9():
220
+ logger.info("Using Flash Attention backend on V1 engine.")
221
+ return ("vllm.v1.attention.backends."
222
+ "rocm_aiter_fa.AiterFlashAttentionBackend")
223
+ else:
224
+ logger.info("Using Triton Attention backend on V1 engine.")
225
+ return ("vllm.v1.attention.backends."
226
+ "triton_attn.TritonAttentionBackend")
227
if selected_backend == _Backend.ROCM_FLASH:
228
if not cls.has_device_capability(90):
229
# not Instinct series GPUs.
vllm/v1/attention/backends/rocm_aiter_fa.py
new
+585
@@ -0,0 +1,585 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+"""Attention layer with AiterFlashAttention."""
3
+from dataclasses import dataclass
4
+from typing import TYPE_CHECKING, Any, Optional
5
+
6
+import torch
7
+
8
+from vllm import _custom_ops as ops
9
+from vllm.attention.backends.abstract import (AttentionBackend, AttentionImpl,
10
+ AttentionMetadata, AttentionType,
11
+ is_quantized_kv_cache)
12
+from vllm.logger import init_logger
13
+from vllm.platforms import current_platform
14
+from vllm.v1.attention.backends.flash_attn import (
15
+ make_local_attention_virtual_batches)
16
+from vllm.v1.attention.backends.utils import CommonAttentionMetadata
17
+from vllm.v1.kv_cache_interface import AttentionSpec
18
+from vllm.v1.worker.block_table import BlockTable
19
+
20
+if TYPE_CHECKING:
21
+ from vllm.v1.core.sched.output import SchedulerOutput
22
+ from vllm.v1.worker.gpu_input_batch import InputBatch
23
+ from vllm.v1.worker.gpu_model_runner import GPUModelRunner
24
+
25
+if current_platform.is_rocm():
26
+ import aiter
27
+
28
+ from vllm.triton_utils import tl, triton
29
+ from vllm.utils import direct_register_custom_op
30
+
31
+ @triton.jit
32
+ def _vllm_layout_trans_kernel(
33
+ k_buffer_ptr,
34
+ v_buffer_ptr,
35
+ k_values_ptr,
36
+ v_values_ptr,
37
+ b_query_lens_loc,
38
+ b_seq_lens_loc,
39
+ block_table,
40
+ block_table_stride_0,
41
+ E_DIM: tl.constexpr,
42
+ BLOCK_SIZE: tl.constexpr,
43
+ ):
44
+ batch_idx = tl.program_id(0)
45
+ block_idx = tl.program_id(1)
46
+ batch_token_indexes = tl.load(b_seq_lens_loc + batch_idx +
47
+ tl.arange(0, 2))
48
+ batch_token_start, batch_token_end = tl.split(batch_token_indexes)
49
+ seq_len = batch_token_end - batch_token_start
50
+
51
+ batch_query_indexes = tl.load(b_query_lens_loc + batch_idx +
52
+ tl.arange(0, 2))
53
+ batch_query_start, batch_query_end = tl.split(batch_query_indexes)
54
+ query_len = batch_query_end - batch_query_start
55
+ if query_len <= 1:
56
+ return
57
+ if block_idx * BLOCK_SIZE < seq_len:
58
+ block_mask = (block_idx * BLOCK_SIZE +
59
+ tl.arange(0, BLOCK_SIZE)[:, None]) < seq_len
60
+
61
+ kv_idx = tl.load(block_table + batch_idx * block_table_stride_0 +
62
+ block_idx)
63
+
64
+ kv_buffer_off = kv_idx * BLOCK_SIZE * E_DIM + tl.arange(
65
+ 0, BLOCK_SIZE)[:, None] * E_DIM + tl.arange(0, E_DIM)[None, :]
66
+ k_vals = tl.load(k_buffer_ptr + kv_buffer_off,
67
+ mask=block_mask,
68
+ other=0.0)
69
+ v_vals = tl.load(v_buffer_ptr + kv_buffer_off,
70
+ mask=block_mask,
71
+ other=0.0)
72
+
73
+ kv_values_off = batch_token_start * E_DIM + \
74
+ block_idx * BLOCK_SIZE * E_DIM + \
75
+ tl.arange(0, BLOCK_SIZE)[:, None] * E_DIM + \
76
+ tl.arange(0, E_DIM)[None, :]
77
+ tl.store(k_values_ptr + kv_values_off, k_vals, mask=block_mask)
78
+ tl.store(v_values_ptr + kv_values_off, v_vals, mask=block_mask)
79
+
80
+ def vllm_layout_trans(b_query_lens_loc, b_seq_lens_loc, block_table,
81
+ k_buffer, v_buffer, max_seq_len, total_tokens):
82
+ H_KV = v_buffer.shape[2]
83
+ D = v_buffer.shape[3]
84
+ BLOCK_SIZE = v_buffer.shape[1]
85
+ dtype = k_buffer.dtype
86
+ k_values = torch.empty((total_tokens, H_KV, D),
87
+ dtype=dtype,
88
+ device="cuda")
89
+ v_values = torch.empty((total_tokens, H_KV, D),
90
+ dtype=dtype,
91
+ device="cuda")
92
+
93
+ grid = (block_table.shape[0],
94
+ (max_seq_len + BLOCK_SIZE - 1) // BLOCK_SIZE)
95
+
96
+ _vllm_layout_trans_kernel[grid](k_buffer,
97
+ v_buffer,
98
+ k_values,
99
+ v_values,
100
+ b_query_lens_loc,
101
+ b_seq_lens_loc,
102
+ block_table,
103
+ block_table.stride(0),
104
+ E_DIM=H_KV * D,
105
+ BLOCK_SIZE=BLOCK_SIZE)
106
+
107
+ return k_values, v_values
108
+
109
+ def flash_attn_varlen_func_impl(
110
+ q: torch.Tensor,
111
+ k_cache: torch.Tensor,
112
+ v_cache: torch.Tensor,
113
+ out: torch.Tensor,
114
+ cu_seqlens_q: torch.Tensor,
115
+ cu_seqlens_k: torch.Tensor,
116
+ total_tokens: int,
117
+ max_seqlen_q: int,
118
+ max_seqlen_k: int,
119
+ softmax_scale: float,
120
+ window_size: Optional[list[int]], # -1 means infinite context window
121
+ alibi_slopes: Optional[list[float]],
122
+ block_table: torch.Tensor,
123
+ ) -> torch.Tensor:
124
+ k, v = vllm_layout_trans(cu_seqlens_q, cu_seqlens_k, block_table,
125
+ k_cache, v_cache, max_seqlen_k, total_tokens)
126
+ output = aiter.flash_attn_varlen_func(
127
+ q=q,
128
+ k=k,
129
+ v=v,
130
+ cu_seqlens_q=cu_seqlens_q,
131
+ max_seqlen_q=max_seqlen_q,
132
+ min_seqlen_q=1,
133
+ cu_seqlens_k=cu_seqlens_k,
134
+ max_seqlen_k=max_seqlen_k,
135
+ softmax_scale=softmax_scale,
136
+ causal=True,
137
+ alibi_slopes=alibi_slopes,
138
+ window_size=window_size,
139
+ out=out,
140
+ )
141
+ return output
142
+
143
+ def flash_attn_varlen_func_fake(
144
+ q: torch.Tensor,
145
+ k_cache: torch.Tensor,
146
+ v_cache: torch.Tensor,
147
+ out: torch.Tensor,
148
+ cu_seqlens_q: torch.Tensor,
149
+ cu_seqlens_k: torch.Tensor,
150
+ total_tokens: int,
151
+ max_seqlen_q: int,
152
+ max_seqlen_k: int,
153
+ softmax_scale: float,
154
+ window_size: Optional[list[int]], # -1 means infinite context window
155
+ alibi_slopes: Optional[list[float]],
156
+ block_table: torch.Tensor,
157
+ ) -> torch.Tensor:
158
+ return torch.empty(q.shape[0],
159
+ q.shape[1],
160
+ v_cache.shape[-2],
161
+ dtype=torch.float8_e4m3fnuz,
162
+ device="cuda")
163
+
164
+ direct_register_custom_op("flash_attn_varlen_func",
165
+ flash_attn_varlen_func_impl, ["out"],
166
+ flash_attn_varlen_func_fake,
167
+ dispatch_key=current_platform.dispatch_key)
168
+
169
+logger = init_logger(__name__)
170
+
171
+
172
+class AiterFlashAttentionMetadataBuilder:
173
+
174
+ def __init__(self, runner: "GPUModelRunner", kv_cache_spec: AttentionSpec,
175
+ block_table: BlockTable):
176
+ model_config = runner.model_config
177
+
178
+ self.runner = runner
179
+ self.num_heads_q = model_config.get_num_attention_heads(
180
+ runner.parallel_config)
181
+ self.num_heads_kv = model_config.get_num_kv_heads(
182
+ runner.parallel_config)
183
+ self.headdim = model_config.get_head_size()
184
+ self.block_size = kv_cache_spec.block_size
185
+ self.kv_cache_spec = kv_cache_spec
186
+ self.block_table = block_table
187
+
188
+ # Sliding window size to be used with the AOT scheduler will be
189
+ # populated on first build() call.
190
+ self.aot_sliding_window: Optional[tuple[int, int]] = None
191
+
192
+ def reorder_batch(self, input_batch: "InputBatch",
193
+ scheduler_output: "SchedulerOutput") -> bool:
194
+ return False
195
+
196
+ def build(self, common_prefix_len: int,
197
+ common_attn_metadata: CommonAttentionMetadata):
198
+
199
+ num_reqs = common_attn_metadata.num_reqs
200
+ num_actual_tokens = common_attn_metadata.num_actual_tokens
201
+ max_query_len = common_attn_metadata.max_query_len
202
+
203
+ max_seq_len = int(self.runner.seq_lens_np[:num_reqs].max())
204
+ total_tokens = int(self.runner.seq_lens_np[:num_reqs].sum())
205
+ query_start_loc = common_attn_metadata.query_start_loc
206
+ seq_lens = common_attn_metadata.seq_lens
207
+ block_table = self.block_table
208
+ block_table_tensor = block_table.get_device_tensor()[:num_reqs]
209
+
210
+ block_table.slot_mapping[:num_actual_tokens].copy_(
211
+ block_table.slot_mapping_cpu[:num_actual_tokens],
212
+ non_blocking=True)
213
+ # Fill unused with -1. Needed for reshape_and_cache in full cuda graph
214
+ # mode.
215
+ block_table.slot_mapping[num_actual_tokens:].fill_(-1)
216
+
217
+ slot_mapping = block_table.slot_mapping[:num_actual_tokens]
218
+
219
+ cu_seq_lens = torch.zeros(seq_lens.shape[0] + 1,
220
+ dtype=torch.int32,
221
+ device="cuda")
222
+ torch.cumsum(seq_lens,
223
+ dim=0,
224
+ dtype=cu_seq_lens.dtype,
225
+ out=cu_seq_lens[1:])
226
+
227
+ def schedule(batch_size, cu_query_lens, max_query_len, seqlens,
228
+ max_seq_len, causal):
229
+ return None
230
+
231
+ # for local attention
232
+ local_attn_metadata = None
233
+ if self.runner.attention_chunk_size is not None:
234
+ seqlens_q_local_np, virt_q_cu_seqlens_np, virt_k_seqlens_np, \
235
+ virt_block_table_tensor = make_local_attention_virtual_batches(
236
+ self.runner.attention_chunk_size,
237
+ self.runner.query_start_loc_np[:num_reqs + 1],
238
+ self.runner.seq_lens_np[:num_reqs],
239
+ block_table_tensor,
240
+ self.block_size,
241
+ )
242
+ local_query_start_loc = torch.from_numpy(virt_q_cu_seqlens_np).to(
243
+ self.runner.device, non_blocking=True)
244
+ local_seqused_k = torch.from_numpy(virt_k_seqlens_np).to(
245
+ self.runner.device, non_blocking=True)
246
+ local_max_query_len = seqlens_q_local_np.max()
247
+ local_max_seq_len = virt_k_seqlens_np.max()
248
+ local_scheduler_metadata = schedule(
249
+ batch_size=local_query_start_loc.shape[0] - 1,
250
+ cu_query_lens=local_query_start_loc,
251
+ max_query_len=local_max_query_len,
252
+ seqlens=local_seqused_k,
253
+ max_seq_len=local_max_seq_len,
254
+ causal=True)
255
+
256
+ local_attn_metadata = \
257
+ AiterFlashAttentionMetadata.LocalAttentionMetadata(
258
+ local_query_start_loc=local_query_start_loc,
259
+ local_seqused_k=local_seqused_k,
260
+ local_block_table=virt_block_table_tensor,
261
+ local_max_query_len=local_max_query_len,
262
+ local_max_seq_len=local_max_seq_len,
263
+ local_scheduler_metadata=local_scheduler_metadata,
264
+ )
265
+
266
+ use_cascade = common_prefix_len > 0
267
+
268
+ cu_prefix_query_lens = None
269
+ prefix_kv_lens = None
270
+ suffix_kv_lens = None
271
+
272
+ attn_metadata = AiterFlashAttentionMetadata(
273
+ num_actual_tokens=num_actual_tokens,
274
+ max_query_len=max_query_len,
275
+ query_start_loc=query_start_loc,
276
+ max_seq_len=max_seq_len,
277
+ seq_lens=seq_lens,
278
+ cu_seq_lens=cu_seq_lens,
279
+ total_tokens=total_tokens,
280
+ block_table=block_table_tensor,
281
+ slot_mapping=slot_mapping,
282
+ use_cascade=use_cascade,
283
+ common_prefix_len=common_prefix_len,
284
+ cu_prefix_query_lens=cu_prefix_query_lens,
285
+ prefix_kv_lens=prefix_kv_lens,
286
+ suffix_kv_lens=suffix_kv_lens,
287
+ local_attn_metadata=local_attn_metadata,
288
+ )
289
+ return attn_metadata
290
+
291
+ def can_run_in_cudagraph(
292
+ self, common_attn_metadata: CommonAttentionMetadata) -> bool:
293
+ # Full CUDA Graph always supported (FA2 support checked separately)
294
+ return True
295
+
296
+ def use_cascade_attention(self, *args, **kwargs) -> bool:
297
+ return False
298
+
299
+
300
+class AiterFlashAttentionBackend(AttentionBackend):
301
+
302
+ accept_output_buffer: bool = True
303
+
304
+ @staticmethod
305
+ def get_supported_head_sizes() -> list[int]:
306
+ return [32, 64, 96, 128, 160, 192, 224, 256]
307
+
308
+ @staticmethod
309
+ def get_name() -> str:
310
+ return "FLASH_ATTN_VLLM_V1"
311
+
312
+ @staticmethod
313
+ def get_impl_cls() -> type["AiterFlashAttentionImpl"]:
314
+ return AiterFlashAttentionImpl
315
+
316
+ @staticmethod
317
+ def get_metadata_cls() -> type["AttentionMetadata"]:
318
+ return AiterFlashAttentionMetadata
319
+
320
+ @staticmethod
321
+ def get_builder_cls() -> type["AiterFlashAttentionMetadataBuilder"]:
322
+ return AiterFlashAttentionMetadataBuilder
323
+
324
+ @staticmethod
325
+ def get_kv_cache_shape(
326
+ num_blocks: int,
327
+ block_size: int,
328
+ num_kv_heads: int,
329
+ head_size: int,
330
+ ) -> tuple[int, ...]:
331
+ if block_size % 16 != 0:
332
+ raise ValueError("Block size must be a multiple of 16.")
333
+ return (2, num_blocks, block_size, num_kv_heads, head_size)
334
+
335
+
336
+@dataclass
337
+class AiterFlashAttentionMetadata:
338
+ # NOTE(sang): Definition of context_len, query_len, and seq_len.
339
+ # |---------- N-1 iteration --------|
340
+ # |---------------- N iteration ---------------------|
341
+ # |- tokenA -|......................|-- newTokens ---|
342
+ # |---------- context_len ----------|
343
+ # |-------------------- seq_len ---------------------|
344
+ # |-- query_len ---|
345
+
346
+ num_actual_tokens: int # Number of tokens excluding padding.
347
+ max_query_len: int
348
+ query_start_loc: torch.Tensor
349
+ max_seq_len: int
350
+ seq_lens: torch.Tensor
351
+ cu_seq_lens: torch.Tensor
352
+ total_tokens: int
353
+ block_table: torch.Tensor
354
+ slot_mapping: torch.Tensor
355
+
356
+ # For cascade attention.
357
+ use_cascade: bool
358
+ common_prefix_len: int
359
+ cu_prefix_query_lens: Optional[torch.Tensor]
360
+ prefix_kv_lens: Optional[torch.Tensor]
361
+ suffix_kv_lens: Optional[torch.Tensor]
362
+
363
+ # for local attention
364
+ @dataclass
365
+ class LocalAttentionMetadata:
366
+ local_query_start_loc: torch.Tensor
367
+ local_seqused_k: torch.Tensor
368
+ local_block_table: torch.Tensor
369
+ local_max_query_len: int
370
+ local_max_seq_len: int
371
+ local_scheduler_metadata: Optional[torch.Tensor]
372
+
373
+ local_attn_metadata: Optional[LocalAttentionMetadata] = None
374
+
375
+
376
+class AiterFlashAttentionImpl(AttentionImpl):
377
+
378
+ def __init__(
379
+ self,
380
+ num_heads: int,
381
+ head_size: int,
382
+ scale: float,
383
+ num_kv_heads: int,
384
+ alibi_slopes: Optional[list[float]],
385
+ sliding_window: Optional[int],
386
+ kv_cache_dtype: str,
387
+ blocksparse_params: Optional[dict[str, Any]] = None,
388
+ logits_soft_cap: Optional[float] = None,
389
+ attn_type: AttentionType = AttentionType.DECODER,
390
+ use_irope: bool = False,
391
+ ) -> None:
392
+ if blocksparse_params is not None:
393
+ raise ValueError(
394
+ "AiterFlashAttention does not support block-sparse attention.")
395
+ self.num_heads = num_heads
396
+ self.head_size = head_size
397
+ self.scale = float(scale)
398
+ self.num_kv_heads = num_kv_heads
399
+ if alibi_slopes is not None:
400
+ alibi_slopes = torch.tensor(alibi_slopes, dtype=torch.float32)
401
+ self.alibi_slopes = alibi_slopes
402
+ if sliding_window is None:
403
+ self.sliding_window = [-1, -1]
404
+ else:
405
+ self.sliding_window = [sliding_window - 1, 0]
406
+ self.kv_cache_dtype = kv_cache_dtype
407
+ if logits_soft_cap is None:
408
+ # In flash-attn, setting logits_soft_cap as 0 means no soft cap.
409
+ logits_soft_cap = 0.
410
+ self.logits_soft_cap = logits_soft_cap
411
+
412
+ assert self.num_heads % self.num_kv_heads == 0
413
+ self.num_queries_per_kv = self.num_heads // self.num_kv_heads
414
+
415
+ support_head_sizes = \
416
+ AiterFlashAttentionBackend.get_supported_head_sizes()
417
+ if head_size not in support_head_sizes:
418
+ raise ValueError(
419
+ f"Head size {head_size} is not supported by "
420
+ "AiterFlashAttention. "
421
+ f"Supported head sizes are: {support_head_sizes}. "
422
+ "Set VLLM_USE_V1=0 to use another attention backend.")
423
+
424
+ if attn_type != AttentionType.DECODER:
425
+ raise NotImplementedError("Encoder self-attention and "
426
+ "encoder/decoder cross-attention "
427
+ "are not implemented for "
428
+ "FlashAttentionImpl")
429
+ self.use_irope = use_irope
430
+ if is_quantized_kv_cache(self.kv_cache_dtype):
431
+ raise NotImplementedError(
432
+ "AiterFlashAttention does not support fp8 kv-cache on this "
433
+ "device.")
434
+
435
+ def forward(
436
+ self,
437
+ layer: torch.nn.Module,
438
+ query: torch.Tensor,
439
+ key: torch.Tensor,
440
+ value: torch.Tensor,
441
+ kv_cache: torch.Tensor,
442
+ attn_metadata: AiterFlashAttentionMetadata,
443
+ output: Optional[torch.Tensor] = None,
444
+ output_scale: Optional[torch.Tensor] = None,
445
+ ) -> torch.Tensor:
446
+ """Forward pass with AiterFlashAttention.
447
+
448
+ Args:
449
+ query: shape = [num_tokens, num_heads, head_size]
450
+ key: shape = [num_tokens, num_kv_heads, head_size]
451
+ value: shape = [num_tokens, num_kv_heads, head_size]
452
+ kv_cache = [2, num_blocks, block_size, num_kv_heads, head_size]
453
+ attn_metadata: Metadata for attention.
454
+ Returns:
455
+ shape = [num_tokens, num_heads * head_size]
456
+ NOTE: FP8 quantization, flash-attn expect the size of
457
+ {q,k,v}_descale to be (num_sequences, num_kv_heads).
458
+ We use torch's .expand() to avoid duplicating values
459
+ """
460
+ assert output is not None, "Output tensor must be provided."
461
+
462
+ if output_scale is not None:
463
+ raise NotImplementedError(
464
+ "fused output quantization is not yet supported"
465
+ " for FlashAttentionImpl")
466
+
467
+ if attn_metadata is None:
468
+ # Profiling run.
469
+ return output
470
+
471
+ # IMPORTANT!
472
+ # NOTE(woosuk): With piece-wise CUDA graphs, this method is executed in
473
+ # eager-mode PyTorch. Thus, we need to be careful about any CPU overhead
474
+ # in this method. For example, `view` and `slice` (or `[:n]`) operations
475
+ # are surprisingly slow even in the case they do not invoke any GPU ops.
476
+ # Minimize the PyTorch ops in this method as much as possible.
477
+ # Whenever making a change in this method, please benchmark the
478
+ # performance to make sure it does not introduce any overhead.
479
+
480
+ num_actual_tokens = attn_metadata.num_actual_tokens
481
+ # Reshape the input keys and values and store them in the cache.
482
+ # NOTE(woosuk): Here, key and value are padded while slot_mapping is
483
+ # not padded. However, we don't need to do key[:num_actual_tokens] and
484
+ # value[:num_actual_tokens] because the reshape_and_cache_flash op uses
485
+ # the slot_mapping's shape to determine the number of actual tokens.
486
+ key_cache, value_cache = kv_cache.unbind(0)
487
+ torch.ops._C_cache_ops.reshape_and_cache_flash(
488
+ key,
489
+ value,
490
+ key_cache,
491
+ value_cache,
492
+ attn_metadata.slot_mapping,
493
+ self.kv_cache_dtype,
494
+ layer._k_scale,
495
+ layer._v_scale,
496
+ )
497
+
498
+ if self.kv_cache_dtype.startswith("fp8"):
499
+ key_cache = key_cache.view(torch.float8_e4m3fnuz)
500
+ value_cache = value_cache.view(torch.float8_e4m3fnuz)
501
+ num_tokens, num_heads, head_size = query.shape
502
+ query, _ = ops.scaled_fp8_quant(
503
+ query.reshape(
504
+ (num_tokens, num_heads * head_size)).contiguous(),
505
+ layer._q_scale)
506
+ query = query.reshape((num_tokens, num_heads, head_size))
507
+
508
+ # Compute attention and update output up to `num_actual_tokens`.
509
+ use_local_attn = \
510
+ (self.use_irope and attn_metadata.local_attn_metadata is not None)
511
+
512
+ if not attn_metadata.use_cascade or use_local_attn:
513
+ if use_local_attn:
514
+ assert attn_metadata.local_attn_metadata is not None
515
+ local_metadata = attn_metadata.local_attn_metadata
516
+ cu_seqlens_q = local_metadata.local_query_start_loc
517
+ seqused_k = local_metadata.local_seqused_k
518
+ max_seqlen_q = local_metadata.local_max_query_len
519
+ max_seqlen_k = local_metadata.local_max_seq_len
520
+ block_table = local_metadata.local_block_table
521
+ else:
522
+ cu_seqlens_q = attn_metadata.query_start_loc
523
+ seqused_k = attn_metadata.seq_lens
524
+ max_seqlen_q = attn_metadata.max_query_len
525
+ max_seqlen_k = attn_metadata.max_seq_len
526
+ block_table = attn_metadata.block_table
527
+
528
+ if max_seqlen_q > 1:
529
+ cu_seq_lens = attn_metadata.cu_seq_lens
530
+ total_tokens = attn_metadata.total_tokens
531
+ torch.ops.vllm.flash_attn_varlen_func(
532
+ query[:num_actual_tokens],
533
+ key_cache,
534
+ value_cache,
535
+ out=output[:num_actual_tokens],
536
+ cu_seqlens_q=cu_seqlens_q,
537
+ max_seqlen_q=max_seqlen_q,
538
+ max_seqlen_k=max_seqlen_k,
539
+ total_tokens=total_tokens,
540
+ softmax_scale=self.scale,
541
+ alibi_slopes=self.alibi_slopes,
542
+ window_size=self.sliding_window,
543
+ block_table=block_table,
544
+ cu_seqlens_k=cu_seq_lens,
545
+ )
546
+
547
+ _, num_heads, head_size = query.shape
548
+ _PARTITION_SIZE_ROCM = 256
549
+ num_seqs = seqused_k.shape[0]
550
+ nbyes_per_qo_elem = torch.finfo(output.dtype).bits // 8
551
+ max_num_partitions = (max_seqlen_k + _PARTITION_SIZE_ROCM -
552
+ 1) // _PARTITION_SIZE_ROCM
553
+
554
+ workspace_buffer = torch.empty(
555
+ (num_seqs * num_heads * max_num_partitions * head_size) *
556
+ nbyes_per_qo_elem + 2 *
557
+ (num_seqs * num_heads * max_num_partitions) * 4,
558
+ dtype=torch.uint8,
559
+ device=output.device,
560
+ )
561
+
562
+ aiter.paged_attention_v1(
563
+ output[:num_actual_tokens],
564
+ workspace_buffer,
565
+ query[:num_actual_tokens],
566
+ key_cache,
567
+ value_cache,
568
+ self.scale,
569
+ block_table,
570
+ cu_seqlens_q,
571
+ seqused_k,
572
+ max_seqlen_k,
573
+ self.alibi_slopes,
574
+ self.kv_cache_dtype,
575
+ "NHD",
576
+ self.logits_soft_cap,
577
+ layer._k_scale,
578
+ layer._v_scale,
579
+ None,
580
+ _PARTITION_SIZE_ROCM,
581
+ )
582
+ return output
583
+ else:
584
+ raise NotImplementedError(
585
+ "Cascade attention is not implemented for ROCM AITER")