@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / dc937175d4
Fork 0

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 13, 19:30 UTC

[ROCm][Perf] New design on ROCm AITER MHA backend Implementation (#25763)

Signed-off-by: ganyi <ygan@amd.com>

Pleaplusone committed Nov 5, 2025 at 02:05 UTC dc937175d496a801f175b5dcce6a8157506bae52
2 files changed +593 -275
vllm/v1/attention/backends/rocm_aiter_fa.py
+526 -275
@@ -13,221 +13,202 @@ from vllm.attention.backends.abstract import (
13 AttentionType,
14 MultipleOf,
15 )
16 +from vllm.attention.ops.merge_attn_states import merge_attn_states
17 from vllm.config import VllmConfig
18 from vllm.logger import init_logger
19 from vllm.platforms import current_platform
20 +from vllm.utils.math_utils import cdiv
21 from vllm.v1.attention.backends.utils import (
22 AttentionCGSupport,
23 AttentionMetadataBuilder,
24 CommonAttentionMetadata,
25 + split_decodes_prefills_and_extends,
26 )
27 from vllm.v1.kv_cache_interface import AttentionSpec
28
29 _PARTITION_SIZE_ROCM = 256
30 +_CP_TOKENS_PER_ITER_ROCM = 32 * 1024
31
32 if current_platform.is_rocm():
33 import aiter
34 + from aiter.ops.triton.utils.device_info import get_num_sms
35
36 from vllm.triton_utils import tl, triton
32 - from vllm.utils.torch_utils import direct_register_custom_op
37 +
38 + def block_size(x, head_dim):
39 + return min(65536 // x.element_size(), triton.next_power_of_2(head_dim))
40 +
41 + def num_programs(head_dim):
42 + return min(head_dim, get_num_sms())
43
44 @triton.jit
35 - def _vllm_layout_trans_kernel(
36 - k_buffer_ptr,
37 - v_buffer_ptr,
38 - k_values_ptr,
39 - v_values_ptr,
40 - b_query_lens_loc,
41 - b_seq_lens_loc,
42 - block_table,
43 - block_table_stride_0,
44 - k_scale,
45 - v_scale,
46 - output_dtype: tl.constexpr,
47 - E_DIM: tl.constexpr,
45 + def cp_mha_gather_cache_kernel(
46 + key_cache_ptr, # [num_blocks, page_size, num_head, head_size]
47 + value_cache_ptr, # [num_blocks, page_size, num_head, head_size]
48 + key_ptr, # [num_tokens, num_heads, head_size]
49 + value_ptr, # [num_tokens, num_heads, head_size]
50 + block_table_ptr, # [num_batches, max_block_num]
51 + cu_seqlens_kv_ptr, # [num_batches + 1]
52 + token_to_batch_ptr, # [max_cum_tokens]
53 + seq_start_ptr, # [num_batches]
54 + k_scale_ptr,
55 + v_scale_ptr,
56 + num_heads,
57 + head_size,
58 + x,
59 + max_block_num,
60 + num_tokens,
61 + DEQUANT: tl.constexpr,
62 + PAGE_SIZE: tl.constexpr,
63 + CACHE_FORMAT: tl.constexpr,
64 BLOCK_SIZE: tl.constexpr,
65 + NUM_PRGMS: tl.constexpr,
66 ):
50 - batch_idx = tl.program_id(0)
51 - block_idx = tl.program_id(1)
52 -
53 - batch_query_indexes = tl.load(b_query_lens_loc + batch_idx + tl.arange(0, 2))
54 - batch_query_start, batch_query_end = tl.split(batch_query_indexes)
55 - query_len = batch_query_end - batch_query_start
56 -
57 - if query_len <= 1:
58 - return
59 -
60 - batch_token_indexes = tl.load(b_seq_lens_loc + batch_idx + tl.arange(0, 2))
61 - batch_token_start, batch_token_end = tl.split(batch_token_indexes)
62 - seq_len = batch_token_end - batch_token_start
63 -
64 - if block_idx * BLOCK_SIZE < seq_len:
65 - block_mask = (
66 - block_idx * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)[:, None]
67 - ) < seq_len
68 -
69 - kv_idx = tl.load(
70 - block_table + batch_idx * block_table_stride_0 + block_idx
71 - ).to(tl.int64)
72 -
73 - kv_buffer_off = (
74 - kv_idx * BLOCK_SIZE * E_DIM
75 - + tl.arange(0, BLOCK_SIZE)[:, None] * E_DIM
76 - + tl.arange(0, E_DIM)[None, :]
67 + bid = tl.program_id(0)
68 + col_offsets = tl.arange(0, BLOCK_SIZE)
69 + if DEQUANT:
70 + k_scale = tl.load(k_scale_ptr)
71 + v_scale = tl.load(v_scale_ptr)
72 +
73 + for token_id in tl.range(bid, num_tokens, NUM_PRGMS):
74 + key_ptr_offset = key_ptr + token_id * head_size * num_heads
75 + value_ptr_offset = value_ptr + token_id * head_size * num_heads
76 + batch_idx = tl.load(token_to_batch_ptr + token_id)
77 + batch_start = tl.load(seq_start_ptr + batch_idx)
78 + token_start = tl.load(cu_seqlens_kv_ptr + batch_idx)
79 + batch_offset = token_id - token_start + batch_start
80 + block_offset = batch_offset // PAGE_SIZE
81 + block_id = tl.load(
82 + block_table_ptr + max_block_num * batch_idx + block_offset
83 )
78 - k_vals = tl.load(k_buffer_ptr + kv_buffer_off, mask=block_mask, other=0.0)
79 - if k_vals.dtype.is_fp8():
80 - k_vals = (k_vals.to(tl.float32) * tl.load(k_scale)).to(output_dtype)
81 - else:
82 - k_vals = k_vals.to(output_dtype)
84 + slot_id = batch_offset % PAGE_SIZE
85 +
86 + if CACHE_FORMAT == "NHD":
87 + # for kv cache layout as
88 + # K: [num_blocks, page_size, num_head, head_dim]
89 + # V: [num_blocks, page_size, num_head, head_dim]
90 + key_cache_ptr_offset = (
91 + key_cache_ptr
92 + + block_id * num_heads * head_size * PAGE_SIZE
93 + + slot_id * num_heads * head_size
94 + )
95 + value_cache_ptr_offset = (
96 + value_cache_ptr
97 + + block_id * num_heads * head_size * PAGE_SIZE
98 + + slot_id * num_heads * head_size
99 + )
100
84 - v_vals = tl.load(v_buffer_ptr + kv_buffer_off, mask=block_mask, other=0.0)
85 - if v_vals.dtype.is_fp8():
86 - v_vals = (v_vals.to(tl.float32) * tl.load(v_scale)).to(output_dtype)
87 - else:
88 - v_vals = v_vals.to(output_dtype)
89 - kv_values_off = (
90 - batch_token_start * E_DIM
91 - + block_idx * BLOCK_SIZE * E_DIM
92 - + tl.arange(0, BLOCK_SIZE)[:, None] * E_DIM
93 - + tl.arange(0, E_DIM)[None, :]
94 - )
95 - tl.store(k_values_ptr + kv_values_off, k_vals, mask=block_mask)
96 - tl.store(v_values_ptr + kv_values_off, v_vals, mask=block_mask)
97 -
98 - def vllm_layout_trans(
99 - b_query_lens_loc,
100 - b_seq_lens_loc,
101 - block_table,
102 - k_cache,
103 - v_cache,
104 - max_seq_len,
105 - k_scale,
106 - v_scale,
107 - output_dtype,
108 - total_tokens,
101 + for i in tl.range(0, head_size * num_heads, BLOCK_SIZE):
102 + mask = (col_offsets + i) < head_size * num_heads
103 + k_reg = tl.load(key_cache_ptr_offset + col_offsets + i, mask=mask)
104 + v_reg = tl.load(value_cache_ptr_offset + col_offsets + i, mask=mask)
105 + if DEQUANT:
106 + k_dtype = k_reg.dtype
107 + v_dtype = v_reg.dtype
108 + k_reg = (k_reg.to(tl.float32) * k_scale).to(k_dtype)
109 + v_reg = (v_reg.to(tl.float32) * v_scale).to(v_dtype)
110 + tl.store(key_ptr_offset + col_offsets + i, k_reg, mask=mask)
111 + tl.store(value_ptr_offset + col_offsets + i, v_reg, mask=mask)
112 +
113 + def cp_mha_gather_cache(
114 + key_cache: torch.Tensor,
115 + value_cache: torch.Tensor,
116 + key: torch.Tensor,
117 + value: torch.Tensor,
118 + block_tables: torch.Tensor,
119 + k_scales: torch.Tensor,
120 + v_scales: torch.Tensor,
121 + cu_seqlens_kv: torch.Tensor,
122 + token_to_batch: torch.Tensor,
123 + seq_starts: torch.Tensor,
124 + dequant: bool,
125 + kv_cache_layout: str,
126 + total_tokens: int,
127 ):
110 - H_KV = v_cache.shape[2]
111 - D = v_cache.shape[3]
112 - BLOCK_SIZE = v_cache.shape[1]
113 -
114 - k_values = torch.empty(
115 - (total_tokens, H_KV, D),
116 - dtype=output_dtype,
117 - device=k_cache.device,
128 + assert kv_cache_layout in ["v0", "NHD", "HND"], (
129 + "kv_cache_layout only support v0, NHD, HND"
130 )
119 - v_values = torch.empty(
120 - (total_tokens, H_KV, D),
121 - dtype=output_dtype,
122 - device=v_cache.device,
131 + head_dim = key.shape[2]
132 + x = 0
133 + # assert dequant is True, "Currently, we only support "\
134 + # "gather cache with dequant"
135 + # For k cache layout: [num_blocks, num_heads, page_size, head_dim]
136 + assert kv_cache_layout == "NHD", (
137 + "ROCM_AITER_FA_BACKEND Only support NHD kv cache layout for now"
138 )
124 -
125 - grid = (block_table.shape[0], (max_seq_len + BLOCK_SIZE - 1) // BLOCK_SIZE)
126 -
127 - if output_dtype == torch.float16:
128 - output_dtype = tl.float16
129 - elif output_dtype == torch.bfloat16:
130 - output_dtype = tl.bfloat16
131 - else:
132 - raise ValueError(f"Unsupported output dtype: {output_dtype}")
133 -
134 - _vllm_layout_trans_kernel[grid](
135 - k_cache,
136 - v_cache,
137 - k_values,
138 - v_values,
139 - b_query_lens_loc,
140 - b_seq_lens_loc,
141 - block_table,
142 - block_table.stride(0),
143 - k_scale,
144 - v_scale,
145 - output_dtype=output_dtype,
146 - E_DIM=H_KV * D,
139 + assert head_dim == key_cache.shape[3], (
140 + "We assume your kv cache layout is [num_blocks, "
141 + "page_size, num_heads, head_dim], but got otherwise"
142 + )
143 + page_size = key_cache.shape[1]
144 + num_heads = key_cache.shape[2]
145 +
146 + NUM_PRGMS = num_programs(total_tokens)
147 + BLOCK_SIZE = block_size(key_cache, head_dim)
148 + grid = lambda meta: (NUM_PRGMS,)
149 + cp_mha_gather_cache_kernel[grid](
150 + key_cache,
151 + value_cache,
152 + key,
153 + value,
154 + block_tables,
155 + cu_seqlens_kv,
156 + token_to_batch,
157 + seq_starts,
158 + k_scales,
159 + v_scales,
160 + num_heads,
161 + head_dim,
162 + x,
163 + block_tables.size(1),
164 + total_tokens,
165 + DEQUANT=dequant,
166 + PAGE_SIZE=page_size,
167 + CACHE_FORMAT=kv_cache_layout,
168 BLOCK_SIZE=BLOCK_SIZE,
169 + NUM_PRGMS=NUM_PRGMS,
170 )
171
150 - return k_values, v_values
172
152 - def flash_attn_varlen_func_impl(
153 - q: torch.Tensor,
154 - k_cache: torch.Tensor,
155 - v_cache: torch.Tensor,
156 - out: torch.Tensor,
157 - cu_seqlens_q: torch.Tensor,
158 - cu_seqlens_k: torch.Tensor,
159 - max_seqlen_q: int,
160 - max_seqlen_k: int,
161 - softmax_scale: float,
162 - window_size: list[int] | None, # -1 means infinite context window
163 - alibi_slopes: list[float] | None,
164 - block_table: torch.Tensor,
165 - k_scale: torch.Tensor,
166 - v_scale: torch.Tensor,
167 - total_tokens: int = 0,
168 - ) -> torch.Tensor:
169 - if total_tokens == 0:
170 - total_tokens = int(cu_seqlens_k[-1].item())
171 - k, v = vllm_layout_trans(
172 - cu_seqlens_q,
173 - cu_seqlens_k,
174 - block_table,
175 - k_cache,
176 - v_cache,
177 - max_seqlen_k,
178 - k_scale,
179 - v_scale,
180 - q.dtype,
181 - total_tokens,
182 - )
173 +logger = init_logger(__name__)
174
184 - output = aiter.flash_attn_varlen_func(
185 - q=q,
186 - k=k,
187 - v=v,
188 - cu_seqlens_q=cu_seqlens_q,
189 - max_seqlen_q=max_seqlen_q,
190 - min_seqlen_q=1,
191 - cu_seqlens_k=cu_seqlens_k,
192 - max_seqlen_k=max_seqlen_k,
193 - softmax_scale=softmax_scale,
194 - causal=True,
195 - alibi_slopes=alibi_slopes,
196 - window_size=window_size,
197 - out=out,
198 - )
199 - return output
175
201 - def flash_attn_varlen_func_fake(
202 - q: torch.Tensor,
203 - k_cache: torch.Tensor,
204 - v_cache: torch.Tensor,
205 - out: torch.Tensor,
206 - cu_seqlens_q: torch.Tensor,
207 - cu_seqlens_k: torch.Tensor,
208 - max_seqlen_q: int,
209 - max_seqlen_k: int,
210 - softmax_scale: float,
211 - window_size: list[int] | None, # -1 means infinite context window
212 - alibi_slopes: list[float] | None,
213 - block_table: torch.Tensor,
214 - k_scale: torch.Tensor,
215 - v_scale: torch.Tensor,
216 - total_tokens: int = 0,
217 - ) -> torch.Tensor:
218 - return torch.empty(
219 - q.shape[0], q.shape[1], v_cache.shape[-2], dtype=q.dtype, device=q.device
220 - )
176 +@dataclass
177 +class AiterFlashAttentionDecodeMetadata:
178 + max_query_len: int
179 + min_query_len: int
180 + max_seq_len: int
181 + query_start_loc: torch.Tensor
182
222 - direct_register_custom_op(
223 - "flash_attn_varlen_func",
224 - flash_attn_varlen_func_impl,
225 - ["out"],
226 - flash_attn_varlen_func_fake,
227 - dispatch_key=current_platform.dispatch_key,
228 - )
183
230 -logger = init_logger(__name__)
184 +@dataclass
185 +class AiterFlashAttentionPrefillMetadata:
186 + max_query_len: int
187 + min_query_len: int
188 + max_seq_len: int
189 + query_start_loc: torch.Tensor
190 +
191 +
192 +@dataclass
193 +class AiterChunkContextMetadata:
194 + workspace: torch.Tensor
195 + cu_seq_lens_chunk: torch.Tensor
196 + chunk_starts: torch.Tensor
197 + token_to_batch: torch.Tensor
198 + seq_tot: list[int]
199 + max_seq_lens: list[int]
200 + seq_lens: torch.Tensor
201 + num_chunks: int
202 + total_token_per_batch: list[int]
203 +
204 +
205 +@dataclass
206 +class AiterFlashAttentionChunkPrefillMetadata:
207 + max_query_len: int
208 + min_query_len: int
209 + max_seq_len: int
210 + query_start_loc: torch.Tensor
211 + chunk_context_metadata: AiterChunkContextMetadata
212
213
214 @dataclass
@@ -248,7 +229,18 @@ class AiterFlashAttentionMetadata:
229 seq_lens: torch.Tensor
230 slot_mapping: torch.Tensor
231 block_table: torch.Tensor
251 - cu_seq_lens: torch.Tensor | None
232 +
233 + # prefill and deocde split
234 + num_decodes: int
235 + num_decode_tokens: int
236 + num_prefills: int
237 + num_prefill_tokens: int
238 + num_extends: int
239 + num_extend_tokens: int
240 +
241 + decode_metadata: AiterFlashAttentionDecodeMetadata | None
242 + prefill_metadata: AiterFlashAttentionPrefillMetadata | None
243 + extend_metadata: AiterFlashAttentionChunkPrefillMetadata | None
244
245 # For cascade attention.
246 use_cascade: bool
@@ -260,6 +252,7 @@ class AiterFlashAttentionMetadataBuilder(
252 AttentionMetadataBuilder[AiterFlashAttentionMetadata]
253 ):
254 cudagraph_support = AttentionCGSupport.UNIFORM_SINGLE_TOKEN_DECODE
255 + reorder_batch_threshold: int = 1
256
257 def __init__(
258 self,
@@ -285,6 +278,12 @@ class AiterFlashAttentionMetadataBuilder(
278 self.aot_sliding_window: tuple[int, int] | None = None
279 self.total_tokens: int = 0
280
281 + self.extend_workspace = torch.empty(
282 + [2, _CP_TOKENS_PER_ITER_ROCM, self.num_heads_kv, self.headdim],
283 + dtype=self.model_config.dtype,
284 + device=device,
285 + )
286 +
287 def build_for_cudagraph_capture(
288 self, common_attn_metadata: CommonAttentionMetadata
289 ):
@@ -302,42 +301,139 @@ class AiterFlashAttentionMetadataBuilder(
301 common_attn_metadata: CommonAttentionMetadata,
302 fast_build: bool = False,
303 ) -> "AiterFlashAttentionMetadata":
305 - num_actual_tokens = common_attn_metadata.num_actual_tokens
306 - max_query_len = common_attn_metadata.max_query_len
307 - max_seq_len = common_attn_metadata.max_seq_len
308 - query_start_loc = common_attn_metadata.query_start_loc
309 - seq_lens = common_attn_metadata.seq_lens
310 - block_table_tensor = common_attn_metadata.block_table_tensor
311 - slot_mapping = common_attn_metadata.slot_mapping
312 - if max_query_len > 1:
313 - # We pre-compute cumulative seq len needed for prefill attention
314 - # here to avoid recomputing it for every layer
304 + split_ret = split_decodes_prefills_and_extends(
305 + common_attn_metadata,
306 + decode_threshold=self.reorder_batch_threshold,
307 + )
308 +
309 + (
310 + num_decodes,
311 + num_extends,
312 + num_prefills,
313 + num_decode_tokens,
314 + num_extend_tokens,
315 + num_prefill_tokens,
316 + ) = split_ret
317 +
318 + query_start_loc_cpu = common_attn_metadata.query_start_loc_cpu
319 +
320 + seq_lens = common_attn_metadata.seq_lens_cpu
321 +
322 + query_lens_cpu = query_start_loc_cpu[1:] - query_start_loc_cpu[:-1]
323 +
324 + decode_metadata = None
325 + if num_decodes > 0:
326 + decode_metadata = AiterFlashAttentionDecodeMetadata(
327 + max_query_len=query_lens_cpu[:num_decodes].max().item(),
328 + min_query_len=query_lens_cpu[:num_decodes].min().item(),
329 + max_seq_len=seq_lens[:num_decodes].max().item(),
330 + query_start_loc=common_attn_metadata.query_start_loc[: num_decodes + 1],
331 + )
332 +
333 + prefill_metadata = None
334 + if num_prefills > 0:
335 + query_lens_for_prefill = query_lens_cpu[num_decodes + num_extends :]
336 + query_start_loc_device = common_attn_metadata.query_start_loc[
337 + num_decodes + num_extends :
338 + ]
339 + prefill_metadata = AiterFlashAttentionPrefillMetadata(
340 + max_query_len=query_lens_for_prefill.max().item(),
341 + min_query_len=query_lens_for_prefill.min().item(),
342 + max_seq_len=seq_lens[num_decodes + num_extends :].max().item(),
343 + query_start_loc=query_start_loc_device - query_start_loc_device[0],
344 + )
345 +
346 + extend_metadata = None
347 + if num_extends > 0:
348 + num_extends_slice = slice(num_decodes, num_decodes + num_extends)
349 + query_lens_for_extend = query_lens_cpu[num_extends_slice]
350 + seq_lens_for_extend = common_attn_metadata.seq_lens_cpu[num_extends_slice]
351 + computed_kv_lens = seq_lens_for_extend - query_lens_for_extend
352 +
353 + # allocate the equal amount of workspace for
354 + # each chunk prefill request
355 + max_context_chunk = _CP_TOKENS_PER_ITER_ROCM // num_extends
356 + num_chunks = cdiv(computed_kv_lens.max().item(), max_context_chunk)
357 +
358 + chunk_starts = (
359 + torch.arange(num_chunks, dtype=torch.int32)
360 + .unsqueeze(1)
361 + .expand(-1, num_extends)
362 + * max_context_chunk
363 + )
364 + chunk_ends = torch.min(
365 + computed_kv_lens.unsqueeze(0), chunk_starts + max_context_chunk
366 + )
367 + chunk_seq_lens = (chunk_ends - chunk_starts).clamp(
368 + min=0
369 + ) # [num_chunks, num_extends]
370 + cu_seq_lens_cpu = torch.zeros(
371 + [num_chunks, num_extends + 1], dtype=torch.int32, pin_memory=True
372 + )
373 + torch.cumsum(
374 + chunk_seq_lens, dim=1, out=cu_seq_lens_cpu[:, 1:], dtype=torch.int32
375 + )
376 + max_cum_tokens = cu_seq_lens_cpu[:, -1].max().item()
377 +
378 + range_idx = torch.arange(max_cum_tokens, dtype=torch.int32)[None, None, :]
379 + idx_to_batch_tensor = range_idx == cu_seq_lens_cpu[:, 1:][:, :, None]
380 + idx_to_batch_tensor = idx_to_batch_tensor.sum(
381 + dim=1
382 + ) # [num_chunks, max_cum_tokens]
383 + token_to_batch_tensor = torch.cumsum(idx_to_batch_tensor, dim=1)
384 +
385 + chunk_context_metadata = AiterChunkContextMetadata(
386 + workspace=self.extend_workspace,
387 + cu_seq_lens_chunk=cu_seq_lens_cpu.to(self.device, non_blocking=True),
388 + chunk_starts=chunk_starts.to(self.device, non_blocking=True),
389 + seq_tot=chunk_seq_lens.sum(dim=1).tolist(),
390 + max_seq_lens=chunk_seq_lens.max(dim=1).values.tolist(),
391 + seq_lens=chunk_seq_lens,
392 + token_to_batch=token_to_batch_tensor.to(self.device, non_blocking=True),
393 + num_chunks=num_chunks,
394 + total_token_per_batch=cu_seq_lens_cpu[:, -1].tolist(),
395 + )
396 +
397 + query_start_loc_device = common_attn_metadata.query_start_loc[
398 + num_decodes : num_decodes + num_extends + 1
399 + ]
400 + seq_lens_device = common_attn_metadata.seq_lens[num_extends_slice]
401 cu_seq_lens = torch.zeros(
316 - seq_lens.shape[0] + 1, dtype=torch.int32, device=seq_lens.device
402 + num_extends + 1, dtype=torch.int32, device=seq_lens_device.device
403 + )
404 + torch.cumsum(
405 + seq_lens_device, dim=0, dtype=cu_seq_lens.dtype, out=cu_seq_lens[1:]
406 + )
407 + extend_metadata = AiterFlashAttentionChunkPrefillMetadata(
408 + max_query_len=query_lens_for_extend.max().item(),
409 + min_query_len=query_lens_for_extend.min().item(),
410 + max_seq_len=seq_lens[num_extends_slice].max().item(),
411 + query_start_loc=query_start_loc_device - query_start_loc_device[0],
412 + chunk_context_metadata=chunk_context_metadata,
413 )
318 - torch.cumsum(seq_lens, dim=0, dtype=cu_seq_lens.dtype, out=cu_seq_lens[1:])
319 - num_actual_kv_tokens = int(cu_seq_lens[-1].item())
320 - else:
321 - cu_seq_lens = None
322 - num_actual_kv_tokens = 0
414
324 - def schedule(
325 - batch_size, cu_query_lens, max_query_len, seqlens, max_seq_len, causal
326 - ):
327 - return None
415 + num_actual_kv_tokens = torch.sum(seq_lens).item()
416
417 use_cascade = common_prefix_len > 0
418
419 attn_metadata = AiterFlashAttentionMetadata(
332 - num_actual_tokens=num_actual_tokens,
420 + num_actual_tokens=common_attn_metadata.num_actual_tokens,
421 num_actual_kv_tokens=num_actual_kv_tokens,
334 - max_query_len=max_query_len,
335 - query_start_loc=query_start_loc,
336 - max_seq_len=max_seq_len,
337 - seq_lens=seq_lens,
338 - block_table=block_table_tensor,
339 - slot_mapping=slot_mapping,
340 - cu_seq_lens=cu_seq_lens,
422 + max_query_len=common_attn_metadata.max_query_len,
423 + query_start_loc=common_attn_metadata.query_start_loc,
424 + max_seq_len=common_attn_metadata.max_seq_len,
425 + seq_lens=common_attn_metadata.seq_lens,
426 + block_table=common_attn_metadata.block_table_tensor,
427 + slot_mapping=common_attn_metadata.slot_mapping,
428 + num_decodes=num_decodes,
429 + num_decode_tokens=num_decode_tokens,
430 + num_prefills=num_prefills,
431 + num_prefill_tokens=num_prefill_tokens,
432 + num_extends=num_extends,
433 + num_extend_tokens=num_extend_tokens,
434 + decode_metadata=decode_metadata,
435 + prefill_metadata=prefill_metadata,
436 + extend_metadata=extend_metadata,
437 use_cascade=use_cascade,
438 common_prefix_len=common_prefix_len,
439 total_tokens=self.total_tokens,
@@ -401,6 +497,7 @@ class AiterFlashAttentionBackend(AttentionBackend):
497 ) -> tuple[int, ...]:
498 if block_size % 16 != 0:
499 raise ValueError("Block size must be a multiple of 16.")
500 +
501 return (2, num_blocks, block_size, num_kv_heads, head_size)
502
503
@@ -449,6 +546,110 @@ class AiterFlashAttentionImpl(AttentionImpl):
546 "FlashAttentionImpl"
547 )
548
549 + def extend_forward(
550 + self,
551 + attn_metadata: AiterFlashAttentionMetadata,
552 + query: torch.Tensor,
553 + key: torch.Tensor,
554 + value: torch.Tensor,
555 + key_cache: torch.Tensor,
556 + value_cache: torch.Tensor,
557 + output: torch.Tensor,
558 + cu_seqlens_q: torch.Tensor,
559 + max_seqlen_q: int,
560 + max_seqlen_k: int,
561 + min_seqlen_q: int,
562 + block_table: torch.Tensor,
563 + slot_mapping: torch.Tensor,
564 + k_scale: float,
565 + v_scale: float,
566 + ):
567 + out, lse = aiter.flash_attn_varlen_func(
568 + q=query,
569 + k=key,
570 + v=value,
571 + cu_seqlens_q=cu_seqlens_q,
572 + cu_seqlens_k=cu_seqlens_q,
573 + max_seqlen_q=max_seqlen_q,
574 + max_seqlen_k=max_seqlen_q,
575 + min_seqlen_q=min_seqlen_q,
576 + dropout_p=0.0,
577 + softmax_scale=self.scale,
578 + causal=True,
579 + window_size=self.sliding_window,
580 + alibi_slopes=self.alibi_slopes,
581 + return_lse=True,
582 + )
583 + assert attn_metadata.extend_metadata is not None
584 + chunk_context_metadata = attn_metadata.extend_metadata.chunk_context_metadata
585 + num_chunks = chunk_context_metadata.num_chunks
586 + workspace = chunk_context_metadata.workspace
587 + cu_seqlens_kv = chunk_context_metadata.cu_seq_lens_chunk
588 + max_seqlens = chunk_context_metadata.max_seq_lens
589 + chunk_starts = chunk_context_metadata.chunk_starts
590 + token_to_batch = chunk_context_metadata.token_to_batch
591 + total_token_per_batch = chunk_context_metadata.total_token_per_batch
592 + key_fetched, value_fetched = workspace[0], workspace[1]
593 + chunked_output = None
594 + chunked_lse = None
595 + for chunk_idx in range(num_chunks):
596 + cp_mha_gather_cache(
597 + key_cache=key_cache,
598 + value_cache=value_cache,
599 + key=key_fetched,
600 + value=value_fetched,
601 + block_tables=block_table,
602 + k_scales=k_scale,
603 + v_scales=v_scale,
604 + cu_seqlens_kv=cu_seqlens_kv[chunk_idx],
605 + token_to_batch=token_to_batch[chunk_idx],
606 + seq_starts=chunk_starts[chunk_idx],
607 + dequant=False,
608 + kv_cache_layout="NHD",
609 + total_tokens=total_token_per_batch[chunk_idx],
610 + )
611 +
612 + suf_out, suf_lse = aiter.flash_attn_varlen_func(
613 + q=query,
614 + k=key_fetched,
615 + v=value_fetched,
616 + cu_seqlens_q=cu_seqlens_q,
617 + cu_seqlens_k=cu_seqlens_kv[chunk_idx],
618 + max_seqlen_q=max_seqlen_q,
619 + max_seqlen_k=max_seqlens[chunk_idx],
620 + min_seqlen_q=min_seqlen_q,
621 + dropout_p=0.0,
622 + softmax_scale=self.scale,
623 + causal=False,
624 + window_size=self.sliding_window,
625 + alibi_slopes=self.alibi_slopes,
626 + return_lse=True,
627 + )
628 + if chunked_output is None:
629 + chunked_output = suf_out
630 + chunked_lse = suf_lse
631 + else:
632 + tmp_output = torch.empty_like(out)
633 + tmp_lse = torch.empty_like(lse)
634 + merge_attn_states(
635 + output=tmp_output,
636 + output_lse=tmp_lse,
637 + prefix_output=chunked_output,
638 + prefix_lse=chunked_lse,
639 + suffix_output=suf_out,
640 + suffix_lse=suf_lse,
641 + )
642 + chunked_output = tmp_output
643 + chunked_lse = tmp_lse
644 +
645 + merge_attn_states(
646 + output=output,
647 + prefix_output=chunked_output,
648 + prefix_lse=chunked_lse,
649 + suffix_output=out,
650 + suffix_lse=lse,
651 + )
652 +
653 def forward(
654 self,
655 layer: torch.nn.Module,
@@ -488,24 +689,25 @@ class AiterFlashAttentionImpl(AttentionImpl):
689 return output.fill_(0)
690
691 # IMPORTANT!
491 - # NOTE(woosuk): With piece-wise CUDA graphs, this method is executed in
492 - # eager-mode PyTorch. Thus, we need to be careful about any CPU overhead
493 - # in this method. For example, `view` and `slice` (or `[:n]`) operations
494 - # are surprisingly slow even in the case they do not invoke any GPU ops.
692 + # NOTE(woosuk): With piece-wise CUDA graphs, this method is
693 + # executed in eager-mode PyTorch. Thus, we need to be careful
694 + # about any CPU overhead in this method. For example, `view`
695 + # and `slice` (or `[:n]`) operations are surprisingly slow even
696 + # in the case they do not invoke any GPU ops.
697 # Minimize the PyTorch ops in this method as much as possible.
698 # Whenever making a change in this method, please benchmark the
699 # performance to make sure it does not introduce any overhead.
498 -
700 num_actual_tokens = attn_metadata.num_actual_tokens
701 key_cache, value_cache = kv_cache.unbind(0)
702 if self.kv_sharing_target_layer_name is None:
703 # Reshape the input keys and values and store them in the cache.
704 # Skip this if sharing KV cache with an earlier attention layer.
504 - # NOTE(woosuk): Here, key and value are padded while slot_mapping is
505 - # not padded. However, we don't need to do key[:num_actual_tokens]
506 - # and value[:num_actual_tokens] because the reshape_and_cache_flash
507 - # op uses the slot_mapping's shape to determine the number of
508 - # actual tokens.
705 + # NOTE(woosuk): Here, key and value are padded while slot_mapping
706 + # is not padded. However, we don't need to do
707 + # key[:num_actual_tokens] and value[:num_actual_tokens] because
708 + # the reshape_and_cache_flash op uses the slot_mapping's shape
709 + # to determine the number of actual tokens.
710 +
711 torch.ops._C_cache_ops.reshape_and_cache_flash(
712 key,
713 value,
@@ -521,69 +723,118 @@ class AiterFlashAttentionImpl(AttentionImpl):
723 key_cache = key_cache.view(current_platform.fp8_dtype())
724 value_cache = value_cache.view(current_platform.fp8_dtype())
725
726 + # decode:extend:prefill
727 + query = query[:num_actual_tokens]
728 + key = key[:num_actual_tokens]
729 + value = value[:num_actual_tokens]
730 +
731 + output_actual_tokens = output[:num_actual_tokens]
732 +
733 + num_decodes = attn_metadata.num_decodes
734 + num_prefills = attn_metadata.num_prefills
735 + num_extends = attn_metadata.num_extends
736 +
737 + num_decode_tokens = attn_metadata.num_decode_tokens
738 + num_extend_tokens = attn_metadata.num_extend_tokens
739 if not attn_metadata.use_cascade:
525 - cu_seqlens_q = attn_metadata.query_start_loc
526 - seqused_k = attn_metadata.seq_lens
527 - max_seqlen_q = attn_metadata.max_query_len
528 - max_seqlen_k = attn_metadata.max_seq_len
529 - block_table = attn_metadata.block_table
530 -
531 - if max_seqlen_q > 1:
532 - torch.ops.vllm.flash_attn_varlen_func(
533 - query[:num_actual_tokens],
534 - key_cache,
535 - value_cache,
536 - out=output[:num_actual_tokens],
537 - cu_seqlens_q=cu_seqlens_q,
538 - max_seqlen_q=max_seqlen_q,
539 - max_seqlen_k=max_seqlen_k,
740 + # calculate for pure prefills
741 + if num_prefills > 0:
742 + assert attn_metadata.prefill_metadata is not None
743 +
744 + prefill_query = query[num_decode_tokens + num_extend_tokens :]
745 + prefill_key = key[num_decode_tokens + num_extend_tokens :]
746 + prefill_value = value[num_decode_tokens + num_extend_tokens :]
747 +
748 + aiter.flash_attn_varlen_func(
749 + q=prefill_query,
750 + k=prefill_key,
751 + v=prefill_value,
752 + cu_seqlens_q=attn_metadata.prefill_metadata.query_start_loc,
753 + cu_seqlens_k=attn_metadata.prefill_metadata.query_start_loc,
754 + max_seqlen_q=attn_metadata.prefill_metadata.max_query_len,
755 + max_seqlen_k=attn_metadata.prefill_metadata.max_seq_len,
756 + min_seqlen_q=attn_metadata.prefill_metadata.min_query_len,
757 + dropout_p=0.0,
758 softmax_scale=self.scale,
541 - alibi_slopes=self.alibi_slopes,
759 + causal=True,
760 window_size=self.sliding_window,
543 - block_table=block_table,
544 - cu_seqlens_k=attn_metadata.cu_seq_lens,
761 + alibi_slopes=self.alibi_slopes,
762 + out=output_actual_tokens[num_decode_tokens + num_extend_tokens :],
763 + )
764 +
765 + # calculate for extends
766 + if num_extends > 0:
767 + assert attn_metadata.extend_metadata is not None
768 + extend_tokens_slice = slice(
769 + num_decode_tokens, num_decode_tokens + num_extend_tokens
770 + )
771 + extend_querys = query[extend_tokens_slice]
772 + extend_keys = key[extend_tokens_slice]
773 + extend_values = value[extend_tokens_slice]
774 + extend_outputs = output[extend_tokens_slice]
775 + self.extend_forward(
776 + attn_metadata=attn_metadata,
777 + query=extend_querys,
778 + key=extend_keys,
779 + value=extend_values,
780 + key_cache=key_cache,
781 + value_cache=value_cache,
782 + output=extend_outputs,
783 + cu_seqlens_q=attn_metadata.extend_metadata.query_start_loc,
784 + max_seqlen_q=attn_metadata.extend_metadata.max_query_len,
785 + max_seqlen_k=attn_metadata.extend_metadata.max_seq_len,
786 + min_seqlen_q=attn_metadata.extend_metadata.min_query_len,
787 + block_table=attn_metadata.block_table[
788 + num_decodes : num_decodes + num_extends
789 + ],
790 + slot_mapping=attn_metadata.slot_mapping[
791 + num_decodes : num_decodes + num_extends
792 + ],
793 k_scale=layer._k_scale,
794 v_scale=layer._v_scale,
547 - total_tokens=attn_metadata.num_actual_kv_tokens,
795 )
796
550 - _, num_heads, head_size = query.shape
551 - nbytes_per_qo_elem = torch.finfo(query.dtype).bits // 8
552 - num_seqs = seqused_k.shape[0]
553 - max_num_partitions = (
554 - max_seqlen_k + _PARTITION_SIZE_ROCM - 1
555 - ) // _PARTITION_SIZE_ROCM
556 -
557 - workspace_buffer = torch.empty(
558 - (num_seqs * num_heads * max_num_partitions * head_size)
559 - * nbytes_per_qo_elem
560 - + 2 * (num_seqs * num_heads * max_num_partitions) * 4,
561 - dtype=torch.uint8,
562 - device=output.device,
563 - )
797 + # calculate for decodes
798 + if num_decodes > 0:
799 + assert attn_metadata.decode_metadata is not None
800 + _, num_heads, head_size = query.shape
801 + nbytes_per_qo_elem = torch.finfo(query.dtype).bits // 8
802 + num_seqs = attn_metadata.seq_lens.shape[0]
803 + max_num_partitions = (
804 + attn_metadata.max_seq_len + _PARTITION_SIZE_ROCM - 1
805 + ) // _PARTITION_SIZE_ROCM
806 +
807 + workspace_buffer = torch.empty(
808 + (num_seqs * num_heads * max_num_partitions * head_size)
809 + * nbytes_per_qo_elem
810 + + 2 * (num_seqs * num_heads * max_num_partitions) * 4,
811 + dtype=torch.uint8,
812 + device=output.device,
813 + )
814
565 - torch.ops.aiter.paged_attention_v1(
566 - output[:num_actual_tokens],
567 - workspace_buffer,
568 - query[:num_actual_tokens],
569 - key_cache,
570 - value_cache,
571 - self.scale,
572 - block_table,
573 - cu_seqlens_q,
574 - seqused_k,
575 - max_seqlen_k,
576 - self.alibi_slopes,
577 - self.kv_cache_dtype,
578 - "NHD",
579 - self.logits_soft_cap,
580 - layer._k_scale,
581 - layer._v_scale,
582 - None,
583 - _PARTITION_SIZE_ROCM,
584 - )
585 - return output
815 + torch.ops.aiter.paged_attention_v1(
816 + output[:num_decode_tokens],
817 + workspace_buffer,
818 + query[:num_decode_tokens],
819 + key_cache,
820 + value_cache,
821 + self.scale,
822 + attn_metadata.block_table[:num_decodes],
823 + attn_metadata.query_start_loc[:num_decodes],
824 + attn_metadata.seq_lens[:num_decodes],
825 + attn_metadata.max_seq_len,
826 + self.alibi_slopes,
827 + self.kv_cache_dtype,
828 + "NHD",
829 + self.logits_soft_cap,
830 + layer._k_scale,
831 + layer._v_scale,
832 + None,
833 + _PARTITION_SIZE_ROCM,
834 + )
835 else:
836 raise NotImplementedError(
837 "Cascade attention is not implemented for ROCM AITER"
838 )
839 +
840 + return output
vllm/v1/attention/backends/utils.py
+67
@@ -728,6 +728,73 @@ def subclass_attention_backend(
728 )
729
730
731 +def split_decodes_prefills_and_extends(
732 + common_attn_metadata: CommonAttentionMetadata,
733 + decode_threshold: int = 1,
734 +) -> tuple[int, int, int, int, int, int]:
735 + """
736 + Assuming a reordered batch, finds the boundary between prefill and decode
737 + requests.
738 +
739 + Args:
740 + common_attn_metadata: CommonAttentionMetadata object containing the
741 + batch metadata.
742 + decode_threshold: The maximum query length to be considered a decode.
743 +
744 + Returns:
745 + num_decodes: The number of decode requests.
746 + num_extends: The number of extend requests.
747 + num_prefills: The number of prefill requests.
748 + num_decode_tokens: The number of tokens in the decode requests.
749 + num_extend_tokens: The number of tokens in the extend requests.
750 + num_prefill_tokens: The number of tokens in the prefill requests.
751 + """
752 + max_query_len = common_attn_metadata.max_query_len
753 + num_reqs = common_attn_metadata.num_reqs
754 + num_tokens = common_attn_metadata.num_actual_tokens
755 + query_start_loc = common_attn_metadata.query_start_loc_cpu
756 + seq_lens = common_attn_metadata.seq_lens_cpu
757 +
758 + if max_query_len <= decode_threshold:
759 + return num_reqs, 0, 0, num_tokens, 0, 0
760 +
761 + query_lens = query_start_loc[1:] - query_start_loc[:-1]
762 + is_prefill_or_extend = query_lens > decode_threshold
763 + is_prefill = (seq_lens == query_lens) & is_prefill_or_extend
764 + first_extend = is_prefill_or_extend.int().argmax(dim=-1).item()
765 + first_prefill = is_prefill.int().argmax(dim=-1).item()
766 + num_decodes = first_extend
767 + num_decode_tokens = query_start_loc[first_extend].item()
768 + if not torch.any(is_prefill_or_extend):
769 + return (num_decodes, 0, 0, num_decode_tokens, 0, 0)
770 +
771 + num_prefills_or_extends = num_reqs - num_decodes
772 + num_prefill_or_extend_tokens = num_tokens - num_decode_tokens
773 + if not torch.any(is_prefill):
774 + return (
775 + num_decodes,
776 + num_prefills_or_extends,
777 + 0,
778 + num_decode_tokens,
779 + num_prefill_or_extend_tokens,
780 + 0,
781 + )
782 +
783 + num_extends = first_prefill - num_decodes
784 + num_prefills = num_reqs - first_prefill
785 +
786 + num_prefill_tokens = num_tokens - query_start_loc[first_prefill]
787 + num_extend_tokens = num_prefill_or_extend_tokens - num_prefill_tokens
788 + return (
789 + num_decodes,
790 + num_extends,
791 + num_prefills,
792 + num_decode_tokens,
793 + num_extend_tokens,
794 + num_prefill_tokens,
795 + )
796 +
797 +
798 def split_decodes_and_prefills(
799 common_attn_metadata: CommonAttentionMetadata,
800 decode_threshold: int = 1,