@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
py 817 lines 22.6 KB
Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4 # Adapted from
5 # https://github.com/sgl-project/sglang/blob/9f635ea50de920aa507f486daafba26a5b837574/python/sglang/srt/layers/attention/triton_ops/decode_attention.py
6 # which was originally adapted from
7 # https://github.com/ModelTC/lightllm/blob/96353e868a840db4d103138caf15ed9dbea8c186/lightllm/models/deepseek2/triton_kernel/gqa_flash_decoding_stage1.py
8 # https://github.com/ModelTC/lightllm/blob/96353e868a840db4d103138caf15ed9dbea8c186/lightllm/models/deepseek2/triton_kernel/gqa_flash_decoding_stage2.py
9
10 # Changes:
11 # - Add support for page size >= 1.
12
13 # Copyright 2025 vLLM Team
14 # Copyright 2023-2024 SGLang Team
15 # Licensed under the Apache License, Version 2.0 (the "License");
16 # you may not use this file except in compliance with the License.
17 # You may obtain a copy of the License at
18 #
19 # http://www.apache.org/licenses/LICENSE-2.0
20 #
21 # Unless required by applicable law or agreed to in writing, software
22 # distributed under the License is distributed on an "AS IS" BASIS,
23 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
24 # See the License for the specific language governing permissions and
25 # limitations under the License.
26 # ==============================================================================
27 """Memory-efficient attention for decoding.
28 It supports page size >= 1.
29 """
30
31 import logging
32
33 import torch
34 from packaging import version
35
36 from vllm.platforms import current_platform
37 from vllm.triton_utils import tl, triton
38
39 is_hip_ = current_platform.is_rocm()
40
41 logger = logging.getLogger(__name__)
42
43 # Only print the following warnings when triton version < 3.2.0.
44 # The issue won't affect performance or accuracy.
45 if version.parse(triton.__version__) < version.parse("3.2.0"):
46 logger.warning(
47 "The following error message 'operation scheduled before its operands' "
48 "can be ignored."
49 )
50
51
52 @triton.jit
53 def tanh(x):
54 # Tanh is just a scaled sigmoid
55 return 2 * tl.sigmoid(2 * x) - 1
56
57
58 def _page_stride(buf, page_size):
59 # Stride between pages. 4D buffers have a page dim; 3D buffers pack pages
60 # along the token dim, so split it out first. Read the real stride (a
61 # cross-layer view has gaps), don't assume PAGE_SIZE * token stride.
62 if buf.ndim == 3:
63 buf = buf.unflatten(-3, (-1, page_size))
64 return buf.stride(-4)
65
66
67 @triton.jit
68 def _fwd_kernel_stage1(
69 Q,
70 K_Buffer,
71 V_Buffer,
72 sm_scale,
73 Req_to_tokens,
74 B_Seqlen,
75 Att_Out,
76 stride_req_to_tokens_b,
77 stride_qbs,
78 stride_qh,
79 stride_buf_kpbs,
80 stride_buf_kbs,
81 stride_buf_kh,
82 stride_buf_vpbs,
83 stride_buf_vbs,
84 stride_buf_vh,
85 stride_mid_ob,
86 stride_mid_oh,
87 stride_mid_os,
88 k_scale,
89 v_scale,
90 kv_group_num: tl.constexpr,
91 BLOCK_DMODEL: tl.constexpr,
92 BLOCK_DV: tl.constexpr,
93 BLOCK_N: tl.constexpr,
94 NUM_KV_SPLITS: tl.constexpr,
95 PAGE_SIZE: tl.constexpr,
96 logit_cap: tl.constexpr,
97 Lk: tl.constexpr,
98 Lv: tl.constexpr,
99 ):
100 cur_batch = tl.program_id(0)
101 cur_head = tl.program_id(1)
102 split_kv_id = tl.program_id(2)
103
104 cur_kv_head = cur_head // kv_group_num
105
106 offs_d = tl.arange(0, BLOCK_DMODEL)
107 offs_dv = tl.arange(0, BLOCK_DV)
108 mask_d = offs_d < Lk
109 mask_dv = offs_dv < Lv
110 cur_batch_seq_len = tl.load(B_Seqlen + cur_batch)
111 cur_batch_req_idx = cur_batch
112
113 off_q = cur_batch * stride_qbs + cur_head * stride_qh + offs_d
114 q = tl.load(Q + off_q, mask=mask_d, other=0.0)
115
116 kv_len_per_split = tl.cdiv(cur_batch_seq_len, NUM_KV_SPLITS)
117 split_kv_start = kv_len_per_split * split_kv_id
118 split_kv_end = tl.minimum(split_kv_start + kv_len_per_split, cur_batch_seq_len)
119
120 e_max = -float("inf")
121 e_sum = 0.0
122 acc = tl.zeros([BLOCK_DV], dtype=tl.float32)
123
124 if split_kv_end > split_kv_start:
125 ks = tl.load(k_scale)
126 vs = tl.load(v_scale)
127 for start_n in range(split_kv_start, split_kv_end, BLOCK_N):
128 offs_n = start_n + tl.arange(0, BLOCK_N)
129 kv_page_number = tl.load(
130 Req_to_tokens
131 + stride_req_to_tokens_b * cur_batch_req_idx
132 + offs_n // PAGE_SIZE,
133 mask=offs_n < split_kv_end,
134 other=0,
135 ).to(tl.int64) # page_number * page stride overflows int32
136 kv_in_page = offs_n % PAGE_SIZE
137 offs_buf_k = (
138 (kv_page_number * stride_buf_kpbs + kv_in_page * stride_buf_kbs)[
139 :, None
140 ]
141 + cur_kv_head * stride_buf_kh
142 + offs_d[None, :]
143 )
144 k = tl.load(
145 K_Buffer + offs_buf_k,
146 mask=(offs_n[:, None] < split_kv_end) & (mask_d[None, :]),
147 other=0.0,
148 )
149 if k.dtype.is_fp8():
150 k = (k.to(tl.float32) * ks).to(q.dtype)
151 qk = tl.sum(q[None, :] * k, 1)
152 qk *= sm_scale
153
154 if logit_cap > 0:
155 qk = logit_cap * tanh(qk / logit_cap)
156
157 qk = tl.where(offs_n < split_kv_end, qk, float("-inf"))
158
159 offs_buf_v = (
160 (kv_page_number * stride_buf_vpbs + kv_in_page * stride_buf_vbs)[
161 :, None
162 ]
163 + cur_kv_head * stride_buf_vh
164 + offs_dv[None, :]
165 )
166 v = tl.load(
167 V_Buffer + offs_buf_v,
168 mask=(offs_n[:, None] < split_kv_end) & (mask_dv[None, :]),
169 other=0.0,
170 )
171 if v.dtype.is_fp8():
172 v = (v.to(tl.float32) * vs).to(q.dtype)
173
174 n_e_max = tl.maximum(tl.max(qk, 0), e_max)
175 re_scale = tl.exp(e_max - n_e_max)
176 p = tl.exp(qk - n_e_max)
177 acc *= re_scale
178 acc += tl.sum(p[:, None] * v, 0)
179
180 e_sum = e_sum * re_scale + tl.sum(p, 0)
181 e_max = n_e_max
182
183 offs_mid_o = (
184 cur_batch * stride_mid_ob
185 + cur_head * stride_mid_oh
186 + split_kv_id * stride_mid_os
187 + offs_dv
188 )
189
190 tl.store(
191 Att_Out + offs_mid_o,
192 acc / e_sum,
193 mask=(mask_dv),
194 )
195
196 offs_mid_o_1 = (
197 cur_batch * stride_mid_ob
198 + cur_head * stride_mid_oh
199 + split_kv_id * stride_mid_os
200 + Lv
201 )
202
203 tl.store(
204 Att_Out + offs_mid_o_1,
205 e_max + tl.log(e_sum),
206 )
207
208
209 def _decode_att_m_fwd(
210 q,
211 k_buffer,
212 v_buffer,
213 att_out,
214 Req_to_tokens,
215 B_Seqlen,
216 num_kv_splits,
217 sm_scale,
218 page_size,
219 logit_cap,
220 k_scale,
221 v_scale,
222 ):
223 BLOCK = 64 if not is_hip_ else 8
224
225 NUM_KV_SPLITS = num_kv_splits
226 Lk = k_buffer.shape[-1]
227 Lv = v_buffer.shape[-1]
228
229 batch, head_num = q.shape[0], q.shape[1]
230
231 grid = (batch, head_num, NUM_KV_SPLITS)
232 kv_group_num = q.shape[1] // k_buffer.shape[-2]
233
234 num_warps = 4
235 if kv_group_num != 1:
236 num_warps = 1 if is_hip_ else 2
237
238 BLOCK_DMODEL = triton.next_power_of_2(Lk)
239 BLOCK_DV = triton.next_power_of_2(Lv)
240
241 _fwd_kernel_stage1[grid](
242 q,
243 k_buffer,
244 v_buffer,
245 sm_scale,
246 Req_to_tokens,
247 B_Seqlen,
248 att_out,
249 Req_to_tokens.stride(0),
250 q.stride(0),
251 q.stride(1),
252 _page_stride(k_buffer, page_size),
253 k_buffer.stride(-3), # Assume (..., PAGE_SIZE, NUM_HEADS, HEAD_DIM)
254 k_buffer.stride(-2), # Assume (..., PAGE_SIZE, NUM_HEADS, HEAD_DIM)
255 _page_stride(v_buffer, page_size),
256 v_buffer.stride(-3), # Assume (..., PAGE_SIZE, NUM_HEADS, HEAD_DIM)
257 v_buffer.stride(-2), # Assume (..., PAGE_SIZE, NUM_HEADS, HEAD_DIM)
258 att_out.stride(0),
259 att_out.stride(1),
260 att_out.stride(2),
261 k_scale,
262 v_scale,
263 kv_group_num=kv_group_num,
264 BLOCK_DMODEL=BLOCK_DMODEL,
265 BLOCK_DV=BLOCK_DV,
266 BLOCK_N=BLOCK,
267 NUM_KV_SPLITS=NUM_KV_SPLITS,
268 PAGE_SIZE=page_size,
269 logit_cap=logit_cap,
270 num_warps=num_warps,
271 num_stages=2,
272 Lk=Lk,
273 Lv=Lv,
274 )
275
276
277 @triton.jit
278 def _fwd_grouped_kernel_stage1(
279 Q,
280 K_Buffer,
281 V_Buffer,
282 sm_scale,
283 Req_to_tokens,
284 B_Seqlen,
285 Att_Out,
286 stride_req_to_tokens_b,
287 stride_qbs,
288 stride_qh,
289 stride_buf_kpbs,
290 stride_buf_kbs,
291 stride_buf_kh,
292 stride_buf_vpbs,
293 stride_buf_vbs,
294 stride_buf_vh,
295 stride_mid_ob,
296 stride_mid_oh,
297 stride_mid_os,
298 k_scale,
299 v_scale,
300 kv_group_num: tl.constexpr,
301 q_head_num: tl.constexpr,
302 BLOCK_DMODEL: tl.constexpr,
303 BLOCK_DPE: tl.constexpr,
304 BLOCK_DV: tl.constexpr,
305 BLOCK_N: tl.constexpr,
306 BLOCK_H: tl.constexpr,
307 NUM_KV_SPLITS: tl.constexpr,
308 PAGE_SIZE: tl.constexpr,
309 logit_cap: tl.constexpr,
310 Lk: tl.constexpr,
311 Lv: tl.constexpr,
312 IS_MLA: tl.constexpr = False,
313 ):
314 cur_batch = tl.program_id(0)
315 cur_head_id = tl.program_id(1)
316 cur_kv_head = cur_head_id // tl.cdiv(kv_group_num, BLOCK_H)
317 split_kv_id = tl.program_id(2)
318
319 VALID_BLOCK_H: tl.constexpr = BLOCK_H if kv_group_num > BLOCK_H else kv_group_num
320 cur_head = cur_head_id * VALID_BLOCK_H + tl.arange(0, BLOCK_H)
321 mask_h = cur_head < (cur_head_id + 1) * VALID_BLOCK_H
322 mask_h = mask_h & (cur_head < q_head_num)
323
324 offs_d = tl.arange(0, BLOCK_DMODEL)
325 offs_dv = tl.arange(0, BLOCK_DV)
326 mask_d = offs_d < Lk
327 mask_dv = offs_dv < Lv
328 cur_batch_seq_len = tl.load(B_Seqlen + cur_batch)
329 cur_batch_req_idx = cur_batch
330
331 offs_q = cur_batch * stride_qbs + cur_head[:, None] * stride_qh + offs_d[None, :]
332 q = tl.load(
333 Q + offs_q,
334 mask=(mask_h[:, None]) & (mask_d[None, :]),
335 other=0.0,
336 cache_modifier=".ca",
337 )
338
339 if BLOCK_DPE > 0:
340 offs_dpe = BLOCK_DMODEL + tl.arange(0, BLOCK_DPE)
341 mask_dpe = offs_dpe < Lk
342 off_qpe = (
343 cur_batch * stride_qbs + cur_head[:, None] * stride_qh + offs_dpe[None, :]
344 )
345 qpe = tl.load(
346 Q + off_qpe,
347 mask=(mask_h[:, None]) & (mask_dpe[None, :]),
348 other=0.0,
349 cache_modifier=".ca",
350 )
351
352 kv_len_per_split = tl.cdiv(cur_batch_seq_len, NUM_KV_SPLITS)
353 split_kv_start = kv_len_per_split * split_kv_id
354 split_kv_end = tl.minimum(split_kv_start + kv_len_per_split, cur_batch_seq_len)
355
356 e_max = tl.zeros([BLOCK_H], dtype=tl.float32) - float("inf")
357 e_sum = tl.zeros([BLOCK_H], dtype=tl.float32)
358 acc = tl.zeros([BLOCK_H, BLOCK_DV], dtype=tl.float32)
359
360 if split_kv_end > split_kv_start:
361 base_offs_k = cur_kv_head * stride_buf_kh + offs_d[:, None]
362 base_offs_v = cur_kv_head * stride_buf_vh + offs_dv[None, :]
363 if BLOCK_DPE > 0:
364 base_offs_kpe = cur_kv_head * stride_buf_kh + offs_dpe[:, None]
365
366 ks = tl.load(k_scale)
367 vs = tl.load(v_scale)
368 for start_n in tl.range(split_kv_start, split_kv_end, BLOCK_N):
369 offs_n = start_n + tl.arange(0, BLOCK_N)
370 kv_page_number = tl.load(
371 Req_to_tokens
372 + stride_req_to_tokens_b * cur_batch_req_idx
373 + offs_n // PAGE_SIZE,
374 mask=offs_n < split_kv_end,
375 other=0,
376 cache_modifier=".ca",
377 ).to(tl.int64) # page_number * page stride overflows int32
378 kv_off_k = (
379 kv_page_number * stride_buf_kpbs + (offs_n % PAGE_SIZE) * stride_buf_kbs
380 )
381
382 # explicitly facilitate overlapping load/compute
383 offs_buf_k = kv_off_k[None, :] + base_offs_k
384 k = tl.load(
385 K_Buffer + offs_buf_k,
386 mask=(offs_n[None, :] < split_kv_end) & (mask_d[:, None]),
387 other=0.0,
388 cache_modifier=".cg",
389 )
390
391 if k.dtype.is_fp8():
392 k = (k.to(tl.float32) * ks).to(q.dtype)
393 qk = tl.dot(q, k.to(q.dtype))
394 if BLOCK_DPE > 0:
395 offs_buf_kpe = kv_off_k[None, :] + base_offs_kpe
396 kpe = tl.load(
397 K_Buffer + offs_buf_kpe,
398 mask=(offs_n[None, :] < split_kv_end) & (mask_dpe[:, None]),
399 other=0.0,
400 cache_modifier=".cg",
401 )
402 if kpe.dtype.is_fp8():
403 kpe = (kpe.to(tl.float32) * ks).to(qpe.dtype)
404 qk += tl.dot(qpe, kpe.to(qpe.dtype))
405 qk *= sm_scale
406
407 if logit_cap > 0:
408 qk = logit_cap * tanh(qk / logit_cap)
409
410 qk = tl.where(
411 mask_h[:, None] & (offs_n[None, :] < split_kv_end), qk, float("-inf")
412 )
413
414 if not IS_MLA:
415 kv_off_v = (
416 kv_page_number * stride_buf_vpbs
417 + (offs_n % PAGE_SIZE) * stride_buf_vbs
418 )
419 offs_buf_v = kv_off_v[:, None] + base_offs_v
420 v = tl.load(
421 V_Buffer + offs_buf_v,
422 mask=(offs_n[:, None] < split_kv_end) & (mask_dv[None, :]),
423 other=0.0,
424 )
425 if v.dtype.is_fp8():
426 v = (v.to(tl.float32) * vs).to(q.dtype)
427 else:
428 # MLA uses a single c_kv.
429 # loading the same c_kv to interpret it as v is not necessary.
430 # transpose the existing c_kv (aka k) for the dot product.
431 v = tl.trans(k)
432
433 n_e_max = tl.maximum(tl.max(qk, 1), e_max)
434 re_scale = tl.exp(e_max - n_e_max)
435 p = tl.exp(qk - n_e_max[:, None])
436 acc *= re_scale[:, None]
437 acc += tl.dot(p.to(v.dtype), v)
438
439 e_sum = e_sum * re_scale + tl.sum(p, 1)
440 e_max = n_e_max
441
442 offs_mid_o = (
443 cur_batch * stride_mid_ob
444 + cur_head[:, None] * stride_mid_oh
445 + split_kv_id * stride_mid_os
446 + offs_dv[None, :]
447 )
448
449 tl.store(
450 Att_Out + offs_mid_o,
451 acc / e_sum[:, None],
452 mask=(mask_h[:, None]) & (mask_dv[None, :]),
453 )
454
455 offs_mid_o_1 = (
456 cur_batch * stride_mid_ob
457 + cur_head * stride_mid_oh
458 + split_kv_id * stride_mid_os
459 + Lv
460 )
461
462 tl.store(
463 Att_Out + offs_mid_o_1,
464 e_max + tl.log(e_sum),
465 mask=mask_h,
466 )
467
468
469 def _decode_grouped_att_m_fwd(
470 q,
471 k_buffer,
472 v_buffer,
473 att_out,
474 Req_to_tokens,
475 B_Seqlen,
476 num_kv_splits,
477 sm_scale,
478 page_size,
479 logit_cap,
480 k_scale,
481 v_scale,
482 is_mla=False,
483 ):
484 # with is_mla there is only a single c_kv in smem.
485 # could increase BLOCK or num_stages.
486 Lk = k_buffer.shape[-1]
487 Lv = v_buffer.shape[-1]
488
489 # Align tile dimensions with latent rank for MLA to avoid shape mismatch.
490 if is_mla:
491 if not is_hip_ and Lk == 576:
492 BLOCK_DMODEL = 512
493 BLOCK_DPE = 64
494 elif not is_hip_ and Lk == 288:
495 BLOCK_DMODEL = 256
496 BLOCK_DPE = 32
497 else:
498 BLOCK_DMODEL = triton.next_power_of_2(Lv)
499 BLOCK_DPE = triton.next_power_of_2(Lk - Lv) if Lk > Lv else 0
500 else:
501 BLOCK_DMODEL = triton.next_power_of_2(Lk)
502 BLOCK_DPE = 0
503 BLOCK_DV = triton.next_power_of_2(Lv)
504
505 BLOCK = 32
506 if is_hip_:
507 BLOCK = 16
508
509 batch, head_num = q.shape[0], q.shape[1]
510 kv_group_num = q.shape[1] // k_buffer.shape[-2]
511
512 BLOCK_H = 16
513 NUM_KV_SPLITS = num_kv_splits
514 grid = (
515 batch,
516 triton.cdiv(head_num, min(BLOCK_H, kv_group_num)),
517 NUM_KV_SPLITS,
518 )
519
520 extra_kargs = {}
521 num_stages = 2
522 if is_hip_:
523 # https://rocm.docs.amd.com/en/latest/how-to/rocm-for-ai/inference-optimization/workload.html#mi300x-triton-kernel-performance-optimization
524 # https://github.com/triton-lang/triton/blob/main/third_party/amd/backend/compiler.py
525 extra_kargs = {"waves_per_eu": 1, "matrix_instr_nonkdim": 16, "kpack": 2}
526 num_stages = 1
527 elif not is_hip_ and BLOCK_DMODEL >= 1024:
528 # Avoid shared memory overflow on NVIDIA when BLOCK_DMODEL is large
529 # like non-MLA D_QK=576, BLOCK_DMODEL=1024, BLOCK_H=16
530 # exceeds 101376 bytes limit
531 num_stages = 1
532
533 _fwd_grouped_kernel_stage1[grid](
534 q,
535 k_buffer,
536 v_buffer,
537 sm_scale,
538 Req_to_tokens,
539 B_Seqlen,
540 att_out,
541 Req_to_tokens.stride(0),
542 q.stride(0),
543 q.stride(1),
544 _page_stride(k_buffer, page_size),
545 k_buffer.stride(-3), # Assume (..., PAGE_SIZE, NUM_HEADS, HEAD_DIM)
546 k_buffer.stride(-2), # Assume (..., PAGE_SIZE, NUM_HEADS, HEAD_DIM)
547 _page_stride(v_buffer, page_size),
548 v_buffer.stride(-3), # Assume (..., PAGE_SIZE, NUM_HEADS, HEAD_DIM)
549 v_buffer.stride(-2), # Assume (..., PAGE_SIZE, NUM_HEADS, HEAD_DIM)
550 att_out.stride(0),
551 att_out.stride(1),
552 att_out.stride(2),
553 k_scale,
554 v_scale,
555 kv_group_num=kv_group_num,
556 q_head_num=head_num,
557 BLOCK_DMODEL=BLOCK_DMODEL,
558 BLOCK_DPE=BLOCK_DPE,
559 BLOCK_DV=BLOCK_DV,
560 BLOCK_N=BLOCK,
561 BLOCK_H=BLOCK_H,
562 NUM_KV_SPLITS=NUM_KV_SPLITS,
563 PAGE_SIZE=page_size,
564 logit_cap=logit_cap,
565 num_warps=4,
566 num_stages=num_stages,
567 Lk=Lk,
568 Lv=Lv,
569 IS_MLA=is_mla,
570 **extra_kargs,
571 )
572
573
574 @triton.jit
575 def _fwd_kernel_stage2(
576 Mid_O,
577 o,
578 lse,
579 B_Seqlen,
580 stride_mid_ob,
581 stride_mid_oh,
582 stride_mid_os,
583 stride_obs,
584 stride_oh,
585 stride_lse_bs,
586 NUM_KV_SPLITS: tl.constexpr,
587 BLOCK_DV: tl.constexpr,
588 Lv: tl.constexpr,
589 OUTPUT_FP16: tl.constexpr = 0,
590 ):
591 cur_batch = tl.program_id(0)
592 cur_head = tl.program_id(1)
593
594 cur_batch_seq_len = tl.load(B_Seqlen + cur_batch)
595
596 offs_d = tl.arange(0, BLOCK_DV)
597 mask_d = offs_d < Lv
598
599 e_sum = 0.0
600 e_max = -float("inf")
601 acc = tl.zeros([BLOCK_DV], dtype=tl.float32)
602
603 offs_v = cur_batch * stride_mid_ob + cur_head * stride_mid_oh + offs_d
604 offs_logic = cur_batch * stride_mid_ob + cur_head * stride_mid_oh + Lv
605
606 for split_kv_id in range(0, NUM_KV_SPLITS):
607 kv_len_per_split = tl.cdiv(cur_batch_seq_len, NUM_KV_SPLITS)
608 split_kv_start = kv_len_per_split * split_kv_id
609 split_kv_end = tl.minimum(split_kv_start + kv_len_per_split, cur_batch_seq_len)
610
611 if split_kv_end > split_kv_start:
612 tv = tl.load(
613 Mid_O + offs_v + split_kv_id * stride_mid_os, mask=mask_d, other=0.0
614 )
615 tlogic = tl.load(Mid_O + offs_logic + split_kv_id * stride_mid_os)
616 n_e_max = tl.maximum(tlogic, e_max)
617
618 old_scale = tl.exp(e_max - n_e_max)
619 acc *= old_scale
620 exp_logic = tl.exp(tlogic - n_e_max)
621 acc += exp_logic * tv
622
623 e_sum = e_sum * old_scale + exp_logic
624 e_max = n_e_max
625
626 result = acc / e_sum
627 if OUTPUT_FP16:
628 result = result.to(tl.float16)
629 tl.store(
630 o + cur_batch * stride_obs + cur_head * stride_oh + offs_d,
631 result,
632 mask=mask_d,
633 )
634 lse_val = e_max + tl.log(e_sum)
635 tl.store(
636 lse + cur_batch * stride_lse_bs + cur_head,
637 lse_val,
638 )
639
640
641 def _decode_softmax_reducev_fwd(
642 logits,
643 q,
644 o,
645 lse,
646 v_buffer,
647 b_seq_len,
648 num_kv_splits,
649 ):
650 batch, head_num = q.shape[0], q.shape[1]
651 Lv = v_buffer.shape[-1]
652 BLOCK_DV = triton.next_power_of_2(Lv)
653
654 NUM_KV_SPLITS = num_kv_splits
655
656 extra_kargs = {}
657 if is_hip_:
658 # https://rocm.docs.amd.com/en/docs-6.2.0/how-to/llm-fine-tuning-optimization/optimizing-triton-kernel.html
659 # https://github.com/triton-lang/triton/blob/main/third_party/amd/backend/compiler.py
660 extra_kargs = {"waves_per_eu": 4, "matrix_instr_nonkdim": 16, "kpack": 2}
661
662 grid = (batch, head_num)
663 _fwd_kernel_stage2[grid](
664 logits,
665 o,
666 lse,
667 b_seq_len,
668 logits.stride(0),
669 logits.stride(1),
670 logits.stride(2),
671 o.stride(0),
672 o.stride(1),
673 lse.stride(0),
674 NUM_KV_SPLITS=NUM_KV_SPLITS,
675 BLOCK_DV=BLOCK_DV,
676 Lv=Lv,
677 num_warps=4,
678 num_stages=2,
679 **extra_kargs,
680 )
681
682
683 def decode_attention_fwd_normal(
684 q,
685 k_buffer,
686 v_buffer,
687 o,
688 lse,
689 req_to_token,
690 b_seq_len,
691 attn_logits,
692 num_kv_splits,
693 sm_scale,
694 page_size,
695 logit_cap=0.0,
696 k_scale=None,
697 v_scale=None,
698 ):
699 _decode_att_m_fwd(
700 q,
701 k_buffer,
702 v_buffer,
703 attn_logits,
704 req_to_token,
705 b_seq_len,
706 num_kv_splits,
707 sm_scale,
708 page_size,
709 logit_cap,
710 k_scale,
711 v_scale,
712 )
713 _decode_softmax_reducev_fwd(
714 attn_logits, q, o, lse, v_buffer, b_seq_len, num_kv_splits
715 )
716
717
718 def decode_attention_fwd_grouped(
719 q,
720 k_buffer,
721 v_buffer,
722 o,
723 lse,
724 req_to_token,
725 b_seq_len,
726 attn_logits,
727 num_kv_splits,
728 sm_scale,
729 page_size,
730 logit_cap=0.0,
731 k_scale=None,
732 v_scale=None,
733 is_mla=False,
734 ):
735 _decode_grouped_att_m_fwd(
736 q,
737 k_buffer,
738 v_buffer,
739 attn_logits,
740 req_to_token,
741 b_seq_len,
742 num_kv_splits,
743 sm_scale,
744 page_size,
745 logit_cap,
746 k_scale,
747 v_scale,
748 is_mla=is_mla,
749 )
750 _decode_softmax_reducev_fwd(
751 attn_logits, q, o, lse, v_buffer, b_seq_len, num_kv_splits
752 )
753
754
755 def decode_attention_fwd(
756 q,
757 k_buffer,
758 v_buffer,
759 o,
760 lse,
761 req_to_token,
762 b_seq_len,
763 attn_logits,
764 num_kv_splits,
765 sm_scale,
766 page_size=1,
767 logit_cap=0.0,
768 k_scale=None,
769 v_scale=None,
770 is_mla=False,
771 ):
772 assert num_kv_splits == attn_logits.shape[2]
773
774 if k_scale is None:
775 k_scale = torch.tensor(1.0, dtype=torch.float32, device=q.device)
776 if v_scale is None:
777 v_scale = torch.tensor(1.0, dtype=torch.float32, device=q.device)
778
779 kv_group_num = q.shape[1] // v_buffer.shape[-2]
780
781 if kv_group_num == 1:
782 # MHA
783 decode_attention_fwd_normal(
784 q,
785 k_buffer,
786 v_buffer,
787 o,
788 lse,
789 req_to_token,
790 b_seq_len,
791 attn_logits,
792 num_kv_splits,
793 sm_scale,
794 page_size,
795 logit_cap,
796 k_scale,
797 v_scale,
798 )
799 else:
800 # GQA/MQA/MLA
801 decode_attention_fwd_grouped(
802 q,
803 k_buffer,
804 v_buffer,
805 o,
806 lse,
807 req_to_token,
808 b_seq_len,
809 attn_logits,
810 num_kv_splits,
811 sm_scale,
812 page_size,
813 logit_cap,
814 k_scale,
815 v_scale,
816 is_mla=is_mla,
817 )