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 12, 09:45 UTC
[V0 deprecation] Clean up legacy paged attention helper functions (#28043)
Signed-off-by: Isotr0py <mozf@mail2.sysu.edu.cn>
Isotr0py committed
Nov 29, 2025 at 00:44 UTC
6f9d81d03b3b60f8a8b51d624c86f99bf26e96a4
2 files changed
-334
vllm/attention/ops/paged_attn.py
-211
@@ -1,58 +1,18 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4
-from dataclasses import dataclass
4
5
import torch
6
7
from vllm.platforms import current_platform
9
-from vllm.triton_utils import HAS_TRITON
8
9
if current_platform.is_cuda_alike():
10
from vllm import _custom_ops as ops
11
elif current_platform.is_xpu():
12
from vllm._ipex_ops import ipex_ops as ops
13
16
-if HAS_TRITON:
17
- from vllm.attention.ops.prefix_prefill import context_attention_fwd
18
-
19
-# Should be the same as PARTITION_SIZE in `paged_attention_v2_launcher`.
20
-_PARTITION_SIZE = 512
21
-
22
-
23
-@dataclass
24
-class PagedAttentionMetadata:
25
- """Metadata for PagedAttention."""
26
-
27
- # (batch_size,). The length of sequences (entire tokens seen so far) per
28
- # sequence.
29
- seq_lens_tensor: torch.Tensor | None
30
- # Maximum sequence length in the batch. 0 if it is prefill-only batch.
31
- max_decode_seq_len: int
32
- # (batch_size, max_blocks_per_seq).
33
- # Block addresses per sequence. (Seq id -> list of physical block)
34
- # E.g., [0, 1, 2] means tokens are stored in 0th, 1st, and 2nd blocks
35
- # in the kv cache. Each block can contain up to block_size tokens.
36
- # 2nd dimensions are padded up to max_blocks_per_seq if it is cuda-graph
37
- # captured.
38
- block_tables: torch.Tensor | None
39
-
14
15
class PagedAttention:
42
- @staticmethod
43
- def get_supported_head_sizes() -> list[int]:
44
- return [32, 64, 80, 96, 112, 120, 128, 192, 256]
45
-
46
- @staticmethod
47
- def get_kv_cache_shape(
48
- num_blocks: int,
49
- block_size: int,
50
- num_kv_heads: int,
51
- head_size: int,
52
- cache_dtype_str: str = "auto",
53
- ) -> tuple[int, ...]:
54
- return (2, num_blocks, block_size * num_kv_heads * head_size)
55
-
16
@staticmethod
17
def split_kv_cache(
18
kv_cache: torch.Tensor,
@@ -89,174 +49,3 @@ class PagedAttention:
49
k_scale,
50
v_scale,
51
)
92
-
93
- @staticmethod
94
- def forward_decode(
95
- query: torch.Tensor,
96
- key_cache: torch.Tensor,
97
- value_cache: torch.Tensor,
98
- block_tables: torch.Tensor,
99
- seq_lens: torch.Tensor,
100
- max_seq_len: int,
101
- kv_cache_dtype: str,
102
- num_kv_heads: int,
103
- scale: float,
104
- alibi_slopes: torch.Tensor | None,
105
- k_scale: torch.Tensor,
106
- v_scale: torch.Tensor,
107
- tp_rank: int = 0,
108
- blocksparse_local_blocks: int = 0,
109
- blocksparse_vert_stride: int = 0,
110
- blocksparse_block_size: int = 64,
111
- blocksparse_head_sliding_step: int = 0,
112
- ) -> torch.Tensor:
113
- if blocksparse_vert_stride is not None and blocksparse_vert_stride > 1:
114
- # use blocksparse paged attention
115
- block_size = value_cache.size(-1)
116
- assert (
117
- blocksparse_block_size > 0 and blocksparse_block_size % block_size == 0
118
- ), (
119
- f"{blocksparse_block_size=} needs to be a multiple of"
120
- f"{block_size=} used in block_tables."
121
- )
122
-
123
- output = torch.empty_like(query)
124
- block_size = value_cache.shape[3]
125
- num_seqs, num_heads, head_size = query.shape
126
- max_num_partitions = (max_seq_len + _PARTITION_SIZE - 1) // _PARTITION_SIZE
127
- # NOTE(woosuk): We use a simple heuristic to decide whether to use
128
- # PagedAttention V1 or V2. If the number of partitions is 1, we use
129
- # V1 to avoid the overhead of reduction. Also, if the number of
130
- # sequences or heads is large, we use V1 since there is enough work
131
- # to parallelize.
132
- # TODO(woosuk): Tune this heuristic.
133
- # For context len > 8192, use V2 kernel to avoid shared memory shortage.
134
- use_v1 = max_seq_len <= 8192 and (
135
- max_num_partitions == 1 or num_seqs * num_heads > 512
136
- )
137
-
138
- if use_v1:
139
- # Run PagedAttention V1.
140
- ops.paged_attention_v1(
141
- output,
142
- query,
143
- key_cache,
144
- value_cache,
145
- num_kv_heads,
146
- scale,
147
- block_tables,
148
- seq_lens,
149
- block_size,
150
- max_seq_len,
151
- alibi_slopes,
152
- kv_cache_dtype,
153
- k_scale,
154
- v_scale,
155
- tp_rank,
156
- blocksparse_local_blocks,
157
- blocksparse_vert_stride,
158
- blocksparse_block_size,
159
- blocksparse_head_sliding_step,
160
- )
161
- else:
162
- # Run PagedAttention V2.
163
- assert _PARTITION_SIZE % block_size == 0
164
- tmp_output = torch.empty(
165
- size=(num_seqs, num_heads, max_num_partitions, head_size),
166
- dtype=output.dtype,
167
- device=output.device,
168
- )
169
- exp_sums = torch.empty(
170
- size=(num_seqs, num_heads, max_num_partitions),
171
- dtype=torch.float32,
172
- device=output.device,
173
- )
174
- max_logits = torch.empty_like(exp_sums)
175
- ops.paged_attention_v2(
176
- output,
177
- exp_sums,
178
- max_logits,
179
- tmp_output,
180
- query,
181
- key_cache,
182
- value_cache,
183
- num_kv_heads,
184
- scale,
185
- block_tables,
186
- seq_lens,
187
- block_size,
188
- max_seq_len,
189
- alibi_slopes,
190
- kv_cache_dtype,
191
- k_scale,
192
- v_scale,
193
- tp_rank,
194
- blocksparse_local_blocks,
195
- blocksparse_vert_stride,
196
- blocksparse_block_size,
197
- blocksparse_head_sliding_step,
198
- )
199
- return output
200
-
201
- @staticmethod
202
- def forward_prefix(
203
- query: torch.Tensor,
204
- key: torch.Tensor,
205
- value: torch.Tensor,
206
- kv_cache_dtype: str,
207
- key_cache: torch.Tensor,
208
- value_cache: torch.Tensor,
209
- block_tables: torch.Tensor,
210
- query_start_loc: torch.Tensor,
211
- seq_lens_tensor: torch.Tensor,
212
- max_query_len: int,
213
- alibi_slopes: torch.Tensor | None,
214
- sliding_window: int | None,
215
- k_scale: torch.Tensor,
216
- v_scale: torch.Tensor,
217
- ) -> torch.Tensor:
218
- output = torch.empty_like(query)
219
- max_seq_len = None
220
- context_attention_fwd(
221
- query,
222
- key,
223
- value,
224
- output,
225
- kv_cache_dtype,
226
- key_cache,
227
- value_cache,
228
- block_tables,
229
- # query_start_loc is (batch_size + 1,)
230
- query_start_loc,
231
- seq_lens_tensor,
232
- max_seq_len,
233
- max_query_len,
234
- k_scale,
235
- v_scale,
236
- alibi_slopes,
237
- sliding_window,
238
- )
239
- return output
240
-
241
- @staticmethod
242
- def swap_blocks(
243
- src_kv_cache: torch.Tensor,
244
- dst_kv_cache: torch.Tensor,
245
- src_to_dst: torch.Tensor,
246
- ) -> None:
247
- src_key_cache = src_kv_cache[0]
248
- dst_key_cache = dst_kv_cache[0]
249
- ops.swap_blocks(src_key_cache, dst_key_cache, src_to_dst)
250
-
251
- src_value_cache = src_kv_cache[1]
252
- dst_value_cache = dst_kv_cache[1]
253
- ops.swap_blocks(src_value_cache, dst_value_cache, src_to_dst)
254
-
255
- @staticmethod
256
- def copy_blocks(
257
- kv_caches: list[torch.Tensor],
258
- src_to_dists: torch.Tensor,
259
- ) -> None:
260
- key_caches = [kv_cache[0] for kv_cache in kv_caches]
261
- value_caches = [kv_cache[1] for kv_cache in kv_caches]
262
- ops.copy_blocks(key_caches, value_caches, src_to_dists)
vllm/attention/ops/rocm_aiter_paged_attn.py
deleted
-123
@@ -1,123 +0,0 @@
1
-# SPDX-License-Identifier: Apache-2.0
2
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
-
4
-import aiter as rocm_aiter
5
-import torch
6
-
7
-from vllm.attention.ops.paged_attn import PagedAttention
8
-from vllm.platforms import current_platform
9
-from vllm.utils.math_utils import cdiv
10
-
11
-FP8_DTYPE = current_platform.fp8_dtype()
12
-
13
-
14
-class AITERPagedAttention(PagedAttention):
15
- @staticmethod
16
- def write_to_paged_cache(
17
- key: torch.Tensor,
18
- value: torch.Tensor,
19
- key_cache: torch.Tensor,
20
- value_cache: torch.Tensor,
21
- slot_mapping: torch.Tensor,
22
- kv_cache_dtype: str,
23
- k_scale: torch.Tensor,
24
- v_scale: torch.Tensor,
25
- ) -> None:
26
- if kv_cache_dtype not in ["int8", "fp8", "fp8_e4m3"]:
27
- PagedAttention.write_to_paged_cache(
28
- key,
29
- value,
30
- key_cache,
31
- value_cache,
32
- slot_mapping,
33
- kv_cache_dtype,
34
- k_scale,
35
- v_scale,
36
- )
37
- else:
38
- kv_cache_torch_dtype = FP8_DTYPE if "fp8" in kv_cache_dtype else torch.int8
39
- key_cache = key_cache.view(kv_cache_torch_dtype)
40
- value_cache = value_cache.view(kv_cache_torch_dtype)
41
-
42
- rocm_aiter.reshape_and_cache_with_pertoken_quant(
43
- key,
44
- value,
45
- key_cache,
46
- value_cache,
47
- k_scale,
48
- v_scale,
49
- slot_mapping.flatten(),
50
- True,
51
- )
52
-
53
- @staticmethod
54
- def forward_decode(
55
- query: torch.Tensor,
56
- key_cache: torch.Tensor,
57
- value_cache: torch.Tensor,
58
- block_tables: torch.Tensor,
59
- seq_lens: torch.Tensor,
60
- max_seq_len: int,
61
- kv_cache_dtype: str,
62
- num_kv_heads: int,
63
- scale: float,
64
- alibi_slopes: torch.Tensor | None,
65
- k_scale: torch.Tensor,
66
- v_scale: torch.Tensor,
67
- tp_rank: int = 0,
68
- blocksparse_local_blocks: int = 0,
69
- blocksparse_vert_stride: int = 0,
70
- blocksparse_block_size: int = 64,
71
- blocksparse_head_sliding_step: int = 0,
72
- ) -> torch.Tensor:
73
- if kv_cache_dtype not in ["int8", "fp8", "fp8_e4m3"]:
74
- return PagedAttention.forward_decode(
75
- query=query,
76
- key_cache=key_cache,
77
- value_cache=value_cache,
78
- block_tables=block_tables,
79
- seq_lens=seq_lens,
80
- max_seq_len=max_seq_len,
81
- kv_cache_dtype=kv_cache_dtype,
82
- num_kv_heads=num_kv_heads,
83
- scale=scale,
84
- alibi_slopes=alibi_slopes,
85
- k_scale=k_scale,
86
- v_scale=v_scale,
87
- tp_rank=tp_rank,
88
- blocksparse_local_blocks=blocksparse_local_blocks,
89
- blocksparse_vert_stride=blocksparse_vert_stride,
90
- blocksparse_block_size=blocksparse_block_size,
91
- blocksparse_head_sliding_step=blocksparse_head_sliding_step,
92
- )
93
-
94
- if "fp8" in kv_cache_dtype:
95
- key_cache = key_cache.view(current_platform.fp8_dtype())
96
- value_cache = value_cache.view(current_platform.fp8_dtype())
97
-
98
- if blocksparse_vert_stride is not None and blocksparse_vert_stride > 1:
99
- # use blocksparse paged attention
100
- block_size = value_cache.size(-1)
101
- assert (
102
- blocksparse_block_size > 0 and blocksparse_block_size % block_size == 0
103
- ), (
104
- f"{blocksparse_block_size=} needs to be a multiple of"
105
- f"{block_size=} used in block_tables."
106
- )
107
-
108
- output = torch.empty_like(query)
109
- block_size = value_cache.shape[3]
110
- max_num_blocks_per_seq = cdiv(max_seq_len, block_size)
111
-
112
- rocm_aiter.pa_fwd_asm(
113
- query,
114
- key_cache,
115
- value_cache,
116
- block_tables,
117
- seq_lens,
118
- max_num_blocks_per_seq,
119
- k_scale,
120
- v_scale,
121
- output,
122
- )
123
- return output