@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
py 298 lines 9.46 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/97cb762bb65ebf05025eb342de03c184660427a3/python/sglang/srt/layers/attention/triton_ops/prefill_attention.py
6 # Changes:
7 # - Add support for sliding window attention
8
9 # Copyright 2023-2024 SGLang Team
10 # Licensed under the Apache License, Version 2.0 (the "License");
11 # you may not use this file except in compliance with the License.
12 # You may obtain a copy of the License at
13 #
14 # http://www.apache.org/licenses/LICENSE-2.0
15 #
16 # Unless required by applicable law or agreed to in writing, software
17 # distributed under the License is distributed on an "AS IS" BASIS,
18 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
19 # See the License for the specific language governing permissions and
20 # limitations under the License.
21 # ==============================================================================
22 """Memory-efficient attention for prefill.
23 It supports page size = 1.
24 """
25
26 # Adapted from
27 # https://github.com/ModelTC/lightllm/blob/f2a54f0912293f683bf1d1695fd12c4098a5bf82/lightllm/models/llama/triton_kernel/context_flashattention_nopad.py#L1
28 import torch
29
30 from vllm.platforms import current_platform
31 from vllm.triton_utils import tl, triton
32 from vllm.utils.math_utils import RCP_LN2
33
34
35 def _prefer_narrow_kv_tile() -> bool:
36 """RDNA3/RDNA4 prefer a narrower KV tile than the shared default."""
37 if not current_platform.is_rocm():
38 return False
39 from vllm.platforms.rocm import on_gfx1x
40
41 # RDNA3/RDNA4 only: on_gfx1x() excludes gfx10xx and gfx1250.
42 return on_gfx1x()
43
44
45 @triton.jit
46 def _fwd_kernel(
47 Q,
48 K,
49 V,
50 Sinks,
51 sm_scale,
52 B_Start_Loc,
53 B_Seqlen,
54 Out,
55 stride_qbs,
56 stride_qh,
57 stride_kbs,
58 stride_kh,
59 stride_vbs,
60 stride_vh,
61 stride_obs,
62 stride_oh,
63 kv_group_num: tl.constexpr,
64 BLOCK_M: tl.constexpr,
65 BLOCK_DMODEL: tl.constexpr,
66 BLOCK_N: tl.constexpr,
67 IS_CAUSAL: tl.constexpr,
68 SLIDING_WINDOW_Q: tl.constexpr,
69 SLIDING_WINDOW_K: tl.constexpr,
70 SINKS_BIAS_KEY0: tl.constexpr,
71 USE_SINKS: tl.constexpr,
72 Lk: tl.constexpr,
73 ):
74 cur_batch = tl.program_id(0)
75 cur_head = tl.program_id(1)
76 start_m = tl.program_id(2)
77
78 cur_kv_head = cur_head // kv_group_num
79
80 cur_batch_seq_len = tl.load(B_Seqlen + cur_batch)
81 cur_batch_in_all_start_index = tl.load(B_Start_Loc + cur_batch)
82
83 block_start_loc = BLOCK_M * start_m
84
85 # initialize offsets
86 offs_n = tl.arange(0, BLOCK_N)
87 offs_d = tl.arange(0, BLOCK_DMODEL)
88 offs_m = start_m * BLOCK_M + tl.arange(0, BLOCK_M)
89 off_q = (
90 (cur_batch_in_all_start_index + offs_m[:, None]) * stride_qbs
91 + cur_head * stride_qh
92 + offs_d[None, :]
93 )
94 off_k = offs_n[None, :] * stride_kbs + cur_kv_head * stride_kh + offs_d[:, None]
95 off_v = offs_n[:, None] * stride_vbs + cur_kv_head * stride_vh + offs_d[None, :]
96
97 mask_d = offs_d < Lk
98
99 q = tl.load(
100 Q + off_q,
101 mask=(offs_m[:, None] < cur_batch_seq_len) & (mask_d[None, :]),
102 other=0.0,
103 )
104
105 k_ptrs = K + off_k
106 v_ptrs = V + off_v
107
108 # initialize pointer to m and l
109 if USE_SINKS:
110 sink = tl.load(Sinks + cur_head) * 1.4426950408889634
111 if SINKS_BIAS_KEY0:
112 # Sinks bias the logit of key 0, so the softmax starts empty and
113 # normalizes over the biased scores as usual.
114 m_i = tl.zeros([BLOCK_M], dtype=tl.float32) - float("inf")
115 l_i = tl.zeros([BLOCK_M], dtype=tl.float32)
116 else:
117 # Sinks are a null logit that only inflates the denominator.
118 m_i = tl.full([BLOCK_M], sink, dtype=tl.float32)
119 l_i = tl.full([BLOCK_M], 1.0, dtype=tl.float32)
120 else:
121 m_i = tl.zeros([BLOCK_M], dtype=tl.float32) - float("inf")
122 l_i = tl.zeros([BLOCK_M], dtype=tl.float32)
123 acc = tl.zeros([BLOCK_M, BLOCK_DMODEL], dtype=tl.float32)
124
125 block_mask = tl.where(block_start_loc < cur_batch_seq_len, 1, 0)
126
127 # Calculate the end position for attention computation
128 end_n = cur_batch_seq_len
129
130 # Apply causal attention pruning and sliding window attention pruning
131 end_n = tl.minimum(end_n, (start_m + 1) * BLOCK_M) if IS_CAUSAL else end_n
132
133 # Calculate the start position for backward sliding window.
134 # Keys outside the window are fully masked below, so skip those blocks
135 # rather than loading them and discarding the result. Without this a
136 # windowed pass still costs O(seq_len^2).
137 start_n_limit = 0
138 if SLIDING_WINDOW_Q > 0:
139 first_needed = start_m * BLOCK_M - SLIDING_WINDOW_Q
140 start_n_limit = tl.maximum(0, (first_needed // BLOCK_N) * BLOCK_N)
141 if SLIDING_WINDOW_K > 0:
142 last_needed = (start_m + 1) * BLOCK_M - 1 + SLIDING_WINDOW_K
143 end_n = tl.minimum(end_n, last_needed + 1)
144 end_n_limit = block_mask * end_n
145
146 for start_n in range(start_n_limit, end_n_limit, BLOCK_N):
147 # -- prepare attention mask ----
148 # Position indices in the sequence
149 pos_q = offs_m[:, None] # Query positions [BLOCK_M, 1]
150 pos_k = start_n + offs_n[None, :] # Key positions [1, BLOCK_N]
151
152 # Valid sequence mask
153 mask = pos_k < cur_batch_seq_len
154 # Causal mask
155 if IS_CAUSAL:
156 mask &= pos_q >= pos_k
157
158 # Bidirectional sliding window masks
159 sliding_mask_q = (
160 pos_q - pos_k <= SLIDING_WINDOW_Q if SLIDING_WINDOW_Q > 0 else None
161 )
162 sliding_mask_k = (
163 pos_k - pos_q <= SLIDING_WINDOW_K if SLIDING_WINDOW_K > 0 else None
164 )
165 if sliding_mask_q is not None:
166 mask &= sliding_mask_q
167 if sliding_mask_k is not None:
168 mask &= sliding_mask_k
169
170 start_n = tl.multiple_of(start_n, BLOCK_N)
171 # -- compute qk ----
172 k = tl.load(
173 k_ptrs + (cur_batch_in_all_start_index + start_n) * stride_kbs,
174 mask=(pos_k < cur_batch_seq_len) & (mask_d[:, None]),
175 other=0.0,
176 )
177
178 qk = tl.dot(q, k)
179 qk = tl.where(mask, qk * sm_scale, -1.0e8)
180 if USE_SINKS and SINKS_BIAS_KEY0:
181 qk = tl.where(mask & (pos_k == 0), qk + sink, qk)
182 m_ij = tl.maximum(m_i, tl.max(qk, 1))
183 qk -= m_ij[:, None]
184 p = tl.math.exp2(qk)
185 l_ij = tl.sum(p, 1)
186
187 # -- update m_i and l_i
188 alpha = tl.math.exp2(m_i - m_ij)
189 l_i = l_i * alpha + l_ij
190 # -- update output accumulator --
191 acc = acc * alpha[:, None]
192 # update acc
193 v = tl.load(
194 v_ptrs + (cur_batch_in_all_start_index + start_n) * stride_vbs,
195 mask=((start_n + offs_n[:, None]) < cur_batch_seq_len) & (mask_d[None, :]),
196 other=0.0,
197 )
198 p = p.to(v.dtype)
199 acc = tl.dot(p, v, acc)
200 # update m_i
201 m_i = m_ij
202
203 acc = acc / l_i[:, None]
204 off_o = (
205 (cur_batch_in_all_start_index + offs_m[:, None]) * stride_obs
206 + cur_head * stride_oh
207 + offs_d[None, :]
208 )
209 out_ptrs = Out + off_o
210 tl.store(
211 out_ptrs, acc, mask=(offs_m[:, None] < cur_batch_seq_len) & (mask_d[None, :])
212 )
213
214
215 def get_block_size(dtype: torch.dtype) -> int:
216 if dtype == torch.float32:
217 return 32
218 elif current_platform.is_cuda_alike() and current_platform.has_device_capability(
219 80
220 ):
221 return 128
222 else:
223 return 64
224
225
226 def context_attention_fwd(
227 q: torch.Tensor,
228 k: torch.Tensor,
229 v: torch.Tensor,
230 o: torch.Tensor,
231 b_start_loc: torch.Tensor,
232 b_seq_len: torch.Tensor,
233 max_input_len: int,
234 is_causal: bool = True,
235 softmax_scale: float | None = None,
236 sliding_window_q: int | None = None,
237 sliding_window_k: int | None = None,
238 sinks: torch.Tensor | None = None,
239 sinks_bias_key0: bool = False,
240 ):
241 """q, k, v: [b * s, head, head_dim]
242 b_start_loc: [b]
243 b_seq_len: [b]
244 out: [b * s, head, head_dim]
245 """
246 BLOCK = get_block_size(q.dtype)
247
248 Lq, Lk, _ = q.shape[-1], k.shape[-1], v.shape[-1]
249
250 sm_scale = 1.0 / (Lq**0.5) if softmax_scale is None else softmax_scale
251 # rescale with 1/ln(2) for triton exp2
252 sm_scale *= RCP_LN2
253
254 batch, head = b_seq_len.shape[0], q.shape[1]
255 kv_group_num = q.shape[1] // k.shape[1]
256 if sinks is not None:
257 assert sinks.shape[0] == head, "Sinks must be num_query_heads size"
258
259 grid = (batch, head, triton.cdiv(max_input_len, BLOCK))
260 num_warps = 4 if Lk <= 64 else 8
261
262 # BLOCK_M, num_warps and num_stages stay at the shared defaults; min()
263 # leaves dtypes whose default tile is already 32, such as float32, alone.
264 BLOCK_N = min(BLOCK, 32) if _prefer_narrow_kv_tile() else BLOCK
265
266 sliding_window_q = sliding_window_q if sliding_window_q is not None else 0
267 sliding_window_k = sliding_window_k if sliding_window_k is not None else 0
268
269 _fwd_kernel[grid](
270 q,
271 k,
272 v,
273 sinks if sinks is not None else q,
274 sm_scale,
275 b_start_loc,
276 b_seq_len,
277 o,
278 q.stride(0),
279 q.stride(1),
280 k.stride(0),
281 k.stride(1),
282 v.stride(0),
283 v.stride(1),
284 o.stride(0),
285 o.stride(1),
286 kv_group_num=kv_group_num,
287 BLOCK_M=BLOCK,
288 BLOCK_DMODEL=triton.next_power_of_2(Lk),
289 BLOCK_N=BLOCK_N,
290 IS_CAUSAL=is_causal,
291 SLIDING_WINDOW_Q=sliding_window_q,
292 SLIDING_WINDOW_K=sliding_window_k,
293 USE_SINKS=sinks is not None,
294 SINKS_BIAS_KEY0=sinks_bias_key0,
295 num_warps=num_warps,
296 num_stages=1,
297 Lk=Lk,
298 )