@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
py 1,017 lines 32.6 KB
Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4 # The kernels in this file are adapted from LightLLM's context_attention_fwd:
5 # https://github.com/ModelTC/lightllm/blob/main/lightllm/models/llama/triton_kernel/context_flashattention_nopad.py
6
7 from typing import Any
8
9 import torch
10
11 from vllm.platforms import current_platform
12 from vllm.triton_utils import tl, triton
13
14 # Static kernels parameters
15 BASE_BLOCK = 128 if current_platform.has_device_capability(80) else 64
16 NUM_WARPS = 4 if current_platform.is_rocm() else 8
17
18 # To check compatibility
19 IS_TURING = current_platform.get_device_capability() == (7, 5)
20 float8_info = torch.finfo(current_platform.fp8_dtype())
21
22
23 # Here's an example autotuner config for this kernel. This config does provide
24 # a performance improvement, but dramatically increases first call latency in
25 # triton 3.2. Because of this tradeoff, it's currently commented out.
26 # @triton.autotune(
27 # configs=[
28 # triton.Config({'BLOCK_M': 128, 'BLOCK_N': 64, \
29 # "num_unroll_cache": 4, \
30 # "num_unroll_request": 1 } | \
31 # ({"kpack": 2, "waves_per_eu": 2} \
32 # if current_platform.is_rocm() else {}), \
33 # num_warps=4, \
34 # num_stages=1)
35 # ],
36 # key=["BLOCK_SIZE", "MAX_Q_LEN", "MAX_CTX_LEN"]
37 # )
38 @triton.jit
39 def _paged_kv_cache_offsets(
40 B_Loc,
41 cur_batch,
42 token_indices,
43 token_valid,
44 offs_d,
45 cur_kv_head,
46 x,
47 stride_b_loc_b,
48 stride_b_loc_s,
49 stride_k_cache_bs,
50 stride_k_cache_h,
51 stride_k_cache_d,
52 stride_k_cache_bl,
53 stride_k_cache_x,
54 stride_v_cache_bs,
55 stride_v_cache_h,
56 stride_v_cache_d,
57 stride_v_cache_bl,
58 PHYSICAL_BLOCK_SIZE: tl.constexpr,
59 MASK_BLOCK_TABLE: tl.constexpr = False,
60 ):
61 """Compute paged K/V cache element offsets for a tile of token positions.
62
63 `token_indices` holds the absolute sequence positions of the tile. The
64 logical block for each token is looked up in `B_Loc` (handling cross-block
65 tiles), then the physical element offsets are built with the same layout the
66 kernel uses everywhere:
67 K cache: [num_blocks, num_kv_heads, head_size // x, block_size, x]
68 V cache: [num_blocks, num_kv_heads, head_size, block_size]
69 Returns `(off_k, off_v)` with shapes [D, N] and [N, D] respectively.
70
71 When `MASK_BLOCK_TABLE` is set, the block-table load is masked with
72 `token_valid` so padded lanes (a tile that extends past the end of the
73 sequence) cannot read a block-table entry past this batch's row. Invalid
74 lanes resolve to block 0 (always in-bounds) and are dropped by the caller's
75 later K/V load mask. When it is not set, `token_valid` is unused and the
76 load is emitted exactly as before.
77 """
78 bn_logical = token_indices // PHYSICAL_BLOCK_SIZE
79 if MASK_BLOCK_TABLE:
80 bn = tl.load(
81 B_Loc + cur_batch * stride_b_loc_b + bn_logical * stride_b_loc_s,
82 mask=token_valid,
83 other=0,
84 ).to(tl.int64)
85 else:
86 bn = tl.load(
87 B_Loc + cur_batch * stride_b_loc_b + bn_logical * stride_b_loc_s
88 ).to(tl.int64)
89 internal = token_indices % PHYSICAL_BLOCK_SIZE
90 off_k = (
91 bn[None, :] * stride_k_cache_bs
92 + cur_kv_head * stride_k_cache_h
93 + (offs_d[:, None] // x) * stride_k_cache_d
94 + internal[None, :] * stride_k_cache_bl
95 + (offs_d[:, None] % x) * stride_k_cache_x
96 )
97 off_v = (
98 bn[:, None] * stride_v_cache_bs
99 + cur_kv_head * stride_v_cache_h
100 + offs_d[None, :] * stride_v_cache_d
101 + internal[:, None] * stride_v_cache_bl
102 )
103 return off_k, off_v
104
105
106 @triton.jit
107 def _fwd_kernel(
108 Q,
109 K,
110 V,
111 K_cache,
112 V_cache,
113 sink_ptr,
114 B_Loc,
115 sm_scale,
116 k_scale,
117 v_scale,
118 out_scale_inv,
119 B_Start_Loc,
120 B_Seqlen,
121 x: tl.constexpr,
122 Out,
123 stride_b_loc_b,
124 stride_b_loc_s,
125 stride_qbs,
126 stride_qh,
127 stride_qd,
128 stride_kbs,
129 stride_kh,
130 stride_kd,
131 stride_vbs,
132 stride_vh,
133 stride_vd,
134 stride_obs,
135 stride_oh,
136 stride_od,
137 stride_k_cache_bs,
138 stride_k_cache_h,
139 stride_k_cache_d,
140 stride_k_cache_bl: tl.constexpr,
141 stride_k_cache_x,
142 stride_v_cache_bs,
143 stride_v_cache_h,
144 stride_v_cache_d,
145 stride_v_cache_bl,
146 num_queries_per_kv: tl.constexpr,
147 IN_PRECISION: tl.constexpr,
148 BLOCK_M: tl.constexpr,
149 BLOCK_DMODEL: tl.constexpr,
150 BLOCK_DMODEL_PADDED: tl.constexpr,
151 BLOCK_SIZE: tl.constexpr,
152 PHYSICAL_BLOCK_SIZE: tl.constexpr,
153 BLOCK_N: tl.constexpr,
154 SLIDING_WINDOW: tl.constexpr,
155 num_unroll_cache: tl.constexpr,
156 num_unroll_request: tl.constexpr,
157 SKIP_DECODE: tl.constexpr,
158 USE_SINKS: tl.constexpr,
159 USE_FP8: tl.constexpr,
160 CAUSAL: tl.constexpr = True,
161 MAX_Q_LEN: tl.constexpr = 0,
162 MAX_CTX_LEN: tl.constexpr = 0,
163 FP8_MIN: tl.constexpr = float8_info.min,
164 FP8_MAX: tl.constexpr = float8_info.max,
165 # When True, the current-chunk K/V are not provided as dense tensors
166 # (`K`/`V` are unused placeholders); instead they are read back from the
167 # paged KV cache. This supports layers that re-attend an already-cached
168 # sequence with query only (e.g. IQuest LoopCoder's `attn(q, None, None)`).
169 KV_FROM_CACHE: tl.constexpr = False,
170 ):
171 cur_batch = tl.program_id(0)
172 cur_head = tl.program_id(1)
173 start_m = tl.program_id(2)
174
175 cur_kv_head = cur_head // num_queries_per_kv
176
177 cur_batch_seq_len = tl.load(B_Seqlen + cur_batch)
178 cur_batch_in_all_start_index = tl.load(B_Start_Loc + cur_batch)
179 cur_batch_in_all_stop_index = tl.load(B_Start_Loc + cur_batch + 1)
180 cur_batch_query_len = cur_batch_in_all_stop_index - cur_batch_in_all_start_index
181 cur_batch_ctx_len = cur_batch_seq_len - cur_batch_query_len
182
183 if SKIP_DECODE and cur_batch_query_len == 1:
184 return
185
186 # start position inside of the query
187 # generally, N goes over kv, while M goes over query_len
188 block_start_loc = BLOCK_M * start_m
189
190 # initialize offsets
191 # [BLOCK_SIZE]; starts at 0
192 offs_bs_n = tl.arange(0, BLOCK_SIZE)
193 # [N]; starts at 0
194 offs_n = tl.arange(0, BLOCK_N)
195 # [D]; starts at 0
196 offs_d = tl.arange(0, BLOCK_DMODEL_PADDED)
197 # [M]; starts at current position in query
198 offs_m = start_m * BLOCK_M + tl.arange(0, BLOCK_M)
199 # [M,D]
200 off_q = (
201 (cur_batch_in_all_start_index + offs_m[:, None]) * stride_qbs
202 + cur_head * stride_qh
203 + offs_d[None, :] * stride_qd
204 )
205
206 dim_mask = tl.where(tl.arange(0, BLOCK_DMODEL_PADDED) < BLOCK_DMODEL, 1, 0).to(
207 tl.int1
208 ) # [D]
209
210 q = tl.load(
211 Q + off_q,
212 mask=dim_mask[None, :] & (offs_m[:, None] < cur_batch_query_len),
213 other=0.0,
214 ) # [M,D]
215
216 # initialize pointer to m and l
217 if not USE_SINKS:
218 m_i = tl.full([BLOCK_M], float("-inf"), dtype=tl.float32)
219 l_i = tl.zeros([BLOCK_M], dtype=tl.float32)
220 else:
221 m_i = tl.load(
222 sink_ptr + tl.full([BLOCK_M], cur_head, dtype=tl.int64),
223 mask=(offs_m < cur_batch_query_len),
224 other=float("-inf"),
225 ).to(dtype=tl.float32)
226 l_i = tl.where(m_i > float("-inf"), 1.0, 0.0)
227
228 acc = tl.zeros([BLOCK_M, BLOCK_DMODEL_PADDED], dtype=tl.float32) # [M,D]
229
230 # compute query against context (no causal mask here)
231 for start_n in tl.range(
232 0, cur_batch_ctx_len, BLOCK_SIZE, loop_unroll_factor=num_unroll_cache
233 ):
234 # Under a block size of 544 (Qwen/Qwen3-Next-80B-A3B-Thinking),
235 # replace one physical block every 17 32-Tile blocks
236 # Calculate the logical block index of each of the 32 tokens
237 # in the current Tile (handling cross-block cases).
238 # Physical-block K/V cache offsets for the context tokens in this tile
239 # (handles cross-block tiles via B_Loc). Same addressing is reused for
240 # the current-chunk cache read below.
241 token_indices = start_n + offs_bs_n
242 # Context tokens are always followed by the current chunk, so the tile
243 # never steps past this batch's block-table row; the block-table load
244 # stays unmasked (MASK_BLOCK_TABLE=False) and offs_bs_n is an unused
245 # placeholder for token_valid.
246 off_k, off_v = _paged_kv_cache_offsets(
247 B_Loc,
248 cur_batch,
249 token_indices,
250 offs_bs_n,
251 offs_d,
252 cur_kv_head,
253 x,
254 stride_b_loc_b,
255 stride_b_loc_s,
256 stride_k_cache_bs,
257 stride_k_cache_h,
258 stride_k_cache_d,
259 stride_k_cache_bl,
260 stride_k_cache_x,
261 stride_v_cache_bs,
262 stride_v_cache_h,
263 stride_v_cache_d,
264 stride_v_cache_bl,
265 PHYSICAL_BLOCK_SIZE,
266 )
267
268 if (
269 start_n + BLOCK_SIZE > cur_batch_ctx_len
270 or BLOCK_DMODEL != BLOCK_DMODEL_PADDED
271 ):
272 k_load = tl.load(
273 K_cache + off_k,
274 mask=dim_mask[:, None]
275 & ((start_n + offs_bs_n[None, :]) < cur_batch_ctx_len),
276 other=0.0,
277 ) # [D,N]
278 else:
279 k_load = tl.load(K_cache + off_k)
280
281 if k_load.dtype.is_fp8():
282 k = (k_load.to(tl.float32) * tl.load(k_scale)).to(q.dtype)
283 else:
284 k = k_load
285
286 # qk = tl.zeros([BLOCK_M, BLOCK_SIZE], dtype=tl.float32) # [M,N]
287 qk = sm_scale * tl.dot(q, k, input_precision=IN_PRECISION)
288 qk = tl.where(
289 (start_n + offs_bs_n[None, :]) < cur_batch_ctx_len, qk, float("-inf")
290 )
291 # qk *= sm_scale
292 if SLIDING_WINDOW > 0:
293 # (cur_batch_ctx_len + offs_m[:, None]) are the positions of
294 # Q entries in sequence
295 # (start_n + offs_bs_n[None, :]) are the positions of
296 # KV entries in sequence
297 # So the condition makes sure each entry in Q only attends
298 # to KV entries not more than SLIDING_WINDOW away.
299 #
300 # We can't use -inf here, because the
301 # sliding window may lead to the entire row being masked.
302 # This then makes m_ij contain -inf, which causes NaNs in
303 # exp().
304 qk = tl.where(
305 (cur_batch_ctx_len + offs_m[:, None]) - (start_n + offs_bs_n[None, :])
306 < SLIDING_WINDOW,
307 qk,
308 float("-inf"),
309 )
310
311 # compute running maximum
312 m_ij = tl.maximum(m_i, tl.max(qk, axis=1))
313 p = tl.exp(qk - m_ij[:, None])
314 p = tl.where(m_ij[:, None] == float("-inf"), 0.0, p)
315 l_ij = tl.sum(p, axis=1)
316 alpha = tl.exp(m_i - m_ij)
317 alpha = tl.where(m_i == float("-inf"), 0.0, alpha)
318 acc = acc * alpha[:, None]
319
320 # update acc
321 if (
322 start_n + BLOCK_SIZE > cur_batch_ctx_len
323 or BLOCK_DMODEL != BLOCK_DMODEL_PADDED
324 ):
325 v_load = tl.load(
326 V_cache + off_v,
327 mask=dim_mask[None, :]
328 & ((start_n + offs_bs_n[:, None]) < cur_batch_ctx_len),
329 other=0.0,
330 ) # [N,D]
331 else:
332 v_load = tl.load(V_cache + off_v)
333
334 if v_load.dtype.is_fp8():
335 v = (v_load.to(tl.float32) * tl.load(v_scale)).to(q.dtype)
336 else:
337 v = v_load
338 p = p.to(v.dtype)
339
340 acc = tl.dot(p, v, acc=acc, input_precision=IN_PRECISION)
341 # # update m_i and l_i
342 l_i = l_i * alpha + l_ij
343 m_i = m_ij
344
345 off_k = (
346 offs_n[None, :] * stride_kbs
347 + cur_kv_head * stride_kh
348 + offs_d[:, None] * stride_kd
349 )
350 off_v = (
351 offs_n[:, None] * stride_vbs
352 + cur_kv_head * stride_vh
353 + offs_d[None, :] * stride_vd
354 )
355 k_ptrs = K + off_k
356 v_ptrs = V + off_v
357
358 # block_mask is 0 when we're already past the current query length
359 block_mask = tl.where(block_start_loc < cur_batch_query_len, 1, 0)
360
361 # compute query against itself (causal among queries by default;
362 # CAUSAL=False for bidirectional attention over query tokens, e.g. DFlash.)
363 if CAUSAL:
364 key_range_upper = block_mask * (start_m + 1) * BLOCK_M
365 else:
366 q_len_pad = (cur_batch_query_len + BLOCK_N - 1) // BLOCK_N * BLOCK_N
367 key_range_upper = block_mask * q_len_pad
368
369 for start_n in tl.range(
370 0,
371 key_range_upper,
372 BLOCK_N,
373 loop_unroll_factor=num_unroll_request,
374 ):
375 start_n = tl.multiple_of(start_n, BLOCK_N)
376 # -- compute qk ----
377 if KV_FROM_CACHE:
378 # Current-chunk K/V are not passed densely; read them from the paged
379 # cache at their absolute positions [ctx_len, ctx_len + query_len),
380 # reusing the same addressing as the context loop above.
381 cache_token_idx = cur_batch_ctx_len + start_n + offs_n
382 # Padded lanes of the last tile can index past the sequence end;
383 # mask the block-table load so they don't read past this batch's
384 # row (they are dropped by the K/V load mask below anyway).
385 cache_token_valid = (start_n + offs_n) < cur_batch_query_len
386 off_k_cur, off_v_cur = _paged_kv_cache_offsets(
387 B_Loc,
388 cur_batch,
389 cache_token_idx,
390 cache_token_valid,
391 offs_d,
392 cur_kv_head,
393 x,
394 stride_b_loc_b,
395 stride_b_loc_s,
396 stride_k_cache_bs,
397 stride_k_cache_h,
398 stride_k_cache_d,
399 stride_k_cache_bl,
400 stride_k_cache_x,
401 stride_v_cache_bs,
402 stride_v_cache_h,
403 stride_v_cache_d,
404 stride_v_cache_bl,
405 PHYSICAL_BLOCK_SIZE,
406 MASK_BLOCK_TABLE=True,
407 )
408 k_cur_load = tl.load(
409 K_cache + off_k_cur,
410 mask=dim_mask[:, None]
411 & ((start_n + offs_n[None, :]) < cur_batch_query_len),
412 other=0.0,
413 )
414 if k_cur_load.dtype.is_fp8():
415 k = (k_cur_load.to(tl.float32) * tl.load(k_scale)).to(q.dtype)
416 else:
417 k = k_cur_load
418 else:
419 k = tl.load(
420 k_ptrs + (cur_batch_in_all_start_index + start_n) * stride_kbs,
421 mask=dim_mask[:, None]
422 & ((start_n + offs_n[None, :]) < cur_batch_query_len),
423 other=0.0,
424 )
425
426 qk = tl.zeros([BLOCK_M, BLOCK_N], dtype=tl.float32)
427 qk = tl.dot(q, k, acc=qk, input_precision=IN_PRECISION)
428 qk *= sm_scale
429
430 valid_kv = (start_n + offs_n[None, :]) < cur_batch_query_len
431 if CAUSAL:
432 attn_mask = valid_kv & (offs_m[:, None] >= (start_n + offs_n[None, :]))
433 else:
434 attn_mask = valid_kv
435 if SLIDING_WINDOW > 0:
436 attn_mask = attn_mask & (
437 offs_m[:, None] - (start_n + offs_n[None, :]) < SLIDING_WINDOW
438 )
439 qk = tl.where(attn_mask, qk, float("-inf"))
440
441 # compute running maximum
442 m_ij = tl.maximum(m_i, tl.max(qk, axis=1))
443 p = tl.exp(qk - m_ij[:, None])
444 p = tl.where(m_ij[:, None] == float("-inf"), 0.0, p)
445 l_ij = tl.sum(p, axis=1)
446 alpha = tl.exp(m_i - m_ij)
447 # To prevent NaN from appearing in the first round
448 alpha = tl.where(m_i == float("-inf"), 0.0, alpha)
449 acc = acc * alpha[:, None]
450
451 # update acc
452 if KV_FROM_CACHE:
453 # off_v_cur was computed together with off_k_cur above.
454 v_cur_load = tl.load(
455 V_cache + off_v_cur,
456 mask=dim_mask[None, :]
457 & ((start_n + offs_n[:, None]) < cur_batch_query_len),
458 other=0.0,
459 )
460 if v_cur_load.dtype.is_fp8():
461 v = (v_cur_load.to(tl.float32) * tl.load(v_scale)).to(q.dtype)
462 else:
463 v = v_cur_load
464 else:
465 v = tl.load(
466 v_ptrs + (cur_batch_in_all_start_index + start_n) * stride_vbs,
467 mask=dim_mask[None, :]
468 & ((start_n + offs_n[:, None]) < cur_batch_query_len),
469 other=0.0,
470 )
471 p = p.to(v.dtype)
472
473 acc = tl.dot(p, v, acc=acc, input_precision=IN_PRECISION)
474 # update m_i and l_i
475 l_i = l_i * alpha + l_ij
476 m_i = m_ij
477
478 acc = acc / (l_i[:, None] + 1e-10)
479
480 # initialize pointers to output
481 off_o = (
482 (cur_batch_in_all_start_index + offs_m[:, None]) * stride_obs
483 + cur_head * stride_oh
484 + offs_d[None, :] * stride_od
485 )
486 out_ptrs = Out + off_o
487 if USE_FP8:
488 acc = acc * tl.load(out_scale_inv)
489 acc = tl.clamp(acc, FP8_MIN, FP8_MAX)
490 tl.store(
491 out_ptrs, acc, mask=dim_mask[None, :] & (offs_m[:, None] < cur_batch_query_len)
492 )
493 return
494
495
496 @triton.jit
497 def _fwd_kernel_alibi(
498 Q,
499 K,
500 V,
501 K_cache,
502 V_cache,
503 B_Loc,
504 sm_scale,
505 k_scale,
506 v_scale,
507 B_Start_Loc,
508 B_Seqlen,
509 Alibi_slopes,
510 block_size,
511 x,
512 Out,
513 stride_b_loc_b,
514 stride_b_loc_s,
515 stride_qbs,
516 stride_qh,
517 stride_qd,
518 stride_kbs,
519 stride_kh,
520 stride_kd,
521 stride_vbs,
522 stride_vh,
523 stride_vd,
524 stride_obs,
525 stride_oh,
526 stride_od,
527 stride_k_cache_bs,
528 stride_k_cache_h,
529 stride_k_cache_d,
530 stride_k_cache_bl,
531 stride_k_cache_x,
532 stride_v_cache_bs,
533 stride_v_cache_h,
534 stride_v_cache_d,
535 stride_v_cache_bl,
536 num_queries_per_kv: int,
537 IN_PRECISION: tl.constexpr,
538 BLOCK_M: tl.constexpr,
539 BLOCK_DMODEL: tl.constexpr, # head size
540 BLOCK_DMODEL_PADDED: tl.constexpr, # head size padded to a power of 2
541 BLOCK_N: tl.constexpr,
542 SKIP_DECODE: tl.constexpr,
543 ):
544 # attn_bias[]
545 cur_batch = tl.program_id(0)
546 cur_head = tl.program_id(1)
547 start_m = tl.program_id(2)
548
549 cur_kv_head = cur_head // num_queries_per_kv
550
551 # cur_batch_seq_len: the length of prompts
552 # cur_batch_ctx_len: the length of prefix
553 # cur_batch_in_all_start_index: the start id of the dim=0
554 cur_batch_seq_len = tl.load(B_Seqlen + cur_batch)
555 cur_batch_in_all_start_index = tl.load(B_Start_Loc + cur_batch)
556 cur_batch_in_all_stop_index = tl.load(B_Start_Loc + cur_batch + 1)
557 cur_batch_query_len = cur_batch_in_all_stop_index - cur_batch_in_all_start_index
558 cur_batch_ctx_len = cur_batch_seq_len - cur_batch_query_len
559
560 if SKIP_DECODE and cur_batch_query_len == 1:
561 return
562
563 block_start_loc = BLOCK_M * start_m
564
565 # initialize offsets
566 offs_n = tl.arange(0, BLOCK_N)
567 offs_d = tl.arange(0, BLOCK_DMODEL_PADDED)
568 offs_m = start_m * BLOCK_M + tl.arange(0, BLOCK_M)
569 off_q = (
570 (cur_batch_in_all_start_index + offs_m[:, None]) * stride_qbs
571 + cur_head * stride_qh
572 + offs_d[None, :] * stride_qd
573 )
574
575 dim_mask = tl.where(tl.arange(0, BLOCK_DMODEL_PADDED) < BLOCK_DMODEL, 1, 0).to(
576 tl.int1
577 )
578
579 q = tl.load(
580 Q + off_q,
581 mask=dim_mask[None, :]
582 & (offs_m[:, None] < cur_batch_seq_len - cur_batch_ctx_len),
583 other=0.0,
584 )
585
586 # # initialize pointer to m and l
587 m_i = tl.zeros([BLOCK_M], dtype=tl.float32) - float("inf")
588 l_i = tl.zeros([BLOCK_M], dtype=tl.float32)
589 acc = tl.zeros([BLOCK_M, BLOCK_DMODEL_PADDED], dtype=tl.float32)
590
591 alibi_slope = tl.load(Alibi_slopes + cur_head)
592 alibi_start_q = tl.arange(0, BLOCK_M) + block_start_loc + cur_batch_ctx_len
593 alibi_start_k = 0
594 for start_n in range(0, cur_batch_ctx_len, BLOCK_N):
595 start_n = tl.multiple_of(start_n, BLOCK_N)
596 # -- compute qk ----
597 bn = tl.load(
598 B_Loc
599 + cur_batch * stride_b_loc_b
600 + ((start_n + offs_n) // block_size) * stride_b_loc_s,
601 mask=(start_n + offs_n) < cur_batch_ctx_len,
602 other=0,
603 ).to(tl.int64)
604 off_k = (
605 bn[None, :] * stride_k_cache_bs
606 + cur_kv_head * stride_k_cache_h
607 + (offs_d[:, None] // x) * stride_k_cache_d
608 + ((start_n + offs_n[None, :]) % block_size) * stride_k_cache_bl
609 + (offs_d[:, None] % x) * stride_k_cache_x
610 )
611 off_v = (
612 bn[:, None] * stride_v_cache_bs
613 + cur_kv_head * stride_v_cache_h
614 + offs_d[None, :] * stride_v_cache_d
615 + (start_n + offs_n[:, None]) % block_size * stride_v_cache_bl
616 )
617 k_load = tl.load(
618 K_cache + off_k,
619 mask=dim_mask[:, None] & ((start_n + offs_n[None, :]) < cur_batch_ctx_len),
620 other=0.0,
621 ) # [D,N]
622
623 if k_load.dtype.is_fp8():
624 k = (k_load.to(tl.float32) * tl.load(k_scale)).to(q.dtype)
625 else:
626 k = k_load
627
628 qk = tl.zeros([BLOCK_M, BLOCK_N], dtype=tl.float32)
629 qk = tl.dot(q, k, acc=qk, input_precision=IN_PRECISION)
630 qk = tl.where(
631 (start_n + offs_n[None, :]) < cur_batch_ctx_len, qk, float("-inf")
632 )
633 qk *= sm_scale
634
635 # load alibi
636 alibi = (
637 tl.arange(0, BLOCK_N)[None, :] + alibi_start_k - alibi_start_q[:, None]
638 ) * alibi_slope
639 alibi = tl.where(
640 (alibi <= 0) & (alibi_start_q[:, None] < cur_batch_seq_len),
641 alibi,
642 float("-inf"),
643 )
644 qk += alibi
645 alibi_start_k += BLOCK_N
646
647 # -- compute m_ij, p, l_ij
648 m_ij = tl.max(qk, 1)
649 m_i_new = tl.maximum(m_i, m_ij)
650 p = tl.math.exp(qk - m_i_new[:, None])
651 l_ij = tl.sum(p, 1)
652 # -- update m_i and l_i
653
654 alpha = tl.math.exp(m_i - m_i_new)
655 l_i_new = alpha * l_i + l_ij
656 # -- update output accumulator --
657 # scale p
658 # scale acc
659 acc_scale = alpha
660 # acc_scale = l_i / l_i_new * alpha
661 acc = acc * acc_scale[:, None]
662 # update acc
663 v_load = tl.load(
664 V_cache + off_v,
665 mask=dim_mask[None, :] & ((start_n + offs_n[:, None]) < cur_batch_ctx_len),
666 other=0.0,
667 )
668 if v_load.dtype.is_fp8():
669 v = (v_load.to(tl.float32) * tl.load(v_scale)).to(q.dtype)
670 else:
671 v = v_load
672 p = p.to(v.dtype)
673
674 acc = tl.dot(p, v, acc=acc, input_precision="ieee")
675 # update m_i and l_i
676 l_i = l_i_new
677 m_i = m_i_new
678
679 off_k = (
680 offs_n[None, :] * stride_kbs
681 + cur_kv_head * stride_kh
682 + offs_d[:, None] * stride_kd
683 )
684 off_v = (
685 offs_n[:, None] * stride_vbs
686 + cur_kv_head * stride_vh
687 + offs_d[None, :] * stride_vd
688 )
689 k_ptrs = K + off_k
690 v_ptrs = V + off_v
691
692 block_mask = tl.where(block_start_loc < cur_batch_seq_len - cur_batch_ctx_len, 1, 0)
693
694 # init alibi
695 alibi_slope = tl.load(Alibi_slopes + cur_head)
696 alibi_start_q = tl.arange(0, BLOCK_M) + block_start_loc + cur_batch_ctx_len
697 alibi_start_k = cur_batch_ctx_len
698 # # init debugger
699 # offset_db_q = tl.arange(0, BLOCK_M) + block_start_loc
700 # offset_db_k = tl.arange(0, BLOCK_N)
701 # calc q[BLOCK_M, BLOCK_MODEL] mul k[prefix_len: , BLOCK_DMODEL]
702 for start_n in range(0, block_mask * (start_m + 1) * BLOCK_M, BLOCK_N):
703 start_n = tl.multiple_of(start_n, BLOCK_N)
704 # -- compute qk ----
705 k = tl.load(
706 k_ptrs + (cur_batch_in_all_start_index + start_n) * stride_kbs,
707 mask=dim_mask[:, None]
708 & ((start_n + offs_n[None, :]) < cur_batch_seq_len - cur_batch_ctx_len),
709 other=0.0,
710 )
711
712 qk = tl.zeros([BLOCK_M, BLOCK_N], dtype=tl.float32)
713 qk = tl.dot(q, k, acc=qk, input_precision="ieee")
714 qk *= sm_scale
715 qk = tl.where(offs_m[:, None] >= (start_n + offs_n[None, :]), qk, float("-inf"))
716
717 # load alibi
718 alibi = (
719 tl.arange(0, BLOCK_N)[None, :] + alibi_start_k - alibi_start_q[:, None]
720 ) * alibi_slope
721 alibi = tl.where(
722 (alibi <= 0) & (alibi_start_q[:, None] < cur_batch_seq_len),
723 alibi,
724 float("-inf"),
725 )
726 qk += alibi
727 alibi_start_k += BLOCK_N
728
729 # -- compute m_ij, p, l_ij
730 m_ij = tl.max(qk, 1)
731 m_i_new = tl.maximum(m_i, m_ij)
732 p = tl.math.exp(qk - m_i_new[:, None])
733 l_ij = tl.sum(p, 1)
734 # -- update m_i and l_i
735
736 alpha = tl.math.exp(m_i - m_i_new)
737 l_i_new = alpha * l_i + l_ij
738 # -- update output accumulator --
739 # scale p
740 # scale acc
741 acc_scale = alpha
742 # acc_scale = l_i / l_i_new * alpha
743 acc = acc * acc_scale[:, None]
744 # update acc
745 v = tl.load(
746 v_ptrs + (cur_batch_in_all_start_index + start_n) * stride_vbs,
747 mask=dim_mask[None, :]
748 & ((start_n + offs_n[:, None]) < cur_batch_seq_len - cur_batch_ctx_len),
749 other=0.0,
750 )
751 p = p.to(v.dtype)
752
753 acc = tl.dot(p, v, acc=acc, input_precision="ieee")
754 # update m_i and l_i
755 l_i = l_i_new
756 m_i = m_i_new
757
758 acc = acc / l_i[:, None]
759
760 # initialize pointers to output
761 off_o = (
762 (cur_batch_in_all_start_index + offs_m[:, None]) * stride_obs
763 + cur_head * stride_oh
764 + offs_d[None, :] * stride_od
765 )
766 out_ptrs = Out + off_o
767 tl.store(
768 out_ptrs,
769 acc,
770 mask=dim_mask[None, :]
771 & (offs_m[:, None] < cur_batch_seq_len - cur_batch_ctx_len),
772 )
773 return
774
775
776 @torch.inference_mode()
777 def context_attention_fwd(
778 q,
779 k,
780 v,
781 o,
782 kv_cache_dtype: str,
783 k_cache,
784 v_cache,
785 b_loc,
786 b_start_loc,
787 b_seq_len,
788 max_seq_len,
789 max_input_len,
790 k_scale: torch.Tensor,
791 v_scale: torch.Tensor,
792 alibi_slopes=None,
793 sliding_window=None,
794 sm_scale=None,
795 skip_decode=False,
796 fp8_out_scale=None,
797 sinks=None,
798 is_block_table_ptr: bool = False,
799 causal: bool = True,
800 ):
801 q_dtype_is_f32 = q.dtype is torch.float32
802
803 # Turing does have tensor core for float32 multiplication
804 # use ieee as fallback for triton kernels work. There is also
805 # warning on vllm/config.py to inform users this fallback
806 # implementation
807 IN_PRECISION = "ieee" if IS_TURING and q_dtype_is_f32 else None
808
809 # Conversion of FP8 Tensor from uint8 storage to
810 # appropriate torch.dtype for interpretation by Triton
811 if "fp8" in kv_cache_dtype:
812 assert k_cache.dtype in [torch.uint8, current_platform.fp8_dtype()]
813 assert v_cache.dtype in [torch.uint8, current_platform.fp8_dtype()]
814
815 if kv_cache_dtype in ("fp8", "fp8_e4m3"):
816 target_dtype = current_platform.fp8_dtype()
817 elif kv_cache_dtype == "fp8_e5m2":
818 target_dtype = torch.float8_e5m2
819 else:
820 raise ValueError("Unsupported FP8 dtype:", kv_cache_dtype)
821
822 k_cache = k_cache.view(target_dtype)
823 v_cache = v_cache.view(target_dtype)
824
825 if (
826 k_cache.dtype == torch.uint8
827 or v_cache.dtype == torch.uint8
828 and kv_cache_dtype == "auto"
829 ):
830 raise ValueError(
831 "kv_cache_dtype='auto' unsupported for\
832 FP8 KV Cache prefill kernel"
833 )
834
835 assert (k is None) == (v is None), (
836 "k and v must both be None (cached-K/V path) or both be tensors"
837 )
838 kv_from_cache = k is None
839 if kv_from_cache:
840 if alibi_slopes is not None:
841 raise NotImplementedError(
842 "context_attention_fwd with cached K/V (key=None) is not "
843 "supported together with ALiBi slopes."
844 )
845 # k_cache: [num_blocks, num_kv_heads, head_size // x, block_size, x]
846 num_kv_heads = k_cache.shape[1]
847 Lq = Lk = Lv = q.shape[-1]
848 # q is a never-dereferenced placeholder for the dense K/V pointer args.
849 k = q
850 v = q
851 else:
852 num_kv_heads = k.shape[1]
853 Lq, Lk, Lv = q.shape[-1], k.shape[-1], v.shape[-1]
854
855 # shape constraints
856 assert Lq == Lk and Lk == Lv
857 # round up Lk to a power of 2 - this is required for Triton block size
858 Lk_padded = triton.next_power_of_2(Lk)
859
860 if sm_scale is None:
861 sm_scale = 1.0 / (Lq**0.5)
862 batch, head = b_seq_len.shape[0], q.shape[1]
863 num_queries_per_kv = q.shape[1] // num_kv_heads
864
865 assert batch + 1 == len(b_start_loc)
866
867 # 0 means "disable"
868 if sliding_window is None or sliding_window <= 0:
869 sliding_window = 0
870
871 if is_block_table_ptr:
872 kv_element_size = k_cache.element_size()
873 block_byte_stride = k_cache.stride(0) * kv_element_size
874 # The physical starting point of the obtained KV Cache Pool
875 base_addr = k_cache.data_ptr()
876
877 mask = b_loc > 0
878 processed_b_loc = torch.where(
879 mask, (b_loc - base_addr) // block_byte_stride, b_loc
880 ).to(torch.int32)
881 else:
882 processed_b_loc = b_loc.to(torch.int32)
883
884 if alibi_slopes is not None:
885 assert causal, "Non-causal prefix attention is not supported with alibi"
886 assert sinks is None, "Sinks arg is not supported with alibi"
887 assert fp8_out_scale is None, "FP8 output not supported with alibi"
888 # need to reduce num. blocks when using fp32
889 # due to increased use of GPU shared memory
890 # if q.dtype is torch.float32:
891 BLOCK = BASE_BLOCK // 2 if q_dtype_is_f32 else BASE_BLOCK
892 # batch, head,
893 grid = (batch, head, triton.cdiv(max_input_len, BLOCK))
894 _fwd_kernel_alibi[grid](
895 q,
896 k,
897 v,
898 k_cache,
899 v_cache,
900 b_loc,
901 sm_scale,
902 k_scale,
903 v_scale,
904 b_start_loc,
905 b_seq_len,
906 alibi_slopes,
907 v_cache.shape[3],
908 k_cache.shape[4],
909 o,
910 b_loc.stride(0),
911 b_loc.stride(1),
912 q.stride(0),
913 q.stride(1),
914 q.stride(2),
915 k.stride(0),
916 k.stride(1),
917 k.stride(2),
918 v.stride(0),
919 v.stride(1),
920 v.stride(2),
921 o.stride(0),
922 o.stride(1),
923 o.stride(2),
924 k_cache.stride(0),
925 k_cache.stride(1),
926 k_cache.stride(2),
927 k_cache.stride(3),
928 k_cache.stride(4), # [num_blocks, num_kv_heads, head_size/x, block_size, x]
929 v_cache.stride(0),
930 v_cache.stride(1),
931 v_cache.stride(2),
932 v_cache.stride(3), # [num_blocks, num_kv_heads, head_size, block_size]
933 num_queries_per_kv=num_queries_per_kv,
934 IN_PRECISION=IN_PRECISION,
935 BLOCK_M=BLOCK,
936 BLOCK_DMODEL=Lk,
937 BLOCK_DMODEL_PADDED=Lk_padded,
938 BLOCK_N=BLOCK,
939 SKIP_DECODE=skip_decode,
940 num_warps=NUM_WARPS,
941 num_stages=1,
942 )
943 return
944
945 max_seq_len = 0 if max_seq_len is None else max_seq_len
946 extra_kargs: dict[str, Any] = {}
947 if current_platform.is_rocm():
948 extra_kargs = {}
949
950 real_block_size = v_cache.shape[3]
951 # _paged_kv_cache_offsets resolves context tokens against PHYSICAL_BLOCK_SIZE
952 # individually, so tiles need not divide the page size.
953 BLOCK_M = 128
954 BLOCK_N = 64
955 TRITON_BLOCK_SIZE = 32
956
957 grid_fn = lambda META: (batch, head, triton.cdiv(max_input_len, META["BLOCK_M"]))
958 _fwd_kernel[grid_fn](
959 q,
960 k,
961 v,
962 k_cache,
963 v_cache,
964 sinks,
965 processed_b_loc,
966 sm_scale,
967 k_scale,
968 v_scale,
969 1.0 / fp8_out_scale if fp8_out_scale is not None else 1.0,
970 b_start_loc,
971 b_seq_len,
972 k_cache.shape[4],
973 o,
974 processed_b_loc.stride(0),
975 processed_b_loc.stride(1),
976 q.stride(0),
977 q.stride(1),
978 q.stride(2),
979 k.stride(0),
980 k.stride(1),
981 k.stride(2),
982 v.stride(0),
983 v.stride(1),
984 v.stride(2),
985 o.stride(0),
986 o.stride(1),
987 o.stride(2),
988 stride_k_cache_bs=k_cache.stride(0),
989 stride_k_cache_h=k_cache.stride(1),
990 stride_k_cache_d=k_cache.stride(2),
991 stride_k_cache_bl=k_cache.stride(3),
992 stride_k_cache_x=k_cache.stride(4),
993 stride_v_cache_bs=v_cache.stride(0),
994 stride_v_cache_h=v_cache.stride(1),
995 stride_v_cache_d=v_cache.stride(2),
996 stride_v_cache_bl=v_cache.stride(3),
997 BLOCK_SIZE=TRITON_BLOCK_SIZE,
998 PHYSICAL_BLOCK_SIZE=real_block_size,
999 num_queries_per_kv=num_queries_per_kv,
1000 IN_PRECISION=IN_PRECISION,
1001 BLOCK_DMODEL=Lk,
1002 BLOCK_DMODEL_PADDED=Lk_padded,
1003 SLIDING_WINDOW=sliding_window,
1004 SKIP_DECODE=skip_decode,
1005 USE_FP8=fp8_out_scale is not None,
1006 BLOCK_M=BLOCK_M,
1007 BLOCK_N=BLOCK_N,
1008 num_unroll_cache=4,
1009 num_unroll_request=1,
1010 num_warps=4,
1011 num_stages=1,
1012 USE_SINKS=sinks is not None,
1013 CAUSAL=causal,
1014 KV_FROM_CACHE=kv_from_cache,
1015 **extra_kargs,
1016 )
1017 return