@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 1b7cfd5a36
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, 08:30 UTC

[ROCm][V0][Attention] Revert to the previous FA triton kernel (#18226)

Signed-off-by: Gregory Shtrasberg <Gregory.Shtrasberg@amd.com>

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