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, 23:45 UTC
[ROCm][AITER] Enable fp8 kv cache on rocm aiter backend. (#20295)
Signed-off-by: fsx950223 <fsx950223@outlook.com> Signed-off-by: amd-ruitang3 <Rui.Tang2@amd.com> Co-authored-by: amd-ruitang3 <Rui.Tang2@amd.com>
who who who committed
Jul 25, 2025 at 21:50 UTC
b3caeb82e7407d5faa30c49aecd951df3dafd42c
2 files changed
+320
-96
tests/kernels/attention/test_aiter_flash_attn.py
new
+191
@@ -0,0 +1,191 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+
4
+from typing import Optional
5
+
6
+import pytest
7
+import torch
8
+
9
+import vllm.v1.attention.backends.rocm_aiter_fa # noqa: F401
10
+from vllm.platforms import current_platform
11
+
12
+NUM_HEADS = [(4, 4), (8, 2), (16, 2)]
13
+HEAD_SIZES = [128, 256]
14
+BLOCK_SIZES = [16, 32]
15
+DTYPES = [torch.float16, torch.bfloat16]
16
+QDTYPES = [None]
17
+# one value large enough to test overflow in index calculation.
18
+# one value small enough to test the schema op check
19
+NUM_BLOCKS = [32768, 2048]
20
+
21
+
22
+def ref_paged_attn(
23
+ query: torch.Tensor,
24
+ key_cache: torch.Tensor,
25
+ value_cache: torch.Tensor,
26
+ query_lens: list[int],
27
+ kv_lens: list[int],
28
+ block_tables: torch.Tensor,
29
+ scale: float,
30
+ sliding_window: Optional[int] = None,
31
+ soft_cap: Optional[float] = None,
32
+) -> torch.Tensor:
33
+ num_seqs = len(query_lens)
34
+ block_tables = block_tables.cpu().numpy()
35
+ _, block_size, num_kv_heads, head_size = key_cache.shape
36
+
37
+ outputs: list[torch.Tensor] = []
38
+ start_idx = 0
39
+ for i in range(num_seqs):
40
+ query_len = query_lens[i]
41
+ kv_len = kv_lens[i]
42
+ q = query[start_idx:start_idx + query_len]
43
+ q *= scale
44
+
45
+ num_kv_blocks = (kv_len + block_size - 1) // block_size
46
+ block_indices = block_tables[i, :num_kv_blocks]
47
+
48
+ k = key_cache[block_indices].view(-1, num_kv_heads, head_size)
49
+ k = k[:kv_len]
50
+ v = value_cache[block_indices].view(-1, num_kv_heads, head_size)
51
+ v = v[:kv_len]
52
+
53
+ if q.shape[1] != k.shape[1]:
54
+ k = torch.repeat_interleave(k, q.shape[1] // k.shape[1], dim=1)
55
+ v = torch.repeat_interleave(v, q.shape[1] // v.shape[1], dim=1)
56
+ attn = torch.einsum("qhd,khd->hqk", q, k).float()
57
+ empty_mask = torch.ones(query_len, kv_len)
58
+ mask = torch.triu(empty_mask, diagonal=kv_len - query_len + 1).bool()
59
+ if sliding_window is not None:
60
+ sliding_window_mask = torch.triu(empty_mask,
61
+ diagonal=kv_len -
62
+ (query_len + sliding_window) +
63
+ 1).bool().logical_not()
64
+ mask |= sliding_window_mask
65
+ if soft_cap is not None:
66
+ attn = soft_cap * torch.tanh(attn / soft_cap)
67
+ attn.masked_fill_(mask, float("-inf"))
68
+ attn = torch.softmax(attn, dim=-1).to(v.dtype)
69
+ out = torch.einsum("hqk,khd->qhd", attn, v)
70
+
71
+ outputs.append(out)
72
+ start_idx += query_len
73
+
74
+ return torch.cat(outputs, dim=0)
75
+
76
+
77
+@pytest.mark.skipif(not current_platform.is_rocm(),
78
+ reason="Only ROCm is supported")
79
+@pytest.mark.parametrize("seq_lens",
80
+ [[(10, 1328), (5, 18),
81
+ (129, 463)], [(8, 523), (24, 37), (3, 2011)]])
82
+@pytest.mark.parametrize("num_heads", NUM_HEADS)
83
+@pytest.mark.parametrize("head_size", HEAD_SIZES)
84
+@pytest.mark.parametrize("block_size", BLOCK_SIZES)
85
+@pytest.mark.parametrize("sliding_window", [None, 256])
86
+@pytest.mark.parametrize("dtype", DTYPES)
87
+@pytest.mark.parametrize("soft_cap", [None])
88
+@pytest.mark.parametrize("num_blocks", NUM_BLOCKS)
89
+@pytest.mark.parametrize("q_dtype", QDTYPES)
90
+@torch.inference_mode()
91
+def test_varlen_with_paged_kv(
92
+ seq_lens: list[tuple[int, int]],
93
+ num_heads: tuple[int, int],
94
+ head_size: int,
95
+ sliding_window: Optional[int],
96
+ dtype: torch.dtype,
97
+ block_size: int,
98
+ soft_cap: Optional[float],
99
+ num_blocks: int,
100
+ q_dtype: Optional[torch.dtype],
101
+) -> None:
102
+ torch.set_default_device("cuda")
103
+ current_platform.seed_everything(0)
104
+ num_seqs = len(seq_lens)
105
+ query_lens = [x[0] for x in seq_lens]
106
+ kv_lens = [x[1] for x in seq_lens]
107
+ num_query_heads = num_heads[0]
108
+ num_kv_heads = num_heads[1]
109
+ assert num_query_heads % num_kv_heads == 0
110
+ max_query_len = max(query_lens)
111
+ max_kv_len = max(kv_lens)
112
+ window_size = ((sliding_window - 1, 0) if sliding_window is not None else
113
+ (-1, -1))
114
+ scale = head_size**-0.5
115
+
116
+ query = torch.randn(sum(query_lens),
117
+ num_query_heads,
118
+ head_size,
119
+ dtype=dtype)
120
+ key_cache = torch.randn(num_blocks,
121
+ block_size,
122
+ num_kv_heads,
123
+ head_size,
124
+ dtype=dtype)
125
+ value_cache = torch.randn_like(key_cache)
126
+ cu_query_lens = torch.tensor([0] + query_lens,
127
+ dtype=torch.int32).cumsum(dim=0,
128
+ dtype=torch.int32)
129
+
130
+ cu_seq_lens = torch.tensor([0] + kv_lens,
131
+ dtype=torch.int32).cumsum(dim=0,
132
+ dtype=torch.int32)
133
+ kv_lens = torch.tensor(kv_lens, dtype=torch.int32)
134
+
135
+ max_num_blocks_per_seq = (max_kv_len + block_size - 1) // block_size
136
+ block_tables = torch.randint(0,
137
+ num_blocks,
138
+ (num_seqs, max_num_blocks_per_seq),
139
+ dtype=torch.int32)
140
+
141
+ output = torch.empty_like(query)
142
+
143
+ maybe_quantized_query = query
144
+ maybe_quantized_key_cache = key_cache
145
+ maybe_quantized_value_cache = value_cache
146
+ k_descale = None
147
+ v_descale = None
148
+ if q_dtype is not None:
149
+ # QKV are drawn from N(0, 1): no need for a fp8 scaling factor
150
+ maybe_quantized_query = query.to(q_dtype)
151
+ maybe_quantized_key_cache = key_cache.to(q_dtype)
152
+ maybe_quantized_value_cache = value_cache.to(q_dtype)
153
+
154
+ scale_shape = (num_seqs, num_kv_heads)
155
+ k_descale = torch.ones(scale_shape, dtype=torch.float32)
156
+ v_descale = torch.ones(scale_shape, dtype=torch.float32)
157
+
158
+ torch.ops.vllm.flash_attn_varlen_func(
159
+ maybe_quantized_query,
160
+ maybe_quantized_key_cache,
161
+ maybe_quantized_value_cache,
162
+ out=output,
163
+ cu_seqlens_q=cu_query_lens,
164
+ max_seqlen_q=max_query_len,
165
+ max_seqlen_k=max_kv_len,
166
+ softmax_scale=scale,
167
+ alibi_slopes=None,
168
+ window_size=window_size,
169
+ block_table=block_tables,
170
+ cu_seqlens_k=cu_seq_lens,
171
+ k_scale=k_descale,
172
+ v_scale=v_descale,
173
+ )
174
+
175
+ ref_output = ref_paged_attn(
176
+ query=query,
177
+ key_cache=key_cache,
178
+ value_cache=value_cache,
179
+ query_lens=query_lens,
180
+ kv_lens=kv_lens,
181
+ block_tables=block_tables,
182
+ scale=scale,
183
+ sliding_window=sliding_window,
184
+ soft_cap=soft_cap,
185
+ )
186
+
187
+ atol, rtol = 2e-2, 2e-2
188
+ if q_dtype is not None:
189
+ atol, rtol = 1.5e-1, 1.5e-1
190
+ torch.testing.assert_close(output, ref_output, atol=atol, rtol=rtol), \
191
+ f"{torch.max(torch.abs(output - ref_output))}"
vllm/v1/attention/backends/rocm_aiter_fa.py
+129
-96
@@ -2,20 +2,21 @@
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
"""Attention layer with AiterFlashAttention."""
4
from dataclasses import dataclass
5
-from typing import Optional
5
+from typing import ClassVar, Optional
6
7
import torch
8
9
-from vllm import _custom_ops as ops
9
from vllm.attention.backends.abstract import (AttentionBackend, AttentionImpl,
11
- AttentionMetadata, AttentionType,
12
- is_quantized_kv_cache)
10
+ AttentionMetadata, AttentionType)
11
from vllm.config import VllmConfig
12
from vllm.logger import init_logger
13
from vllm.platforms import current_platform
16
-from vllm.v1.attention.backends.utils import CommonAttentionMetadata
14
+from vllm.v1.attention.backends.utils import (AttentionMetadataBuilder,
15
+ CommonAttentionMetadata)
16
from vllm.v1.kv_cache_interface import AttentionSpec
17
18
+_PARTITION_SIZE_ROCM = 256
19
+
20
if current_platform.is_rocm():
21
import aiter
22
@@ -32,38 +33,54 @@ if current_platform.is_rocm():
33
b_seq_lens_loc,
34
block_table,
35
block_table_stride_0,
36
+ k_scale,
37
+ v_scale,
38
+ output_dtype: tl.constexpr,
39
E_DIM: tl.constexpr,
40
BLOCK_SIZE: tl.constexpr,
41
):
42
batch_idx = tl.program_id(0)
43
block_idx = tl.program_id(1)
40
- batch_token_indexes = tl.load(b_seq_lens_loc + batch_idx +
41
- tl.arange(0, 2))
42
- batch_token_start, batch_token_end = tl.split(batch_token_indexes)
43
- seq_len = batch_token_end - batch_token_start
44
45
batch_query_indexes = tl.load(b_query_lens_loc + batch_idx +
46
tl.arange(0, 2))
47
batch_query_start, batch_query_end = tl.split(batch_query_indexes)
48
query_len = batch_query_end - batch_query_start
49
+
50
if query_len <= 1:
51
return
52
+
53
+ batch_token_indexes = tl.load(b_seq_lens_loc + batch_idx +
54
+ tl.arange(0, 2))
55
+ batch_token_start, batch_token_end = tl.split(batch_token_indexes)
56
+ seq_len = batch_token_end - batch_token_start
57
+
58
if block_idx * BLOCK_SIZE < seq_len:
59
block_mask = (block_idx * BLOCK_SIZE +
60
tl.arange(0, BLOCK_SIZE)[:, None]) < seq_len
61
62
kv_idx = tl.load(block_table + batch_idx * block_table_stride_0 +
56
- block_idx)
63
+ block_idx).to(tl.int64)
64
65
kv_buffer_off = kv_idx * BLOCK_SIZE * E_DIM + tl.arange(
66
0, BLOCK_SIZE)[:, None] * E_DIM + tl.arange(0, E_DIM)[None, :]
67
k_vals = tl.load(k_buffer_ptr + kv_buffer_off,
68
mask=block_mask,
69
other=0.0)
70
+ if k_vals.dtype.is_fp8():
71
+ k_vals = (k_vals.to(tl.float32) *
72
+ tl.load(k_scale)).to(output_dtype)
73
+ else:
74
+ k_vals = k_vals.to(output_dtype)
75
+
76
v_vals = tl.load(v_buffer_ptr + kv_buffer_off,
77
mask=block_mask,
78
other=0.0)
66
-
79
+ if v_vals.dtype.is_fp8():
80
+ v_vals = (v_vals.to(tl.float32) *
81
+ tl.load(v_scale)).to(output_dtype)
82
+ else:
83
+ v_vals = v_vals.to(output_dtype)
84
kv_values_off = batch_token_start * E_DIM + \
85
block_idx * BLOCK_SIZE * E_DIM + \
86
tl.arange(0, BLOCK_SIZE)[:, None] * E_DIM + \
@@ -72,29 +89,44 @@ if current_platform.is_rocm():
89
tl.store(v_values_ptr + kv_values_off, v_vals, mask=block_mask)
90
91
def vllm_layout_trans(b_query_lens_loc, b_seq_lens_loc, block_table,
75
- k_buffer, v_buffer, max_seq_len, total_tokens):
76
- H_KV = v_buffer.shape[2]
77
- D = v_buffer.shape[3]
78
- BLOCK_SIZE = v_buffer.shape[1]
79
- dtype = k_buffer.dtype
80
- k_values = torch.empty((total_tokens, H_KV, D),
81
- dtype=dtype,
82
- device="cuda")
83
- v_values = torch.empty((total_tokens, H_KV, D),
84
- dtype=dtype,
85
- device="cuda")
92
+ k_cache, v_cache, max_seq_len, k_scale, v_scale,
93
+ output_dtype, total_tokens):
94
+ H_KV = v_cache.shape[2]
95
+ D = v_cache.shape[3]
96
+ BLOCK_SIZE = v_cache.shape[1]
97
+
98
+ k_values = torch.empty(
99
+ (total_tokens, H_KV, D),
100
+ dtype=output_dtype,
101
+ device=k_cache.device,
102
+ )
103
+ v_values = torch.empty(
104
+ (total_tokens, H_KV, D),
105
+ dtype=output_dtype,
106
+ device=v_cache.device,
107
+ )
108
109
grid = (block_table.shape[0],
110
(max_seq_len + BLOCK_SIZE - 1) // BLOCK_SIZE)
111
90
- _vllm_layout_trans_kernel[grid](k_buffer,
91
- v_buffer,
112
+ if output_dtype == torch.float16:
113
+ output_dtype = tl.float16
114
+ elif output_dtype == torch.bfloat16:
115
+ output_dtype = tl.bfloat16
116
+ else:
117
+ raise ValueError(f"Unsupported output dtype: {output_dtype}")
118
+
119
+ _vllm_layout_trans_kernel[grid](k_cache,
120
+ v_cache,
121
k_values,
122
v_values,
123
b_query_lens_loc,
124
b_seq_lens_loc,
125
block_table,
126
block_table.stride(0),
127
+ k_scale,
128
+ v_scale,
129
+ output_dtype=output_dtype,
130
E_DIM=H_KV * D,
131
BLOCK_SIZE=BLOCK_SIZE)
132
@@ -107,16 +139,22 @@ if current_platform.is_rocm():
139
out: torch.Tensor,
140
cu_seqlens_q: torch.Tensor,
141
cu_seqlens_k: torch.Tensor,
110
- total_tokens: int,
142
max_seqlen_q: int,
143
max_seqlen_k: int,
144
softmax_scale: float,
145
window_size: Optional[list[int]], # -1 means infinite context window
146
alibi_slopes: Optional[list[float]],
147
block_table: torch.Tensor,
148
+ k_scale: torch.Tensor,
149
+ v_scale: torch.Tensor,
150
+ total_tokens: int = 0,
151
) -> torch.Tensor:
152
+ if total_tokens == 0:
153
+ total_tokens = int(cu_seqlens_k[-1].item())
154
k, v = vllm_layout_trans(cu_seqlens_q, cu_seqlens_k, block_table,
119
- k_cache, v_cache, max_seqlen_k, total_tokens)
155
+ k_cache, v_cache, max_seqlen_k, k_scale,
156
+ v_scale, q.dtype, total_tokens)
157
+
158
output = aiter.flash_attn_varlen_func(
159
q=q,
160
k=k,
@@ -141,19 +179,21 @@ if current_platform.is_rocm():
179
out: torch.Tensor,
180
cu_seqlens_q: torch.Tensor,
181
cu_seqlens_k: torch.Tensor,
144
- total_tokens: int,
182
max_seqlen_q: int,
183
max_seqlen_k: int,
184
softmax_scale: float,
185
window_size: Optional[list[int]], # -1 means infinite context window
186
alibi_slopes: Optional[list[float]],
187
block_table: torch.Tensor,
188
+ k_scale: torch.Tensor,
189
+ v_scale: torch.Tensor,
190
+ total_tokens: int = 0,
191
) -> torch.Tensor:
192
return torch.empty(q.shape[0],
193
q.shape[1],
194
v_cache.shape[-2],
155
- dtype=torch.float8_e4m3fnuz,
156
- device="cuda")
195
+ dtype=q.dtype,
196
+ device=q.device)
197
198
direct_register_custom_op("flash_attn_varlen_func",
199
flash_attn_varlen_func_impl, ["out"],
@@ -163,7 +203,33 @@ if current_platform.is_rocm():
203
logger = init_logger(__name__)
204
205
166
-class AiterFlashAttentionMetadataBuilder:
206
+@dataclass
207
+class AiterFlashAttentionMetadata:
208
+ # NOTE(sang): Definition of context_len, query_len, and seq_len.
209
+ # |---------- N-1 iteration --------|
210
+ # |---------------- N iteration ---------------------|
211
+ # |- tokenA -|......................|-- newTokens ---|
212
+ # |---------- context_len ----------|
213
+ # |-------------------- seq_len ---------------------|
214
+ # |-- query_len ---|
215
+
216
+ num_actual_tokens: int # Number of tokens excluding padding.
217
+ max_query_len: int
218
+ query_start_loc: torch.Tensor
219
+ max_seq_len: int
220
+ seq_lens: torch.Tensor
221
+ slot_mapping: torch.Tensor
222
+ block_table: torch.Tensor
223
+
224
+ # For cascade attention.
225
+ use_cascade: bool
226
+ common_prefix_len: int
227
+ total_tokens: int
228
+
229
+
230
+class AiterFlashAttentionMetadataBuilder(
231
+ AttentionMetadataBuilder[AiterFlashAttentionMetadata]):
232
+ full_cudagraph_supported: ClassVar[bool] = True
233
234
def __init__(self, kv_cache_spec: AttentionSpec, vllm_config: VllmConfig,
235
device: torch.device):
@@ -180,14 +246,23 @@ class AiterFlashAttentionMetadataBuilder:
246
self.headdim = self.model_config.get_head_size()
247
self.block_size = kv_cache_spec.block_size
248
self.kv_cache_spec = kv_cache_spec
183
-
249
# Sliding window size to be used with the AOT scheduler will be
250
# populated on first build() call.
251
self.aot_sliding_window: Optional[tuple[int, int]] = None
252
+ self.total_tokens: int = 0
253
254
def reorder_batch(self, input_batch, scheduler_output) -> bool:
255
return False
256
257
+ def build_for_cudagraph_capture(
258
+ self, common_attn_metadata: CommonAttentionMetadata):
259
+ self.total_tokens = self.model_config.max_model_len \
260
+ * self.vllm_config.scheduler_config.max_num_partial_prefills
261
+ res = self.build(common_prefix_len=0,
262
+ common_attn_metadata=common_attn_metadata)
263
+ self.total_tokens = 0
264
+ return res
265
+
266
def build(self,
267
common_prefix_len: int,
268
common_attn_metadata: CommonAttentionMetadata,
@@ -195,43 +270,29 @@ class AiterFlashAttentionMetadataBuilder:
270
271
num_actual_tokens = common_attn_metadata.num_actual_tokens
272
max_query_len = common_attn_metadata.max_query_len
198
-
273
max_seq_len = int(common_attn_metadata.seq_lens_cpu.max())
200
- total_tokens = int(common_attn_metadata.seq_lens_cpu.sum())
274
query_start_loc = common_attn_metadata.query_start_loc
275
seq_lens = common_attn_metadata.seq_lens
276
block_table_tensor = common_attn_metadata.block_table_tensor
277
slot_mapping = common_attn_metadata.slot_mapping
278
206
- cu_seq_lens = torch.zeros(seq_lens.shape[0] + 1,
207
- dtype=torch.int32,
208
- device=self.device)
209
- torch.cumsum(seq_lens,
210
- dim=0,
211
- dtype=cu_seq_lens.dtype,
212
- out=cu_seq_lens[1:])
279
+ def schedule(batch_size, cu_query_lens, max_query_len, seqlens,
280
+ max_seq_len, causal):
281
+ return None
282
283
use_cascade = common_prefix_len > 0
284
216
- cu_prefix_query_lens = None
217
- prefix_kv_lens = None
218
- suffix_kv_lens = None
219
-
285
attn_metadata = AiterFlashAttentionMetadata(
286
num_actual_tokens=num_actual_tokens,
287
max_query_len=max_query_len,
288
query_start_loc=query_start_loc,
289
max_seq_len=max_seq_len,
290
seq_lens=seq_lens,
226
- cu_seq_lens=cu_seq_lens,
227
- total_tokens=total_tokens,
291
block_table=block_table_tensor,
292
slot_mapping=slot_mapping,
293
use_cascade=use_cascade,
294
common_prefix_len=common_prefix_len,
232
- cu_prefix_query_lens=cu_prefix_query_lens,
233
- prefix_kv_lens=prefix_kv_lens,
234
- suffix_kv_lens=suffix_kv_lens,
295
+ total_tokens=self.total_tokens,
296
)
297
return attn_metadata
298
@@ -254,7 +315,7 @@ class AiterFlashAttentionBackend(AttentionBackend):
315
316
@classmethod
317
def get_supported_head_sizes(cls) -> list[int]:
257
- return [32, 64, 96, 128, 160, 192, 224, 256]
318
+ return [64, 128, 256]
319
320
@classmethod
321
def validate_head_size(cls, head_size: int) -> None:
@@ -295,34 +356,6 @@ class AiterFlashAttentionBackend(AttentionBackend):
356
return (2, num_blocks, block_size, num_kv_heads, head_size)
357
358
298
-@dataclass
299
-class AiterFlashAttentionMetadata:
300
- # NOTE(sang): Definition of context_len, query_len, and seq_len.
301
- # |---------- N-1 iteration --------|
302
- # |---------------- N iteration ---------------------|
303
- # |- tokenA -|......................|-- newTokens ---|
304
- # |---------- context_len ----------|
305
- # |-------------------- seq_len ---------------------|
306
- # |-- query_len ---|
307
-
308
- num_actual_tokens: int # Number of tokens excluding padding.
309
- max_query_len: int
310
- query_start_loc: torch.Tensor
311
- max_seq_len: int
312
- seq_lens: torch.Tensor
313
- cu_seq_lens: torch.Tensor
314
- total_tokens: int
315
- block_table: torch.Tensor
316
- slot_mapping: torch.Tensor
317
-
318
- # For cascade attention.
319
- use_cascade: bool
320
- common_prefix_len: int
321
- cu_prefix_query_lens: Optional[torch.Tensor]
322
- prefix_kv_lens: Optional[torch.Tensor]
323
- suffix_kv_lens: Optional[torch.Tensor]
324
-
325
-
359
class AiterFlashAttentionImpl(AttentionImpl):
360
361
def __init__(
@@ -366,10 +399,6 @@ class AiterFlashAttentionImpl(AttentionImpl):
399
"encoder/decoder cross-attention "
400
"are not implemented for "
401
"FlashAttentionImpl")
369
- if is_quantized_kv_cache(self.kv_cache_dtype):
370
- raise NotImplementedError(
371
- "AiterFlashAttention does not support fp8 kv-cache on this "
372
- "device.")
402
403
def forward(
404
self,
@@ -440,12 +469,6 @@ class AiterFlashAttentionImpl(AttentionImpl):
469
if self.kv_cache_dtype.startswith("fp8"):
470
key_cache = key_cache.view(torch.float8_e4m3fnuz)
471
value_cache = value_cache.view(torch.float8_e4m3fnuz)
443
- num_tokens, num_heads, head_size = query.shape
444
- query, _ = ops.scaled_fp8_quant(
445
- query.reshape(
446
- (num_tokens, num_heads * head_size)).contiguous(),
447
- layer._q_scale)
448
- query = query.reshape((num_tokens, num_heads, head_size))
472
473
if not attn_metadata.use_cascade:
474
cu_seqlens_q = attn_metadata.query_start_loc
@@ -455,8 +478,16 @@ class AiterFlashAttentionImpl(AttentionImpl):
478
block_table = attn_metadata.block_table
479
480
if max_seqlen_q > 1:
458
- cu_seq_lens = attn_metadata.cu_seq_lens
459
- total_tokens = attn_metadata.total_tokens
481
+
482
+ cu_seq_lens = torch.zeros(seqused_k.shape[0] + 1,
483
+ dtype=torch.int32,
484
+ device=query.device)
485
+
486
+ torch.cumsum(seqused_k,
487
+ dim=0,
488
+ dtype=cu_seq_lens.dtype,
489
+ out=cu_seq_lens[1:])
490
+
491
torch.ops.vllm.flash_attn_varlen_func(
492
query[:num_actual_tokens],
493
key_cache,
@@ -465,29 +496,31 @@ class AiterFlashAttentionImpl(AttentionImpl):
496
cu_seqlens_q=cu_seqlens_q,
497
max_seqlen_q=max_seqlen_q,
498
max_seqlen_k=max_seqlen_k,
468
- total_tokens=total_tokens,
499
softmax_scale=self.scale,
500
alibi_slopes=self.alibi_slopes,
501
window_size=self.sliding_window,
502
block_table=block_table,
473
- cu_seqlens_k=cu_seq_lens)
503
+ cu_seqlens_k=cu_seq_lens,
504
+ k_scale=layer._k_scale,
505
+ v_scale=layer._v_scale,
506
+ total_tokens=attn_metadata.total_tokens,
507
+ )
508
509
_, num_heads, head_size = query.shape
476
- _PARTITION_SIZE_ROCM = 256
510
+ nbytes_per_qo_elem = torch.finfo(query.dtype).bits // 8
511
num_seqs = seqused_k.shape[0]
478
- nbyes_per_qo_elem = torch.finfo(output.dtype).bits // 8
512
max_num_partitions = (max_seqlen_k + _PARTITION_SIZE_ROCM -
513
1) // _PARTITION_SIZE_ROCM
514
515
workspace_buffer = torch.empty(
516
(num_seqs * num_heads * max_num_partitions * head_size) *
484
- nbyes_per_qo_elem + 2 *
517
+ nbytes_per_qo_elem + 2 *
518
(num_seqs * num_heads * max_num_partitions) * 4,
519
dtype=torch.uint8,
520
device=output.device,
521
)
522
490
- aiter.paged_attention_v1(
523
+ torch.ops.aiter.paged_attention_v1(
524
output[:num_actual_tokens],
525
workspace_buffer,
526
query[:num_actual_tokens],