@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / bc7c4d206b
Fork 0

A high-throughput and memory-efficient inference and serving engine for LLMs

Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git

Synced Nov 15, 13:30 UTC

[Kernel][ROCM] Upstream prefix prefill speed up for vLLM V1 (#13305)

Signed-off-by: Sage Moore <sage@neuralmagic.com> Signed-off-by: root <root@banff-cyxtera-s73-5.ctr.dcgpu> Signed-off-by: Aleksandr Malyshev <maleksan@amd.com> Signed-off-by: root <root@banff-cyxtera-s65-4.amd.com> Signed-off-by: maleksan85 <maleksan@amd.com> Signed-off-by: <> Co-authored-by: Sage Moore <sage@neuralmagic.com> Co-authored-by: root <root@banff-cyxtera-s73-5.ctr.dcgpu> Co-authored-by: Aleksandr Malyshev <maleksan@amd.com> Co-authored-by: qli88 <qiang.li2@amd.com> Co-authored-by: root <root@banff-cyxtera-s65-4.amd.com>

Aleksandr Malyshev committed Apr 22, 2025 at 19:11 UTC bc7c4d206bbfb56b06d218b6c2971e8ca191db36
2 files changed +824 -816
tests/core/block/e2e/test_correctness.py
+3 -3
@@ -195,15 +195,15 @@ def test_lookahead_greedy_equality_with_preemption(baseline_llm_generator,
195 ])
196 @pytest.mark.parametrize("per_test_common_llm_kwargs",
197 [{
198 - "block_size": 8,
198 + "block_size": 16,
199 "max_num_batched_tokens": 2,
200 "max_num_seqs": 2,
201 }, {
202 - "block_size": 8,
202 + "block_size": 16,
203 "max_num_batched_tokens": 3,
204 "max_num_seqs": 2,
205 }, {
206 - "block_size": 8,
206 + "block_size": 16,
207 "max_num_batched_tokens": 256,
208 "max_num_seqs": 10,
209 }])
vllm/attention/ops/prefix_prefill.py
+821 -813
@@ -16,831 +16,778 @@ NUM_WARPS = 4 if current_platform.is_rocm() else 8
16 # To check compatibility
17 IS_TURING = current_platform.get_device_capability() == (7, 5)
18
19 -if triton.__version__ >= "2.1.0":
20 -
21 - @triton.jit
22 - def _fwd_kernel(
23 - Q,
24 - K,
25 - V,
26 - K_cache,
27 - V_cache,
28 - B_Loc,
29 - sm_scale,
30 - k_scale,
31 - v_scale,
32 - B_Start_Loc,
33 - B_Seqlen,
34 - block_size,
35 - x,
36 - Out,
37 - stride_b_loc_b,
38 - stride_b_loc_s,
39 - stride_qbs,
40 - stride_qh,
41 - stride_qd,
42 - stride_kbs,
43 - stride_kh,
44 - stride_kd,
45 - stride_vbs,
46 - stride_vh,
47 - stride_vd,
48 - stride_obs,
49 - stride_oh,
50 - stride_od,
51 - stride_k_cache_bs,
52 - stride_k_cache_h,
53 - stride_k_cache_d,
54 - stride_k_cache_bl,
55 - stride_k_cache_x,
56 - stride_v_cache_bs,
57 - stride_v_cache_h,
58 - stride_v_cache_d,
59 - stride_v_cache_bl,
60 - num_queries_per_kv: int,
61 - IN_PRECISION: tl.constexpr,
62 - BLOCK_M: tl.constexpr,
63 - BLOCK_DMODEL: tl.constexpr, # head size
64 - BLOCK_DMODEL_PADDED: tl.constexpr, # head size padded to a power of 2
65 - BLOCK_N: tl.constexpr,
66 - SLIDING_WINDOW: tl.constexpr,
67 - SKIP_DECODE: tl.constexpr,
68 - ):
69 -
70 - cur_batch = tl.program_id(0)
71 - cur_head = tl.program_id(1)
72 - start_m = tl.program_id(2)
73 -
74 - cur_kv_head = cur_head // num_queries_per_kv
75 -
76 - cur_batch_seq_len = tl.load(B_Seqlen + cur_batch)
77 - cur_batch_in_all_start_index = tl.load(B_Start_Loc + cur_batch)
78 - cur_batch_in_all_stop_index = tl.load(B_Start_Loc + cur_batch + 1)
79 - cur_batch_query_len = (cur_batch_in_all_stop_index -
80 - cur_batch_in_all_start_index)
81 - cur_batch_ctx_len = cur_batch_seq_len - cur_batch_query_len
82 -
83 - if SKIP_DECODE and cur_batch_query_len == 1:
84 - return
85 -
86 - # start position inside of the query
87 - # generally, N goes over kv, while M goes over query_len
88 - block_start_loc = BLOCK_M * start_m
89 -
90 - # initialize offsets
91 - # [N]; starts at 0
92 - offs_n = tl.arange(0, BLOCK_N)
93 - # [D]; starts at 0
94 - offs_d = tl.arange(0, BLOCK_DMODEL_PADDED)
95 - # [M]; starts at current position in query
96 - offs_m = start_m * BLOCK_M + tl.arange(0, BLOCK_M)
97 - # [M,D]
98 - off_q = (
99 - (cur_batch_in_all_start_index + offs_m[:, None]) * stride_qbs +
100 - cur_head * stride_qh + offs_d[None, :] * stride_qd)
101 -
102 - dim_mask = tl.where(
103 - tl.arange(0, BLOCK_DMODEL_PADDED) < BLOCK_DMODEL, 1,
104 - 0).to(tl.int1) # [D]
105 -
106 - q = tl.load(Q + off_q,
107 - mask=dim_mask[None, :] &
108 - (offs_m[:, None] < cur_batch_query_len),
109 - other=0.0) # [M,D]
110 -
111 - # initialize pointer to m and l
112 - m_i = tl.zeros([BLOCK_M], dtype=tl.float32) - float("inf") # [M]
113 - l_i = tl.zeros([BLOCK_M], dtype=tl.float32) # [M]
114 - acc = tl.zeros([BLOCK_M, BLOCK_DMODEL_PADDED],
115 - dtype=tl.float32) # [M,D]
116 -
117 - # compute query against context (no causal mask here)
118 - for start_n in range(0, cur_batch_ctx_len, BLOCK_N):
119 - start_n = tl.multiple_of(start_n, BLOCK_N)
120 - # -- compute qk ----
121 - bn = tl.load(B_Loc + cur_batch * stride_b_loc_b +
122 - ((start_n + offs_n) // block_size) * stride_b_loc_s,
123 - mask=(start_n + offs_n) < cur_batch_ctx_len,
124 - other=0) # [N]
125 - # [D,N]
126 - off_k = (bn[None, :] * stride_k_cache_bs +
127 - cur_kv_head * stride_k_cache_h +
128 - (offs_d[:, None] // x) * stride_k_cache_d +
129 - ((start_n + offs_n[None, :]) % block_size) *
130 - stride_k_cache_bl +
131 - (offs_d[:, None] % x) * stride_k_cache_x)
132 - # [N,D]
133 - off_v = (
134 - bn[:, None] * stride_v_cache_bs +
135 - cur_kv_head * stride_v_cache_h +
136 - offs_d[None, :] * stride_v_cache_d +
137 - (start_n + offs_n[:, None]) % block_size * stride_v_cache_bl)
138 - k_load = tl.load(K_cache + off_k,
139 - mask=dim_mask[:, None] &
140 - ((start_n + offs_n[None, :]) < cur_batch_ctx_len),
141 - other=0.0) # [D,N]
142 -
143 - if k_load.dtype.is_fp8():
144 - k = (k_load.to(tl.float32) * tl.load(k_scale)).to(q.dtype)
145 - else:
146 - k = k_load
147 -
148 - qk = tl.zeros([BLOCK_M, BLOCK_N], dtype=tl.float32) # [M,N]
149 - qk = tl.dot(q, k, acc=qk, input_precision=IN_PRECISION)
150 - qk = tl.where((start_n + offs_n[None, :]) < cur_batch_ctx_len, qk,
151 - float("-inf"))
152 - qk *= sm_scale
153 - if SLIDING_WINDOW > 0:
154 - # (cur_batch_ctx_len + offs_m[:, None]) are the positions of
155 - # Q entries in sequence
156 - # (start_n + offs_n[None, :]) are the positions of
157 - # KV entries in sequence
158 - # So the condition makes sure each entry in Q only attends
159 - # to KV entries not more than SLIDING_WINDOW away.
160 - #
161 - # We can't use -inf here, because the
162 - # sliding window may lead to the entire row being masked.
163 - # This then makes m_ij contain -inf, which causes NaNs in
164 - # exp().
165 - qk = tl.where((cur_batch_ctx_len + offs_m[:, None]) -
166 - (start_n + offs_n[None, :]) < SLIDING_WINDOW, qk,
167 - -10000)
168 -
169 - # -- compute m_ij, p, l_ij
170 - m_ij = tl.max(qk, 1) # [M]
171 - p = tl.exp(qk - m_ij[:, None]) # [M,N]
172 - l_ij = tl.sum(p, 1) # [M]
173 - # -- update m_i and l_i
174 - m_i_new = tl.maximum(m_i, m_ij) # [M]
175 - alpha = tl.exp(m_i - m_i_new) # [M]
176 - beta = tl.exp(m_ij - m_i_new) # [M]
177 - l_i_new = alpha * l_i + beta * l_ij # [M]
178 -
179 - # -- update output accumulator --
180 - # scale p
181 - p_scale = beta / l_i_new
182 - p = p * p_scale[:, None]
183 - # scale acc
184 - acc_scale = l_i / l_i_new * alpha
185 - acc = acc * acc_scale[:, None]
186 - # update acc
187 - v_load = tl.load(V_cache + off_v,
188 - mask=dim_mask[None, :] &
189 - ((start_n + offs_n[:, None]) < cur_batch_ctx_len),
190 - other=0.0) # [N,D]
191 - if v_load.dtype.is_fp8():
192 - v = (v_load.to(tl.float32) * tl.load(v_scale)).to(q.dtype)
193 - else:
194 - v = v_load
195 - p = p.to(v.dtype)
196 -
197 - acc = tl.dot(p, v, acc=acc, input_precision=IN_PRECISION)
198 - # # update m_i and l_i
199 - l_i = l_i_new
200 - m_i = m_i_new
201 -
202 - off_k = (offs_n[None, :] * stride_kbs + cur_kv_head * stride_kh +
203 - offs_d[:, None] * stride_kd)
204 - off_v = (offs_n[:, None] * stride_vbs + cur_kv_head * stride_vh +
205 - offs_d[None, :] * stride_vd)
206 - k_ptrs = K + off_k
207 - v_ptrs = V + off_v
208 -
209 - # block_mask is 0 when we're already past the current query length
210 - block_mask = tl.where(block_start_loc < cur_batch_query_len, 1, 0)
211 -
212 - # compute query against itself (with causal mask)
213 - for start_n in range(0, block_mask * (start_m + 1) * BLOCK_M, BLOCK_N):
214 - start_n = tl.multiple_of(start_n, BLOCK_N)
215 - # -- compute qk ----
216 - k = tl.load(k_ptrs +
217 - (cur_batch_in_all_start_index + start_n) * stride_kbs,
218 - mask=dim_mask[:, None] &
219 - ((start_n + offs_n[None, :]) < cur_batch_query_len),
220 - other=0.0)
221 -
222 - qk = tl.zeros([BLOCK_M, BLOCK_N], dtype=tl.float32)
223 - qk = tl.dot(q, k, acc=qk, input_precision=IN_PRECISION)
224 - qk *= sm_scale
225 - # apply causal mask
226 - qk = tl.where(offs_m[:, None] >= (start_n + offs_n[None, :]), qk,
227 - float("-inf"))
228 - if SLIDING_WINDOW > 0:
229 - qk = tl.where(
230 - offs_m[:, None] - (start_n + offs_n[None, :])
231 - < SLIDING_WINDOW, qk, -10000)
232 -
233 - # -- compute m_ij, p, l_ij
234 - m_ij = tl.max(qk, 1)
235 - p = tl.exp(qk - m_ij[:, None])
236 - l_ij = tl.sum(p, 1)
237 - # -- update m_i and l_i
238 - m_i_new = tl.maximum(m_i, m_ij)
239 - alpha = tl.exp(m_i - m_i_new)
240 - beta = tl.exp(m_ij - m_i_new)
241 - l_i_new = alpha * l_i + beta * l_ij
242 - # -- update output accumulator --
243 - # scale p
244 - p_scale = beta / l_i_new
245 - p = p * p_scale[:, None]
246 - # scale acc
247 - acc_scale = l_i / l_i_new * alpha
248 - acc = acc * acc_scale[:, None]
249 - # update acc
250 - v = tl.load(v_ptrs +
251 - (cur_batch_in_all_start_index + start_n) * stride_vbs,
252 - mask=dim_mask[None, :] &
253 - ((start_n + offs_n[:, None]) < cur_batch_query_len),
254 - other=0.0)
255 - p = p.to(v.dtype)
256 -
257 - acc = tl.dot(p, v, acc=acc, input_precision=IN_PRECISION)
258 - # update m_i and l_i
259 - l_i = l_i_new
260 - m_i = m_i_new
261 - # initialize pointers to output
262 - off_o = (
263 - (cur_batch_in_all_start_index + offs_m[:, None]) * stride_obs +
264 - cur_head * stride_oh + offs_d[None, :] * stride_od)
265 - out_ptrs = Out + off_o
266 - tl.store(out_ptrs,
267 - acc,
268 - mask=dim_mask[None, :] &
269 - (offs_m[:, None] < cur_batch_query_len))
19 +
20 +# Here's an example autotuner config for this kernel. This config does provide
21 +# a performance improvement, but dramatically increases first call latency in
22 +# triton 3.2. Because of this tradeoff, it's currently commented out.
23 +# @triton.autotune(
24 +# configs=[
25 +# triton.Config({'BLOCK_M': 128, 'BLOCK_N': 64, \
26 +# "num_unroll_cache": 4, \
27 +# "num_unroll_request": 1 } | \
28 +# ({"kpack": 2, "waves_per_eu": 2} \
29 +# if current_platform.is_rocm() else {}), \
30 +# num_warps=4, \
31 +# num_stages=1)
32 +# ],
33 +# key=["BLOCK_SIZE", "MAX_Q_LEN", "MAX_CTX_LEN"]
34 +# )
35 +@triton.jit
36 +def _fwd_kernel(Q,
37 + K,
38 + V,
39 + K_cache,
40 + V_cache,
41 + B_Loc,
42 + sm_scale,
43 + k_scale,
44 + v_scale,
45 + B_Start_Loc,
46 + B_Seqlen,
47 + x: tl.constexpr,
48 + Out,
49 + stride_b_loc_b,
50 + stride_b_loc_s,
51 + stride_qbs,
52 + stride_qh,
53 + stride_qd,
54 + stride_kbs,
55 + stride_kh,
56 + stride_kd,
57 + stride_vbs,
58 + stride_vh,
59 + stride_vd,
60 + stride_obs,
61 + stride_oh,
62 + stride_od,
63 + stride_k_cache_bs,
64 + stride_k_cache_h,
65 + stride_k_cache_d,
66 + stride_k_cache_bl: tl.constexpr,
67 + stride_k_cache_x,
68 + stride_v_cache_bs,
69 + stride_v_cache_h,
70 + stride_v_cache_d,
71 + stride_v_cache_bl,
72 + num_queries_per_kv: tl.constexpr,
73 + IN_PRECISION: tl.constexpr,
74 + BLOCK_M: tl.constexpr,
75 + BLOCK_DMODEL: tl.constexpr,
76 + BLOCK_DMODEL_PADDED: tl.constexpr,
77 + BLOCK_SIZE: tl.constexpr,
78 + BLOCK_N: tl.constexpr,
79 + SLIDING_WINDOW: tl.constexpr,
80 + num_unroll_cache: tl.constexpr,
81 + num_unroll_request: tl.constexpr,
82 + SKIP_DECODE: tl.constexpr,
83 + MAX_Q_LEN: tl.constexpr = 0,
84 + MAX_CTX_LEN: tl.constexpr = 0):
85 +
86 + cur_batch = tl.program_id(0)
87 + cur_head = tl.program_id(1)
88 + start_m = tl.program_id(2)
89 +
90 + cur_kv_head = cur_head // num_queries_per_kv
91 +
92 + cur_batch_seq_len = tl.load(B_Seqlen + cur_batch)
93 + cur_batch_in_all_start_index = tl.load(B_Start_Loc + cur_batch)
94 + cur_batch_in_all_stop_index = tl.load(B_Start_Loc + cur_batch + 1)
95 + cur_batch_query_len = (cur_batch_in_all_stop_index -
96 + cur_batch_in_all_start_index)
97 + cur_batch_ctx_len = cur_batch_seq_len - cur_batch_query_len
98 +
99 + if SKIP_DECODE and cur_batch_query_len == 1:
100 return
101
272 - @triton.jit
273 - def _fwd_kernel_flash_attn_v2(
274 - Q,
275 - K,
276 - V,
277 - K_cache,
278 - V_cache,
279 - B_Loc,
280 - sm_scale,
281 - B_Start_Loc,
282 - B_Seqlen,
283 - B_Ctxlen,
284 - block_size,
285 - x,
286 - Out,
287 - stride_b_loc_b,
288 - stride_b_loc_s,
289 - stride_qbs,
290 - stride_qh,
291 - stride_qd,
292 - stride_kbs,
293 - stride_kh,
294 - stride_kd,
295 - stride_vbs,
296 - stride_vh,
297 - stride_vd,
298 - stride_obs,
299 - stride_oh,
300 - stride_od,
301 - stride_k_cache_bs,
302 - stride_k_cache_h,
303 - stride_k_cache_d,
304 - stride_k_cache_bl,
305 - stride_k_cache_x,
306 - stride_v_cache_bs,
307 - stride_v_cache_h,
308 - stride_v_cache_d,
309 - stride_v_cache_bl,
310 - num_queries_per_kv: int,
311 - BLOCK_M: tl.constexpr,
312 - BLOCK_DMODEL: tl.constexpr,
313 - BLOCK_N: tl.constexpr,
314 - ):
315 - cur_batch = tl.program_id(0)
316 - cur_head = tl.program_id(1)
317 - start_m = tl.program_id(2)
318 -
319 - cur_kv_head = cur_head // num_queries_per_kv
320 -
321 - cur_batch_ctx_len = tl.load(B_Ctxlen + cur_batch)
322 - cur_batch_seq_len = tl.load(B_Seqlen + cur_batch)
323 - cur_batch_in_all_start_index = tl.load(B_Start_Loc + cur_batch)
324 -
325 - block_start_loc = BLOCK_M * start_m
326 -
327 - # initialize offsets
328 - offs_n = tl.arange(0, BLOCK_N)
329 - offs_d = tl.arange(0, BLOCK_DMODEL)
330 - offs_m = start_m * BLOCK_M + tl.arange(0, BLOCK_M)
331 - off_q = (
332 - (cur_batch_in_all_start_index + offs_m[:, None]) * stride_qbs +
333 - cur_head * stride_qh + offs_d[None, :] * stride_qd)
334 -
335 - q = tl.load(Q + off_q,
336 - mask=offs_m[:, None]
337 - < cur_batch_seq_len - cur_batch_ctx_len,
102 + # start position inside of the query
103 + # generally, N goes over kv, while M goes over query_len
104 + block_start_loc = BLOCK_M * start_m
105 +
106 + # initialize offsets
107 + # [BLOCK_SIZE]; starts at 0
108 + offs_bs_n = tl.arange(0, BLOCK_SIZE)
109 + # [N]; starts at 0
110 + offs_n = tl.arange(0, BLOCK_N)
111 + # [D]; starts at 0
112 + offs_d = tl.arange(0, BLOCK_DMODEL_PADDED)
113 + # [M]; starts at current position in query
114 + offs_m = start_m * BLOCK_M + tl.arange(0, BLOCK_M)
115 + # [M,D]
116 + off_q = ((cur_batch_in_all_start_index + offs_m[:, None]) * stride_qbs +
117 + cur_head * stride_qh + offs_d[None, :] * stride_qd)
118 +
119 + dim_mask = tl.where(
120 + tl.arange(0, BLOCK_DMODEL_PADDED) < BLOCK_DMODEL, 1,
121 + 0).to(tl.int1) # [D]
122 +
123 + q = tl.load(Q + off_q,
124 + mask=dim_mask[None, :] &
125 + (offs_m[:, None] < cur_batch_query_len),
126 + other=0.0) # [M,D]
127 +
128 + # initialize pointer to m and l
129 + m_i = tl.full([BLOCK_M], float("-inf"), dtype=tl.float32)
130 + l_i = tl.full([BLOCK_M], 1.0, dtype=tl.float32)
131 + acc = tl.zeros([BLOCK_M, BLOCK_DMODEL_PADDED], dtype=tl.float32) # [M,D]
132 +
133 + # compute query against context (no causal mask here)
134 + for start_n in tl.range(0, cur_batch_ctx_len, BLOCK_SIZE, \
135 + loop_unroll_factor=num_unroll_cache):
136 + start_n = tl.multiple_of(start_n, BLOCK_SIZE)
137 + # -- compute qk ----
138 + bn = tl.load(B_Loc + cur_batch * stride_b_loc_b +
139 + (start_n // BLOCK_SIZE) * stride_b_loc_s)
140 + # [D,BLOCK_SIZE]
141 + off_k = (
142 + bn[None, :] * stride_k_cache_bs + cur_kv_head * stride_k_cache_h +
143 + (offs_d[:, None] // x) * stride_k_cache_d +
144 + ((start_n + offs_bs_n[None, :]) % BLOCK_SIZE) * stride_k_cache_bl +
145 + (offs_d[:, None] % x) * stride_k_cache_x)
146 +
147 + # [BLOCK_SIZE,D]
148 + off_v = (bn[:, None] * stride_v_cache_bs +
149 + cur_kv_head * stride_v_cache_h +
150 + offs_d[None, :] * stride_v_cache_d +
151 + offs_bs_n[:, None] * stride_v_cache_bl)
152 +
153 + if start_n + BLOCK_SIZE > cur_batch_ctx_len or \
154 + BLOCK_DMODEL != BLOCK_DMODEL_PADDED:
155 + k_load = tl.load(
156 + K_cache + off_k,
157 + mask=dim_mask[:, None] &
158 + ((start_n + offs_bs_n[None, :]) < cur_batch_ctx_len),
159 + other=0.0) # [D,N]
160 + else:
161 + k_load = tl.load(K_cache + off_k)
162 +
163 + if k_load.dtype.is_fp8():
164 + k = (k_load.to(tl.float32) * tl.load(k_scale)).to(q.dtype)
165 + else:
166 + k = k_load
167 +
168 + qk = tl.zeros([BLOCK_M, BLOCK_SIZE], dtype=tl.float32) # [M,N]
169 + qk = tl.dot(q, k, acc=qk, input_precision=IN_PRECISION)
170 + qk = tl.where((start_n + offs_bs_n[None, :]) < cur_batch_ctx_len, qk,
171 + float("-inf"))
172 + qk *= sm_scale
173 + if SLIDING_WINDOW > 0:
174 + # (cur_batch_ctx_len + offs_m[:, None]) are the positions of
175 + # Q entries in sequence
176 + # (start_n + offs_bs_n[None, :]) are the positions of
177 + # KV entries in sequence
178 + # So the condition makes sure each entry in Q only attends
179 + # to KV entries not more than SLIDING_WINDOW away.
180 + #
181 + # We can't use -inf here, because the
182 + # sliding window may lead to the entire row being masked.
183 + # This then makes m_ij contain -inf, which causes NaNs in
184 + # exp().
185 + qk = tl.where((cur_batch_ctx_len + offs_m[:, None]) -
186 + (start_n + offs_bs_n[None, :]) < SLIDING_WINDOW, qk,
187 + -10000)
188 +
189 + # compute running maximum
190 + m_ij = tl.maximum(m_i, tl.max(qk, axis=1))
191 + p = tl.exp(qk - m_ij[:, None])
192 + l_ij = tl.sum(p, axis=1)
193 + alpha = tl.exp(m_i - m_ij)
194 + acc = acc * alpha[:, None]
195 +
196 + # update acc
197 + if start_n + BLOCK_SIZE > cur_batch_ctx_len or \
198 + BLOCK_DMODEL != BLOCK_DMODEL_PADDED:
199 + v_load = tl.load(
200 + V_cache + off_v,
201 + mask=dim_mask[None, :] &
202 + ((start_n + offs_bs_n[:, None]) < cur_batch_ctx_len),
203 + other=0.0) # [N,D]
204 + else:
205 + v_load = tl.load(V_cache + off_v)
206 +
207 + if v_load.dtype.is_fp8():
208 + v = (v_load.to(tl.float32) * tl.load(v_scale)).to(q.dtype)
209 + else:
210 + v = v_load
211 + p = p.to(v.dtype)
212 +
213 + acc = tl.dot(p, v, acc=acc, input_precision=IN_PRECISION)
214 + # # update m_i and l_i
215 + l_i = l_i * alpha + l_ij
216 + m_i = m_ij
217 +
218 + off_k = (offs_n[None, :] * stride_kbs + cur_kv_head * stride_kh +
219 + offs_d[:, None] * stride_kd)
220 + off_v = (offs_n[:, None] * stride_vbs + cur_kv_head * stride_vh +
221 + offs_d[None, :] * stride_vd)
222 + k_ptrs = K + off_k
223 + v_ptrs = V + off_v
224 +
225 + # block_mask is 0 when we're already past the current query length
226 + block_mask = tl.where(block_start_loc < cur_batch_query_len, 1, 0)
227 +
228 + # compute query against itself (with causal mask)
229 + for start_n in tl.range(0, \
230 + block_mask * (start_m + 1) * BLOCK_M, BLOCK_N, \
231 + loop_unroll_factor=num_unroll_request):
232 + start_n = tl.multiple_of(start_n, BLOCK_N)
233 + # -- compute qk ----
234 + k = tl.load(k_ptrs +
235 + (cur_batch_in_all_start_index + start_n) * stride_kbs,
236 + mask=dim_mask[:, None] &
237 + ((start_n + offs_n[None, :]) < cur_batch_query_len),
238 other=0.0)
239
340 - # # initialize pointer to m and l
341 - m_i = tl.zeros([BLOCK_M], dtype=tl.float32) - float("inf")
342 - l_i = tl.zeros([BLOCK_M], dtype=tl.float32)
343 - acc = tl.zeros([BLOCK_M, BLOCK_DMODEL], dtype=tl.float32)
344 -
345 - for start_n in range(0, cur_batch_ctx_len, BLOCK_N):
346 - start_n = tl.multiple_of(start_n, BLOCK_N)
347 - # -- compute qk ----
348 - bn = tl.load(B_Loc + cur_batch * stride_b_loc_b +
349 - ((start_n + offs_n) // block_size) * stride_b_loc_s,
350 - mask=(start_n + offs_n) < cur_batch_ctx_len,
351 - other=0)
352 - off_k = (bn[None, :] * stride_k_cache_bs +
353 - cur_kv_head * stride_k_cache_h +
354 - (offs_d[:, None] // x) * stride_k_cache_d +
355 - ((start_n + offs_n[None, :]) % block_size) *
356 - stride_k_cache_bl +
357 - (offs_d[:, None] % x) * stride_k_cache_x)
358 - off_v = (
359 - bn[:, None] * stride_v_cache_bs +
360 - cur_kv_head * stride_v_cache_h +
361 - offs_d[None, :] * stride_v_cache_d +
362 - (start_n + offs_n[:, None]) % block_size * stride_v_cache_bl)
363 - k = tl.load(K_cache + off_k,
364 - mask=(start_n + offs_n[None, :]) < cur_batch_ctx_len,
365 - other=0.0)
366 - qk = tl.zeros([BLOCK_M, BLOCK_N], dtype=tl.float32)
367 - qk += tl.dot(q, k)
368 - qk = tl.where((start_n + offs_n[None, :]) < cur_batch_ctx_len, qk,
369 - float("-inf"))
370 - qk *= sm_scale
371 -
372 - # -- compute m_ij, p, l_ij
373 - m_ij = tl.max(qk, 1)
374 - m_i_new = tl.maximum(m_i, m_ij)
375 - p = tl.math.exp(qk - m_i_new[:, None])
376 - l_ij = tl.sum(p, 1)
377 - # -- update m_i and l_i
378 -
379 - alpha = tl.math.exp(m_i - m_i_new)
380 - l_i_new = alpha * l_i + l_ij
381 - # -- update output accumulator --
382 - # scale p
383 - # scale acc
384 - acc_scale = alpha
385 - # acc_scale = l_i / l_i_new * alpha
386 - acc = acc * acc_scale[:, None]
387 - # update acc
388 - v = tl.load(V_cache + off_v,
389 - mask=(start_n + offs_n[:, None]) < cur_batch_ctx_len,
390 - other=0.0)
391 -
392 - p = p.to(v.dtype)
393 - acc += tl.dot(p, v)
394 - # update m_i and l_i
395 - l_i = l_i_new
396 - m_i = m_i_new
397 -
398 - off_k = (offs_n[None, :] * stride_kbs + cur_kv_head * stride_kh +
399 - offs_d[:, None] * stride_kd)
400 - off_v = (offs_n[:, None] * stride_vbs + cur_kv_head * stride_vh +
401 - offs_d[None, :] * stride_vd)
402 - k_ptrs = K + off_k
403 - v_ptrs = V + off_v
404 -
405 - block_mask = tl.where(
406 - block_start_loc < cur_batch_seq_len - cur_batch_ctx_len, 1, 0)
407 -
408 - for start_n in range(0, block_mask * (start_m + 1) * BLOCK_M, BLOCK_N):
409 - start_n = tl.multiple_of(start_n, BLOCK_N)
410 - # -- compute qk ----
411 - k = tl.load(k_ptrs +
412 - (cur_batch_in_all_start_index + start_n) * stride_kbs,
413 - mask=(start_n + offs_n[None, :])
414 - < cur_batch_seq_len - cur_batch_ctx_len,
415 - other=0.0)
416 -
417 - qk = tl.zeros([BLOCK_M, BLOCK_N], dtype=tl.float32)
418 - qk += tl.dot(q, k)
419 - qk *= sm_scale
420 - qk = tl.where(offs_m[:, None] >= (start_n + offs_n[None, :]), qk,
421 - float("-inf"))
422 -
423 - # -- compute m_ij, p, l_ij
424 - m_ij = tl.max(qk, 1)
425 - m_i_new = tl.maximum(m_i, m_ij)
426 - p = tl.math.exp(qk - m_i_new[:, None])
427 - l_ij = tl.sum(p, 1)
428 - # -- update m_i and l_i
429 -
430 - alpha = tl.math.exp(m_i - m_i_new)
431 - l_i_new = alpha * l_i + l_ij
432 - # -- update output accumulator --
433 - # scale p
434 - # scale acc
435 - acc_scale = alpha
436 - # acc_scale = l_i / l_i_new * alpha
437 - acc = acc * acc_scale[:, None]
438 - # update acc
439 - v = tl.load(v_ptrs +
440 - (cur_batch_in_all_start_index + start_n) * stride_vbs,
441 - mask=(start_n + offs_n[:, None])
442 - < cur_batch_seq_len - cur_batch_ctx_len,
443 - other=0.0)
444 -
445 - p = p.to(v.dtype)
446 - acc += tl.dot(p, v)
447 - # update m_i and l_i
448 - l_i = l_i_new
449 - m_i = m_i_new
450 -
451 - # acc /= l_i[:, None]
452 - # initialize pointers to output
453 - off_o = (
454 - (cur_batch_in_all_start_index + offs_m[:, None]) * stride_obs +
455 - cur_head * stride_oh + offs_d[None, :] * stride_od)
456 - out_ptrs = Out + off_o
457 - tl.store(out_ptrs,
458 - acc,
459 - mask=offs_m[:, None] < cur_batch_seq_len - cur_batch_ctx_len)
460 - return
461 -
462 - @triton.jit
463 - def _fwd_kernel_alibi(
464 - Q,
465 - K,
466 - V,
467 - K_cache,
468 - V_cache,
469 - B_Loc,
470 - sm_scale,
471 - k_scale,
472 - v_scale,
473 - B_Start_Loc,
474 - B_Seqlen,
475 - Alibi_slopes,
476 - block_size,
477 - x,
478 - Out,
479 - stride_b_loc_b,
480 - stride_b_loc_s,
481 - stride_qbs,
482 - stride_qh,
483 - stride_qd,
484 - stride_kbs,
485 - stride_kh,
486 - stride_kd,
487 - stride_vbs,
488 - stride_vh,
489 - stride_vd,
490 - stride_obs,
491 - stride_oh,
492 - stride_od,
493 - stride_k_cache_bs,
494 - stride_k_cache_h,
495 - stride_k_cache_d,
496 - stride_k_cache_bl,
497 - stride_k_cache_x,
498 - stride_v_cache_bs,
499 - stride_v_cache_h,
500 - stride_v_cache_d,
501 - stride_v_cache_bl,
502 - num_queries_per_kv: int,
503 - IN_PRECISION: tl.constexpr,
504 - BLOCK_M: tl.constexpr,
505 - BLOCK_DMODEL: tl.constexpr, # head size
506 - BLOCK_DMODEL_PADDED: tl.constexpr, # head size padded to a power of 2
507 - BLOCK_N: tl.constexpr,
508 - SKIP_DECODE: tl.constexpr,
509 - ):
510 - # attn_bias[]
511 - cur_batch = tl.program_id(0)
512 - cur_head = tl.program_id(1)
513 - start_m = tl.program_id(2)
514 -
515 - cur_kv_head = cur_head // num_queries_per_kv
516 -
517 - # cur_batch_seq_len: the length of prompts
518 - # cur_batch_ctx_len: the length of prefix
519 - # cur_batch_in_all_start_index: the start id of the dim=0
520 - cur_batch_seq_len = tl.load(B_Seqlen + cur_batch)
521 - cur_batch_in_all_start_index = tl.load(B_Start_Loc + cur_batch)
522 - cur_batch_in_all_stop_index = tl.load(B_Start_Loc + cur_batch + 1)
523 - cur_batch_query_len = (cur_batch_in_all_stop_index -
524 - cur_batch_in_all_start_index)
525 - cur_batch_ctx_len = cur_batch_seq_len - cur_batch_query_len
526 -
527 - if SKIP_DECODE and cur_batch_query_len == 1:
528 - return
529 -
530 - block_start_loc = BLOCK_M * start_m
531 -
532 - # initialize offsets
533 - offs_n = tl.arange(0, BLOCK_N)
534 - offs_d = tl.arange(0, BLOCK_DMODEL_PADDED)
535 - offs_m = start_m * BLOCK_M + tl.arange(0, BLOCK_M)
536 - off_q = (
537 - (cur_batch_in_all_start_index + offs_m[:, None]) * stride_qbs +
538 - cur_head * stride_qh + offs_d[None, :] * stride_qd)
539 -
540 - dim_mask = tl.where(
541 - tl.arange(0, BLOCK_DMODEL_PADDED) < BLOCK_DMODEL, 1, 0).to(tl.int1)
542 -
543 - q = tl.load(Q + off_q,
240 + qk = tl.zeros([BLOCK_M, BLOCK_N], dtype=tl.float32)
241 + qk = tl.dot(q, k, acc=qk, input_precision=IN_PRECISION)
242 + qk *= sm_scale
243 + # apply causal mask
244 + qk = tl.where(offs_m[:, None] >= (start_n + offs_n[None, :]), qk,
245 + float("-inf"))
246 + if SLIDING_WINDOW > 0:
247 + qk = tl.where(
248 + offs_m[:, None] - (start_n + offs_n[None, :]) < SLIDING_WINDOW,
249 + qk, -10000)
250 +
251 + # compute running maximum
252 + m_ij = tl.maximum(m_i, tl.max(qk, axis=1))
253 + p = tl.exp(qk - m_ij[:, None])
254 + l_ij = tl.sum(p, axis=1)
255 + alpha = tl.exp(m_i - m_ij)
256 + acc = acc * alpha[:, None]
257 +
258 + # update acc
259 + v = tl.load(v_ptrs +
260 + (cur_batch_in_all_start_index + start_n) * stride_vbs,
261 mask=dim_mask[None, :] &
545 - (offs_m[:, None] < cur_batch_seq_len - cur_batch_ctx_len),
262 + ((start_n + offs_n[:, None]) < cur_batch_query_len),
263 + other=0.0)
264 + p = p.to(v.dtype)
265 +
266 + acc = tl.dot(p, v, acc=acc, input_precision=IN_PRECISION)
267 + # update m_i and l_i
268 + l_i = l_i * alpha + l_ij
269 + m_i = m_ij
270 +
271 + acc = acc / l_i[:, None]
272 +
273 + # initialize pointers to output
274 + off_o = ((cur_batch_in_all_start_index + offs_m[:, None]) * stride_obs +
275 + cur_head * stride_oh + offs_d[None, :] * stride_od)
276 + out_ptrs = Out + off_o
277 + tl.store(out_ptrs,
278 + acc,
279 + mask=dim_mask[None, :] & (offs_m[:, None] < cur_batch_query_len))
280 + return
281 +
282 +
283 +@triton.jit
284 +def _fwd_kernel_flash_attn_v2(
285 + Q,
286 + K,
287 + V,
288 + K_cache,
289 + V_cache,
290 + B_Loc,
291 + sm_scale,
292 + B_Start_Loc,
293 + B_Seqlen,
294 + B_Ctxlen,
295 + block_size,
296 + x,
297 + Out,
298 + stride_b_loc_b,
299 + stride_b_loc_s,
300 + stride_qbs,
301 + stride_qh,
302 + stride_qd,
303 + stride_kbs,
304 + stride_kh,
305 + stride_kd,
306 + stride_vbs,
307 + stride_vh,
308 + stride_vd,
309 + stride_obs,
310 + stride_oh,
311 + stride_od,
312 + stride_k_cache_bs,
313 + stride_k_cache_h,
314 + stride_k_cache_d,
315 + stride_k_cache_bl,
316 + stride_k_cache_x,
317 + stride_v_cache_bs,
318 + stride_v_cache_h,
319 + stride_v_cache_d,
320 + stride_v_cache_bl,
321 + num_queries_per_kv: int,
322 + BLOCK_M: tl.constexpr,
323 + BLOCK_DMODEL: tl.constexpr,
324 + BLOCK_N: tl.constexpr,
325 +):
326 + cur_batch = tl.program_id(0)
327 + cur_head = tl.program_id(1)
328 + start_m = tl.program_id(2)
329 +
330 + cur_kv_head = cur_head // num_queries_per_kv
331 +
332 + cur_batch_ctx_len = tl.load(B_Ctxlen + cur_batch)
333 + cur_batch_seq_len = tl.load(B_Seqlen + cur_batch)
334 + cur_batch_in_all_start_index = tl.load(B_Start_Loc + cur_batch)
335 +
336 + block_start_loc = BLOCK_M * start_m
337 +
338 + # initialize offsets
339 + offs_n = tl.arange(0, BLOCK_N)
340 + offs_d = tl.arange(0, BLOCK_DMODEL)
341 + offs_m = start_m * BLOCK_M + tl.arange(0, BLOCK_M)
342 + off_q = ((cur_batch_in_all_start_index + offs_m[:, None]) * stride_qbs +
343 + cur_head * stride_qh + offs_d[None, :] * stride_qd)
344 +
345 + q = tl.load(Q + off_q,
346 + mask=offs_m[:, None] < cur_batch_seq_len - cur_batch_ctx_len,
347 + other=0.0)
348 +
349 + # # initialize pointer to m and l
350 + m_i = tl.zeros([BLOCK_M], dtype=tl.float32) - float("inf")
351 + l_i = tl.zeros([BLOCK_M], dtype=tl.float32)
352 + acc = tl.zeros([BLOCK_M, BLOCK_DMODEL], dtype=tl.float32)
353 +
354 + for start_n in range(0, cur_batch_ctx_len, BLOCK_N):
355 + start_n = tl.multiple_of(start_n, BLOCK_N)
356 + # -- compute qk ----
357 + bn = tl.load(B_Loc + cur_batch * stride_b_loc_b +
358 + ((start_n + offs_n) // block_size) * stride_b_loc_s,
359 + mask=(start_n + offs_n) < cur_batch_ctx_len,
360 + other=0)
361 + off_k = (
362 + bn[None, :] * stride_k_cache_bs + cur_kv_head * stride_k_cache_h +
363 + (offs_d[:, None] // x) * stride_k_cache_d +
364 + ((start_n + offs_n[None, :]) % block_size) * stride_k_cache_bl +
365 + (offs_d[:, None] % x) * stride_k_cache_x)
366 + off_v = (bn[:, None] * stride_v_cache_bs +
367 + cur_kv_head * stride_v_cache_h +
368 + offs_d[None, :] * stride_v_cache_d +
369 + (start_n + offs_n[:, None]) % block_size * stride_v_cache_bl)
370 + k = tl.load(K_cache + off_k,
371 + mask=(start_n + offs_n[None, :]) < cur_batch_ctx_len,
372 + other=0.0)
373 + qk = tl.zeros([BLOCK_M, BLOCK_N], dtype=tl.float32)
374 + qk += tl.dot(q, k)
375 + qk = tl.where((start_n + offs_n[None, :]) < cur_batch_ctx_len, qk,
376 + float("-inf"))
377 + qk *= sm_scale
378 +
379 + # -- compute m_ij, p, l_ij
380 + m_ij = tl.max(qk, 1)
381 + m_i_new = tl.maximum(m_i, m_ij)
382 + p = tl.math.exp(qk - m_i_new[:, None])
383 + l_ij = tl.sum(p, 1)
384 + # -- update m_i and l_i
385 +
386 + alpha = tl.math.exp(m_i - m_i_new)
387 + l_i_new = alpha * l_i + l_ij
388 + # -- update output accumulator --
389 + # scale p
390 + # scale acc
391 + acc_scale = alpha
392 + # acc_scale = l_i / l_i_new * alpha
393 + acc = acc * acc_scale[:, None]
394 + # update acc
395 + v = tl.load(V_cache + off_v,
396 + mask=(start_n + offs_n[:, None]) < cur_batch_ctx_len,
397 + other=0.0)
398 +
399 + p = p.to(v.dtype)
400 + acc += tl.dot(p, v)
401 + # update m_i and l_i
402 + l_i = l_i_new
403 + m_i = m_i_new
404 +
405 + off_k = (offs_n[None, :] * stride_kbs + cur_kv_head * stride_kh +
406 + offs_d[:, None] * stride_kd)
407 + off_v = (offs_n[:, None] * stride_vbs + cur_kv_head * stride_vh +
408 + offs_d[None, :] * stride_vd)
409 + k_ptrs = K + off_k
410 + v_ptrs = V + off_v
411 +
412 + block_mask = tl.where(
413 + block_start_loc < cur_batch_seq_len - cur_batch_ctx_len, 1, 0)
414 +
415 + for start_n in range(0, block_mask * (start_m + 1) * BLOCK_M, BLOCK_N):
416 + start_n = tl.multiple_of(start_n, BLOCK_N)
417 + # -- compute qk ----
418 + k = tl.load(k_ptrs +
419 + (cur_batch_in_all_start_index + start_n) * stride_kbs,
420 + mask=(start_n + offs_n[None, :])
421 + < cur_batch_seq_len - cur_batch_ctx_len,
422 other=0.0)
423
548 - # # initialize pointer to m and l
549 - m_i = tl.zeros([BLOCK_M], dtype=tl.float32) - float("inf")
550 - l_i = tl.zeros([BLOCK_M], dtype=tl.float32)
551 - acc = tl.zeros([BLOCK_M, BLOCK_DMODEL_PADDED], dtype=tl.float32)
552 -
553 - alibi_slope = tl.load(Alibi_slopes + cur_head)
554 - alibi_start_q = tl.arange(
555 - 0, BLOCK_M) + block_start_loc + cur_batch_ctx_len
556 - alibi_start_k = 0
557 - for start_n in range(0, cur_batch_ctx_len, BLOCK_N):
558 - start_n = tl.multiple_of(start_n, BLOCK_N)
559 - # -- compute qk ----
560 - bn = tl.load(B_Loc + cur_batch * stride_b_loc_b +
561 - ((start_n + offs_n) // block_size) * stride_b_loc_s,
562 - mask=(start_n + offs_n) < cur_batch_ctx_len,
563 - other=0)
564 - off_k = (bn[None, :] * stride_k_cache_bs +
565 - cur_kv_head * stride_k_cache_h +
566 - (offs_d[:, None] // x) * stride_k_cache_d +
567 - ((start_n + offs_n[None, :]) % block_size) *
568 - stride_k_cache_bl +
569 - (offs_d[:, None] % x) * stride_k_cache_x)
570 - off_v = (
571 - bn[:, None] * stride_v_cache_bs +
572 - cur_kv_head * stride_v_cache_h +
573 - offs_d[None, :] * stride_v_cache_d +
574 - (start_n + offs_n[:, None]) % block_size * stride_v_cache_bl)
575 - k_load = tl.load(K_cache + off_k,
576 - mask=dim_mask[:, None] &
577 - ((start_n + offs_n[None, :]) < cur_batch_ctx_len),
578 - other=0.0) # [D,N]
579 -
580 - if k_load.dtype.is_fp8():
581 - k = (k_load.to(tl.float32) * tl.load(k_scale)).to(q.dtype)
582 - else:
583 - k = k_load
584 -
585 - qk = tl.zeros([BLOCK_M, BLOCK_N], dtype=tl.float32)
586 - qk = tl.dot(q, k, acc=qk, input_precision=IN_PRECISION)
587 - qk = tl.where((start_n + offs_n[None, :]) < cur_batch_ctx_len, qk,
588 - float("-inf"))
589 - qk *= sm_scale
590 -
591 - # load alibi
592 - alibi = (tl.arange(0, BLOCK_N)[None, :] + alibi_start_k -
593 - alibi_start_q[:, None]) * alibi_slope
594 - alibi = tl.where(
595 - (alibi <= 0) & (alibi_start_q[:, None] < cur_batch_seq_len),
596 - alibi, float("-inf"))
597 - qk += alibi
598 - alibi_start_k += BLOCK_N
599 -
600 - # -- compute m_ij, p, l_ij
601 - m_ij = tl.max(qk, 1)
602 - m_i_new = tl.maximum(m_i, m_ij)
603 - p = tl.math.exp(qk - m_i_new[:, None])
604 - l_ij = tl.sum(p, 1)
605 - # -- update m_i and l_i
606 -
607 - alpha = tl.math.exp(m_i - m_i_new)
608 - l_i_new = alpha * l_i + l_ij
609 - # -- update output accumulator --
610 - # scale p
611 - # scale acc
612 - acc_scale = alpha
613 - # acc_scale = l_i / l_i_new * alpha
614 - acc = acc * acc_scale[:, None]
615 - # update acc
616 - v_load = tl.load(V_cache + off_v,
617 - mask=dim_mask[None, :] &
618 - ((start_n + offs_n[:, None]) < cur_batch_ctx_len),
619 - other=0.0)
620 - if v_load.dtype.is_fp8():
621 - v = (v_load.to(tl.float32) * tl.load(v_scale)).to(q.dtype)
622 - else:
623 - v = v_load
624 - p = p.to(v.dtype)
625 -
626 - acc = tl.dot(p, v, acc=acc, input_precision='ieee')
627 - # update m_i and l_i
628 - l_i = l_i_new
629 - m_i = m_i_new
630 -
631 - off_k = (offs_n[None, :] * stride_kbs + cur_kv_head * stride_kh +
632 - offs_d[:, None] * stride_kd)
633 - off_v = (offs_n[:, None] * stride_vbs + cur_kv_head * stride_vh +
634 - offs_d[None, :] * stride_vd)
635 - k_ptrs = K + off_k
636 - v_ptrs = V + off_v
637 -
638 - block_mask = tl.where(
639 - block_start_loc < cur_batch_seq_len - cur_batch_ctx_len, 1, 0)
640 -
641 - # init alibi
642 - alibi_slope = tl.load(Alibi_slopes + cur_head)
643 - alibi_start_q = tl.arange(
644 - 0, BLOCK_M) + block_start_loc + cur_batch_ctx_len
645 - alibi_start_k = cur_batch_ctx_len
646 - # # init debugger
647 - # offset_db_q = tl.arange(0, BLOCK_M) + block_start_loc
648 - # offset_db_k = tl.arange(0, BLOCK_N)
649 - # calc q[BLOCK_M, BLOCK_MODEL] mul k[prefix_len: , BLOCK_DMODEL]
650 - for start_n in range(0, block_mask * (start_m + 1) * BLOCK_M, BLOCK_N):
651 - start_n = tl.multiple_of(start_n, BLOCK_N)
652 - # -- compute qk ----
653 - k = tl.load(k_ptrs +
654 - (cur_batch_in_all_start_index + start_n) * stride_kbs,
655 - mask=dim_mask[:, None] &
656 - ((start_n + offs_n[None, :])
657 - < cur_batch_seq_len - cur_batch_ctx_len),
658 - other=0.0)
659 -
660 - qk = tl.zeros([BLOCK_M, BLOCK_N], dtype=tl.float32)
661 - qk = tl.dot(q, k, acc=qk, input_precision='ieee')
662 - qk *= sm_scale
663 - qk = tl.where(offs_m[:, None] >= (start_n + offs_n[None, :]), qk,
664 - float("-inf"))
665 -
666 - # load alibi
667 - alibi = (tl.arange(0, BLOCK_N)[None, :] + alibi_start_k -
668 - alibi_start_q[:, None]) * alibi_slope
669 - alibi = tl.where(
670 - (alibi <= 0) & (alibi_start_q[:, None] < cur_batch_seq_len),
671 - alibi, float("-inf"))
672 - qk += alibi
673 - alibi_start_k += BLOCK_N
674 -
675 - # -- compute m_ij, p, l_ij
676 - m_ij = tl.max(qk, 1)
677 - m_i_new = tl.maximum(m_i, m_ij)
678 - p = tl.math.exp(qk - m_i_new[:, None])
679 - l_ij = tl.sum(p, 1)
680 - # -- update m_i and l_i
681 -
682 - alpha = tl.math.exp(m_i - m_i_new)
683 - l_i_new = alpha * l_i + l_ij
684 - # -- update output accumulator --
685 - # scale p
686 - # scale acc
687 - acc_scale = alpha
688 - # acc_scale = l_i / l_i_new * alpha
689 - acc = acc * acc_scale[:, None]
690 - # update acc
691 - v = tl.load(v_ptrs +
692 - (cur_batch_in_all_start_index + start_n) * stride_vbs,
693 - mask=dim_mask[None, :] &
694 - ((start_n + offs_n[:, None])
695 - < cur_batch_seq_len - cur_batch_ctx_len),
696 - other=0.0)
697 - p = p.to(v.dtype)
698 -
699 - acc = tl.dot(p, v, acc=acc, input_precision='ieee')
700 - # update m_i and l_i
701 - l_i = l_i_new
702 - m_i = m_i_new
703 -
704 - acc = acc / l_i[:, None]
705 -
706 - # initialize pointers to output
707 - off_o = (
708 - (cur_batch_in_all_start_index + offs_m[:, None]) * stride_obs +
709 - cur_head * stride_oh + offs_d[None, :] * stride_od)
710 - out_ptrs = Out + off_o
711 - tl.store(out_ptrs,
712 - acc,
713 - mask=dim_mask[None, :] &
714 - (offs_m[:, None] < cur_batch_seq_len - cur_batch_ctx_len))
424 + qk = tl.zeros([BLOCK_M, BLOCK_N], dtype=tl.float32)
425 + qk += tl.dot(q, k)
426 + qk *= sm_scale
427 + qk = tl.where(offs_m[:, None] >= (start_n + offs_n[None, :]), qk,
428 + float("-inf"))
429 +
430 + # -- compute m_ij, p, l_ij
431 + m_ij = tl.max(qk, 1)
432 + m_i_new = tl.maximum(m_i, m_ij)
433 + p = tl.math.exp(qk - m_i_new[:, None])
434 + l_ij = tl.sum(p, 1)
435 + # -- update m_i and l_i
436 +
437 + alpha = tl.math.exp(m_i - m_i_new)
438 + l_i_new = alpha * l_i + l_ij
439 + # -- update output accumulator --
440 + # scale p
441 + # scale acc
442 + acc_scale = alpha
443 + # acc_scale = l_i / l_i_new * alpha
444 + acc = acc * acc_scale[:, None]
445 + # update acc
446 + v = tl.load(v_ptrs +
447 + (cur_batch_in_all_start_index + start_n) * stride_vbs,
448 + mask=(start_n + offs_n[:, None])
449 + < cur_batch_seq_len - cur_batch_ctx_len,
450 + other=0.0)
451 +
452 + p = p.to(v.dtype)
453 + acc += tl.dot(p, v)
454 + # update m_i and l_i
455 + l_i = l_i_new
456 + m_i = m_i_new
457 +
458 + # acc /= l_i[:, None]
459 + # initialize pointers to output
460 + off_o = ((cur_batch_in_all_start_index + offs_m[:, None]) * stride_obs +
461 + cur_head * stride_oh + offs_d[None, :] * stride_od)
462 + out_ptrs = Out + off_o
463 + tl.store(out_ptrs,
464 + acc,
465 + mask=offs_m[:, None] < cur_batch_seq_len - cur_batch_ctx_len)
466 + return
467 +
468 +
469 +@triton.jit
470 +def _fwd_kernel_alibi(
471 + Q,
472 + K,
473 + V,
474 + K_cache,
475 + V_cache,
476 + B_Loc,
477 + sm_scale,
478 + k_scale,
479 + v_scale,
480 + B_Start_Loc,
481 + B_Seqlen,
482 + Alibi_slopes,
483 + block_size,
484 + x,
485 + Out,
486 + stride_b_loc_b,
487 + stride_b_loc_s,
488 + stride_qbs,
489 + stride_qh,
490 + stride_qd,
491 + stride_kbs,
492 + stride_kh,
493 + stride_kd,
494 + stride_vbs,
495 + stride_vh,
496 + stride_vd,
497 + stride_obs,
498 + stride_oh,
499 + stride_od,
500 + stride_k_cache_bs,
501 + stride_k_cache_h,
502 + stride_k_cache_d,
503 + stride_k_cache_bl,
504 + stride_k_cache_x,
505 + stride_v_cache_bs,
506 + stride_v_cache_h,
507 + stride_v_cache_d,
508 + stride_v_cache_bl,
509 + num_queries_per_kv: int,
510 + IN_PRECISION: tl.constexpr,
511 + BLOCK_M: tl.constexpr,
512 + BLOCK_DMODEL: tl.constexpr, # head size
513 + BLOCK_DMODEL_PADDED: tl.constexpr, # head size padded to a power of 2
514 + BLOCK_N: tl.constexpr,
515 + SKIP_DECODE: tl.constexpr,
516 +):
517 + # attn_bias[]
518 + cur_batch = tl.program_id(0)
519 + cur_head = tl.program_id(1)
520 + start_m = tl.program_id(2)
521 +
522 + cur_kv_head = cur_head // num_queries_per_kv
523 +
524 + # cur_batch_seq_len: the length of prompts
525 + # cur_batch_ctx_len: the length of prefix
526 + # cur_batch_in_all_start_index: the start id of the dim=0
527 + cur_batch_seq_len = tl.load(B_Seqlen + cur_batch)
528 + cur_batch_in_all_start_index = tl.load(B_Start_Loc + cur_batch)
529 + cur_batch_in_all_stop_index = tl.load(B_Start_Loc + cur_batch + 1)
530 + cur_batch_query_len = (cur_batch_in_all_stop_index -
531 + cur_batch_in_all_start_index)
532 + cur_batch_ctx_len = cur_batch_seq_len - cur_batch_query_len
533 +
534 + if SKIP_DECODE and cur_batch_query_len == 1:
535 return
536
717 - @torch.inference_mode()
718 - def context_attention_fwd(q,
719 - k,
720 - v,
721 - o,
722 - kv_cache_dtype: str,
723 - k_cache,
724 - v_cache,
725 - b_loc,
726 - b_start_loc,
727 - b_seq_len,
728 - max_seq_len,
729 - max_input_len,
730 - k_scale: torch.Tensor,
731 - v_scale: torch.Tensor,
732 - alibi_slopes=None,
733 - sliding_window=None,
734 - sm_scale=None,
735 - skip_decode=False):
736 -
737 - q_dtype_is_f32 = q.dtype is torch.float32
537 + block_start_loc = BLOCK_M * start_m
538 +
539 + # initialize offsets
540 + offs_n = tl.arange(0, BLOCK_N)
541 + offs_d = tl.arange(0, BLOCK_DMODEL_PADDED)
542 + offs_m = start_m * BLOCK_M + tl.arange(0, BLOCK_M)
543 + off_q = ((cur_batch_in_all_start_index + offs_m[:, None]) * stride_qbs +
544 + cur_head * stride_qh + offs_d[None, :] * stride_qd)
545 +
546 + dim_mask = tl.where(
547 + tl.arange(0, BLOCK_DMODEL_PADDED) < BLOCK_DMODEL, 1, 0).to(tl.int1)
548 +
549 + q = tl.load(Q + off_q,
550 + mask=dim_mask[None, :] &
551 + (offs_m[:, None] < cur_batch_seq_len - cur_batch_ctx_len),
552 + other=0.0)
553 +
554 + # # initialize pointer to m and l
555 + m_i = tl.zeros([BLOCK_M], dtype=tl.float32) - float("inf")
556 + l_i = tl.zeros([BLOCK_M], dtype=tl.float32)
557 + acc = tl.zeros([BLOCK_M, BLOCK_DMODEL_PADDED], dtype=tl.float32)
558 +
559 + alibi_slope = tl.load(Alibi_slopes + cur_head)
560 + alibi_start_q = tl.arange(0, BLOCK_M) + block_start_loc + cur_batch_ctx_len
561 + alibi_start_k = 0
562 + for start_n in range(0, cur_batch_ctx_len, BLOCK_N):
563 + start_n = tl.multiple_of(start_n, BLOCK_N)
564 + # -- compute qk ----
565 + bn = tl.load(B_Loc + cur_batch * stride_b_loc_b +
566 + ((start_n + offs_n) // block_size) * stride_b_loc_s,
567 + mask=(start_n + offs_n) < cur_batch_ctx_len,
568 + other=0)
569 + off_k = (
570 + bn[None, :] * stride_k_cache_bs + cur_kv_head * stride_k_cache_h +
571 + (offs_d[:, None] // x) * stride_k_cache_d +
572 + ((start_n + offs_n[None, :]) % block_size) * stride_k_cache_bl +
573 + (offs_d[:, None] % x) * stride_k_cache_x)
574 + off_v = (bn[:, None] * stride_v_cache_bs +
575 + cur_kv_head * stride_v_cache_h +
576 + offs_d[None, :] * stride_v_cache_d +
577 + (start_n + offs_n[:, None]) % block_size * stride_v_cache_bl)
578 + k_load = tl.load(K_cache + off_k,
579 + mask=dim_mask[:, None] &
580 + ((start_n + offs_n[None, :]) < cur_batch_ctx_len),
581 + other=0.0) # [D,N]
582 +
583 + if k_load.dtype.is_fp8():
584 + k = (k_load.to(tl.float32) * tl.load(k_scale)).to(q.dtype)
585 + else:
586 + k = k_load
587 +
588 + qk = tl.zeros([BLOCK_M, BLOCK_N], dtype=tl.float32)
589 + qk = tl.dot(q, k, acc=qk, input_precision=IN_PRECISION)
590 + qk = tl.where((start_n + offs_n[None, :]) < cur_batch_ctx_len, qk,
591 + float("-inf"))
592 + qk *= sm_scale
593 +
594 + # load alibi
595 + alibi = (tl.arange(0, BLOCK_N)[None, :] + alibi_start_k -
596 + alibi_start_q[:, None]) * alibi_slope
597 + alibi = tl.where(
598 + (alibi <= 0) & (alibi_start_q[:, None] < cur_batch_seq_len), alibi,
599 + float("-inf"))
600 + qk += alibi
601 + alibi_start_k += BLOCK_N
602 +
603 + # -- compute m_ij, p, l_ij
604 + m_ij = tl.max(qk, 1)
605 + m_i_new = tl.maximum(m_i, m_ij)
606 + p = tl.math.exp(qk - m_i_new[:, None])
607 + l_ij = tl.sum(p, 1)
608 + # -- update m_i and l_i
609 +
610 + alpha = tl.math.exp(m_i - m_i_new)
611 + l_i_new = alpha * l_i + l_ij
612 + # -- update output accumulator --
613 + # scale p
614 + # scale acc
615 + acc_scale = alpha
616 + # acc_scale = l_i / l_i_new * alpha
617 + acc = acc * acc_scale[:, None]
618 + # update acc
619 + v_load = tl.load(V_cache + off_v,
620 + mask=dim_mask[None, :] &
621 + ((start_n + offs_n[:, None]) < cur_batch_ctx_len),
622 + other=0.0)
623 + if v_load.dtype.is_fp8():
624 + v = (v_load.to(tl.float32) * tl.load(v_scale)).to(q.dtype)
625 + else:
626 + v = v_load
627 + p = p.to(v.dtype)
628 +
629 + acc = tl.dot(p, v, acc=acc, input_precision='ieee')
630 + # update m_i and l_i
631 + l_i = l_i_new
632 + m_i = m_i_new
633 +
634 + off_k = (offs_n[None, :] * stride_kbs + cur_kv_head * stride_kh +
635 + offs_d[:, None] * stride_kd)
636 + off_v = (offs_n[:, None] * stride_vbs + cur_kv_head * stride_vh +
637 + offs_d[None, :] * stride_vd)
638 + k_ptrs = K + off_k
639 + v_ptrs = V + off_v
640 +
641 + block_mask = tl.where(
642 + block_start_loc < cur_batch_seq_len - cur_batch_ctx_len, 1, 0)
643 +
644 + # init alibi
645 + alibi_slope = tl.load(Alibi_slopes + cur_head)
646 + alibi_start_q = tl.arange(0, BLOCK_M) + block_start_loc + cur_batch_ctx_len
647 + alibi_start_k = cur_batch_ctx_len
648 + # # init debugger
649 + # offset_db_q = tl.arange(0, BLOCK_M) + block_start_loc
650 + # offset_db_k = tl.arange(0, BLOCK_N)
651 + # calc q[BLOCK_M, BLOCK_MODEL] mul k[prefix_len: , BLOCK_DMODEL]
652 + for start_n in range(0, block_mask * (start_m + 1) * BLOCK_M, BLOCK_N):
653 + start_n = tl.multiple_of(start_n, BLOCK_N)
654 + # -- compute qk ----
655 + k = tl.load(
656 + k_ptrs + (cur_batch_in_all_start_index + start_n) * stride_kbs,
657 + mask=dim_mask[:, None] & ((start_n + offs_n[None, :])
658 + < cur_batch_seq_len - cur_batch_ctx_len),
659 + other=0.0)
660 +
661 + qk = tl.zeros([BLOCK_M, BLOCK_N], dtype=tl.float32)
662 + qk = tl.dot(q, k, acc=qk, input_precision='ieee')
663 + qk *= sm_scale
664 + qk = tl.where(offs_m[:, None] >= (start_n + offs_n[None, :]), qk,
665 + float("-inf"))
666 +
667 + # load alibi
668 + alibi = (tl.arange(0, BLOCK_N)[None, :] + alibi_start_k -
669 + alibi_start_q[:, None]) * alibi_slope
670 + alibi = tl.where(
671 + (alibi <= 0) & (alibi_start_q[:, None] < cur_batch_seq_len), alibi,
672 + float("-inf"))
673 + qk += alibi
674 + alibi_start_k += BLOCK_N
675 +
676 + # -- compute m_ij, p, l_ij
677 + m_ij = tl.max(qk, 1)
678 + m_i_new = tl.maximum(m_i, m_ij)
679 + p = tl.math.exp(qk - m_i_new[:, None])
680 + l_ij = tl.sum(p, 1)
681 + # -- update m_i and l_i
682 +
683 + alpha = tl.math.exp(m_i - m_i_new)
684 + l_i_new = alpha * l_i + l_ij
685 + # -- update output accumulator --
686 + # scale p
687 + # scale acc
688 + acc_scale = alpha
689 + # acc_scale = l_i / l_i_new * alpha
690 + acc = acc * acc_scale[:, None]
691 + # update acc
692 + v = tl.load(
693 + v_ptrs + (cur_batch_in_all_start_index + start_n) * stride_vbs,
694 + mask=dim_mask[None, :] & ((start_n + offs_n[:, None])
695 + < cur_batch_seq_len - cur_batch_ctx_len),
696 + other=0.0)
697 + p = p.to(v.dtype)
698 +
699 + acc = tl.dot(p, v, acc=acc, input_precision='ieee')
700 + # update m_i and l_i
701 + l_i = l_i_new
702 + m_i = m_i_new
703 +
704 + acc = acc / l_i[:, None]
705 +
706 + # initialize pointers to output
707 + off_o = ((cur_batch_in_all_start_index + offs_m[:, None]) * stride_obs +
708 + cur_head * stride_oh + offs_d[None, :] * stride_od)
709 + out_ptrs = Out + off_o
710 + tl.store(out_ptrs,
711 + acc,
712 + mask=dim_mask[None, :] &
713 + (offs_m[:, None] < cur_batch_seq_len - cur_batch_ctx_len))
714 + return
715 +
716 +
717 +@torch.inference_mode()
718 +def context_attention_fwd(q,
719 + k,
720 + v,
721 + o,
722 + kv_cache_dtype: str,
723 + k_cache,
724 + v_cache,
725 + b_loc,
726 + b_start_loc,
727 + b_seq_len,
728 + max_seq_len,
729 + max_input_len,
730 + k_scale: torch.Tensor,
731 + v_scale: torch.Tensor,
732 + alibi_slopes=None,
733 + sliding_window=None,
734 + sm_scale=None,
735 + skip_decode=False):
736 +
737 + q_dtype_is_f32 = q.dtype is torch.float32
738 +
739 + # Turing does have tensor core for float32 multiplication
740 + # use ieee as fallback for triton kernels work. There is also
741 + # warning on vllm/config.py to inform users this fallback
742 + # implementation
743 + IN_PRECISION = 'ieee' if IS_TURING and q_dtype_is_f32 else None
744 +
745 + # Conversion of FP8 Tensor from uint8 storage to
746 + # appropriate torch.dtype for interpretation by Triton
747 + if "fp8" in kv_cache_dtype:
748 + assert (k_cache.dtype == torch.uint8)
749 + assert (v_cache.dtype == torch.uint8)
750 +
751 + if kv_cache_dtype in ("fp8", "fp8_e4m3"):
752 + target_dtype = current_platform.fp8_dtype()
753 + elif kv_cache_dtype == "fp8_e5m2":
754 + target_dtype = torch.float8_e5m2
755 + else:
756 + raise ValueError("Unsupported FP8 dtype:", kv_cache_dtype)
757 +
758 + k_cache = k_cache.view(target_dtype)
759 + v_cache = v_cache.view(target_dtype)
760 +
761 + if (k_cache.dtype == torch.uint8
762 + or v_cache.dtype == torch.uint8 and kv_cache_dtype == "auto"):
763 + raise ValueError("kv_cache_dtype='auto' unsupported for\
764 + FP8 KV Cache prefill kernel")
765 +
766 + # shape constraints
767 + Lq, Lk, Lv = q.shape[-1], k.shape[-1], v.shape[-1]
768 + assert Lq == Lk and Lk == Lv
769 + # round up Lk to a power of 2 - this is required for Triton block size
770 + Lk_padded = triton.next_power_of_2(Lk)
771 +
772 + if sm_scale is None:
773 + sm_scale = 1.0 / (Lq**0.5)
774 + batch, head = b_seq_len.shape[0], q.shape[1]
775 + num_queries_per_kv = q.shape[1] // k.shape[1]
776 +
777 + assert batch + 1 == len(b_start_loc)
778 +
779 + # 0 means "disable"
780 + if sliding_window is None or sliding_window <= 0:
781 + sliding_window = 0
782 +
783 + if alibi_slopes is not None:
784 # need to reduce num. blocks when using fp32
785 # due to increased use of GPU shared memory
786 # if q.dtype is torch.float32:
787 BLOCK = BASE_BLOCK // 2 if q_dtype_is_f32 else BASE_BLOCK
742 -
743 - # Turing does have tensor core for float32 multiplication
744 - # use ieee as fallback for triton kernels work. There is also
745 - # warning on vllm/config.py to inform users this fallback
746 - # implementation
747 - IN_PRECISION = 'ieee' if IS_TURING and q_dtype_is_f32 else None
748 -
749 - # Conversion of FP8 Tensor from uint8 storage to
750 - # appropriate torch.dtype for interpretation by Triton
751 - if "fp8" in kv_cache_dtype:
752 - assert (k_cache.dtype == torch.uint8)
753 - assert (v_cache.dtype == torch.uint8)
754 -
755 - if kv_cache_dtype in ("fp8", "fp8_e4m3"):
756 - target_dtype = current_platform.fp8_dtype()
757 - elif kv_cache_dtype == "fp8_e5m2":
758 - target_dtype = torch.float8_e5m2
759 - else:
760 - raise ValueError("Unsupported FP8 dtype:", kv_cache_dtype)
761 -
762 - k_cache = k_cache.view(target_dtype)
763 - v_cache = v_cache.view(target_dtype)
764 -
765 - if (k_cache.dtype == torch.uint8
766 - or v_cache.dtype == torch.uint8 and kv_cache_dtype == "auto"):
767 - raise ValueError("kv_cache_dtype='auto' unsupported for\
768 - FP8 KV Cache prefill kernel")
769 -
770 - # shape constraints
771 - Lq, Lk, Lv = q.shape[-1], k.shape[-1], v.shape[-1]
772 - assert Lq == Lk and Lk == Lv
773 - # round up Lk to a power of 2 - this is required for Triton block size
774 - Lk_padded = triton.next_power_of_2(Lk)
775 -
776 - if sm_scale is None:
777 - sm_scale = 1.0 / (Lq**0.5)
778 - batch, head = b_seq_len.shape[0], q.shape[1]
779 - num_queries_per_kv = q.shape[1] // k.shape[1]
780 -
781 - assert batch + 1 == len(b_start_loc)
782 - grid = (batch, head, triton.cdiv(max_input_len, BLOCK)) # batch, head,
783 -
784 - # 0 means "disable"
785 - if sliding_window is None or sliding_window <= 0:
786 - sliding_window = 0
787 -
788 - if alibi_slopes is not None:
789 - _fwd_kernel_alibi[grid](
790 - q,
791 - k,
792 - v,
793 - k_cache,
794 - v_cache,
795 - b_loc,
796 - sm_scale,
797 - k_scale,
798 - v_scale,
799 - b_start_loc,
800 - b_seq_len,
801 - alibi_slopes,
802 - v_cache.shape[3],
803 - k_cache.shape[4],
804 - o,
805 - b_loc.stride(0),
806 - b_loc.stride(1),
807 - q.stride(0),
808 - q.stride(1),
809 - q.stride(2),
810 - k.stride(0),
811 - k.stride(1),
812 - k.stride(2),
813 - v.stride(0),
814 - v.stride(1),
815 - v.stride(2),
816 - o.stride(0),
817 - o.stride(1),
818 - o.stride(2),
819 - k_cache.stride(0),
820 - k_cache.stride(1),
821 - k_cache.stride(2),
822 - k_cache.stride(3),
823 - k_cache.stride(
824 - 4
825 - ), #[num_blocks, num_kv_heads, head_size/x, block_size, x]
826 - v_cache.stride(0),
827 - v_cache.stride(1),
828 - v_cache.stride(2),
829 - v_cache.stride(
830 - 3), #[num_blocks, num_kv_heads, head_size, block_size]
831 - num_queries_per_kv=num_queries_per_kv,
832 - IN_PRECISION=IN_PRECISION,
833 - BLOCK_M=BLOCK,
834 - BLOCK_DMODEL=Lk,
835 - BLOCK_DMODEL_PADDED=Lk_padded,
836 - BLOCK_N=BLOCK,
837 - SKIP_DECODE=skip_decode,
838 - num_warps=NUM_WARPS,
839 - num_stages=1,
840 - )
841 - return
842 -
843 - _fwd_kernel[grid](
788 + # batch, head,
789 + grid = (batch, head, triton.cdiv(max_input_len, BLOCK))
790 + _fwd_kernel_alibi[grid](
791 q,
792 k,
793 v,
@@ -852,6 +799,7 @@ if triton.__version__ >= "2.1.0":
799 v_scale,
800 b_start_loc,
801 b_seq_len,
802 + alibi_slopes,
803 v_cache.shape[3],
804 k_cache.shape[4],
805 o,
@@ -886,9 +834,69 @@ if triton.__version__ >= "2.1.0":
834 BLOCK_DMODEL=Lk,
835 BLOCK_DMODEL_PADDED=Lk_padded,
836 BLOCK_N=BLOCK,
889 - SLIDING_WINDOW=sliding_window,
837 SKIP_DECODE=skip_decode,
838 num_warps=NUM_WARPS,
839 num_stages=1,
840 )
841 return
842 +
843 + max_seq_len = 0 if max_seq_len is None else max_seq_len
844 + extra_kargs = {}
845 + if current_platform.is_rocm():
846 + extra_kargs = {"kpack": 2, "waves_per_eu": 2}
847 +
848 + grid = lambda META: (batch, head,
849 + triton.cdiv(max_input_len, META["BLOCK_M"]))
850 + _fwd_kernel[grid](
851 + q,
852 + k,
853 + v,
854 + k_cache,
855 + v_cache,
856 + b_loc,
857 + sm_scale,
858 + k_scale,
859 + v_scale,
860 + b_start_loc,
861 + b_seq_len,
862 + k_cache.shape[4],
863 + o,
864 + b_loc.stride(0),
865 + b_loc.stride(1),
866 + q.stride(0),
867 + q.stride(1),
868 + q.stride(2),
869 + k.stride(0),
870 + k.stride(1),
871 + k.stride(2),
872 + v.stride(0),
873 + v.stride(1),
874 + v.stride(2),
875 + o.stride(0),
876 + o.stride(1),
877 + o.stride(2),
878 + k_cache.stride(0),
879 + k_cache.stride(1),
880 + k_cache.stride(2),
881 + k_cache.stride(3),
882 + k_cache.stride(
883 + 4), #[num_blocks, num_kv_heads, head_size/x, block_size, x]
884 + v_cache.stride(0),
885 + v_cache.stride(1),
886 + v_cache.stride(2),
887 + v_cache.stride(3), #[num_blocks, num_kv_heads, head_size, block_size]
888 + BLOCK_SIZE=v_cache.shape[3],
889 + num_queries_per_kv=num_queries_per_kv,
890 + IN_PRECISION=IN_PRECISION,
891 + BLOCK_DMODEL=Lk,
892 + BLOCK_DMODEL_PADDED=Lk_padded,
893 + SLIDING_WINDOW=sliding_window,
894 + SKIP_DECODE=skip_decode,
895 + BLOCK_M=128,
896 + BLOCK_N=64,
897 + num_unroll_cache=4,
898 + num_unroll_request=1,
899 + num_warps=4,
900 + num_stages=1,
901 + **extra_kargs)
902 + return