@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 8e4b351a0c
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][Triton][FP8] Adding fp8 and variable length sequence support to Triton FAv2 kernel (#12591)

Signed-off-by: Randall Smith <Randall.Smith@amd.com>

rasmith committed Apr 26, 2025 at 19:35 UTC 8e4b351a0c9e414b0c56c32cbdef51a21d1ea1be
2 files changed +1657 -604
tests/kernels/test_triton_flash_attention.py new
+499
@@ -0,0 +1,499 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +"""Tests for the triton_flash_attention kernel
3 +
4 +Run `pytest tests/kernels/test_triton_flash_attention.py`.
5 +"""
6 +import pytest
7 +import torch
8 +
9 +from vllm.attention.ops.triton_flash_attention import (SUPPORTED_LAYOUTS,
10 + MetaData,
11 + compute_alibi_tensor,
12 + scale_fp8,
13 + triton_attention_rocm)
14 +from vllm.platforms import current_platform
15 +
16 +
17 +class ReferenceAttention:
18 +
19 + def __init__(self, Z, HQ, HK, N_CTX_Q, N_CTX_K, D_HEAD, use_alibi, dtype,
20 + input_metadata):
21 + self.Z = Z
22 + self.HQ = HQ
23 + self.HK = HK
24 + self.N_CTX_Q = N_CTX_Q
25 + self.N_CTX_K = N_CTX_K
26 + self.D_HEAD = D_HEAD
27 + self.use_alibi = use_alibi
28 + self.dtype = dtype
29 + self.input_metadata = input_metadata
30 +
31 + def fwd(self, q, k, v):
32 + scores = torch.einsum('bhqd,bhkd->bhqk', q,
33 + k).float() * self.input_metadata.sm_scale
34 + if self.input_metadata.causal:
35 + mask = torch.tril(torch.ones(self.N_CTX_Q,
36 + self.N_CTX_K,
37 + device="cuda"),
38 + diagonal=self.N_CTX_K - self.N_CTX_Q)
39 + scores[:, :, mask == 0] = float("-inf")
40 +
41 + if self.input_metadata.bias is not None:
42 + scores += self.input_metadata.bias
43 +
44 + if self.use_alibi:
45 + scores += compute_alibi_tensor(self.input_metadata.alibi_slopes,
46 + self.N_CTX_Q, self.N_CTX_K)
47 +
48 + p = torch.softmax(scores, dim=-1)
49 + if self.input_metadata.causal:
50 + # If N_CTX_Q > N_CTX_K, there's at least one row of all -infs going
51 + # into softmax. This creates a row of NaNs as -inf - -inf == NaN.
52 + # So we fix this by converting the NaNs to 0s, which is what they
53 + # should be out of the softmax.
54 + nan_mask = torch.isnan(p)
55 + p[nan_mask == 1] = 0
56 + ref_out = torch.einsum('bhqk,bhkd->bhqd', p.to(self.dtype), v)
57 + # compare
58 + if self.input_metadata.layout == 'bshd':
59 + ref_out = ref_out.transpose(1, 2).clone()
60 + return ref_out
61 +
62 + def fwd_fp8(self, q_quantized, k_quantized, v_quantized):
63 + q = (q_quantized.to(torch.float16) * self.input_metadata.q_descale).to(
64 + self.dtype)
65 + k = (k_quantized.to(torch.float16) * self.input_metadata.k_descale).to(
66 + self.dtype)
67 + v = (v_quantized.to(torch.float16) * self.input_metadata.v_descale).to(
68 + self.dtype)
69 + result = self.fwd(q, k, v)
70 + if self.input_metadata.o_scale is not None:
71 + result, _ = scale_fp8(result, self.input_metadata.o_scale)
72 + return result
73 +
74 + def fwd_fp8_kv(self, q, k_quantized, v_quantized):
75 + k_descale, v_descale = (self.input_metadata.k_descale,
76 + self.input_metadata.v_descale)
77 + k_dequantized = (k_quantized.to(torch.float32) *
78 + k_descale.to(torch.float32)).to(self.dtype)
79 + v_dequantized = (v_quantized.to(torch.float32) *
80 + v_descale.to(torch.float32)).to(self.dtype)
81 + return self.fwd(q, k_dequantized, v_dequantized)
82 +
83 + def varlen_fwd(self, q, k, v, is_mqa=False):
84 + ref_out = torch.empty_like(q)
85 + if is_mqa:
86 + # Make KV look like HQ/HK "groups" of HK. Later, we will reshape so
87 + # the size aligns with Q.
88 + k_ref = k.view(k.shape[0], k.shape[1], 1,
89 + k.shape[2]).expand(-1, -1, self.HQ // self.HK, -1)
90 + v_ref = v.view(v.shape[0], v.shape[1], 1,
91 + v.shape[2]).expand(-1, -1, self.HQ // self.HK, -1)
92 + else:
93 + k_ref = k
94 + v_ref = v
95 +
96 + for i in range(0, self.input_metadata.num_contexts):
97 + start_q, start_k = self.input_metadata.cu_seqlens_q[
98 + i], self.input_metadata.cu_seqlens_k[i]
99 + end_q, end_k = self.input_metadata.cu_seqlens_q[
100 + i + 1], self.input_metadata.cu_seqlens_k[i + 1]
101 + k_curr = k_ref[start_k:end_k]
102 + v_curr = v_ref[start_k:end_k]
103 + if is_mqa:
104 + k_curr = k_curr.reshape(k_curr.shape[0], -1, k_curr.shape[3])
105 + v_curr = v_curr.reshape(v_curr.shape[0], -1, v_curr.shape[3])
106 + scores = torch.einsum('qhd,khd->qhk', q[start_q:end_q],
107 + k_curr).float()
108 + p = torch.softmax(scores * self.input_metadata.sm_scale,
109 + dim=-1).half()
110 + ref_out[start_q:end_q] = torch.einsum('qhk,khd->qhd', p, v_curr)
111 + return ref_out
112 +
113 +
114 +def quantize_input(q, k, v, fp8_kv=False, use_o_scale=False):
115 + q_descale = None
116 + if not fp8_kv:
117 + q, q_descale = scale_fp8(q)
118 + k, k_descale = scale_fp8(k)
119 + v, v_descale = scale_fp8(v)
120 +
121 + # In real world use case, the p scale would be a parameter trained by the
122 + # model.
123 + p_scale = None
124 +
125 + o_scale = torch.rand(1, device="cuda",
126 + requires_grad=False) if use_o_scale else None
127 +
128 + return q, k, v, q_descale, k_descale, v_descale, p_scale, o_scale
129 +
130 +
131 +def input_helper(
132 + Z,
133 + HQ,
134 + HK,
135 + N_CTX_Q,
136 + N_CTX_K,
137 + D_HEAD,
138 + dtype,
139 + layout=None,
140 + use_alibi=None,
141 + causal=None,
142 + is_fp8=False,
143 + fp8_kv=False,
144 + use_o_scale=False,
145 + use_bias=False,
146 +):
147 + assert layout in SUPPORTED_LAYOUTS, "Got unsupported layout."
148 +
149 + current_platform.seed_everything(0)
150 +
151 + # Initialize q, k, v
152 + if layout == 'bhsd':
153 + q_tensor_shape = (Z, HQ, N_CTX_Q, D_HEAD)
154 + k_tensor_shape = (Z, HK, N_CTX_K, D_HEAD)
155 + elif layout == 'bshd':
156 + q_tensor_shape = (Z, N_CTX_Q, HQ, D_HEAD)
157 + k_tensor_shape = (Z, N_CTX_K, HK, D_HEAD)
158 +
159 + if use_alibi:
160 + # for n heads the set of slopes is the geometric sequence that starts
161 + # 2^(-8/n)
162 + alibi_slopes = torch.tensor(
163 + [2**(-8 / HQ * i) for i in range(1, HQ + 1)],
164 + dtype=torch.float32,
165 + device="cuda").repeat(Z, 1)
166 + else:
167 + alibi_slopes = None
168 +
169 + if use_bias:
170 + bias = torch.randn((1, HQ, N_CTX_Q, N_CTX_K),
171 + dtype=dtype,
172 + device="cuda",
173 + requires_grad=False)
174 + else:
175 + bias = None
176 +
177 + q = torch.randn(q_tensor_shape,
178 + dtype=dtype,
179 + device="cuda",
180 + requires_grad=False)
181 + k = torch.randn(k_tensor_shape,
182 + dtype=dtype,
183 + device="cuda",
184 + requires_grad=False)
185 + v = torch.randn(k_tensor_shape,
186 + dtype=dtype,
187 + device="cuda",
188 + requires_grad=False)
189 +
190 + if is_fp8:
191 + (q, k, v, q_descale, k_descale, v_descale, p_scale,
192 + o_scale) = quantize_input(q,
193 + k,
194 + v,
195 + use_o_scale=use_o_scale,
196 + fp8_kv=fp8_kv)
197 + else:
198 + q_descale = k_descale = v_descale = p_scale = o_scale = None
199 +
200 + input_metadata = MetaData(sm_scale=D_HEAD**-0.5,
201 + max_seqlens_q=N_CTX_Q,
202 + max_seqlens_k=N_CTX_K,
203 + layout=layout,
204 + alibi_slopes=alibi_slopes,
205 + alibi_batch=Z,
206 + alibi_nheads=HQ,
207 + q_descale=q_descale,
208 + k_descale=k_descale,
209 + v_descale=v_descale,
210 + p_scale=p_scale,
211 + o_scale=o_scale,
212 + bias=bias,
213 + seqlen_q=N_CTX_Q,
214 + seqlen_k=N_CTX_K)
215 + return q, k, v, input_metadata
216 +
217 +
218 +def varlen_input_helper(Z,
219 + HQ,
220 + HK,
221 + N_CTX_Q,
222 + N_CTX_K,
223 + D_HEAD,
224 + dtype,
225 + equal_seqlens=False):
226 + current_platform.seed_everything(0)
227 +
228 + # Random sequence lengths. Using N_CTX as kind of max of sum of individual
229 + # seqs
230 + if not equal_seqlens:
231 + max_seqlens_q = N_CTX_Q // Z
232 + max_seqlens_k = N_CTX_K // Z
233 + seqlens_q = torch.randint(1,
234 + max_seqlens_q + 1, (Z, ),
235 + dtype=torch.int32)
236 + seqlens_k = torch.randint(1,
237 + max_seqlens_k + 1, (Z, ),
238 + dtype=torch.int32)
239 + else:
240 + seqlens_q = torch.full((Z, ), N_CTX_Q // Z)
241 + seqlens_k = torch.full((Z, ), N_CTX_K // Z)
242 +
243 + # Calculate cumulative sequence lengths
244 + cu_seqlens_q = torch.cat([
245 + torch.tensor([0], dtype=torch.int32),
246 + seqlens_q.cumsum(dim=0, dtype=torch.int32)
247 + ])
248 + cu_seqlens_k = torch.cat([
249 + torch.tensor([0], dtype=torch.int32),
250 + seqlens_k.cumsum(dim=0, dtype=torch.int32)
251 + ])
252 + cu_seqlens_q = cu_seqlens_q.to(device="cuda")
253 + cu_seqlens_k = cu_seqlens_k.to(device="cuda")
254 +
255 + # Initialize q, k, v with variable lengths
256 + total_q = cu_seqlens_q[-1].item()
257 + total_k = cu_seqlens_k[-1].item()
258 + q = torch.randn((total_q, HQ, D_HEAD), dtype=dtype,
259 + device="cuda").normal_(mean=0., std=0.5).requires_grad_()
260 + k = torch.randn((total_k, HK, D_HEAD), dtype=dtype,
261 + device="cuda").normal_(mean=0., std=0.5).requires_grad_()
262 + v = torch.randn((total_k, HK, D_HEAD), dtype=dtype,
263 + device="cuda").normal_(mean=0., std=0.5).requires_grad_()
264 + sm_scale = D_HEAD**-0.5
265 + input_metadata = MetaData(sm_scale=sm_scale)
266 + input_metadata.set_varlen_params(cu_seqlens_q, cu_seqlens_k)
267 + return q, k, v, input_metadata
268 +
269 +
270 +@pytest.mark.parametrize('Z, HQ, HK, N_CTX_Q, N_CTX_K, D_HEAD', [
271 + (1, 48, 12, 1, 1, 64),
272 + (4, 4, 4, 128, 128, 65),
273 + (16, 48, 48, 1, 1, 128),
274 + (64, 48, 24, 3, 3, 128),
275 + (4, 4, 4, 113, 123, 1),
276 +])
277 +@pytest.mark.parametrize('causal', [True, False])
278 +@pytest.mark.parametrize('use_alibi', [True, False])
279 +@pytest.mark.parametrize('layout', ['bshd'])
280 +def test_op_fwd(Z,
281 + HQ,
282 + HK,
283 + N_CTX_Q,
284 + N_CTX_K,
285 + D_HEAD,
286 + causal,
287 + use_alibi,
288 + layout,
289 + dtype=torch.float16):
290 + current_platform.seed_everything(0)
291 + q, k, v, input_metadata = input_helper(Z, HQ, HK, N_CTX_Q, N_CTX_K, D_HEAD,
292 + dtype, layout, use_alibi, causal)
293 +
294 + o = torch.empty_like(q)
295 +
296 + # triton implementation
297 + tri_out, _ = triton_attention_rocm(q, k, v, o, input_metadata)
298 +
299 + # Transpose here if layout is bshd so we have same reference code for all
300 + # layouts
301 + if layout == 'bshd':
302 + q = q.transpose(1, 2).clone()
303 + k = k.transpose(1, 2).clone()
304 + v = v.transpose(1, 2).clone()
305 + # Replicate K and V if using MQA/GQA
306 + if HQ != HK:
307 + k = k.view(k.shape[0], k.shape[1], -1, k.shape[2],
308 + k.shape[3]).expand(-1, -1, HQ // HK, -1,
309 + -1).reshape(k.shape[0], -1, k.shape[2],
310 + k.shape[3])
311 + v = v.view(v.shape[0], v.shape[1], -1, v.shape[2],
312 + v.shape[3]).expand(-1, -1, HQ // HK, -1,
313 + -1).reshape(v.shape[0], -1, v.shape[2],
314 + v.shape[3])
315 +
316 + ref_impl = ReferenceAttention(Z, HQ, HK, N_CTX_Q, N_CTX_K, D_HEAD,
317 + use_alibi, dtype, input_metadata)
318 + ref_out = ref_impl.fwd(q, k, v)
319 +
320 + torch.testing.assert_close(ref_out, tri_out, atol=2e-2, rtol=2e-2)
321 +
322 +
323 +@pytest.mark.parametrize('Z, H, N_CTX_Q, N_CTX_K, D_HEAD', [
324 + (4, 48, 1, 1, 64),
325 + (4, 48, 1, 1, 128),
326 + (4, 48, 3, 3, 128),
327 + (4, 4, 128, 128, 65),
328 +])
329 +@pytest.mark.parametrize('causal', [True, False])
330 +@pytest.mark.parametrize('layout', ['bhsd'])
331 +@pytest.mark.parametrize('use_o_scale', [True, False])
332 +@pytest.mark.skipif(torch.cuda.get_device_capability() < (9, 0),
333 + reason="Triton FP8 requires CUDA 9.0 or higher")
334 +def test_op_fwd_fp8(Z,
335 + H,
336 + N_CTX_Q,
337 + N_CTX_K,
338 + D_HEAD,
339 + causal,
340 + layout,
341 + use_o_scale,
342 + dtype=torch.float32):
343 + current_platform.seed_everything(0)
344 +
345 + # Disable grad to save memory it won't run into OOM on CI machine.
346 + # q, k, v, input_metadata = input_helper(Z, H, H, N_CTX_Q, N_CTX_K, D_HEAD,
347 + # dtype, layout)
348 +
349 + q_quantized, k_quantized, v_quantized, input_metadata = input_helper(
350 + Z,
351 + H,
352 + H,
353 + N_CTX_Q,
354 + N_CTX_K,
355 + D_HEAD,
356 + dtype,
357 + causal=causal,
358 + layout=layout,
359 + is_fp8=True,
360 + use_o_scale=use_o_scale)
361 +
362 + o = torch.empty_like(q_quantized) if use_o_scale else None
363 +
364 + tri_out, _ = triton_attention_rocm(q_quantized, k_quantized, v_quantized,
365 + o, input_metadata)
366 +
367 + ref_impl = ReferenceAttention(Z, H, H, N_CTX_Q, N_CTX_K, D_HEAD, False,
368 + dtype, input_metadata)
369 + ref_out = ref_impl.fwd_fp8(q_quantized, k_quantized, v_quantized)
370 +
371 + # compare
372 + torch.testing.assert_close(ref_out.to(torch.float32),
373 + tri_out.to(torch.float32),
374 + atol=7e-2,
375 + rtol=2e-1)
376 +
377 +
378 +@pytest.mark.parametrize('Z, H, N_CTX_Q, N_CTX_K, D_HEAD', [
379 + (4, 48, 1, 1, 64),
380 + (4, 48, 1, 1, 128),
381 + (4, 48, 3, 3, 128),
382 + (4, 4, 128, 128, 65),
383 + (4, 4, 113, 123, 1),
384 +])
385 +@pytest.mark.parametrize('causal', [True, False])
386 +@pytest.mark.parametrize('layout', ['bhsd'])
387 +def test_op_fwd_fp8_kv(Z,
388 + H,
389 + N_CTX_Q,
390 + N_CTX_K,
391 + D_HEAD,
392 + causal,
393 + layout,
394 + dtype=torch.float32):
395 + current_platform.seed_everything(0)
396 +
397 + q, k_quantized, v_quantized, input_metadata = input_helper(Z,
398 + H,
399 + H,
400 + N_CTX_Q,
401 + N_CTX_K,
402 + D_HEAD,
403 + dtype,
404 + causal=causal,
405 + layout=layout,
406 + is_fp8=True,
407 + fp8_kv=True)
408 +
409 + o = torch.empty_like(q)
410 +
411 + tri_out, _ = triton_attention_rocm(q, k_quantized, v_quantized, o,
412 + input_metadata)
413 +
414 + ref_impl = ReferenceAttention(Z, H, H, N_CTX_Q, N_CTX_K, D_HEAD, False,
415 + dtype, input_metadata)
416 + ref_out = ref_impl.fwd_fp8_kv(q, k_quantized, v_quantized)
417 +
418 + torch.testing.assert_close(ref_out, tri_out, atol=3e-2, rtol=8e-1)
419 +
420 +
421 +@pytest.mark.parametrize('Z, H, N_CTX_Q, N_CTX_K, D_HEAD', [
422 + (4, 48, 1, 1, 64),
423 + (4, 48, 1, 1, 128),
424 + (4, 48, 3, 3, 128),
425 + (4, 4, 128, 128, 65),
426 +])
427 +@pytest.mark.parametrize('causal', [True, False])
428 +@pytest.mark.parametrize('use_bias', [True])
429 +@pytest.mark.parametrize('dtype', [torch.bfloat16])
430 +def test_op_fwd_bias(Z, H, N_CTX_Q, N_CTX_K, D_HEAD, causal, use_bias, dtype):
431 + current_platform.seed_everything(0)
432 + q, k, v, input_metadata = input_helper(Z,
433 + H,
434 + H,
435 + N_CTX_Q,
436 + N_CTX_K,
437 + D_HEAD,
438 + dtype,
439 + layout='bhsd',
440 + causal=causal,
441 + use_bias=use_bias)
442 + o = torch.empty_like(q)
443 +
444 + # triton implementation
445 + tri_out, _ = triton_attention_rocm(q, k, v, o, input_metadata)
446 +
447 + ref_impl = ReferenceAttention(Z, H, H, N_CTX_Q, N_CTX_K, D_HEAD, False,
448 + dtype, input_metadata)
449 + ref_out = ref_impl.fwd(q, k, v)
450 +
451 + # compare
452 + torch.testing.assert_close(ref_out, tri_out, atol=2e-2, rtol=2e-2)
453 +
454 +
455 +# NOTE: Uses thd layout, so also tests thd.
456 +@pytest.mark.parametrize('Z, H, N_CTX, D_HEAD', [(1, 48, 256, 64),
457 + (4, 48, 512, 64),
458 + (16, 48, 512, 64),
459 + (64, 48, 128, 128)])
460 +@pytest.mark.parametrize('causal', [True, False])
461 +def test_op_varlen_fwd(Z, H, N_CTX, D_HEAD, causal, dtype=torch.float16):
462 +
463 + q, k, v, input_metadata = varlen_input_helper(Z, H, H, N_CTX, N_CTX,
464 + D_HEAD, dtype)
465 +
466 + tri_out = torch.empty_like(q)
467 + triton_attention_rocm(q, k, v, tri_out, input_metadata)
468 +
469 + ref_impl = ReferenceAttention(Z, H, H, N_CTX, N_CTX, D_HEAD, False, dtype,
470 + input_metadata)
471 + ref_out = ref_impl.varlen_fwd(q, k, v, is_mqa=False)
472 +
473 + torch.testing.assert_close(ref_out, tri_out, atol=2e-2, rtol=2e-2)
474 +
475 +
476 +# NOTE: Uses thd layout, so also tests thd.
477 +@pytest.mark.parametrize('Z, HQ, HK, N_CTX, D_HEAD', [(2, 48, 24, 128, 64),
478 + (4, 48, 12, 256, 64),
479 + (4, 48, 4, 512, 64),
480 + (4, 64, 16, 128, 128)])
481 +@pytest.mark.parametrize('causal', [False])
482 +def test_op_varlen_mqa_fwd(Z,
483 + HQ,
484 + HK,
485 + N_CTX,
486 + D_HEAD,
487 + causal,
488 + dtype=torch.float16):
489 + q, k, v, input_metadata = varlen_input_helper(Z, HQ, HK, N_CTX, N_CTX,
490 + D_HEAD, dtype)
491 +
492 + tri_out = torch.empty_like(q)
493 + triton_attention_rocm(q, k, v, tri_out, input_metadata)
494 +
495 + ref_impl = ReferenceAttention(Z, HQ, HK, N_CTX, N_CTX, D_HEAD, False,
496 + dtype, input_metadata)
497 + ref_out = ref_impl.varlen_fwd(q, k, v, is_mqa=True)
498 +
499 + torch.testing.assert_close(ref_out, tri_out, atol=2e-2, rtol=2e-2)
vllm/attention/ops/triton_flash_attention.py
+1158 -604
@@ -1,31 +1,237 @@
1 -#!/usr/bin/env python
1 # SPDX-License-Identifier: Apache-2.0
2 """
3 Fused Attention
4 ===============
5
7 -This is a Triton implementation of the Flash Attention v2 algorithm from Tri Dao
8 -(https://tridao.me/publications/flash2/flash2.pdf)
9 -Credits: OpenAI kernel team, AMD ML Frameworks Triton team
6 +This is a Triton implementation of the Flash Attention v2 algorithm
7 +See https://tridao.me/publications/flash2/flash2.pdf
8
11 -Features supported:
9 +Credits:
10 +AMD Triton kernels team
11 +OpenAI kernel team
12
13 -1) Fwd with causal masking
14 -2) Any sequence lengths without padding (currently fwd kernel only)
15 -3) Support for different sequence lengths for q and k
16 -4) Nested tensor API currently does not support dropout or bias.
17 -
18 -Not currently supported:
13 +Currently only the forward kernel is supported, and contains these features:
14
20 -1) Non power of two head dims
15 +1) Fwd with causal masking
16 +2) Arbitrary Q and KV sequence lengths
17 +3) Arbitrary head sizes
18 +4) Multi and grouped query attention
19 +5) Variable sequence lengths
20 +6) ALiBi and matrix bias
21 +7) FP8 support
22
23 """
24
25 +from typing import Optional
26 +
27 import torch
28 import triton
29 import triton.language as tl
30
28 -torch_dtype: tl.constexpr = torch.float16
31 +from vllm import _custom_ops as ops
32 +from vllm.platforms import current_platform
33 +
34 +SUPPORTED_LAYOUTS = ['thd', 'bhsd', 'bshd']
35 +
36 +default_eight_bit_dtype_triton = tl.float8e4b8
37 +default_eight_bit_dtype_torch = current_platform.fp8_dtype()
38 +default_float8_info = torch.finfo(default_eight_bit_dtype_torch)
39 +
40 +FP8_MIN = triton.language.constexpr(default_float8_info.min)
41 +
42 +# According to https://github.com/vllm-project/vllm/blob/main
43 +# /csrc/quantization/utils.cuh#L31,
44 +# need to make the max for the uz datatype be 224.0 for accuracy reasons.
45 +FP8_MAX = triton.language.constexpr(
46 + default_float8_info.max if default_eight_bit_dtype_torch !=
47 + torch.float8_e4m3fnuz else 224.0)
48 +
49 +
50 +class MetaData:
51 + cu_seqlens_q = None
52 + cu_seqlens_k = None
53 + max_seqlens_q = 0
54 + max_seqlens_k = 0
55 + bias = None
56 + alibi_slopes = None
57 + causal = False
58 + num_contexts = 0
59 + varlen = False
60 + eight_bit = False
61 + layout = None
62 + return_encoded_softmax = False
63 + eight_bit_dtype_triton = default_eight_bit_dtype_triton
64 + eight_bit_dtype_torch = default_eight_bit_dtype_torch
65 + output_dtype = None
66 +
67 + # Note about layouts:
68 + #
69 + # thd - [num_tokens, num_heads, head_size]
70 + # bshd - [batch_size, seq_len, num_heads, head_size]
71 + # bhsd - [batch_size, num_heads, seq_len, head_size]
72 + #
73 + # This is for each tensor, all tensors must have same layout.
74 + # Q can have num_heads and seq_len differ from from K and V,
75 + # however K and V must agree on this.
76 + #
77 + # Notes about varlen and bias:
78 + # Only one or the other is implemented, meaning can't combine
79 + # both varlen and bias right now.
80 + #
81 + # Note about quantization:
82 + # Only 8-bit quantization supported (for now) and specifically fp8.
83 + # Scales must be tensors.
84 + # o_scale: This is 'output scaling', but comes from parameter called
85 + # 'input_scale', this is applied to the output from the kernel.
86 + # o_scale should be None if none of the other quantization parameters
87 + # are used.
88 + #
89 + # NOTE: Object is in a tentatively good state after initialized, however,
90 + # to verify, call check_args(q,k,v,o) where o is the output tensor.
91 + def __init__(
92 + self,
93 + sm_scale=1.0,
94 + layout=None, # layout can be 'bshd', 'bhsd', or 'thd'
95 + output_dtype=None,
96 + max_seqlens_q=0,
97 + max_seqlens_k=0,
98 + # varlen params
99 + cu_seqlens_q=None, # only 'thd' layout supported for varlen
100 + cu_seqlens_k=None,
101 + # quant params
102 + q_descale=None,
103 + k_descale=None,
104 + v_descale=None,
105 + p_scale=None,
106 + o_scale=None,
107 + # bias params
108 + bias=None, # varlen not implemented for bias
109 + seqlen_q=None,
110 + seqlen_k=None,
111 + # alibi params
112 + alibi_slopes=None,
113 + alibi_batch=None,
114 + alibi_nheads=None,
115 + # causal
116 + causal=None,
117 + ):
118 + self.sm_scale = sm_scale
119 + self.output_dtype = output_dtype
120 + self.max_seqlens_q = max_seqlens_q
121 + self.max_seqlens_k = max_seqlens_k
122 + self.layout = layout
123 + if cu_seqlens_q is not None or cu_seqlens_k is not None:
124 + assert cu_seqlens_q is not None and cu_seqlens_k is not None
125 + assert layout is None or layout not in [
126 + 'bshd', 'bhsd'
127 + ], "Varlen only implemented for thd layout"
128 + self.set_varlen_params(cu_seqlens_q, cu_seqlens_k)
129 + quant_params = [q_descale, k_descale, v_descale, p_scale, o_scale]
130 + if any(x is not None for x in quant_params):
131 + p_descale = 1.0 / p_scale if p_scale is not None else None
132 + self.set_eight_bit_params(q_descale, k_descale, v_descale, p_scale,
133 + p_descale, o_scale)
134 + if bias is not None:
135 + self.need_bias(bias, seqlen_q, seqlen_k)
136 + if alibi_slopes is not None:
137 + self.need_alibi(alibi_slopes, alibi_batch, alibi_nheads)
138 + if causal is not None and causal:
139 + self.need_causal()
140 +
141 + def set_varlen_params(self, cu_seqlens_q, cu_seqlens_k):
142 + self.varlen = True
143 + self.layout = 'thd'
144 + self.cu_seqlens_q = cu_seqlens_q
145 + self.cu_seqlens_k = cu_seqlens_k
146 + # Without "varlen", there should still be one sequence.
147 + assert len(cu_seqlens_q) >= 2
148 + assert len(cu_seqlens_q) == len(cu_seqlens_k)
149 + self.num_contexts = len(cu_seqlens_q) - 1
150 + for i in range(0, self.num_contexts):
151 + self.max_seqlens_q = max(
152 + cu_seqlens_q[i + 1].item() - cu_seqlens_q[i].item(),
153 + self.max_seqlens_q)
154 + self.max_seqlens_k = max(
155 + cu_seqlens_k[i + 1].item() - cu_seqlens_k[i].item(),
156 + self.max_seqlens_k)
157 +
158 + def set_eight_bit_params(self, q_descale, k_descale, v_descale, p_scale,
159 + p_descale, o_scale):
160 + self.eight_bit = True
161 + self.q_descale = q_descale
162 + self.k_descale = k_descale
163 + self.v_descale = v_descale
164 + self.p_scale = p_scale
165 + self.p_descale = p_descale
166 + self.o_scale = o_scale
167 + self.use_p_scale = (p_scale is not None) and (
168 + p_descale is not None) and (v_descale is not None)
169 + self.eight_bit_kv = ((q_descale is None) and (k_descale is not None)
170 + and (v_descale is not None))
171 + self.eight_bit_dtype_torch = default_eight_bit_dtype_torch
172 +
173 + def need_bias(self, bias, seqlen_q, seqlen_k):
174 + assert bias is not None
175 + assert bias.is_cuda
176 + assert bias.dim() == 4
177 + assert bias.shape[0] == 1
178 + assert bias.shape[2:] == (seqlen_q, seqlen_k)
179 + self.bias = bias
180 +
181 + def need_alibi(self, alibi_slopes, batch, nheads):
182 + assert alibi_slopes.is_cuda
183 + assert alibi_slopes.dim() == 2
184 + assert alibi_slopes.shape[0] == batch
185 + assert alibi_slopes.shape[1] == nheads
186 + self.alibi_slopes = alibi_slopes
187 +
188 + def need_causal(self):
189 + self.causal = True
190 +
191 + def check_args(self, q, k, v, o):
192 + assert q.dim() == k.dim() and q.dim() == v.dim()
193 +
194 + batch, nheads_q, nheads_k, head_size = get_shape_from_layout(
195 + q, k, self)
196 + if self.varlen:
197 + assert q.dim() == 3
198 + assert self.cu_seqlens_q is not None
199 + assert self.cu_seqlens_k is not None
200 + assert len(self.cu_seqlens_q) == len(self.cu_seqlens_k)
201 + # TODO: Remove once bias is supported with varlen
202 + assert self.bias is None
203 + assert not self.return_encoded_softmax
204 + else:
205 + assert q.dim() == 4
206 + assert self.max_seqlens_q > 0 and self.max_seqlens_k > 0
207 + assert self.cu_seqlens_q is None and self.cu_seqlens_k is None
208 + assert k.shape == v.shape
209 + assert q.shape[-1] == k.shape[-1] and q.shape[-1] == v.shape[-1]
210 + # TODO: Change assert if we support qkl f8 and v f16
211 + if self.eight_bit:
212 + if self.eight_bit_kv:
213 + assert (v.dtype == k.dtype
214 + and k.dtype == self.eight_bit_dtype_torch)
215 + assert q.dtype != k.dtype
216 + assert (self.v_descale is not None) and (self.k_descale
217 + is not None)
218 + else:
219 + assert (q.dtype == k.dtype and q.dtype == v.dtype
220 + and q.dtype == self.eight_bit_dtype_torch)
221 + assert (self.q_descale
222 + is not None) and (self.k_descale
223 + is not None) and (self.v_descale
224 + is not None)
225 + if self.use_p_scale:
226 + assert (self.p_scale is not None) and (self.p_descale
227 + is not None)
228 + else:
229 + assert (q.dtype == k.dtype) and (q.dtype == v.dtype)
230 + assert head_size <= 256
231 + assert o.shape == q.shape
232 + assert (nheads_q % nheads_k) == 0
233 + assert self.layout is not None
234 + assert self.layout == 'thd' or not self.varlen
235
236
237 @triton.jit
@@ -38,40 +244,85 @@ def max_fn(x, y):
244 return tl.math.max(x, y)
245
246
247 +# Convenience function to load with optional boundary checks.
248 +# "First" is the major dim, "second" is the minor dim.
249 @triton.jit
42 -def dropout_offsets(philox_seed, philox_offset, dropout_p, m, n, stride):
43 - ms = tl.arange(0, m)
44 - ns = tl.arange(0, n)
45 - return philox_offset + ms[:, None] * stride + ns[None, :]
250 +def masked_load(ptrs, offset_first, offset_second, boundary_first,
251 + boundary_second):
252 + if offset_first is not None and offset_second is not None:
253 + mask = (offset_first[:, None] < boundary_first) & \
254 + (offset_second[None, :] < boundary_second)
255 + tensor = tl.load(ptrs, mask=mask, other=0.0)
256 + elif offset_first is not None:
257 + mask = offset_first[:, None] < boundary_first
258 + tensor = tl.load(ptrs, mask=mask, other=0.0)
259 + elif offset_second is not None:
260 + mask = offset_second[None, :] < boundary_second
261 + tensor = tl.load(ptrs, mask=mask, other=0.0)
262 + else:
263 + tensor = tl.load(ptrs)
264 + return tensor
265
266
267 @triton.jit
49 -def dropout_rng(philox_seed, philox_offset, dropout_p, m, n, stride):
50 - rng_offsets = dropout_offsets(philox_seed, philox_offset, dropout_p, m, n,
51 - stride).to(tl.uint32)
52 - # TODO: use tl.randint for better performance
53 - return tl.rand(philox_seed, rng_offsets)
268 +def compute_alibi_block(alibi_slope,
269 + seqlen_q,
270 + seqlen_k,
271 + offs_m,
272 + offs_n,
273 + transpose=False):
274 + # when seqlen_k and seqlen_q are different we want the diagonal to stick to
275 + # the bottom right of the attention matrix
276 + # for casual mask we want something like this where (1 is kept and 0 is
277 + # masked)
278 + # seqlen_q = 2 and seqlen_k = 5
279 + # 1 1 1 1 0
280 + # 1 1 1 1 1
281 + # seqlen_q = 5 and seqlen_k = 2
282 + # 0 0
283 + # 0 0
284 + # 0 0
285 + # 1 0
286 + # 1 1
287 + # for alibi the diagonal is 0 indicating no penalty for attending to that
288 + # spot and increasing penalty for attending further from the diagonal
289 + # e.g. alibi_slope = 1, seqlen_q = 2, seqlen_k = 5,
290 + # offs_m = [0, 1, 2, 3], offs_n = [0, 1, 2, 3, 4], transpose = False
291 + # 1. offs_m[:,None] = [[0],
292 + # [1],
293 + # 2. offs_m[:,None] + seqlen_k = [[5],
294 + # [6],
295 + # 3. offs_m[:,None] + seqlen_k - seqlen_q = [[3],
296 + # [4],
297 + # 4. offs_m[:,None] + seqlen_k - seqlen_q - offs_n[None,:] =
298 + # [[3], - [[0, 1, 2, 3, 4]] = [[ 3, 2, 1, 0,-1], [4], [ 4, 3, 2, 1, 0]]
299 + # 5. -1 * alibi_slope * tl.abs(relative_pos_block) = [[ -3, -2, -1, 0,-1],
300 + # [ -4, -3, -2, -1, 0]],
301 + relative_pos_block = (offs_m[:, None] + seqlen_k - seqlen_q -
302 + offs_n[None, :])
303 + alibi_block = -1 * alibi_slope * tl.abs(relative_pos_block)
304 + if transpose:
305 + return alibi_block.T
306 + else:
307 + return alibi_block
308
309
56 -@triton.jit
57 -def dropout_mask(philox_seed, philox_offset, dropout_p, m, n, stride):
58 - rng_output = dropout_rng(philox_seed, philox_offset, dropout_p, m, n,
59 - stride)
60 - rng_keep = rng_output > dropout_p
61 - return rng_keep
310 +def compute_alibi_tensor(alibi_slopes, seqlen_q, seqlen_k):
311 + q_idx = torch.arange(seqlen_q, dtype=torch.int32,
312 + device="cuda").unsqueeze(-1) # (N_CTX_Q, 1)
313 + k_idx = torch.arange(seqlen_k, dtype=torch.int32,
314 + device="cuda").unsqueeze(0) # (1, N_CTX_K)
315 + relative_pos = torch.abs(q_idx + seqlen_k - seqlen_q -
316 + k_idx) # (N_CTX_Q, N_CTX_K)
317 + return -1 * alibi_slopes.unsqueeze(-1).unsqueeze(
318 + -1) * relative_pos # (Z, H, N_CTX_Q, N_CTX_K)
319
320
321 @triton.jit
65 -def load_fn(block_ptr, first, second, pad):
66 - if first and second:
67 - tensor = tl.load(block_ptr, boundary_check=(0, 1), padding_option=pad)
68 - elif first:
69 - tensor = tl.load(block_ptr, boundary_check=(0, ), padding_option=pad)
70 - elif second:
71 - tensor = tl.load(block_ptr, boundary_check=(1, ), padding_option=pad)
72 - else:
73 - tensor = tl.load(block_ptr)
74 - return tensor
322 +def quant_fp8(x, scale):
323 + x *= scale
324 + x = tl.clamp(x, FP8_MIN, FP8_MAX)
325 + return x
326
327
328 @triton.jit
@@ -80,58 +331,68 @@ def _attn_fwd_inner(
331 l_i,
332 m_i,
333 q,
83 - K_block_ptr,
84 - V_block_ptr,
334 + k_ptrs,
335 + v_ptrs,
336 + bias_ptrs,
337 + stride_kn,
338 + stride_vk,
339 + stride_bn,
340 start_m,
341 actual_seqlen_k,
87 - dropout_p,
342 + actual_seqlen_q,
343 philox_seed,
344 batch_philox_offset,
90 - encoded_softmax_block_ptr,
345 + encoded_sm_ptrs,
346 block_min,
347 block_max,
348 offs_n_causal,
349 masked_blocks,
350 n_extra_tokens,
96 - bias_ptr,
351 + alibi_slope,
352 + q_descale,
353 + k_descale,
354 + v_descale,
355 + p_scale,
356 IS_CAUSAL: tl.constexpr,
357 BLOCK_M: tl.constexpr,
358 BLOCK_DMODEL: tl.constexpr,
359 BLOCK_N: tl.constexpr,
360 OFFS_M: tl.constexpr,
361 OFFS_N: tl.constexpr,
103 - PRE_LOAD_V: tl.constexpr,
104 - MASK_STEPS: tl.constexpr,
105 - ENABLE_DROPOUT: tl.constexpr,
106 - RETURN_ENCODED_SOFTMAX: tl.constexpr,
107 - PADDED_HEAD: tl.constexpr,
362 + SHOULD_PRE_LOAD_V: tl.constexpr,
363 + SHOULD_MASK_STEPS: tl.constexpr,
364 + SHOULD_RETURN_ENCODED_SOFTMAX: tl.constexpr,
365 + USE_PADDED_HEAD: tl.constexpr,
366 + IS_ACTUAL_BLOCK_DMODEL: tl.constexpr,
367 + QK_SCALE: tl.constexpr,
368 + IS_EIGHT_BIT_GEMM: tl.constexpr,
369 + USE_P_SCALE: tl.constexpr,
370 + IS_EIGHT_BIT_KV: tl.constexpr,
371 + QUANT_DTYPE: tl.constexpr = default_eight_bit_dtype_triton,
372 ):
373 +
374 # loop over k, v, and update accumulator
375 for start_n in range(block_min, block_max, BLOCK_N):
376 # For padded blocks, we will overrun the tensor size if
377 # we load all BLOCK_N. For others, the blocks are all within range.
113 - k = load_fn(
114 - K_block_ptr,
115 - PADDED_HEAD,
116 - MASK_STEPS and (n_extra_tokens != 0),
117 - "zero",
118 - )
119 - if PRE_LOAD_V:
120 - v = load_fn(
121 - V_block_ptr,
122 - MASK_STEPS and (n_extra_tokens != 0),
123 - PADDED_HEAD,
124 - "zero",
125 - )
378 + k_offs_n = start_n + tl.arange(0,
379 + BLOCK_N) if SHOULD_MASK_STEPS else None
380 + k_offs_k = None if not USE_PADDED_HEAD else tl.arange(0, BLOCK_DMODEL)
381 + k = masked_load(k_ptrs, k_offs_k, k_offs_n, IS_ACTUAL_BLOCK_DMODEL,
382 + actual_seqlen_k)
383 + if SHOULD_PRE_LOAD_V:
384 + # We can use the same offsets as k, just with dims transposed.
385 + v = masked_load(v_ptrs, k_offs_n, k_offs_k, actual_seqlen_k,
386 + IS_ACTUAL_BLOCK_DMODEL)
387 qk = tl.zeros([BLOCK_M, BLOCK_N], dtype=tl.float32)
388 # We start from end of seqlen_k so only the first iteration would need
389 # to be checked for padding if it is not a multiple of block_n
390 # TODO: This can be optimized to only be true for the padded block.
130 - if MASK_STEPS: # noqa: SIM102
391 + if SHOULD_MASK_STEPS: # noqa: SIM102
392 # If this is the last block / iteration, we want to
393 # mask if the sequence length is not a multiple of block size
133 - # a solution is to always do BLOCK_M // BLOCK_N + 1 steps
134 - # if not is_modulo_mn. last step might get wasted but that is okay.
394 + # a solution is to always do BLOCK_M // BLOCK_N + 1 steps if not
395 + # is_modulo_mn. last step might get wasted but that is okay.
396 # check if this masking works for that case.
397 if (start_n + BLOCK_N == block_max) and (n_extra_tokens != 0):
398 boundary_m = tl.full([BLOCK_M],
@@ -144,167 +405,276 @@ def _attn_fwd_inner(
405 causal_boundary = start_n + offs_n_causal
406 causal_mask = OFFS_M[:, None] >= causal_boundary[None, :]
407 qk = tl.where(causal_mask, qk, float("-inf"))
408 +
409 # -- compute qk ----
148 - qk += tl.dot(q, k)
149 - if bias_ptr is not None:
150 - bias = load_fn(bias_ptr, False, MASK_STEPS
151 - and (n_extra_tokens != 0), "zero")
152 - # While bias is added after multiplying qk with sm_scale, our
153 - # optimization to use 2^x instead of e^x results in an additional
154 - # scale factor of log2(e) which we must also multiply the bias with.
155 - qk += bias * 1.44269504089
410 + if IS_EIGHT_BIT_GEMM:
411 + qk += ((((tl.dot(q, k).to(tl.float32) * q_descale)) * k_descale) *
412 + QK_SCALE)
413 + else:
414 + if IS_EIGHT_BIT_KV:
415 + k = (k * k_descale).to(q.type.element_ty)
416 + qk += (tl.dot(q, k) * QK_SCALE)
417 +
418 + if bias_ptrs is not None:
419 + bias_offs_n = start_n + tl.arange(
420 + 0, BLOCK_N) if SHOULD_MASK_STEPS else None
421 + bias = masked_load(bias_ptrs, OFFS_M, bias_offs_n, actual_seqlen_q,
422 + actual_seqlen_k)
423 + # While bias is added after multiplying qk with sm_scale,
424 + # our optimization to use 2^x instead of e^x results in an
425 + # additional scale factor of log2(e) which we must also multiply
426 + # the bias with.
427 + qk += (bias * 1.44269504089)
428 +
429 + if alibi_slope is not None:
430 + # Compute the global position of each token within the sequence
431 + global_m_positions = start_m * BLOCK_M + tl.arange(0, BLOCK_M)
432 + global_n_positions = start_n + tl.arange(0, BLOCK_N)
433 + alibi_block = compute_alibi_block(alibi_slope, actual_seqlen_q,
434 + actual_seqlen_k,
435 + global_m_positions,
436 + global_n_positions)
437 + qk += (alibi_block * 1.44269504089) # scale factor of log2(e)
438 +
439 + # softmax
440 m_ij = tl.maximum(m_i, tl.max(qk, 1))
441 qk = qk - m_ij[:, None]
442 p = tl.math.exp2(qk)
443
444 # CAVEAT: Must update l_ij before applying dropout
445 l_ij = tl.sum(p, 1)
162 - if ENABLE_DROPOUT:
163 - philox_offset = (batch_philox_offset +
164 - start_m * BLOCK_M * actual_seqlen_k + start_n -
165 - BLOCK_N)
166 - keep = dropout_mask(
167 - philox_seed,
168 - philox_offset,
169 - dropout_p,
170 - BLOCK_M,
171 - BLOCK_N,
172 - actual_seqlen_k,
173 - )
174 - if RETURN_ENCODED_SOFTMAX:
175 - tl.store(
176 - encoded_softmax_block_ptr,
177 - tl.where(keep, p,
178 - -p).to(encoded_softmax_block_ptr.type.element_ty),
179 - )
180 - p = tl.where(keep, p, 0.0)
181 - elif RETURN_ENCODED_SOFTMAX:
182 - tl.store(
183 - encoded_softmax_block_ptr,
184 - p.to(encoded_softmax_block_ptr.type.element_ty),
185 - )
446 + if SHOULD_RETURN_ENCODED_SOFTMAX:
447 + tl.store(encoded_sm_ptrs, p.to(encoded_sm_ptrs.type.element_ty))
448 # -- update output accumulator --
449 alpha = tl.math.exp2(m_i - m_ij)
450 acc = acc * alpha[:, None]
189 - if not PRE_LOAD_V:
190 - v = load_fn(
191 - V_block_ptr,
192 - MASK_STEPS and (n_extra_tokens != 0),
193 - PADDED_HEAD,
194 - "zero",
195 - )
451 + if not SHOULD_PRE_LOAD_V:
452 + v = masked_load(v_ptrs, k_offs_n, k_offs_k, actual_seqlen_k,
453 + IS_ACTUAL_BLOCK_DMODEL)
454 # -- update m_i and l_i
455 l_i = l_i * alpha + l_ij
456 # update m_i and l_i
457 m_i = m_ij
200 - acc += tl.dot(p.to(V_block_ptr.type.element_ty), v)
201 - V_block_ptr = tl.advance(V_block_ptr, (BLOCK_N, 0))
202 - K_block_ptr = tl.advance(K_block_ptr, (0, BLOCK_N))
203 - if bias_ptr is not None:
204 - bias_ptr = tl.advance(bias_ptr, (0, BLOCK_N))
205 - if RETURN_ENCODED_SOFTMAX:
206 - encoded_softmax_block_ptr = tl.advance(encoded_softmax_block_ptr,
207 - (0, BLOCK_N))
458 +
459 + if IS_EIGHT_BIT_GEMM:
460 + if USE_P_SCALE:
461 + p = quant_fp8(p, p_scale).to(QUANT_DTYPE)
462 + acc += tl.dot(p, v)
463 + else:
464 + # v is in eight_bit but p is not, we want the gemm in p's type
465 + acc += tl.dot(p, v.to(p.type.element_ty))
466 + else:
467 + if IS_EIGHT_BIT_KV:
468 + v = (v * v_descale).to(p.type.element_ty)
469 + acc += tl.dot(p.to(v.type.element_ty), v)
470 +
471 + k_ptrs += BLOCK_N * stride_kn
472 + v_ptrs += BLOCK_N * stride_vk
473 + if bias_ptrs is not None:
474 + bias_ptrs += BLOCK_N * stride_bn
475 + if SHOULD_RETURN_ENCODED_SOFTMAX:
476 + encoded_sm_ptrs += BLOCK_N
477 return acc, l_i, m_i
478
479
211 -@triton.autotune(
212 - configs=[
480 +def get_cdna_autotune_configs():
481 + return [
482 + triton.Config(
483 + {
484 + 'BLOCK_M': 128,
485 + 'BLOCK_N': 128,
486 + 'waves_per_eu': 2,
487 + 'SHOULD_PRE_LOAD_V': False,
488 + 'GRID_CU_MULTIP': 2
489 + },
490 + num_stages=1,
491 + num_warps=4),
492 + triton.Config(
493 + {
494 + 'BLOCK_M': 128,
495 + 'BLOCK_N': 64,
496 + 'waves_per_eu': 2,
497 + 'SHOULD_PRE_LOAD_V': False,
498 + 'GRID_CU_MULTIP': 2
499 + },
500 + num_stages=1,
501 + num_warps=4),
502 + triton.Config(
503 + {
504 + 'BLOCK_M': 128,
505 + 'BLOCK_N': 64,
506 + 'waves_per_eu': 3,
507 + 'SHOULD_PRE_LOAD_V': False,
508 + 'GRID_CU_MULTIP': 2
509 + },
510 + num_stages=1,
511 + num_warps=4),
512 triton.Config(
513 {
215 - "BLOCK_M": 256,
216 - "BLOCK_N": 64,
217 - "waves_per_eu": 2,
218 - "PRE_LOAD_V": False,
514 + 'BLOCK_M': 128,
515 + 'BLOCK_N': 64,
516 + 'waves_per_eu': 1,
517 + 'SHOULD_PRE_LOAD_V': False,
518 + 'GRID_CU_MULTIP': 2
519 },
520 num_stages=1,
221 - num_warps=8,
222 - ),
521 + num_warps=4),
522 triton.Config(
523 {
225 - "BLOCK_M": 128,
226 - "BLOCK_N": 128,
227 - "waves_per_eu": 2,
228 - "PRE_LOAD_V": False,
524 + 'BLOCK_M': 128,
525 + 'BLOCK_N': 32,
526 + 'waves_per_eu': 2,
527 + 'SHOULD_PRE_LOAD_V': False,
528 + 'GRID_CU_MULTIP': 2
529 },
530 num_stages=1,
231 - num_warps=4,
232 - ),
531 + num_warps=4),
532 + ], [
533 + 'IS_CAUSAL', 'MAX_SEQLENS_Q', 'MAX_SEQLENS_K',
534 + 'IS_ACTUAL_BLOCK_DMODEL', 'VARLEN', 'HQ', 'HK'
535 + ]
536 +
537 +
538 +def get_rdna_autotune_configs():
539 + return [
540 triton.Config(
541 {
235 - "BLOCK_M": 256,
236 - "BLOCK_N": 128,
237 - "waves_per_eu": 2,
238 - "PRE_LOAD_V": False,
542 + 'BLOCK_M': 32,
543 + 'BLOCK_N': 32,
544 + 'waves_per_eu': 4,
545 + 'SHOULD_PRE_LOAD_V': False,
546 + 'GRID_CU_MULTIP': 2
547 },
548 num_stages=1,
241 - num_warps=8,
242 - ),
549 + num_warps=2),
550 triton.Config(
551 {
245 - "BLOCK_M": 128,
246 - "BLOCK_N": 64,
247 - "waves_per_eu": 1,
248 - "PRE_LOAD_V": False,
552 + 'BLOCK_M': 32,
553 + 'BLOCK_N': 32,
554 + 'waves_per_eu': 2,
555 + 'SHOULD_PRE_LOAD_V': False,
556 + 'GRID_CU_MULTIP': 2
557 },
558 num_stages=1,
251 - num_warps=4,
252 - ),
559 + num_warps=2),
560 triton.Config(
561 {
255 - "BLOCK_M": 128,
256 - "BLOCK_N": 64,
257 - "waves_per_eu": 3,
258 - "PRE_LOAD_V": True,
562 + 'BLOCK_M': 32,
563 + 'BLOCK_N': 16,
564 + 'waves_per_eu': 4,
565 + 'SHOULD_PRE_LOAD_V': False,
566 + 'GRID_CU_MULTIP': 2
567 },
568 num_stages=1,
261 - num_warps=4,
262 - ),
569 + num_warps=2),
570 triton.Config(
571 {
265 - "BLOCK_M": 128,
266 - "BLOCK_N": 64,
267 - "waves_per_eu": 3,
268 - "PRE_LOAD_V": False,
572 + 'BLOCK_M': 32,
573 + 'BLOCK_N': 16,
574 + 'waves_per_eu': 2,
575 + 'SHOULD_PRE_LOAD_V': False,
576 + 'GRID_CU_MULTIP': 2
577 },
578 num_stages=1,
271 - num_warps=4,
272 - ),
579 + num_warps=2),
580 triton.Config(
581 {
275 - "BLOCK_M": 64,
276 - "BLOCK_N": 64,
277 - "waves_per_eu": 4,
278 - "PRE_LOAD_V": False,
582 + 'BLOCK_M': 16,
583 + 'BLOCK_N': 16,
584 + 'waves_per_eu': 4,
585 + 'SHOULD_PRE_LOAD_V': False,
586 + 'GRID_CU_MULTIP': 2
587 },
588 num_stages=1,
281 - num_warps=8,
282 - ),
589 + num_warps=2),
590 triton.Config(
591 {
285 - "BLOCK_M": 32,
286 - "BLOCK_N": 32,
287 - "waves_per_eu": 4,
288 - "PRE_LOAD_V": False,
592 + 'BLOCK_M': 16,
593 + 'BLOCK_N': 16,
594 + 'waves_per_eu': 2,
595 + 'SHOULD_PRE_LOAD_V': False,
596 + 'GRID_CU_MULTIP': 2
597 },
598 num_stages=1,
291 - num_warps=8,
292 - ),
293 - # TODO: This config fails with head_size not pow2 with data mismatches.
294 - # triton.Config({'BLOCK_M': 32, 'BLOCK_N': 16, 'waves_per_eu': 1,
295 - # 'PRE_LOAD_V': False}, num_stages=1, num_warps=4),
599 + num_warps=2),
600 + # Fall-back config.
601 triton.Config(
602 {
298 - "BLOCK_M": 16,
299 - "BLOCK_N": 16,
300 - "waves_per_eu": 1,
301 - "PRE_LOAD_V": False,
603 + 'BLOCK_M': 16,
604 + 'BLOCK_N': 16,
605 + 'waves_per_eu': 1,
606 + 'SHOULD_PRE_LOAD_V': False,
607 + 'GRID_CU_MULTIP': 2
608 },
609 num_stages=1,
304 - num_warps=4,
305 - ),
306 - ],
307 - key=['IS_CAUSAL', 'dropout_p', 'BLOCK_DMODEL'],
610 + num_warps=2),
611 + ], [
612 + 'IS_CAUSAL', 'MAX_SEQLENS_Q', 'MAX_SEQLENS_K',
613 + 'IS_ACTUAL_BLOCK_DMODEL', 'VARLEN', 'HQ', 'HK'
614 + ]
615 +
616 +
617 +def get_general_autotune_configs():
618 + return [
619 + triton.Config(
620 + {
621 + 'BLOCK_M': 128,
622 + 'BLOCK_N': 128,
623 + 'SHOULD_PRE_LOAD_V': False,
624 + 'GRID_CU_MULTIP': 2
625 + },
626 + num_stages=1,
627 + num_warps=4),
628 + triton.Config(
629 + {
630 + 'BLOCK_M': 128,
631 + 'BLOCK_N': 64,
632 + 'SHOULD_PRE_LOAD_V': False,
633 + 'GRID_CU_MULTIP': 2
634 + },
635 + num_stages=1,
636 + num_warps=4),
637 + triton.Config(
638 + {
639 + 'BLOCK_M': 128,
640 + 'BLOCK_N': 32,
641 + 'SHOULD_PRE_LOAD_V': False,
642 + 'GRID_CU_MULTIP': 2
643 + },
644 + num_stages=1,
645 + num_warps=4),
646 + ], [
647 + 'IS_CAUSAL', 'MAX_SEQLENS_Q', 'MAX_SEQLENS_K',
648 + 'IS_ACTUAL_BLOCK_DMODEL', 'VARLEN', 'HQ', 'HK'
649 + ]
650 +
651 +
652 +def has_cdna_target():
653 + ROCM_CDNA_TARGETS = ["gfx940", "gfx941", "gfx942", "gfx90a", "gfx908"]
654 + return triton.runtime.driver.active.get_current_target(
655 + ).arch in ROCM_CDNA_TARGETS
656 +
657 +
658 +def is_rocm_cdna():
659 + return current_platform.is_rocm() and has_cdna_target()
660 +
661 +
662 +def get_autotune_configs():
663 + if is_rocm_cdna():
664 + return get_cdna_autotune_configs()
665 + elif current_platform.is_rocm():
666 + return get_rdna_autotune_configs()
667 + else:
668 + return get_general_autotune_configs()
669 +
670 +
671 +autotune_configs, autotune_keys = get_autotune_configs()
672 +
673 +
674 +@triton.autotune(
675 + configs=autotune_configs,
676 + key=autotune_keys,
677 + use_cuda_graph=True,
678 )
679 @triton.jit
680 def attn_fwd(
@@ -312,38 +682,53 @@ def attn_fwd(
682 K,
683 V,
684 bias,
315 - sm_scale,
685 + SM_SCALE: tl.constexpr,
686 L,
687 Out,
318 - stride_qz,
319 - stride_qh,
320 - stride_qm,
321 - stride_qk,
322 - stride_kz,
323 - stride_kh,
324 - stride_kn,
325 - stride_kk,
326 - stride_vz,
327 - stride_vh,
328 - stride_vk,
329 - stride_vn,
330 - stride_oz,
331 - stride_oh,
332 - stride_om,
333 - stride_on,
334 - stride_bz,
335 - stride_bh,
336 - stride_bm,
337 - stride_bn,
688 + stride_qz: tl.int64,
689 + stride_qh: tl.int64,
690 + stride_qm: tl.int64,
691 + stride_qk: tl.int64,
692 + stride_kz: tl.int64,
693 + stride_kh: tl.int64,
694 + stride_kn: tl.int64,
695 + stride_kk: tl.int64,
696 + stride_vz: tl.int64,
697 + stride_vh: tl.int64,
698 + stride_vk: tl.int64,
699 + stride_vn: tl.int64,
700 + stride_oz: tl.int64,
701 + stride_oh: tl.int64,
702 + stride_om: tl.int64,
703 + stride_on: tl.int64,
704 + stride_bz: tl.int64,
705 + stride_bh: tl.int64,
706 + stride_bm: tl.int64,
707 + stride_bn: tl.int64,
708 + stride_az: tl.int64,
709 + stride_ah: tl.int64,
710 + q_descale_ptr,
711 + k_descale_ptr,
712 + p_scale_ptr,
713 + p_descale_ptr,
714 + o_descale_ptr,
715 + v_descale_ptr,
716 + q_descale_has_singleton: tl.constexpr,
717 + k_descale_has_singleton: tl.constexpr,
718 + p_descale_has_singleton: tl.constexpr,
719 + v_descale_has_singleton: tl.constexpr,
720 cu_seqlens_q,
721 cu_seqlens_k,
340 - dropout_p,
722 philox_seed,
723 + NUM_CU: tl.constexpr,
724 + GRID_CU_MULTIP: tl.constexpr,
725 + B: tl.constexpr,
726 philox_offset_base,
727 encoded_softmax,
728 + alibi_slopes,
729 HQ: tl.constexpr,
730 HK: tl.constexpr,
346 - ACTUAL_BLOCK_DMODEL: tl.constexpr,
731 + IS_ACTUAL_BLOCK_DMODEL: tl.constexpr,
732 MAX_SEQLENS_Q: tl.constexpr,
733 MAX_SEQLENS_K: tl.constexpr,
734 VARLEN: tl.constexpr,
@@ -351,24 +736,39 @@ def attn_fwd(
736 BLOCK_M: tl.constexpr,
737 BLOCK_DMODEL: tl.constexpr,
738 BLOCK_N: tl.constexpr,
354 - PRE_LOAD_V: tl.constexpr,
355 - BIAS_TYPE: tl.constexpr,
356 - ENABLE_DROPOUT: tl.constexpr,
357 - RETURN_ENCODED_SOFTMAX: tl.constexpr,
739 + SHOULD_PRE_LOAD_V: tl.constexpr,
740 + USE_BIAS: tl.constexpr,
741 + SHOULD_RETURN_ENCODED_SOFTMAX: tl.constexpr,
742 + USE_ALIBI: tl.constexpr,
743 + IS_EIGHT_BIT: tl.constexpr,
744 + USE_P_SCALE: tl.constexpr,
745 + IS_EIGHT_BIT_KV: tl.constexpr,
746 + QUANT_DTYPE: tl.constexpr = default_eight_bit_dtype_triton,
747 ):
359 - start_m = tl.program_id(0)
360 - off_h_q = tl.program_id(1)
361 - off_z = tl.program_id(2)
362 - offs_m = start_m * BLOCK_M + tl.arange(0, BLOCK_M)
363 - offs_n = tl.arange(0, BLOCK_N)
748 +
749 + if o_descale_ptr is not None:
750 + o_descale = tl.load(o_descale_ptr)
751 +
752 + start_m: tl.int64 = tl.program_id(0)
753 + off_h_q: tl.int64 = tl.program_id(1)
754 + off_z: tl.int64 = tl.program_id(2)
755 +
756 + offs_m = start_m * BLOCK_M + tl.arange(0, BLOCK_M).to(tl.int64)
757 + offs_n = tl.arange(0, BLOCK_N).to(tl.int64)
758 + offs_d = tl.arange(0, BLOCK_DMODEL).to(tl.int64)
759 +
760 + # as we can't have return statements inside while loop in Triton
761 + continue_condition = True
762 +
763 if VARLEN:
764 cu_seqlens_q_start = tl.load(cu_seqlens_q + off_z)
765 cu_seqlens_q_end = tl.load(cu_seqlens_q + off_z + 1)
766 seqlen_q = cu_seqlens_q_end - cu_seqlens_q_start
368 - # We have a one-size-fits-all grid in id(0). Some seqlens might be too
369 - # small for all start_m so for those we return early.
767 + # We have a one-size-fits-all grid in id(0). Some seqlens might be
768 + # too small for all start_m so for those we return early.
769 if start_m * BLOCK_M > seqlen_q:
371 - return
770 + continue_condition = False
771 + # return
772 cu_seqlens_k_start = tl.load(cu_seqlens_k + off_z)
773 cu_seqlens_k_end = tl.load(cu_seqlens_k + off_z + 1)
774 seqlen_k = cu_seqlens_k_end - cu_seqlens_k_start
@@ -378,444 +778,598 @@ def attn_fwd(
778 seqlen_q = MAX_SEQLENS_Q
779 seqlen_k = MAX_SEQLENS_K
780
381 - # Now we compute whether we need to exit early due to causal masking.
382 - # This is because for seqlen_q > seqlen_k, M rows of the attn scores
383 - # are completely masked, resulting in 0s written to the output, and
384 - # inf written to LSE. We don't need to do any GEMMs in this case.
385 - # This block of code determines what N is, and if this WG is operating
386 - # on those M rows.
387 - n_blocks = cdiv_fn(seqlen_k, BLOCK_N)
388 - if IS_CAUSAL:
389 - # If seqlen_q == seqlen_k, the attn scores are a square matrix.
390 - # If seqlen_q != seqlen_k, attn scores are rectangular which means
391 - # the causal mask boundary is bottom right aligned, and ends at either
392 - # the top edge (seqlen_q < seqlen_k) or left edge.
393 - # This captures the decrease in n_blocks if we have a rectangular attn
394 - # matrix
395 - n_blocks_seqlen = cdiv_fn(
396 - (start_m + 1) * BLOCK_M + seqlen_k - seqlen_q, BLOCK_N)
397 - # This is what adjusts the block_max for the current WG, only
398 - # if IS_CAUSAL. Otherwise we want to always iterate through all n_blocks
399 - n_blocks = min(n_blocks, n_blocks_seqlen)
400 - # If we have no blocks after adjusting for seqlen deltas, this WG is
401 - # part of the blocks that are all 0. We exit early.
402 - if n_blocks <= 0:
403 - o_offset = (off_z * stride_oz + cu_seqlens_q_start * stride_om +
404 - off_h_q * stride_oh)
405 - O_block_ptr = tl.make_block_ptr(
406 - base=Out + o_offset,
407 - shape=(seqlen_q, BLOCK_DMODEL),
408 - strides=(stride_om, stride_on),
409 - offsets=(start_m * BLOCK_M, 0),
410 - block_shape=(BLOCK_M, BLOCK_DMODEL),
411 - order=(1, 0),
412 - )
413 - acc = tl.zeros([BLOCK_M, BLOCK_DMODEL], dtype=Out.type.element_ty)
414 - # We still need to write 0s to the result
415 - # tl.store(O_block_ptr,
416 - # acc.to(Out.type.element_ty), boundary_check=(0,1))
417 - # l_ptrs = L + off_z * HQ * MAX_SEQLENS_Q + off_h_q * MAX_SEQLENS_Q
418 - # + offs_m
419 - # We store inf to LSE, not -inf because in the bwd pass,
420 - # we subtract this
421 - # from qk which makes it -inf, such that exp(qk - inf) = 0
422 - # for these masked blocks.
423 - # l = tl.full([BLOCK_M], value=float("inf"), dtype=tl.float32)
424 - # tl.store(l_ptrs, l)
425 - # TODO: Should dropout and return encoded softmax be handled here?
426 - return
427 -
428 - # If MQA / GQA, set the K and V head offsets appropriately.
429 - GROUP_SIZE: tl.constexpr = HQ // HK
430 - off_h_k = off_h_q // GROUP_SIZE if GROUP_SIZE != 1 else off_h_q
431 -
432 - n_extra_tokens = 0
433 - if seqlen_k < BLOCK_N:
434 - n_extra_tokens = BLOCK_N - seqlen_k
435 - elif seqlen_k % BLOCK_N:
436 - n_extra_tokens = seqlen_k % BLOCK_N
437 - padded_head = ACTUAL_BLOCK_DMODEL != BLOCK_DMODEL
438 -
439 - # Compute pointers for all the tensors used in this kernel.
440 - q_offset = (off_z * stride_qz + off_h_q * stride_qh +
441 - cu_seqlens_q_start * stride_qm)
442 - Q_block_ptr = tl.make_block_ptr(
443 - base=Q + q_offset,
444 - shape=(seqlen_q, ACTUAL_BLOCK_DMODEL),
445 - strides=(stride_qm, stride_qk),
446 - offsets=(start_m * BLOCK_M, 0),
447 - block_shape=(BLOCK_M, BLOCK_DMODEL),
448 - order=(1, 0),
449 - )
450 - k_offset = (off_z * stride_kz + off_h_k * stride_kh +
451 - cu_seqlens_k_start * stride_kn)
452 - K_block_ptr = tl.make_block_ptr(
453 - base=K + k_offset,
454 - shape=(ACTUAL_BLOCK_DMODEL, seqlen_k),
455 - strides=(stride_kk, stride_kn),
456 - offsets=(0, 0),
457 - block_shape=(BLOCK_DMODEL, BLOCK_N),
458 - order=(0, 1),
459 - )
460 - v_offset = (off_z * stride_vz + off_h_k * stride_vh +
461 - cu_seqlens_k_start * stride_vk)
462 - V_block_ptr = tl.make_block_ptr(
463 - base=V + v_offset,
464 - shape=(seqlen_k, ACTUAL_BLOCK_DMODEL),
465 - strides=(stride_vk, stride_vn),
466 - offsets=(0, 0),
467 - block_shape=(BLOCK_N, BLOCK_DMODEL),
468 - order=(1, 0),
469 - )
470 - if BIAS_TYPE != 0:
471 - bias_ptr = tl.make_block_ptr(
472 - base=bias + off_h_q * stride_bh,
473 - shape=(seqlen_q, seqlen_k),
474 - strides=(stride_bm, stride_bn),
475 - offsets=(start_m * BLOCK_M, 0),
476 - block_shape=(BLOCK_M, BLOCK_N),
477 - order=(1, 0),
478 - )
479 - else:
480 - bias_ptr = None
481 - if ENABLE_DROPOUT:
482 - batch_philox_offset = philox_offset_base \
483 - + (off_z * HQ + off_h_q) \
484 - * seqlen_q * seqlen_k
485 - else:
486 - batch_philox_offset = 0
487 - # We can ask to return the dropout mask without actually doing any dropout.
488 - # In this case, we return an invalid pointer so indicate the mask is not i
489 - # valid.
490 - # TODO: Fix encoded softmax. It currently uses just h_q in the base offset.
491 - if RETURN_ENCODED_SOFTMAX:
492 - encoded_softmax_block_ptr = tl.make_block_ptr(
493 - base=encoded_softmax + off_h_q * seqlen_q * seqlen_k,
494 - shape=(seqlen_q, seqlen_k),
495 - strides=(seqlen_k, 1),
496 - offsets=(start_m * BLOCK_M, 0),
497 - block_shape=(BLOCK_M, BLOCK_N),
498 - order=(1, 0),
499 - )
500 - else:
501 - encoded_softmax_block_ptr = 0
502 - # initialize pointer to m and l
503 - m_i = tl.full([BLOCK_M], float("-inf"), dtype=tl.float32)
504 - l_i = tl.full([BLOCK_M], 1.0, dtype=tl.float32)
505 - acc = tl.zeros([BLOCK_M, BLOCK_DMODEL], dtype=tl.float32)
506 - # scale sm_scale by log_2(e) and use 2^x in the loop as we do not
507 - # have native e^x support in HW.
508 - qk_scale = sm_scale * 1.44269504089
509 - # Q is loaded once at the beginning and shared by all N blocks.
510 - q = load_fn(Q_block_ptr, True, padded_head, "zero")
511 - q = (q * qk_scale).to(Q_block_ptr.type.element_ty)
512 -
513 - # Here we compute how many full and masked blocks we have.
514 - padded_block_k = n_extra_tokens != 0
515 - is_modulo_mn = not padded_block_k and (seqlen_q % BLOCK_M == 0)
516 - if IS_CAUSAL:
517 - # There are always at least BLOCK_M // BLOCK_N masked blocks.
518 - # Additionally there might be one more due to dissimilar seqlens.
519 - masked_blocks = BLOCK_M // BLOCK_N + (not is_modulo_mn)
520 - else:
521 - # Padding on Q does not need to be masked in the FA loop.
522 - masked_blocks = padded_block_k
523 - # if IS_CAUSAL, not is_modulo_mn does not always result in an additional
524 - # block. In this case we might exceed n_blocks so pick the min.
525 - masked_blocks = min(masked_blocks, n_blocks)
526 - n_full_blocks = n_blocks - masked_blocks
527 - block_min = 0
528 - block_max = n_blocks * BLOCK_N
529 - # Compute for full blocks. Here we set causal to false regardless of its
530 - # value because there is no masking. Similarly we do not need padding.
531 - if n_full_blocks > 0:
532 - block_max = (n_blocks - masked_blocks) * BLOCK_N
533 - acc, l_i, m_i = _attn_fwd_inner(
534 - acc,
535 - l_i,
536 - m_i,
537 - q,
538 - K_block_ptr,
539 - V_block_ptr,
540 - start_m,
541 - seqlen_k,
542 - dropout_p,
543 - philox_seed,
544 - batch_philox_offset,
545 - encoded_softmax_block_ptr,
546 - # _, _, offs_n_causal, masked_blocks, n_extra_tokens, _
547 - block_min,
548 - block_max,
549 - 0,
550 - 0,
551 - 0,
552 - bias_ptr,
553 - # IS_CAUSAL, ....
554 - False,
555 - BLOCK_M,
556 - BLOCK_DMODEL,
557 - BLOCK_N,
558 - offs_m,
559 - offs_n,
560 - # _, MASK_STEPS, ...
561 - PRE_LOAD_V,
562 - False,
563 - ENABLE_DROPOUT,
564 - RETURN_ENCODED_SOFTMAX,
565 - padded_head,
566 - )
567 - block_min = block_max
568 - block_max = n_blocks * BLOCK_N
569 -
570 - tl.debug_barrier()
571 - # Remaining blocks, if any, are full / not masked.
572 - if masked_blocks > 0:
573 - offs_n_causal = offs_n + (seqlen_q - seqlen_k) if IS_CAUSAL else 0
574 - K_block_ptr = tl.advance(K_block_ptr, (0, n_full_blocks * BLOCK_N))
575 - V_block_ptr = tl.advance(V_block_ptr, (n_full_blocks * BLOCK_N, 0))
576 - if bias_ptr is not None:
577 - bias_ptr = tl.advance(bias_ptr, (0, n_full_blocks * BLOCK_N))
578 - if RETURN_ENCODED_SOFTMAX:
579 - encoded_softmax_block_ptr = tl.advance(encoded_softmax_block_ptr,
580 - (0, n_full_blocks))
581 - acc, l_i, m_i = _attn_fwd_inner(
582 - acc,
583 - l_i,
584 - m_i,
585 - q,
586 - K_block_ptr,
587 - V_block_ptr,
588 - start_m,
589 - seqlen_k,
590 - dropout_p,
591 - philox_seed,
592 - batch_philox_offset,
593 - encoded_softmax_block_ptr,
594 - block_min,
595 - block_max,
596 - offs_n_causal,
597 - masked_blocks,
598 - n_extra_tokens,
599 - bias_ptr,
600 - IS_CAUSAL,
601 - BLOCK_M,
602 - BLOCK_DMODEL,
603 - BLOCK_N,
604 - offs_m,
605 - offs_n,
606 - # _, MASK_STEPS, ...
607 - PRE_LOAD_V,
608 - True,
609 - ENABLE_DROPOUT,
610 - RETURN_ENCODED_SOFTMAX,
611 - padded_head,
612 - )
613 - # epilogue
614 - acc = acc / l_i[:, None]
615 - if ENABLE_DROPOUT:
616 - acc = acc / (1 - dropout_p)
617 - # If seqlen_q > seqlen_k but the delta is not a multiple of BLOCK_M,
618 - # then we have one block with a row of all NaNs which come from computing
619 - # softmax over a row of all -infs (-inf - inf = NaN). We check for that here
620 - # and store 0s where there are NaNs as these rows should've been zeroed out.
621 - end_m_idx = (start_m + 1) * BLOCK_M
622 - start_m_idx = start_m * BLOCK_M
623 - causal_start_idx = seqlen_q - seqlen_k
624 - acc = acc.to(Out.type.element_ty)
625 - if IS_CAUSAL: # noqa: SIM102
626 - if causal_start_idx > start_m_idx and causal_start_idx < end_m_idx:
627 - out_mask_boundary = tl.full((BLOCK_DMODEL, ),
628 - causal_start_idx,
629 - dtype=tl.int32)
630 - mask_m_offsets = start_m_idx + tl.arange(0, BLOCK_M)
631 - out_ptrs_mask = (mask_m_offsets[:, None]
632 - >= out_mask_boundary[None, :])
633 - z = 0.0
634 - acc = tl.where(out_ptrs_mask, acc, z.to(acc.type.element_ty))
635 - # write back LSE
636 - # l_ptrs = L + off_z * HQ * MAX_SEQLENS_Q + off_h_q * MAX_SEQLENS_Q + offs_m
637 - # If seqlen_q not multiple of BLOCK_M, we need to mask out the last
638 - # few rows. This is only true for the last M block. For others,
639 - # overflow_size will be -ve
640 - # overflow_size = end_m_idx - seqlen_q
641 - # if overflow_size > 0:
642 - # boundary = tl.full((BLOCK_M,), BLOCK_M - overflow_size, dtype=tl.int32)
643 - # # This is a > check because mask being 0 blocks the store.
644 - # l_ptrs_mask = boundary > tl.arange(0, BLOCK_M)
645 - # tl.store(l_ptrs, m_i + tl.math.log2(l_i), mask=l_ptrs_mask)
646 - # else:
647 - # tl.store(l_ptrs, m_i + tl.math.log2(l_i))
648 -
649 - # write back O
650 - o_offset = (off_z * stride_oz + cu_seqlens_q_start * stride_om +
651 - off_h_q * stride_oh)
652 - O_block_ptr = tl.make_block_ptr(
653 - base=Out + o_offset,
654 - shape=(seqlen_q, ACTUAL_BLOCK_DMODEL),
655 - strides=(stride_om, stride_on),
656 - offsets=(start_m * BLOCK_M, 0),
657 - block_shape=(BLOCK_M, BLOCK_DMODEL),
658 - order=(1, 0),
659 - )
660 - # Need boundary check on this to make sure the padding from the
661 - # Q and KV tensors in both dims are not part of what we store back.
662 - # TODO: Do the boundary check optionally.
663 - tl.store(O_block_ptr, acc, boundary_check=(0, 1))
664 -
665 -
666 -def check_args(
667 - q,
668 - k,
669 - v,
670 - o,
671 - varlen=True,
672 - max_seqlens=None,
673 - cu_seqlens_q=None,
674 - cu_seqlens_k=None,
675 -):
676 - assert q.dim() == k.dim() and q.dim() == v.dim()
677 - if varlen:
678 - assert q.dim() == 3
679 - total_q, nheads_q, head_size = q.shape
680 - total_k, nheads_k, _ = k.shape
681 - assert cu_seqlens_q is not None
682 - assert cu_seqlens_k is not None
683 - assert len(cu_seqlens_q) == len(cu_seqlens_k)
684 - else:
685 - assert q.dim() == 4
686 - batch, nheads_q, seqlen_q, head_size = q.shape
687 - _, nheads_k, seqlen_k, _ = k.shape
688 - assert max_seqlens > 0
689 - assert k.shape == v.shape
690 - assert q.shape[-1] == k.shape[-1] and q.shape[-1] == v.shape[-1]
691 - # TODO: Change assert if we support qkl f8 and v f16
692 - assert q.dtype == k.dtype and q.dtype == v.dtype
693 - assert head_size <= 256
694 - assert o.shape == q.shape
695 - assert (nheads_q % nheads_k) == 0
781 + if continue_condition:
782 + # Now we compute whether we need to exit early due to causal
783 + # masking. This is because for seqlen_q > seqlen_k, M rows of the
784 + # attn scores are completely masked, resulting in 0s written to the
785 + # output, and inf written to LSE. We don't need to do any GEMMs in
786 + # this case. This block of code determines what N is, and if this
787 + # WG is operating on those M rows.
788 + n_blocks = cdiv_fn(seqlen_k, BLOCK_N)
789 + if (IS_CAUSAL):
790 + # If seqlen_q == seqlen_k, the attn scores are a square matrix.
791 + # If seqlen_q != seqlen_k, attn scores are rectangular which
792 + # means the causal mask boundary is bottom right aligned, and
793 + # ends at either the top edge (seqlen_q < seqlen_k) or left
794 + # edge. This captures the decrease in n_blocks if we have a
795 + # rectangular attn matrix
796 + n_blocks_seqlen = cdiv_fn(
797 + (start_m + 1) * BLOCK_M + seqlen_k - seqlen_q, BLOCK_N)
798 + # This is what adjusts the block_max for the current WG, only
799 + # if IS_CAUSAL. Otherwise we want to always iterate through all
800 + # n_blocks
801 + n_blocks = min(n_blocks, n_blocks_seqlen)
802 + # If we have no blocks after adjusting for seqlen deltas, this
803 + # WG is part of the blocks that are all 0. We exit early.
804 + if n_blocks <= 0:
805 + o_offset = (Out + off_z * stride_oz + off_h_q * stride_oh +
806 + cu_seqlens_q_start * stride_om)
807 + o_ptrs = (o_offset + offs_m[:, None] * stride_om +
808 + offs_d[None, :] * stride_on)
809 + acc = tl.zeros([BLOCK_M, BLOCK_DMODEL], dtype=tl.float32)
810 + o_ptrs_mask = (offs_m[:, None] < seqlen_q).broadcast_to(
811 + [BLOCK_M, BLOCK_DMODEL])
812 + # We still need to write 0s to the result
813 + tl.store(o_ptrs, acc, mask=o_ptrs_mask)
814 + # The tensor allocated for L is based on MAX_SEQLENS_Q as
815 + # that is statically known.
816 + l_ptrs = (L + off_z * HQ * MAX_SEQLENS_Q +
817 + off_h_q * MAX_SEQLENS_Q + offs_m)
818 + # We store inf to LSE, not -inf because in the bwd pass,
819 + # we subtract this from qk which makes it -inf, such that
820 + # exp(qk - inf) = 0 for these masked blocks.
821 + l_value = tl.full([BLOCK_M],
822 + value=float("inf"),
823 + dtype=tl.float32)
824 + l_ptrs_mask = offs_m < MAX_SEQLENS_Q
825 + tl.store(l_ptrs, l_value, mask=l_ptrs_mask)
826 + # TODO: Should dropout and return encoded softmax be
827 + # handled here too?
828 + continue_condition = False
829 + # return
830 +
831 + if continue_condition:
832 + # If MQA / GQA, set the K and V head offsets appropriately.
833 + GROUP_SIZE: tl.constexpr = HQ // HK
834 + off_h_k = off_h_q // GROUP_SIZE if GROUP_SIZE != 1 else off_h_q
835 + n_extra_tokens = 0
836 + if seqlen_k < BLOCK_N:
837 + n_extra_tokens = BLOCK_N - seqlen_k
838 + elif seqlen_k % BLOCK_N:
839 + n_extra_tokens = seqlen_k % BLOCK_N
840 + USE_PADDED_HEAD: tl.constexpr = (IS_ACTUAL_BLOCK_DMODEL
841 + != BLOCK_DMODEL)
842 +
843 + # Compute pointers for all the tensors used in this kernel.
844 + q_offset = (Q + off_z * stride_qz + off_h_q * stride_qh +
845 + cu_seqlens_q_start * stride_qm)
846 + q_ptrs = (q_offset + offs_m[:, None] * stride_qm +
847 + offs_d[None, :] * stride_qk)
848 + k_offset = (K + off_z * stride_kz + off_h_k * stride_kh +
849 + cu_seqlens_k_start * stride_kn)
850 + k_ptrs = (k_offset + offs_d[:, None] * stride_kk +
851 + offs_n[None, :] * stride_kn)
852 + v_offset = (V + off_z * stride_vz + off_h_k * stride_vh +
853 + cu_seqlens_k_start * stride_vk)
854 + v_ptrs = (v_offset + offs_n[:, None] * stride_vk +
855 + offs_d[None, :] * stride_vn)
856 + # Compute pointers for all scale tensors used in this kernel.
857 +
858 + IS_EIGHT_BIT_GEMM: tl.constexpr = IS_EIGHT_BIT & (
859 + not IS_EIGHT_BIT_KV)
860 + if IS_EIGHT_BIT:
861 + if k_descale_has_singleton:
862 + k_descale_ptrs = k_descale_ptr
863 + else:
864 + k_descale_ptrs = k_descale_ptr + off_h_k
865 +
866 + if v_descale_has_singleton:
867 + v_descale_ptrs = v_descale_ptr
868 + else:
869 + v_descale_ptrs = v_descale_ptr + off_h_k
870 +
871 + if not IS_EIGHT_BIT_KV:
872 + if q_descale_has_singleton:
873 + q_descale_ptrs = q_descale_ptr
874 + else:
875 + q_descale_ptrs = q_descale_ptr + off_h_q
876 + if USE_P_SCALE:
877 + if p_descale_has_singleton:
878 + p_scale_ptrs = p_scale_ptr
879 + p_descale_ptrs = p_descale_ptr
880 + else:
881 + p_scale_ptrs = p_scale_ptr + off_h_q
882 + p_descale_ptrs = p_descale_ptr + off_h_q
883 +
884 + if USE_BIAS:
885 + bias_offset = off_h_q * stride_bh
886 + bias_ptrs = (bias + bias_offset + offs_m[:, None] * stride_bm +
887 + offs_n[None, :] * stride_bn)
888 + else:
889 + bias_ptrs = None
890 +
891 + if USE_ALIBI:
892 + a_offset = off_z * stride_az + off_h_q * stride_ah
893 + alibi_slope = tl.load(alibi_slopes + a_offset)
894 + else:
895 + alibi_slope = None
896 +
897 + batch_philox_offset = 0
898 + # We can ask to return the dropout mask without doing any
899 + # dropout. In this case, we return an invalid pointer so
900 + # indicate the mask is not valid.
901 + if SHOULD_RETURN_ENCODED_SOFTMAX:
902 + encoded_sm_base = (encoded_softmax +
903 + off_h_q * seqlen_q * seqlen_k)
904 + encoded_sm_ptrs = (encoded_sm_base +
905 + offs_m[:, None] * seqlen_k +
906 + offs_n[None, :])
907 + else:
908 + encoded_sm_ptrs = None
909 + # initialize pointer to m and l
910 + m_i = tl.full([BLOCK_M], float("-inf"), dtype=tl.float32)
911 + l_i = tl.full([BLOCK_M], 1.0, dtype=tl.float32)
912 + acc = tl.zeros([BLOCK_M, BLOCK_DMODEL], dtype=tl.float32)
913 + # scale sm_scale by log_2(e) and use 2^x in the loop as we do
914 + # not have native e^x support in HW.
915 + QK_SCALE: tl.constexpr = SM_SCALE * 1.44269504089
916 + # Q is loaded once at the beginning and shared by all N blocks.
917 + q_ptrs_mask = offs_m[:, None] < seqlen_q
918 + if USE_PADDED_HEAD:
919 + q_ptrs_mask = q_ptrs_mask & (offs_d[None, :]
920 + < IS_ACTUAL_BLOCK_DMODEL)
921 + q = tl.load(q_ptrs, mask=q_ptrs_mask, other=0.0)
922 +
923 + if IS_EIGHT_BIT:
924 + k_descale = tl.load(k_descale_ptrs)
925 + v_descale = tl.load(v_descale_ptrs)
926 + q_descale = None if IS_EIGHT_BIT_KV else tl.load(
927 + q_descale_ptrs)
928 + if USE_P_SCALE:
929 + p_scale = tl.load(p_scale_ptrs)
930 + p_descale = tl.load(p_descale_ptrs)
931 + else:
932 + p_scale = None
933 + p_descale = None
934 + else:
935 + q_descale = None
936 + k_descale = None
937 + v_descale = None
938 + p_scale = None
939 + p_descale = None
940 + # Here we compute how many full and masked blocks we have.
941 + padded_block_k = n_extra_tokens != 0
942 + is_modulo_mn = not padded_block_k and (seqlen_q % BLOCK_M == 0)
943 + if IS_CAUSAL:
944 + # There are always at least BLOCK_M // BLOCK_N masked
945 + # blocks. Additionally there might be one more due to
946 + # dissimilar seqlens.
947 + masked_blocks = BLOCK_M // BLOCK_N + (not is_modulo_mn)
948 + else:
949 + # Padding on Q does not need to be masked in the FA loop.
950 + masked_blocks = padded_block_k
951 + # if IS_CAUSAL, not is_modulo_mn does not always result in an
952 + # additional block. In this case we might exceed n_blocks so
953 + # pick the min.
954 + masked_blocks = min(masked_blocks, n_blocks)
955 + n_full_blocks = n_blocks - masked_blocks
956 + block_min = 0
957 + block_max = n_blocks * BLOCK_N
958 + # Compute for full blocks. Here we set causal to false
959 + # regardless of its actual value because there is no masking.
960 + # Similarly we do not need padding.
961 + if n_full_blocks > 0:
962 + block_max = (n_blocks - masked_blocks) * BLOCK_N
963 + acc, l_i, m_i = _attn_fwd_inner(
964 + acc,
965 + l_i,
966 + m_i,
967 + q,
968 + k_ptrs,
969 + v_ptrs,
970 + bias_ptrs,
971 + stride_kn,
972 + stride_vk,
973 + stride_bn,
974 + start_m,
975 + seqlen_k,
976 + seqlen_q,
977 + philox_seed,
978 + batch_philox_offset,
979 + encoded_sm_ptrs,
980 + # _, _, offs_n_causal, masked_blocks, n_extra_tokens, _
981 + block_min,
982 + block_max,
983 + 0,
984 + 0,
985 + 0,
986 + alibi_slope,
987 + q_descale,
988 + k_descale,
989 + v_descale,
990 + p_scale,
991 + # IS_CAUSAL, ....
992 + False,
993 + BLOCK_M,
994 + BLOCK_DMODEL,
995 + BLOCK_N,
996 + offs_m,
997 + offs_n,
998 + # _, SHOULD_MASK_STEPS, ...
999 + SHOULD_PRE_LOAD_V,
1000 + False,
1001 + SHOULD_RETURN_ENCODED_SOFTMAX,
1002 + USE_PADDED_HEAD,
1003 + IS_ACTUAL_BLOCK_DMODEL,
1004 + QK_SCALE,
1005 + IS_EIGHT_BIT_GEMM,
1006 + USE_P_SCALE,
1007 + IS_EIGHT_BIT_KV,
1008 + QUANT_DTYPE)
1009 + block_min = block_max
1010 + block_max = n_blocks * BLOCK_N
1011 +
1012 + tl.debug_barrier()
1013 + # Remaining blocks, if any, are full / not masked.
1014 + if (masked_blocks > 0):
1015 + if IS_CAUSAL:
1016 + offs_n_causal = offs_n + (seqlen_q - seqlen_k)
1017 + else:
1018 + offs_n_causal = 0
1019 + k_ptrs += n_full_blocks * BLOCK_N * stride_kn
1020 + v_ptrs += n_full_blocks * BLOCK_N * stride_vk
1021 + if USE_BIAS:
1022 + bias_ptrs += n_full_blocks * BLOCK_N * stride_bn
1023 + if SHOULD_RETURN_ENCODED_SOFTMAX:
1024 + encoded_sm_ptrs += n_full_blocks * BLOCK_N
1025 + acc, l_i, m_i = _attn_fwd_inner(
1026 + acc,
1027 + l_i,
1028 + m_i,
1029 + q,
1030 + k_ptrs,
1031 + v_ptrs,
1032 + bias_ptrs,
1033 + stride_kn,
1034 + stride_vk,
1035 + stride_bn,
1036 + start_m,
1037 + seqlen_k,
1038 + seqlen_q,
1039 + philox_seed,
1040 + batch_philox_offset,
1041 + encoded_sm_ptrs,
1042 + block_min,
1043 + block_max,
1044 + offs_n_causal,
1045 + masked_blocks,
1046 + n_extra_tokens,
1047 + alibi_slope,
1048 + q_descale,
1049 + k_descale,
1050 + v_descale,
1051 + p_scale,
1052 + IS_CAUSAL,
1053 + BLOCK_M,
1054 + BLOCK_DMODEL,
1055 + BLOCK_N,
1056 + offs_m,
1057 + offs_n,
1058 + # _, SHOULD_MASK_STEPS, ...
1059 + SHOULD_PRE_LOAD_V,
1060 + True,
1061 + SHOULD_RETURN_ENCODED_SOFTMAX,
1062 + USE_PADDED_HEAD,
1063 + IS_ACTUAL_BLOCK_DMODEL,
1064 + QK_SCALE,
1065 + IS_EIGHT_BIT_GEMM,
1066 + USE_P_SCALE,
1067 + IS_EIGHT_BIT_KV,
1068 + QUANT_DTYPE)
1069 +
1070 + if IS_EIGHT_BIT and not IS_EIGHT_BIT_KV:
1071 + if USE_P_SCALE:
1072 + acc *= p_descale
1073 + acc *= v_descale
1074 +
1075 + # epilogue
1076 + # This helps the compiler do Newton Raphson on l_i vs on acc
1077 + # which is much larger.
1078 + l_recip = 1 / l_i[:, None]
1079 + acc = acc * l_recip
1080 +
1081 + # If seqlen_q > seqlen_k but the delta is not a multiple of
1082 + # BLOCK_M, then we have one block with a row of all NaNs which
1083 + # come from computing softmax over a row of all
1084 + # -infs (-inf - inf = NaN). We check for that here and store 0s
1085 + # where there are NaNs as these rows should've been zeroed out.
1086 + end_m_idx = (start_m + 1) * BLOCK_M
1087 + start_m_idx = start_m * BLOCK_M
1088 + causal_start_idx = seqlen_q - seqlen_k
1089 + if IS_EIGHT_BIT and not IS_EIGHT_BIT_KV: # noqa: SIM102
1090 + if o_descale_ptr is not None:
1091 + acc = quant_fp8(acc, o_descale)
1092 +
1093 + acc = acc.to(Out.type.element_ty)
1094 + if IS_CAUSAL: # noqa: SIM102
1095 + if (causal_start_idx > start_m_idx
1096 + and causal_start_idx < end_m_idx):
1097 + out_mask_boundary = tl.full((BLOCK_DMODEL, ),
1098 + causal_start_idx,
1099 + dtype=tl.int32)
1100 + mask_m_offsets = start_m_idx + tl.arange(0, BLOCK_M)
1101 + out_ptrs_mask = (mask_m_offsets[:, None]
1102 + >= out_mask_boundary[None, :])
1103 + z = tl.zeros((1, ), tl.float32)
1104 + acc = tl.where(out_ptrs_mask, acc,
1105 + z.to(acc.type.element_ty))
1106 + # write back LSE
1107 + l_ptrs = (L + off_z * HQ * MAX_SEQLENS_Q +
1108 + off_h_q * MAX_SEQLENS_Q + offs_m)
1109 + # If seqlen_q not multiple of BLOCK_M, we need to mask out the
1110 + # last few rows. This is only true for the last M block.
1111 + # For others, overflow_size will be -ve
1112 + overflow_size = end_m_idx - seqlen_q
1113 + if overflow_size > 0:
1114 + boundary = tl.full((BLOCK_M, ),
1115 + BLOCK_M - overflow_size,
1116 + dtype=tl.int32)
1117 + l_ptrs_mask = tl.arange(0, BLOCK_M) < boundary
1118 + tl.store(l_ptrs, m_i + tl.math.log2(l_i), mask=l_ptrs_mask)
1119 + else:
1120 + tl.store(l_ptrs, m_i + tl.math.log2(l_i))
1121 +
1122 + # write back O
1123 + o_offset = (Out + off_z * stride_oz + off_h_q * stride_oh +
1124 + cu_seqlens_q_start * stride_om)
1125 + o_ptrs = (o_offset + offs_m[:, None] * stride_om +
1126 + offs_d[None, :] * stride_on)
1127 + o_ptrs_mask = tl.full([BLOCK_M, BLOCK_DMODEL], 1, dtype=tl.int1)
1128 + if overflow_size > 0:
1129 + o_ptrs_mask = o_ptrs_mask & (offs_m[:, None] < seqlen_q)
1130 + if USE_PADDED_HEAD:
1131 + o_ptrs_mask = o_ptrs_mask & (offs_d[None, :]
1132 + < IS_ACTUAL_BLOCK_DMODEL)
1133 + tl.store(o_ptrs, acc.to(Out.dtype.element_ty), mask=o_ptrs_mask)
1134 +
1135 +
1136 +def get_shape_from_layout(q, k, metadata):
1137 + assert metadata.layout in SUPPORTED_LAYOUTS, "Got unsupported layout."
1138 +
1139 + if metadata.layout == 'thd':
1140 + nheads_q, nheads_k = q.shape[1], k.shape[1]
1141 + head_size = q.shape[-1]
1142 + batch = metadata.num_contexts
1143 + elif metadata.layout == 'bhsd':
1144 + batch, nheads_q, _, head_size = q.shape
1145 + nheads_k = k.shape[1]
1146 + elif metadata.layout == 'bshd':
1147 + batch, _, nheads_q, head_size = q.shape
1148 + nheads_k = k.shape[2]
1149 + return batch, nheads_q, nheads_k, head_size
1150 +
1151 +
1152 +def get_strides_from_layout(q, k, v, o, metadata):
1153 + assert metadata.layout in SUPPORTED_LAYOUTS, "Got unsupported layout."
1154 +
1155 + STRIDE_PERMUTATIONS = {
1156 + 'thd': (None, 1, 0, 2),
1157 + 'bhsd': (0, 1, 2, 3),
1158 + 'bshd': (0, 2, 1, 3),
1159 + }
1160 +
1161 + perm = STRIDE_PERMUTATIONS[metadata.layout]
1162 + stride = lambda x, p: (0 if p is None else x.stride(p))
1163 + strides = lambda x: (stride(x, p) for p in perm)
1164 +
1165 + return tuple(strides(x) for x in [q, k, v, o])
1166
1167
1168 class _attention(torch.autograd.Function):
1169
1170 @staticmethod
701 - def forward(
702 - ctx,
703 - q,
704 - k,
705 - v,
706 - o,
707 - cu_seqlens_q,
708 - cu_seqlens_k,
709 - max_seqlens_q,
710 - max_seqlens_k,
711 - causal=False,
712 - sm_scale=1.0,
713 - bias=None,
714 - ):
1171 + def forward(ctx, q, k, v, o, metadata: MetaData):
1172 + # NOTE: a large bias tensor leads to overflow during pointer arithmetic
1173 + if (metadata.bias is not None):
1174 + assert (metadata.bias.numel() < 2**31)
1175 +
1176 if o is None:
716 - o = torch.empty_like(q, dtype=v.dtype)
1177 + if metadata.eight_bit:
1178 + o = torch.empty_like(
1179 + q,
1180 + dtype=metadata.output_dtype if metadata.output_dtype
1181 + is not None else metadata.eight_bit_dtype_torch)
1182 + else:
1183 + o = torch.empty_like(q, dtype=q.dtype)
1184
718 - check_args(
719 - q,
720 - k,
721 - v,
722 - o,
723 - varlen=True,
724 - cu_seqlens_q=cu_seqlens_q,
725 - cu_seqlens_k=cu_seqlens_k,
726 - )
727 - if True: # varlen
728 - total_q, nheads_q, head_size = q.shape
729 - total_k, nheads_k, _ = k.shape
730 - batch = len(cu_seqlens_q) - 1
731 - q_strides = (0, q.stride(1), q.stride(0), q.stride(2))
732 - k_strides = (0, k.stride(1), k.stride(0), k.stride(2))
733 - v_strides = (0, v.stride(1), v.stride(0), v.stride(2))
734 - o_strides = (0, o.stride(1), o.stride(0), o.stride(2))
735 - else:
736 - batch, seqlen_q, nheads_q, head_size = q.shape
737 - _, seqlen_k, nheads_k, _ = k.shape
738 - q_strides = (q.stride(0), q.stride(2), q.stride(1), q.stride(3))
739 - k_strides = (k.stride(0), k.stride(2), k.stride(1), k.stride(3))
740 - v_strides = (v.stride(0), v.stride(2), v.stride(1), v.stride(3))
741 - o_strides = (o.stride(0), o.stride(2), o.stride(1), o.stride(3))
1185 + metadata.check_args(q, k, v, o)
1186 +
1187 + batch, nheads_q, nheads_k, head_size = get_shape_from_layout(
1188 + q, k, metadata)
1189 + q_strides, k_strides, v_strides, o_strides = get_strides_from_layout(
1190 + q, k, v, o, metadata)
1191
1192 # Get closest power of 2 over or equal to 32.
744 - unpadded_head_dims = {32, 64, 128, 256}
745 - if head_size not in unpadded_head_dims:
746 - padded_d_model = None
747 - for i in unpadded_head_dims:
748 - if i > head_size:
749 - padded_d_model = i
750 - break
751 - assert padded_d_model is not None
752 - else:
753 - padded_d_model = head_size
1193 + padded_d_model = 1 << (head_size - 1).bit_length()
1194 + # Smallest head_dim supported is 16. If smaller, the tile in the
1195 + # kernel is padded - there is no padding in memory for any dims.
1196 + padded_d_model = max(padded_d_model, 16)
1197
755 - grid = lambda META: (
756 - triton.cdiv(max_seqlens_q, META["BLOCK_M"]),
757 - nheads_q,
758 - batch,
759 - )
1198 + # encoded_softmax is used to validate dropout behavior vs the
1199 + # PyTorch SDPA math backend reference. We zero this out to give a
1200 + # consistent starting point and then populate it with the output of
1201 + # softmax with the sign bit set according to the dropout mask.
1202 + # The resulting return allows this mask to be fed into the reference
1203 + # implementation for testing only. This return holds no useful output
1204 + # aside from debugging.
1205 + if metadata.return_encoded_softmax:
1206 + encoded_softmax = torch.zeros(
1207 + (q.shape[0], q.shape[1], q.shape[2], k.shape[2]),
1208 + device=q.device,
1209 + dtype=torch.float32)
1210 + else:
1211 + encoded_softmax = None
1212
761 - encoded_softmax = None
1213 + M = torch.empty((batch, nheads_q, metadata.max_seqlens_q),
1214 + device=q.device,
1215 + dtype=torch.float32)
1216
1217 # Seed the RNG so we get reproducible results for testing.
1218 philox_seed = 0x1BF52
1219 philox_offset = 0x1D4B42
1220
767 - if bias is not None:
768 - bias_strides = (
769 - bias.stride(0),
770 - bias.stride(1),
771 - bias.stride(2),
772 - bias.stride(3),
773 - )
1221 + if metadata.bias is not None:
1222 + bias_strides = (metadata.bias.stride(0), metadata.bias.stride(1),
1223 + metadata.bias.stride(2), metadata.bias.stride(3))
1224 else:
1225 bias_strides = (0, 0, 0, 0)
1226
1227 + if metadata.alibi_slopes is not None:
1228 + alibi_strides = (metadata.alibi_slopes.stride(0),
1229 + metadata.alibi_slopes.stride(1))
1230 + else:
1231 + alibi_strides = (0, 0)
1232 +
1233 + if metadata.eight_bit:
1234 + q_descale, k_descale, p_scale, p_descale, v_descale, o_scale = (
1235 + metadata.q_descale, metadata.k_descale, metadata.p_scale,
1236 + metadata.p_descale, metadata.v_descale, metadata.o_scale)
1237 + o_descale = 1.0 / o_scale if o_scale is not None else None
1238 + else:
1239 + q_descale = k_descale = p_scale = None
1240 + p_descale = v_descale = o_descale = None
1241 +
1242 + # number of compute units available
1243 + NUM_CU = torch.cuda.get_device_properties("cuda").multi_processor_count
1244 +
1245 + grid = lambda META: (triton.cdiv(metadata.max_seqlens_q, META[
1246 + 'BLOCK_M']), nheads_q, batch)
1247 +
1248 attn_fwd[grid](
1249 q,
1250 k,
1251 v,
781 - bias,
782 - sm_scale,
783 - None,
1252 + metadata.bias,
1253 + metadata.sm_scale,
1254 + M,
1255 o,
1256 *q_strides,
1257 *k_strides,
1258 *v_strides,
1259 *o_strides,
1260 *bias_strides,
790 - cu_seqlens_q,
791 - cu_seqlens_k,
792 - dropout_p=0.0,
1261 + *alibi_strides,
1262 + q_descale,
1263 + k_descale,
1264 + p_scale,
1265 + p_descale,
1266 + o_descale,
1267 + v_descale,
1268 + q_descale.numel() == 1 if q_descale is not None else False,
1269 + k_descale.numel() == 1 if k_descale is not None else False,
1270 + p_descale.numel() == 1 if p_descale is not None else False,
1271 + v_descale.numel() == 1 if v_descale is not None else False,
1272 + metadata.cu_seqlens_q,
1273 + metadata.cu_seqlens_k,
1274 philox_seed=philox_seed,
1275 philox_offset_base=philox_offset,
1276 encoded_softmax=encoded_softmax,
1277 + alibi_slopes=metadata.alibi_slopes,
1278 HQ=nheads_q,
1279 HK=nheads_k,
798 - ACTUAL_BLOCK_DMODEL=head_size,
799 - MAX_SEQLENS_Q=max_seqlens_q,
800 - MAX_SEQLENS_K=max_seqlens_k,
801 - IS_CAUSAL=causal,
802 - VARLEN=True,
1280 + IS_ACTUAL_BLOCK_DMODEL=head_size,
1281 + MAX_SEQLENS_Q=metadata.max_seqlens_q,
1282 + MAX_SEQLENS_K=metadata.max_seqlens_k,
1283 + IS_CAUSAL=metadata.causal,
1284 + VARLEN=metadata.varlen,
1285 BLOCK_DMODEL=padded_d_model,
804 - BIAS_TYPE=0 if bias is None else 1,
805 - ENABLE_DROPOUT=False,
806 - RETURN_ENCODED_SOFTMAX=False,
807 - )
1286 + USE_BIAS=metadata.bias is not None,
1287 + USE_ALIBI=metadata.alibi_slopes is not None,
1288 + SHOULD_RETURN_ENCODED_SOFTMAX=metadata.return_encoded_softmax,
1289 + IS_EIGHT_BIT=metadata.eight_bit,
1290 + USE_P_SCALE=metadata.eight_bit and metadata.use_p_scale,
1291 + IS_EIGHT_BIT_KV=metadata.eight_bit and metadata.eight_bit_kv,
1292 + NUM_CU=NUM_CU,
1293 + B=batch,
1294 + QUANT_DTYPE=metadata.eight_bit_dtype_triton)
1295
1296 ctx.grid = grid
810 - ctx.sm_scale = sm_scale
1297 + ctx.sm_scale = metadata.sm_scale
1298 ctx.BLOCK_DMODEL = head_size
812 - ctx.causal = causal
813 - ctx.dropout_p = 0.0
1299 + ctx.causal = metadata.causal
1300 + ctx.alibi_slopes = metadata.alibi_slopes
1301 ctx.philox_seed = philox_seed
1302 ctx.philox_offset = philox_offset
1303 ctx.encoded_softmax = encoded_softmax
817 - ctx.return_encoded_softmax = False
1304 + ctx.return_encoded_softmax = metadata.return_encoded_softmax
1305 return o, encoded_softmax
1306
1307
821 -triton_attention = _attention.apply
1308 +triton_attention_rocm = _attention.apply
1309 +
1310 +
1311 +def scale_fp8(t, scale=None):
1312 + t_scaled, scale_out = ops.scaled_fp8_quant(t.reshape(-1, t.shape[-1]),
1313 + scale)
1314 + return t_scaled.reshape(t.shape), scale_out
1315 +
1316 +
1317 +def maybe_quantize_fp8(t, scale):
1318 + eight_bit_dtype = current_platform.fp8_dtype()
1319 + if t.dtype != eight_bit_dtype:
1320 + t, _ = scale_fp8(t, scale)
1321 + return t
1322 +
1323 +
1324 +def check_and_maybe_quantize_qkv(q, k, v, fp8_scales):
1325 + (q_scale, k_scale, v_scale, p_scale) = fp8_scales
1326 +
1327 + q = maybe_quantize_fp8(q, q_scale)
1328 + k = maybe_quantize_fp8(k, k_scale)
1329 + v = maybe_quantize_fp8(v, v_scale)
1330 +
1331 + return q, k, v
1332 +
1333 +
1334 +# query - [num_tokens, num_heads, head_size]
1335 +# key - [num_tokens, num_kv_heads, head_size]
1336 +# value - [num_tokens, num_kv_heads, head_size
1337 +# output - [num_tokens, num_heads, head_size]
1338 +def triton_attention(
1339 + q: torch.Tensor,
1340 + k: torch.Tensor,
1341 + v: torch.Tensor,
1342 + o: torch.Tensor,
1343 + cu_seqlens_q: torch.Tensor,
1344 + cu_seqlens_k: torch.Tensor,
1345 + max_seqlens_q: int,
1346 + max_seqlens_k: int,
1347 + causal: bool = False,
1348 + sm_scale: float = 1.0,
1349 + bias: Optional[torch.Tensor] = None,
1350 + fp8_scales: Optional[tuple[float, ...]] = None,
1351 + input_scale: Optional[torch.Tensor] = None,
1352 +) -> torch.Tensor:
1353 + if fp8_scales is not None:
1354 + q_descale, k_descale, v_descale, p_scale = fp8_scales
1355 + else:
1356 + q_descale = k_descale = v_descale = p_scale = None
1357 +
1358 + attn_metadata = MetaData(sm_scale=sm_scale,
1359 + max_seqlens_q=max_seqlens_q,
1360 + max_seqlens_k=max_seqlens_k,
1361 + causal=causal,
1362 + bias=bias,
1363 + output_dtype=q.dtype,
1364 + cu_seqlens_q=cu_seqlens_q,
1365 + cu_seqlens_k=cu_seqlens_k,
1366 + q_descale=q_descale,
1367 + k_descale=k_descale,
1368 + v_descale=v_descale,
1369 + p_scale=p_scale,
1370 + o_scale=input_scale)
1371 +
1372 + if fp8_scales is not None:
1373 + q, k, v = check_and_maybe_quantize_qkv(q, k, v, fp8_scales)
1374 +
1375 + return triton_attention_rocm(q, k, v, o, attn_metadata)