@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
py 249 lines 8.27 KB
Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 # SPDX-FileCopyrightText: Songlin Yang, Yu Zhang, Zhiyuan Li
4 #
5 # This file contains code adapted from the flash-linear-attention project.
6 # The original source code was licensed under the MIT license and included
7 # the following copyright notice:
8 # Copyright (c) 2023-2026, Songlin Yang, Yu Zhang, Zhiyuan Li
9
10
11 import torch
12
13 from vllm import _custom_ops as ops
14 from vllm.platforms import current_platform
15 from vllm.triton_utils import tl, triton
16
17
18 # Consumed by kimi_k3_triton_warmup.py during kernel_warmup().
19 def get_attn_res_triton_warmup_profiles(
20 max_blocks: int,
21 ) -> tuple[tuple[int, bool, int, bool], ...]:
22 """Return the small-batch profiles that bypass the native kernel."""
23 profiles = [
24 (num_blocks, False, -1, True) for num_blocks in range(2, max_blocks + 1)
25 ]
26 profiles.extend(
27 (block_write_idx, True, block_write_idx, True)
28 for block_write_idx in range(2, max_blocks)
29 )
30 profiles.append((max_blocks, True, -1, False))
31 return tuple(profiles)
32
33
34 @triton.jit
35 def _attn_res_kernel(
36 prefix_ptr,
37 delta_ptr,
38 blocks_ptr,
39 norm_weight_ptr,
40 qk_weight_ptr,
41 output_norm_weight_ptr,
42 output_ptr,
43 stride_prefix_m: tl.constexpr,
44 stride_delta_m: tl.constexpr,
45 stride_block_m: tl.constexpr,
46 stride_block_r: tl.constexpr,
47 stride_output_m: tl.constexpr,
48 num_blocks: tl.constexpr,
49 hidden_size: tl.constexpr,
50 block_write_idx: tl.constexpr,
51 eps: tl.constexpr,
52 output_norm_eps: tl.constexpr,
53 HAS_DELTA: tl.constexpr,
54 WRITE_BLOCK: tl.constexpr,
55 APPLY_OUTPUT_NORM: tl.constexpr,
56 BLOCK_L: tl.constexpr,
57 BLOCK_D: tl.constexpr,
58 launch_pdl: tl.constexpr,
59 ):
60 row_idx = tl.program_id(0).to(tl.int64)
61 d_offsets = tl.max_contiguous(tl.arange(0, BLOCK_D), BLOCK_D)
62 d_mask = d_offsets < hidden_size
63
64 if launch_pdl:
65 tl.extra.cuda.gdc_wait()
66
67 updated_prefix = tl.load(
68 prefix_ptr + row_idx * stride_prefix_m + d_offsets,
69 mask=d_mask,
70 other=0.0,
71 ).to(tl.float32)
72 if HAS_DELTA:
73 delta = tl.load(
74 delta_ptr + row_idx * stride_delta_m + d_offsets,
75 mask=d_mask,
76 other=0.0,
77 ).to(tl.float32)
78 updated_prefix += delta
79 # Match the BF16 prefix-add result before using it as a residual source.
80 updated_prefix = updated_prefix.to(prefix_ptr.dtype.element_ty).to(tl.float32)
81 tl.store(
82 prefix_ptr + row_idx * stride_prefix_m + d_offsets,
83 updated_prefix,
84 mask=d_mask,
85 )
86 if WRITE_BLOCK:
87 tl.store(
88 blocks_ptr
89 + row_idx * stride_block_m
90 + block_write_idx * stride_block_r
91 + d_offsets,
92 updated_prefix,
93 mask=d_mask,
94 )
95 # With only the prefix source, the AttnRes softmax is exactly one.
96 if num_blocks == 0:
97 mixed = updated_prefix
98 else:
99 # Reloading avoids keeping the full prefix vector live across the loop.
100 if HAS_DELTA:
101 tl.debug_barrier()
102 input_qk_weight = tl.load(
103 norm_weight_ptr + d_offsets, mask=d_mask, other=0.0
104 ).to(tl.float32) * tl.load(
105 qk_weight_ptr + d_offsets, mask=d_mask, other=0.0
106 ).to(tl.float32)
107 max_logit = tl.full((), -float("inf"), tl.float32)
108 denominator = tl.zeros((), tl.float32)
109 mixed = tl.zeros((BLOCK_D,), tl.float32)
110
111 num_sources = num_blocks + 1
112 for source_tile in range(tl.cdiv(num_sources, BLOCK_L)):
113 source_offsets = source_tile * BLOCK_L + tl.arange(0, BLOCK_L)
114 source_mask = source_offsets < num_sources
115 is_prefix = source_offsets == num_blocks
116 block_ptrs = (
117 blocks_ptr
118 + row_idx * stride_block_m
119 + source_offsets[:, None] * stride_block_r
120 + d_offsets[None, :]
121 )
122 prefix_ptrs = (
123 prefix_ptr
124 + row_idx * stride_prefix_m
125 + source_offsets[:, None] * 0
126 + d_offsets[None, :]
127 )
128 value_ptrs = tl.where(is_prefix[:, None], prefix_ptrs, block_ptrs)
129 values = tl.load(
130 value_ptrs,
131 mask=source_mask[:, None] & d_mask[None, :],
132 other=0.0,
133 eviction_policy="evict_first",
134 ).to(tl.float32)
135 reciprocal_std = tl.rsqrt(
136 tl.sum(values * values, axis=1) * (1.0 / hidden_size) + eps
137 )
138 logits = tl.sum(values * input_qk_weight[None, :], axis=1) * reciprocal_std
139 scores = tl.where(source_mask, logits, -float("inf"))
140
141 new_max_logit = tl.maximum(max_logit, tl.max(scores, axis=0))
142 old_scale = tl.exp(max_logit - new_max_logit)
143 block_scales = tl.exp(scores - new_max_logit)
144 denominator = denominator * old_scale + tl.sum(block_scales, axis=0)
145 mixed = mixed * old_scale + tl.sum(block_scales[:, None] * values, axis=0)
146 max_logit = new_max_logit
147
148 mixed /= denominator
149 output = mixed
150
151 if launch_pdl:
152 tl.extra.cuda.gdc_launch_dependents()
153
154 if APPLY_OUTPUT_NORM:
155 output_reciprocal_std = tl.rsqrt(
156 tl.sum(tl.where(d_mask, mixed * mixed, 0.0), axis=0) * (1.0 / hidden_size)
157 + output_norm_eps
158 )
159 output_norm_weight = tl.load(
160 output_norm_weight_ptr + d_offsets, mask=d_mask, other=0.0
161 ).to(tl.float32)
162 output = mixed * output_reciprocal_std * output_norm_weight
163 tl.store(
164 output_ptr + row_idx * stride_output_m + d_offsets,
165 output,
166 mask=d_mask,
167 )
168
169
170 def attn_res(
171 prefix: torch.Tensor,
172 delta: torch.Tensor | None,
173 blocks: torch.Tensor,
174 norm_weight: torch.Tensor,
175 qk_weight: torch.Tensor,
176 output_norm_weight: torch.Tensor | None,
177 num_blocks: int,
178 block_write_idx: int,
179 eps: float,
180 output_norm_eps: float,
181 ) -> torch.Tensor:
182 num_tokens, hidden_size = prefix.shape
183 assert prefix.stride(-1) == 1
184 assert delta is None or delta.stride(-1) == 1
185 assert blocks.stride(-1) == 1
186 assert norm_weight.stride(-1) == 1
187 assert qk_weight.stride(-1) == 1
188 assert output_norm_weight is None or output_norm_weight.stride(-1) == 1
189 # The native kernel covers every Kimi-K3 AttnRes variant on dense SM100
190 # inputs. The op is only compiled under CUDA >= 13, so a device check alone
191 # is not enough to know it exists.
192 if (
193 hidden_size == 7168
194 and prefix.stride(0) == hidden_size
195 and (delta is None or delta.stride(0) == hidden_size)
196 and 0 <= num_blocks <= 8
197 and num_blocks <= blocks.shape[1]
198 and -1 <= block_write_idx < blocks.shape[1]
199 and current_platform.is_device_capability_family(100)
200 and hasattr(torch.ops._C, "kimi_k3_attn_res")
201 ):
202 return ops.kimi_k3_attn_res(
203 prefix,
204 delta,
205 blocks,
206 norm_weight,
207 qk_weight,
208 output_norm_weight,
209 num_blocks,
210 block_write_idx,
211 eps,
212 output_norm_eps,
213 )
214 output = prefix.new_empty(prefix.shape)
215 # Tuned on GB300: source tiling helps decode, while one-source tiles scale
216 # better for prefill.
217 # Keep get_attn_res_triton_warmup_profiles in sync with these fallbacks.
218 if num_tokens >= 256 or num_blocks <= 1:
219 block_l, num_warps = 1, 4
220 else:
221 block_l, num_warps = 4, 8
222 _attn_res_kernel[(num_tokens,)](
223 prefix,
224 delta,
225 blocks,
226 norm_weight,
227 qk_weight,
228 output_norm_weight,
229 output,
230 prefix.stride(0),
231 0 if delta is None else delta.stride(0),
232 blocks.stride(0),
233 blocks.stride(1),
234 output.stride(0),
235 num_blocks,
236 hidden_size,
237 block_write_idx,
238 eps,
239 output_norm_eps,
240 HAS_DELTA=delta is not None,
241 WRITE_BLOCK=block_write_idx >= 0,
242 APPLY_OUTPUT_NORM=output_norm_weight is not None,
243 BLOCK_L=block_l,
244 BLOCK_D=triton.next_power_of_2(hidden_size),
245 num_warps=num_warps,
246 num_stages=2,
247 launch_pdl=current_platform.is_arch_support_pdl(),
248 )
249 return output