@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
py 507 lines 17.5 KB
Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4 # Authors:
5 # - Burkhard Ringlein <ngl@zurich.ibm.com>
6 # - Jan van Lunteren <jvl@zurich.ibm.com>
7 # - Chih-Chieh Yang <chih.chieh.yang@ibm.com>
8 # - Thomas Parnell <tpa@zurich.ibm.com>
9
10 import torch
11
12 from vllm import _custom_ops as ops
13 from vllm.logger import init_logger
14 from vllm.platforms import current_platform
15 from vllm.triton_utils import tl, triton
16
17 from .prefix_prefill import context_attention_fwd
18
19 logger = init_logger(__name__)
20
21 float8_info = torch.finfo(current_platform.fp8_dtype())
22
23
24 def has_native_kv_cache_layout(
25 key_cache: torch.Tensor,
26 value_cache: torch.Tensor,
27 ) -> bool:
28 """Return whether KV cache blocks can use the native ROCm pairing.
29
30 The native reshape_and_cache writer assumes packed blocks. If cache update
31 needs reshape_and_cache_flash for a stride-padded hybrid layout, decode
32 should use the matching Triton path too.
33 """
34 return (
35 key_cache.stride(0) == key_cache.shape[1:].numel()
36 and value_cache.stride(0) == value_cache.shape[1:].numel()
37 )
38
39
40 @triton.jit
41 def cdiv_fn(x, y):
42 return (x + y - 1) // y
43
44
45 @triton.jit
46 def kernel_paged_attention_2d(
47 output_ptr, # [num_tokens, num_query_heads, head_size]
48 query_ptr, # [num_tokens, num_query_heads, head_size]
49 key_cache_ptr, # [num_blks, num_kv_heads, head_size // x, blk_size, x]
50 value_cache_ptr, # [num_blks, num_kv_heads, head_size, blk_size]
51 sink_ptr, # [num_query_heads]
52 block_tables_ptr, # [num_seqs, max_num_blocks_per_seq]
53 seq_lens_ptr, # [num_seqs]
54 alibi_slopes_ptr, # [num_query_heads]
55 scale, # float32
56 k_scale, # float32
57 v_scale, # float32
58 out_scale_inv,
59 num_query_heads: tl.constexpr, # int
60 num_queries_per_kv: tl.constexpr, # int
61 num_queries_per_kv_padded: tl.constexpr, # int
62 block_table_stride: tl.int64, # int
63 query_stride_0: tl.int64, # int
64 query_stride_1: tl.int64, # int, should be equal to head_size
65 output_stride_0: tl.int64, # int
66 output_stride_1: tl.int64, # int, should be equal to head_size
67 BLOCK_SIZE: tl.constexpr, # int
68 PHYSICAL_BLOCK_SIZE: tl.constexpr, # int
69 HEAD_SIZE: tl.constexpr, # int
70 HEAD_SIZE_PADDED: tl.constexpr, # int, must be power of 2
71 USE_ALIBI_SLOPES: tl.constexpr, # bool
72 SLIDING_WINDOW: tl.constexpr, # int
73 x: tl.constexpr, # int
74 stride_k_cache_0: tl.int64, # int
75 stride_k_cache_1: tl.int64, # int
76 stride_k_cache_2: tl.int64, # int
77 stride_k_cache_3: tl.int64, # int
78 stride_k_cache_4: tl.int64, # int
79 stride_v_cache_0: tl.int64, # int
80 stride_v_cache_1: tl.int64, # int
81 stride_v_cache_2: tl.int64, # int
82 stride_v_cache_3: tl.int64, # int
83 filter_by_query_len: tl.constexpr, # bool
84 query_start_len_ptr, # [num_seqs+1]
85 USE_SINKS: tl.constexpr, # bool
86 USE_FP8: tl.constexpr,
87 FP8_MIN: tl.constexpr = float8_info.min,
88 FP8_MAX: tl.constexpr = float8_info.max,
89 ):
90 seq_idx = tl.program_id(0)
91 kv_head_idx = tl.program_id(1)
92
93 if filter_by_query_len:
94 cur_batch_in_all_start_index = tl.load(query_start_len_ptr + seq_idx)
95 cur_batch_in_all_stop_index = tl.load(query_start_len_ptr + seq_idx + 1)
96 cur_batch_query_len = cur_batch_in_all_stop_index - cur_batch_in_all_start_index
97 if cur_batch_query_len > 1:
98 return
99 else:
100 cur_batch_in_all_start_index = seq_idx
101
102 query_head_idx = kv_head_idx * num_queries_per_kv + tl.arange(
103 0, num_queries_per_kv_padded
104 )
105
106 query_offset = (
107 cur_batch_in_all_start_index * query_stride_0
108 + query_head_idx[:, None] * query_stride_1
109 )
110
111 head_mask = query_head_idx < (kv_head_idx + 1) * num_queries_per_kv
112 head_mask = head_mask & (query_head_idx < num_query_heads)
113
114 dim_mask = tl.where(tl.arange(0, HEAD_SIZE_PADDED) < HEAD_SIZE, 1, 0).to(tl.int1)
115
116 # Q : (num_queries_per_kv, HEAD_SIZE,)
117 Q = tl.load(
118 query_ptr + query_offset + tl.arange(0, HEAD_SIZE_PADDED)[None, :],
119 mask=dim_mask[None, :] & head_mask[:, None],
120 other=0.0,
121 )
122
123 block_table_offset = seq_idx * block_table_stride
124
125 if not USE_SINKS:
126 M = tl.full([num_queries_per_kv_padded], float("-inf"), dtype=tl.float32)
127 L = tl.zeros([num_queries_per_kv_padded], dtype=tl.float32)
128 else:
129 M = tl.load(
130 sink_ptr + query_head_idx,
131 mask=head_mask,
132 other=float("-inf"),
133 ).to(dtype=tl.float32)
134 L = tl.where(float("-inf") < M, 1.0, 0.0)
135
136 acc = tl.zeros([num_queries_per_kv_padded, HEAD_SIZE_PADDED], dtype=tl.float32)
137
138 # sequence len for this particular sequence
139 seq_len = tl.load(seq_lens_ptr + seq_idx)
140
141 # alibi slope for this head
142 if USE_ALIBI_SLOPES:
143 alibi_slope = tl.load(
144 alibi_slopes_ptr + query_head_idx, mask=head_mask, other=0.0
145 )
146
147 num_blocks = cdiv_fn(seq_len, BLOCK_SIZE)
148
149 offs_n = tl.arange(0, BLOCK_SIZE)
150 offs_d = tl.arange(0, HEAD_SIZE_PADDED)
151 # iterate through tiles
152 for j in range(0, num_blocks):
153 start_n = j * BLOCK_SIZE
154 # Calculate the logical location within a non-standard physical block,
155 # such as 544 in Qwen/Qwen3-Next-80B-A3B-Thinking.
156 # Supports non-contiguous mapping
157 # from logical blocks to physical blocks
158 abs_token_idx = start_n + offs_n
159 l_block_idx = abs_token_idx // PHYSICAL_BLOCK_SIZE
160 # Vectorized loading of physical block IDs
161 p_block_idx = tl.load(block_tables_ptr + block_table_offset + l_block_idx)
162 internal_offsets = abs_token_idx % PHYSICAL_BLOCK_SIZE
163
164 # 5D addressing logic of K
165 k_offset = (
166 p_block_idx[None, :] * stride_k_cache_0
167 + kv_head_idx * stride_k_cache_1
168 + (offs_d[:, None] // x) * stride_k_cache_2
169 + internal_offsets[None, :] * stride_k_cache_3
170 + (offs_d[:, None] % x) * stride_k_cache_4
171 )
172
173 # 4D addressing logic of V (Slot is innermost)
174 v_offset = (
175 p_block_idx[:, None] * stride_v_cache_0
176 + kv_head_idx * stride_v_cache_1
177 + offs_d[None, :] * stride_v_cache_2
178 + internal_offsets[:, None] * stride_v_cache_3
179 )
180
181 # Only the final tile can straddle seq_len. Slots >= seq_len are
182 # unwritten KV cache that may hold NaN/garbage; they are score-masked
183 # below, but 0 * NaN = NaN would still poison the output, so mask them
184 # out of the K/V loads too. Earlier tiles are fully written, so they
185 # use the cheaper token-uniform dim_mask (matching the pre-0.25.0 fast
186 # path) and skip the per-token predicate entirely.
187 # K : (HEAD_SIZE, BLOCK_SIZE), V : (BLOCK_SIZE, HEAD_SIZE)
188 if j == num_blocks - 1:
189 kv_load_mask = abs_token_idx < seq_len
190 K_load = tl.load(
191 key_cache_ptr + k_offset,
192 mask=dim_mask[:, None] & kv_load_mask[None, :],
193 other=0.0,
194 eviction_policy="evict_last",
195 )
196 V_load = tl.load(
197 value_cache_ptr + v_offset,
198 mask=dim_mask[None, :] & kv_load_mask[:, None],
199 other=0.0,
200 eviction_policy="evict_last",
201 )
202 else:
203 K_load = tl.load(
204 key_cache_ptr + k_offset,
205 mask=dim_mask[:, None],
206 other=0.0,
207 eviction_policy="evict_last",
208 )
209 V_load = tl.load(
210 value_cache_ptr + v_offset,
211 mask=dim_mask[None, :],
212 other=0.0,
213 eviction_policy="evict_last",
214 )
215
216 if K_load.dtype.is_fp8():
217 K = (K_load.to(tl.float32) * tl.load(k_scale)).to(Q.dtype)
218 else:
219 K = K_load
220
221 if V_load.dtype.is_fp8():
222 V = (V_load.to(tl.float32) * tl.load(v_scale)).to(Q.dtype)
223 else:
224 V = V_load
225
226 seq_offset = j * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
227 boundary = tl.full([BLOCK_SIZE], seq_len, dtype=tl.int32)
228 seq_mask = seq_offset[None, :] < boundary
229
230 # First calculate the dot, then apply the mask.
231 qk = scale * tl.dot(Q, K)
232 S = tl.where(head_mask[:, None] & seq_mask, qk, float("-inf"))
233
234 context_len = seq_len - 1
235
236 if SLIDING_WINDOW > 0:
237 S = tl.where((context_len - seq_offset) < SLIDING_WINDOW, S, -10000)
238
239 if USE_ALIBI_SLOPES:
240 S += alibi_slope[:, None] * (seq_offset - context_len)
241
242 # compute running maximum
243 # m_j : (num_queries_per_kv,)
244 m_j = tl.maximum(M, tl.max(S, axis=1))
245
246 # P : (num_queries_per_kv, BLOCK_SIZE,)
247 p = tl.exp(S - m_j[:, None])
248 p = tl.where(m_j[:, None] == float("-inf"), 0.0, p)
249
250 # l_j : (num_queries_per_kv,)
251 l_j = tl.sum(p, axis=1)
252
253 # alpha : (num_queries_per_kv, )
254 alpha = tl.exp(M - m_j)
255 alpha = tl.where(float("-inf") == M, 0.0, alpha)
256
257 # acc : (num_queries_per_kv, BLOCK_SIZE,)
258 acc = acc * alpha[:, None]
259
260 # update constants
261 L = L * alpha + l_j
262 M = m_j
263
264 # acc : (num_queries_per_kv, BLOCK_SIZE,)
265 acc += tl.dot(p.to(V.dtype), V)
266
267 # epilogue
268 acc = acc / (L[:, None] + 1e-10)
269 if USE_FP8:
270 acc = acc * tl.load(out_scale_inv)
271 acc = tl.clamp(acc, FP8_MIN, FP8_MAX)
272
273 output_offset = (
274 cur_batch_in_all_start_index * output_stride_0
275 + query_head_idx * output_stride_1
276 )
277
278 tl.store(
279 output_ptr + output_offset[:, None] + tl.arange(0, HEAD_SIZE_PADDED)[None, :],
280 acc,
281 mask=dim_mask[None, :] & head_mask[:, None],
282 )
283
284
285 def chunked_prefill_paged_decode(
286 query,
287 key,
288 value,
289 output,
290 kv_cache_dtype,
291 key_cache,
292 value_cache,
293 block_table,
294 query_start_loc,
295 seq_lens,
296 max_seq_len,
297 max_query_len,
298 k_scale,
299 v_scale,
300 alibi_slopes=None,
301 sliding_window=None,
302 sm_scale=None,
303 output_scale=None,
304 # Optional tensor for sinks
305 sinks=None,
306 is_block_table_ptr: bool = False,
307 causal: bool = True,
308 ):
309 if sm_scale is None:
310 sm_scale = 1.0 / (query.shape[2] ** 0.5)
311
312 use_alibi_slopes = alibi_slopes is not None
313
314 if sliding_window is None or sliding_window <= 0:
315 sliding_window = 0
316
317 if max_query_len > 1:
318 context_attention_fwd(
319 q=query,
320 k=key,
321 v=value,
322 o=output,
323 kv_cache_dtype=kv_cache_dtype,
324 k_cache=key_cache,
325 v_cache=value_cache,
326 b_loc=block_table,
327 b_start_loc=query_start_loc,
328 b_seq_len=seq_lens,
329 max_seq_len=max_seq_len,
330 max_input_len=max_query_len,
331 k_scale=k_scale,
332 v_scale=v_scale,
333 alibi_slopes=alibi_slopes,
334 sliding_window=sliding_window,
335 sm_scale=sm_scale,
336 skip_decode=True,
337 fp8_out_scale=output_scale,
338 sinks=sinks,
339 causal=causal,
340 )
341
342 block_size = value_cache.shape[3]
343 num_seqs = len(seq_lens)
344 num_query_heads = query.shape[1]
345 # key may be None in cross-attention decode (already cached from encoder)
346 num_kv_heads = key.shape[1] if key is not None else key_cache.shape[1]
347 num_queries_per_kv = num_query_heads // num_kv_heads
348 head_size = query.shape[2]
349
350 # Conversion of FP8 Tensor from uint8 storage to
351 # appropriate torch.dtype for interpretation by Triton
352 if "fp8" in kv_cache_dtype:
353 assert key_cache.dtype in [torch.uint8, current_platform.fp8_dtype()]
354 assert value_cache.dtype in [torch.uint8, current_platform.fp8_dtype()]
355
356 if kv_cache_dtype in ("fp8", "fp8_e4m3"):
357 target_dtype = current_platform.fp8_dtype()
358 elif kv_cache_dtype == "fp8_e5m2":
359 target_dtype = torch.float8_e5m2
360 else:
361 raise ValueError(
362 f"Unsupported FP8 kv_cache_dtype {kv_cache_dtype}: "
363 f"should be one of 'fp8', 'fp8_e4m3', 'fp8_e5m2'."
364 )
365
366 key_cache = key_cache.view(target_dtype)
367 value_cache = value_cache.view(target_dtype)
368
369 num_queries_per_kv_padded = max(triton.next_power_of_2(num_queries_per_kv), 16)
370
371 from vllm.platforms.rocm import use_rocm_custom_paged_attention
372
373 use_custom = use_rocm_custom_paged_attention(
374 query.dtype,
375 head_size,
376 block_size,
377 num_queries_per_kv,
378 max_seq_len,
379 sliding_window,
380 kv_cache_dtype,
381 alibi_slopes,
382 sinks,
383 )
384 has_native_layout = has_native_kv_cache_layout(key_cache, value_cache)
385 # Force Triton for non-standard blocks like Qwen3's 544 and for
386 # stride-padded hybrid layouts. The latter use reshape_and_cache_flash
387 # during cache update, so keep decode on the matching stride-aware path.
388 is_pow2 = block_size > 0 and (block_size & (block_size - 1) == 0)
389 if not is_pow2 or not has_native_layout:
390 use_custom = False
391
392 if use_custom:
393 _PARTITION_SIZE_ROCM = 256
394 max_num_partitions = (
395 max_seq_len + _PARTITION_SIZE_ROCM - 1
396 ) // _PARTITION_SIZE_ROCM
397 assert _PARTITION_SIZE_ROCM % block_size == 0
398 total_num_seq = block_table.shape[0]
399 tmp_output = torch.empty(
400 size=(total_num_seq, num_query_heads, max_num_partitions, head_size),
401 dtype=query.dtype,
402 device=output.device,
403 )
404 exp_sums = torch.empty(
405 size=(total_num_seq, num_query_heads, max_num_partitions),
406 dtype=torch.float32,
407 device=output.device,
408 )
409 max_logits = torch.empty_like(exp_sums)
410
411 ops.paged_attention_rocm(
412 output,
413 exp_sums,
414 max_logits,
415 tmp_output,
416 query,
417 key_cache,
418 value_cache,
419 num_kv_heads,
420 scale=sm_scale,
421 block_tables=block_table,
422 seq_lens=seq_lens,
423 query_start_loc=query_start_loc,
424 block_size=block_size,
425 max_seq_len=max_seq_len,
426 alibi_slopes=alibi_slopes,
427 kv_cache_dtype=kv_cache_dtype,
428 k_scale=k_scale,
429 v_scale=v_scale,
430 fp8_out_scale=output_scale,
431 )
432 else:
433 logger.warning_once(
434 "Cannot use ROCm custom paged attention kernel,"
435 " falling back to Triton implementation."
436 )
437 real_block_size = value_cache.shape[3]
438 # The standard model directly uses the original block_size.
439 # Non-standard 544 uses 32 to accommodate integer division logic.
440 # Cap at 128 to avoid exceeding GPU shared memory limits
441 # (e.g. hybrid Mamba models inflate block_size to 2048).
442 # The kernel handles TRITON_BLOCK_SIZE != PHYSICAL_BLOCK_SIZE
443 # via the l_block_idx/internal_offsets addressing logic.
444 MAX_TRITON_BLOCK_SIZE = 128
445 TRITON_BLOCK_SIZE = min(block_size, MAX_TRITON_BLOCK_SIZE) if is_pow2 else 32
446 if is_block_table_ptr:
447 # Using the physical base address of tensors
448 kv_element_size = key_cache.element_size()
449 block_byte_stride = key_cache.stride(0) * kv_element_size
450 # Get the starting physical address of the KV Cache
451 base_addr = key_cache.data_ptr()
452
453 # Normalization: Directly calculate the block offset
454 # of the pointer relative to the base address
455 processed_block_table = ((block_table - base_addr) // block_byte_stride).to(
456 torch.int32
457 )
458 else:
459 processed_block_table = block_table.to(torch.int32)
460
461 kernel_paged_attention_2d[
462 (
463 num_seqs,
464 num_kv_heads,
465 )
466 ](
467 output_ptr=output,
468 query_ptr=query,
469 key_cache_ptr=key_cache,
470 value_cache_ptr=value_cache,
471 sink_ptr=sinks,
472 block_tables_ptr=processed_block_table,
473 seq_lens_ptr=seq_lens,
474 alibi_slopes_ptr=alibi_slopes,
475 scale=sm_scale,
476 k_scale=k_scale,
477 v_scale=v_scale,
478 out_scale_inv=1.0 / output_scale if output_scale is not None else 1.0,
479 num_query_heads=num_query_heads,
480 num_queries_per_kv=num_queries_per_kv,
481 num_queries_per_kv_padded=num_queries_per_kv_padded,
482 block_table_stride=processed_block_table.stride(0),
483 query_stride_0=query.stride(0),
484 query_stride_1=query.stride(1),
485 output_stride_0=output.stride(0),
486 output_stride_1=output.stride(1),
487 BLOCK_SIZE=TRITON_BLOCK_SIZE,
488 PHYSICAL_BLOCK_SIZE=real_block_size,
489 HEAD_SIZE=head_size,
490 HEAD_SIZE_PADDED=triton.next_power_of_2(head_size),
491 USE_ALIBI_SLOPES=use_alibi_slopes,
492 SLIDING_WINDOW=sliding_window,
493 x=key_cache.shape[4],
494 stride_k_cache_0=key_cache.stride(0),
495 stride_k_cache_1=key_cache.stride(1),
496 stride_k_cache_2=key_cache.stride(2),
497 stride_k_cache_3=key_cache.stride(3),
498 stride_k_cache_4=key_cache.stride(4),
499 stride_v_cache_0=value_cache.stride(0),
500 stride_v_cache_1=value_cache.stride(1),
501 stride_v_cache_2=value_cache.stride(2),
502 stride_v_cache_3=value_cache.stride(3),
503 filter_by_query_len=True,
504 query_start_len_ptr=query_start_loc,
505 USE_SINKS=sinks is not None,
506 USE_FP8=output_scale is not None,
507 )