diff --git a/flash_attn/flash_attn_triton_amd/bwd_ref.py b/flash_attn/flash_attn_triton_amd/bwd_ref.py index 8bdccb1d329..cb1637157a0 100644 --- a/flash_attn/flash_attn_triton_amd/bwd_ref.py +++ b/flash_attn/flash_attn_triton_amd/bwd_ref.py @@ -7,7 +7,8 @@ DEBUG_CORE = False def attention_backward_core_ref_impl( - do, q, k, v, o, softmax_lse, sm_scale, causal, dropout_p, philox_seed, philox_offset, alibi_slopes, use_exp2 + do, q, k, v, o, softmax_lse, sm_scale, causal, window_size_left, window_size_right, + dropout_p, philox_seed, philox_offset, alibi_slopes, use_exp2 ): if DEBUG_CORE: print() @@ -16,10 +17,12 @@ def attention_backward_core_ref_impl( print("q:", q, q.shape) print("k:", k, k.shape) print("v:", v, v.shape) - print("o:", o, o.shape) # is a bad number + print("o:", o, o.shape) print("softmax_lse:", softmax_lse, softmax_lse.shape) print("sm_scale:", sm_scale) print("causal:", causal) + print("window_size_left:", window_size_left) + print("window_size_right:", window_size_right) print("dropout_p:", dropout_p) print("philox_seed:", philox_seed) print("philox_offset:", philox_offset) @@ -33,7 +36,6 @@ def attention_backward_core_ref_impl( o = o.to(torch.float32) softmax_lse = softmax_lse.to(torch.float32) - # recompute attention_scores. Make sure it matches the forward impl. i.e. It use float32 attention_scores = torch.matmul(q, k.transpose(-2, -1)) if DEBUG_CORE: @@ -50,51 +52,95 @@ def attention_backward_core_ref_impl( print("alibi_slopes:", alibi_slopes, alibi_slopes.shape) alibi_bias = compute_alibi_tensor_ref(alibi_slopes, L_q, L_k) alibi_bias = alibi_bias.reshape(-1, L_q, L_k) - if True: + if DEBUG_CORE: print("alibi_bias:", alibi_bias, alibi_bias.shape) attention_scaled_scores = attention_scaled_scores + alibi_bias if DEBUG_CORE: print("attention_scaled_scores after alibi:", attention_scaled_scores, attention_scaled_scores.shape) - # Apply causal mask if necessary - if causal: - L_q, L_k = q.shape[1], k.shape[1] - row_idx = torch.arange(L_q, device=q.device).unsqueeze(1) - col_idx = torch.arange(L_k, device=q.device).unsqueeze(0) - col_offset = L_q-L_k - causal_mask = row_idx >= (col_offset + col_idx) + # Apply masks + L_q, L_k = q.shape[1], k.shape[1] + row_idx = torch.arange(L_q, device=q.device).unsqueeze(1) + col_idx = torch.arange(L_k, device=q.device).unsqueeze(0) + col_offset = L_k - L_q + + mask_applied = False + if causal and (window_size_left, window_size_right) == (-1, -1): + # Pure causal: ensure query doesn't attend to future keys + mask = row_idx >= (col_idx - col_offset) + mask_applied = True + if DEBUG_CORE: + print("causal_mask:", mask) + elif (window_size_left, window_size_right) != (-1, -1): + # Handle the case where window sizes exceed sequence length + if window_size_left >= L_k: + window_size_left = -1 # No left limit + if window_size_right >= L_k: + window_size_right = -1 # No right limit + + if causal: + # Causal + sliding window: ensure we don't attend to future + window_size_right = min(window_size_right, 0) if window_size_right != -1 else 0 + + # Create sliding window mask + # Each query at position i attends to keys in [i + offset - left, i + offset + right] + if window_size_left == -1 and window_size_right == -1: + # No window restriction + mask = torch.ones((L_q, L_k), dtype=torch.bool, device=q.device) + else: + mask = torch.ones((L_q, L_k), dtype=torch.bool, device=q.device) + if window_size_left != -1: + # Each query at position i attends to keys from position (i - left) accounting for offset + mask = mask & (col_idx >= (row_idx + col_offset - window_size_left)) + if window_size_right != -1: + # Each query at position i attends to keys up to position (i + right) accounting for offset + mask = mask & (col_idx <= (row_idx + col_offset + window_size_right)) + + # Apply causal constraint + if causal: + causal_mask = row_idx >= (col_idx - col_offset) + mask = mask & causal_mask + + mask_applied = True if DEBUG_CORE: - print("causal_mask:", causal_mask) - # set -inf to places the causal mask is false + print(f"sliding_window_mask (left={window_size_left}, right={window_size_right}):", mask) + + # Apply the mask if created + if mask_applied: attention_scaled_scores = attention_scaled_scores.masked_fill( - torch.logical_not(causal_mask.unsqueeze(0)), float('-inf') + torch.logical_not(mask.unsqueeze(0)), float('-inf') ) if DEBUG_CORE: - print("attention_scaled_scores after causal:", attention_scaled_scores, attention_scaled_scores.shape) + print("attention_scaled_scores after masking:", attention_scaled_scores, attention_scaled_scores.shape) # compute probabilities using softmax_lse if use_exp2: RCP_LN = 1 / math.log(2) attention_scaled_scores_base2 = attention_scaled_scores * RCP_LN softmax_lse_base2 = softmax_lse * RCP_LN - softmax_lse_3d = softmax_lse_base2.unsqueeze(-1) + softmax_lse_3d = softmax_lse_base2.unsqueeze(-1) p = torch.exp2(attention_scaled_scores_base2 - softmax_lse_3d) else: - softmax_lse_3d = softmax_lse.unsqueeze(-1) + softmax_lse_3d = softmax_lse.unsqueeze(-1) p = torch.exp(attention_scaled_scores - softmax_lse_3d) + + # Zero out positions outside the mask + if mask_applied: + p = p.masked_fill(torch.logical_not(mask.unsqueeze(0)), 0.0) + if DEBUG_CORE: print("softmax_lse_3d:", softmax_lse_3d, softmax_lse_3d.shape) print("p:", p, p.shape) if dropout_p > 0.0: rand_vals = torch.rand(p.shape, generator=torch.Generator(device=p.device).manual_seed(philox_seed), device=p.device, dtype=p.dtype) - dropout_mask, dropout_scale = rand_vals > dropout_p, (1.0 / (1 - dropout_p)) - if DEBUG: + dropout_mask, dropout_scale = rand_vals > dropout_p, (1.0 / (1 - dropout_p)) + if DEBUG_CORE: print("dropout_scale:", dropout_scale) print("dropout_mask:", dropout_mask) p_drop = torch.where(dropout_mask, p, torch.zeros_like(p)) - p_drop_scaled = p_drop * dropout_scale + p_drop_scaled = p_drop * dropout_scale if DEBUG_CORE: print("dropout_scale:", dropout_scale) print("p_drop:", p_drop, p_drop.shape) @@ -107,7 +153,7 @@ def attention_backward_core_ref_impl( # compute dp dp_dropout = torch.matmul(do, v.transpose(-2, -1)) - dp = torch.where(dropout_mask, dp_dropout , torch.zeros_like(dp_dropout)) * dropout_scale + dp = torch.where(dropout_mask, dp_dropout, torch.zeros_like(dp_dropout)) * dropout_scale if DEBUG_CORE: print("dp_dropout:", dp_dropout, dp_dropout.shape) print("dp:", dp, dp.shape) @@ -127,9 +173,14 @@ def attention_backward_core_ref_impl( delta = torch.sum(o * do, axis=-1).unsqueeze(-1) else: delta = torch.sum(p * dp, axis=-1).unsqueeze(-1) - if DEBUG: + if DEBUG_CORE: print("delta:", delta, delta.shape) dscores_scaled = p * (dp - delta) + + # Zero out gradients for positions outside the mask + if mask_applied: + dscores_scaled = dscores_scaled.masked_fill(torch.logical_not(mask.unsqueeze(0)), 0.0) + ds = dscores_scaled * sm_scale if DEBUG_CORE: print("dscores_scaled:", dscores_scaled, dscores_scaled.shape) @@ -167,6 +218,8 @@ def attention_varlen_backward_pytorch_ref_impl( softmax_lse, sm_scale, causal, + window_size_left, + window_size_right, layout, cu_seqlens_q, cu_seqlens_k, @@ -255,6 +308,8 @@ def attention_varlen_backward_pytorch_ref_impl( softmax_lse_i, sm_scale, causal, + window_size_left, + window_size_right, dropout_p, philox_seed, philox_offset, @@ -270,7 +325,8 @@ def attention_varlen_backward_pytorch_ref_impl( if group_size != 1: # Reshape dq_i and delta_i back to original shape dq_i = dq_i.view(dq_i.shape[0], nheads_k, group_size, head_dim) - delta_i = delta_i.view(delta_i.shape[0], nheads_k, group_size) + L_q_i = delta_i.shape[1] + delta_i = delta_i.view(nheads_k, group_size, L_q_i) # Sum dk_i and dv_i over group dimension dk_i = dk_i.view(dk_i.shape[0], nheads_k, group_size, head_dim) dv_i = dv_i.view(dv_i.shape[0], nheads_k, group_size, head_dim) @@ -278,7 +334,7 @@ def attention_varlen_backward_pytorch_ref_impl( dv_i = dv_i.sum(dim=2) # Reshape dq_i back to [L_q_i, nheads_q, head_dim] dq_i = dq_i.reshape(dq_i.shape[0], nheads_q, head_dim) - delta_i = delta_i.reshape(delta_i.shape[0], nheads_q) + delta_i = delta_i.reshape(nheads_q, L_q_i) else: # No need to reshape pass @@ -300,6 +356,8 @@ def attention_vanilla_backward_pytorch_ref_impl( softmax_lse, sm_scale, causal, + window_size_left, + window_size_right, layout, dropout_p, philox_seed, @@ -366,6 +424,8 @@ def attention_vanilla_backward_pytorch_ref_impl( softmax_lse, sm_scale, causal, + window_size_left, + window_size_right, dropout_p, philox_seed, philox_offset, @@ -421,6 +481,8 @@ def attention_backward_pytorch_ref_impl( sm_scale: float, alibi_slopes: Optional[torch.Tensor], causal: bool, + window_size_left: int, + window_size_right: int, layout: Literal["bshd", "bhsd", "thd"], cu_seqlens_q: Optional[torch.Tensor], cu_seqlens_k: Optional[torch.Tensor], @@ -441,6 +503,8 @@ def attention_backward_pytorch_ref_impl( softmax_lse, sm_scale, causal, + window_size_left, + window_size_right, layout, cu_seqlens_q, cu_seqlens_k, @@ -462,6 +526,8 @@ def attention_backward_pytorch_ref_impl( softmax_lse, sm_scale, causal, + window_size_left, + window_size_right, layout, dropout_p, philox_seed, @@ -476,4 +542,4 @@ def attention_backward_pytorch_ref_impl( dk.copy_(dk_ref.to(dk.dtype)) dq.copy_(dq_ref.to(dq.dtype)) - return delta + return delta \ No newline at end of file diff --git a/flash_attn/flash_attn_triton_amd/fwd_decode.py b/flash_attn/flash_attn_triton_amd/fwd_decode.py index e165d714876..d3f7f9c32b9 100644 --- a/flash_attn/flash_attn_triton_amd/fwd_decode.py +++ b/flash_attn/flash_attn_triton_amd/fwd_decode.py @@ -133,6 +133,9 @@ def _fwd_kernel_splitK( USE_ALIBI: tl.constexpr, PADDED_HEAD: tl.constexpr, GROUP_SIZE: tl.constexpr, + USE_SLIDING_WINDOW: tl.constexpr, + WINDOW_SIZE_LEFT: tl.constexpr, + WINDOW_SIZE_RIGHT: tl.constexpr, ): # get program ids pid_m = tl.program_id(0) @@ -297,35 +300,62 @@ def _fwd_kernel_splitK( alibi_bias = -1 * alibi_slope * relative_pos qk += (alibi_bias * 1.44269504) - # Apply causal mask if IS_CAUSAL is True - if IS_CAUSAL: - row_idx = pid_m * BLOCK_M + tl.arange(0, BLOCK_M) - col_idx = start_n + tl.arange(0, BLOCK_N) - - # create a N_CTX_Q x kv_len causal mask - col_offset = N_CTX_Q - N_CTX_K_FINAL - causal_mask = row_idx[:, None] >= (col_offset + col_idx[None, :]) + # ------------------------------------------------------------------ + # masking + # ------------------------------------------------------------------ + if USE_SLIDING_WINDOW: + row_idx = pid_m * BLOCK_M + tl.arange(0, BLOCK_M) # q positions + col_idx = start_n + tl.arange(0, BLOCK_N) # k positions + row = row_idx[:, None] # [M,1] + col = col_idx[None, :] # [1,N] + + if IS_CAUSAL: + # -------- causal + window -------- + diag = N_CTX_K_FINAL - N_CTX_Q # sk-sq + causal_ok = col <= row + diag + if WINDOW_SIZE_LEFT < 0: # only right window + win_ok = col <= row + diag + WINDOW_SIZE_RIGHT + else: # both sides + win_ok = ((col >= row + diag - WINDOW_SIZE_LEFT) & + (col <= row + diag + WINDOW_SIZE_RIGHT)) + mask = ~(causal_ok & win_ok) # True ⇒ -inf + else: + # -------- non-causal window -------- + sk, sq = N_CTX_K_FINAL, N_CTX_Q + if WINDOW_SIZE_LEFT < 0: + mask = col > row + (sk - sq) + WINDOW_SIZE_RIGHT + else: + right = tl.minimum(row + (sk - sq) + WINDOW_SIZE_RIGHT, sk) + left = row + (sk - sq) - WINDOW_SIZE_LEFT + mask = (col > right) | (col < left) + qk = tl.where(mask, float("-inf"), qk) + else: + if IS_CAUSAL: + row_idx = pid_m * BLOCK_M + tl.arange(0, BLOCK_M) + col_idx = start_n + tl.arange(0, BLOCK_N) - # Apply the mask - qk = tl.where(causal_mask, qk, float("-inf")) + # create a N_CTX_Q x kv_len causal mask + col_offset = N_CTX_Q - N_CTX_K_FINAL + causal_mask = row_idx[:, None] >= (col_offset + col_idx[None, :]) + + # Apply the mask + qk = tl.where(causal_mask, qk, float("-inf")) # TODO: This is slow, and only needed at the last iteration. # Maybe we can unroll the last iteration instead? if BOUNDS_CHECKS_N: qk = tl.where(tl.arange(0, BLOCK_N) < hi - start_n, qk, float("-inf")) - # -- compute scaling constant --- - m_i_new = tl.maximum(m_i, tl.max(qk, 1)) - if IS_CAUSAL: - alpha = tl.math.exp2(tl.where(m_i > float("-inf"), m_i - m_i_new, float("-inf"))) - else: - alpha = tl.math.exp2(m_i - m_i_new) - # cause of nan because subtracting infs - if IS_CAUSAL: - qk = tl.where(qk > float("-inf"), qk - m_i_new[:, None], float("-inf")) - else: - qk = qk - m_i_new[:, None] - + m_i_new = tl.maximum(m_i, tl.max(qk, 1)) # per-row max so far + + # rows that are *all* -inf after masking + valid = m_i_new > float("-inf") + + # scale previous partial sums safely + alpha = tl.where(valid, tl.math.exp2(m_i - m_i_new), 0.0) + + # subtract the row max only on valid rows + qk = tl.where(valid[:, None], qk - m_i_new[:, None], float("-inf")) p = tl.math.exp2(qk) # -- update m_i and l_i -- @@ -387,7 +417,6 @@ def _splitK_reduce( split_k: tl.constexpr, splitK_pow2: tl.constexpr, MASK_SPLITK: tl.constexpr, - IS_CAUSAL: tl.constexpr, PADDED_HEAD: tl.constexpr, ): # get pids @@ -426,23 +455,15 @@ def _splitK_reduce( g_m = tl.max(l_m, axis=0) - if IS_CAUSAL: - l_m_offset = l_m - g_m - alpha = tl.where(l_m_offset > float("-inf"), tl.math.exp2(l_m_offset), 0.0) - else: - alpha = tl.math.exp2(l_m - g_m) + alpha = tl.where(l_m > float("-inf"), tl.math.exp2(l_m - g_m), 0.0) # read sum l_sum *= alpha g_sum = tl.sum(l_sum, axis=0) acc = acc * alpha[:, None] - if IS_CAUSAL: - # Avoid division by zero - g_sum_safe = tl.where(g_sum > 0, g_sum, 1.0) - acc_out = tl.sum(acc, axis=0) / g_sum_safe - else: - acc_out = tl.sum(acc, axis=0) / g_sum + g_sum_safe = tl.where(g_sum > 0, g_sum, 1.0) + acc_out = tl.sum(acc, axis=0) / g_sum_safe # Store output z_id = pid_zhg // (H * G) @@ -454,11 +475,10 @@ def _splitK_reduce( # Store lse l_ptrs = LSE + pid_zhg * stride_lse_zhg + pid_m - if IS_CAUSAL: - lse = tl.where(g_sum > 0, (g_m + tl.math.log2(g_sum)) / 1.44269504, g_m) - tl.store(l_ptrs, lse) - else: - tl.store(l_ptrs, (g_m + tl.math.log2(g_sum)) / 1.44269504) + lse_val = tl.where(g_sum > 0, + (g_m + tl.math.log2(g_sum)) / 1.44269504, + g_m) + tl.store(l_ptrs, lse_val) @triton.jit @@ -571,10 +591,12 @@ def attention_decode_forward_triton_impl( v_new: Optional[torch.Tensor], out: torch.Tensor, sm_scale: float, - causal: bool, + causal: bool, + window_size_left: int, + window_size_right: int, alibi_slopes: Optional[torch.Tensor], layout: Literal["bshd"], - cache_seqlens: Optional[Union[(int, torch.Tensor)]], + cache_seqlens: Optional[torch.Tensor], cache_batch_idx: Optional[torch.Tensor], ): # triton configs @@ -586,8 +608,9 @@ def attention_decode_forward_triton_impl( # kernel_configs is_new_kv = True if k_new is not None and v_new is not None else False - use_alibi = False if alibi_slopes is None else True + use_alibi, (stride_az, stride_ah) = True if alibi_slopes is not None else False, alibi_slopes.stride() if alibi_slopes is not None else (None, None) use_cache_seqlens = cache_seqlens is not None + use_sliding_window = window_size_left != -1 or window_size_right != -1 SPLIT_K = None NUM_QUANT_GROUPS = 1 @@ -602,11 +625,6 @@ def attention_decode_forward_triton_impl( ( _, seqlen_kn, nheads_kn, dim_kn), (stride_kn_z, stride_kn_h, stride_kn_n, stride_kn_d) = (None, None, None, None), (None, None, None, None) (_, seqlen_vn, nheads_vn, dim_vn), (stride_vn_z, stride_vn_h, stride_vn_n, stride_vn_d) = (None, None, None, None), (None, None, None, None) (_, seqlen_o, nheads_o, dim_o), (stride_oz, stride_oh, stride_om, stride_od) = get_shape_and_strides_from_layout(out, layout) - if use_alibi: - stride_az, stride_ah = alibi_slopes.stride() - else: - stride_az, stride_ah = (None, None) - assert dim_q == dim_kc == dim_vc, f"Dimensions must match: {dim_q}, {dim_kc}, {dim_vc}" # add extra information needed by the kernels @@ -656,7 +674,7 @@ def attention_decode_forward_triton_impl( stride_mzhg, stride_m2, stride_ms, stride_mm = metadata.stride() stride_lse_zhg, stride_lse_m = lse.stride() - if False: + if DEBUG: print("batch_size, seqlen_q, nheads_q, dim_q", (batch_size, seqlen_q, nheads_q, dim_q)) print("_, seqlen_kc, nheads_kc, dim_kc", (_, seqlen_kc, nheads_kc, dim_kc)) print("dim_padded:", dim_padded) @@ -748,6 +766,9 @@ def attention_decode_forward_triton_impl( USE_ALIBI=use_alibi, PADDED_HEAD=is_padded_head, GROUP_SIZE=group_size, + USE_SLIDING_WINDOW=use_sliding_window, + WINDOW_SIZE_LEFT=window_size_left, + WINDOW_SIZE_RIGHT=window_size_right, num_warps=num_warps_fwd, num_stages=num_stages, ) @@ -809,8 +830,7 @@ def attention_decode_forward_triton_impl( split_k=split_k, splitK_pow2=splitK_pow2, MASK_SPLITK=mask_split_k, - IS_CAUSAL=causal, PADDED_HEAD=is_padded_head, num_warps=num_warps_reduce) - return lse + return lse \ No newline at end of file diff --git a/flash_attn/flash_attn_triton_amd/fwd_prefill.py b/flash_attn/flash_attn_triton_amd/fwd_prefill.py index e33982bb6a7..59fe8bfaf4e 100644 --- a/flash_attn/flash_attn_triton_amd/fwd_prefill.py +++ b/flash_attn/flash_attn_triton_amd/fwd_prefill.py @@ -1,8 +1,9 @@ +import os import torch import triton import triton.language as tl from typing import Literal, Optional, Union -from .utils import DROPOUT_USE_PYTORCH, DROPOUT_DUMP, AUTOTUNE, compute_alibi_block, compute_fp8_scaling_factors, get_arch, is_cdna, is_fp8, is_rdna, create_dropout_mask +from .utils import DROPOUT_USE_PYTORCH, DROPOUT_DUMP, AUTOTUNE, compute_alibi_block, compute_fp8_scaling_factors, get_arch, is_cdna, is_fp8, is_rdna, create_dropout_mask, get_fwd_prefill_autotune_configs DEBUG = False @@ -10,70 +11,192 @@ tl_DROPOUT_USE_PYTORCH: tl.constexpr = triton.language.constexpr(DROPOUT_USE_PYTORCH) tl_DROPOUT_DUMP: tl.constexpr = triton.language.constexpr(DROPOUT_DUMP) -# Convenience function to load with optional boundary checks. -# "First" is the major dim, "second" is the minor dim. +fwd_prefill_autotune_configs, fwd_prefill_autotune_keys = get_fwd_prefill_autotune_configs() + @triton.jit -def load_fn(ptrs, offset_first, offset_second, boundary_first, boundary_second): - if offset_first is not None and offset_second is not None: - mask = (offset_first[:, None] < boundary_first) & \ - (offset_second[None, :] < boundary_second) - tensor = tl.load(ptrs, mask=mask, other=0.0) - elif offset_first is not None: - mask = offset_first[:, None] < boundary_first - tensor = tl.load(ptrs, mask=mask, other=0.0) - elif offset_second is not None: - mask = offset_second[None, :] < boundary_second - tensor = tl.load(ptrs, mask=mask, other=0.0) - else: - tensor = tl.load(ptrs) - return tensor +def _attn_fwd_no_mask(acc, l_i, m_i, + q, k_base_ptrs, v_base_ptrs, bias_base_ptrs, + stride_kn, stride_vk, stride_bn, stride_sn, + start_m, seqlen_k, seqlen_q, + dropout_p, philox_seed, philox_base_ptrs, + sd_mask_base_ptrs, dropout_mask_base_ptrs, + offs_m, offs_n, offs_d, + block_min, block_max, alibi_slope, + descale_q, descale_k, descale_v, IS_FP8: tl.constexpr, FP8_MAX: tl.constexpr, + BLOCK_M: tl.constexpr, BLOCK_DMODEL: tl.constexpr, BLOCK_N: tl.constexpr, + PRE_LOAD_V: tl.constexpr, + ENABLE_DROPOUT: tl.constexpr, PADDED_HEAD: tl.constexpr, + ACTUAL_BLOCK_DMODEL: tl.constexpr, SM_SCALE: tl.constexpr, USE_ALIBI: tl.constexpr, USE_EXP2: tl.constexpr, + RETURN_SCORES: tl.constexpr, ACCUMULATOR_TYPE): + if USE_EXP2: + RCP_LN2: tl.constexpr = 1.4426950408889634 + + # loop over k, v, and update accumulator + for start_n in range(block_min, block_max, BLOCK_N): + # get ptrs + k_ptrs = k_base_ptrs + start_n * stride_kn + v_ptrs = v_base_ptrs + start_n * stride_vk + + kv_offs_n = start_n + tl.arange(0, BLOCK_N) + if PADDED_HEAD: + k_mask, k_mask_other = (offs_d[:, None] < ACTUAL_BLOCK_DMODEL), 0.0 + v_mask, v_mask_other = (offs_d[None, :] < ACTUAL_BLOCK_DMODEL), 0.0 + + # load k and if preload_v then v + k = tl.load(k_ptrs, mask=k_mask, other=k_mask_other) if PADDED_HEAD else tl.load(k_ptrs) + if PRE_LOAD_V: + v = tl.load(v_ptrs, mask=v_mask, other=v_mask_other) if PADDED_HEAD else tl.load(v_ptrs) + + # setup qk accumlator + qk = tl.zeros([BLOCK_M, BLOCK_N], dtype=ACCUMULATOR_TYPE) + + # -- compute qk ---- + if IS_FP8 : + qk += (tl.dot(q, k) * descale_q * descale_k) + else: + qk += tl.dot(q, k) + qk_scaled = qk * SM_SCALE + + if USE_ALIBI: + # compute the global position of each token within the sequence + q_offs_m = start_m * BLOCK_M + tl.arange(0, BLOCK_M) + alibi_block = compute_alibi_block(alibi_slope, seqlen_q, seqlen_k, q_offs_m, + kv_offs_n) + qk_scaled += alibi_block + + # compute qk mask + qk_mask = (offs_m[:, None] < seqlen_q) & (kv_offs_n[None, :] < seqlen_k) + + # compute bias + if bias_base_ptrs is not None: + bias_ptrs = bias_base_ptrs + start_n * stride_bn + bias = tl.load(bias_ptrs, mask=qk_mask, other=0.0) + qk_scaled += bias + + # get max scores so far + m_ij = tl.maximum(m_i, tl.max(qk_scaled, 1)) + + # scale and subtract max + q_shifted = tl.where(m_ij[:, None] == float("-inf"), + float("-inf"), + qk_scaled - m_ij[:, None]) + + # Compute scaled QK and softmax probabilities + if USE_EXP2: + p = tl.math.exp2(q_shifted * RCP_LN2) + else: + p = tl.math.exp(q_shifted) + + # CAVEAT: Must update l_ij before applying dropout + l_ij = tl.sum(p, 1) + if ENABLE_DROPOUT: + dropout_mask_ptrs = dropout_mask_base_ptrs + start_n * stride_sn + sd_mask_ptrs = sd_mask_base_ptrs + start_n * stride_sn + philox_ptrs = philox_base_ptrs + start_n * stride_sn + if tl_DROPOUT_USE_PYTORCH: + dropout_mask = tl.load(dropout_mask_ptrs, mask=qk_mask) + else: + rng_output = tl.rand(philox_seed, philox_ptrs) # TODO: use tl.randint for better performance + dropout_mask = rng_output > dropout_p + if tl_DROPOUT_DUMP: + tl.store(dropout_mask_ptrs, dropout_mask, mask=qk_mask) + + # return scores with negative values for dropped vals + sd_mask = tl.where(dropout_mask, p, -p) + tl.store(sd_mask_ptrs, sd_mask, mask=qk_mask) + + # apply dropout mask in place + p = tl.where(dropout_mask, p, 0.0) + elif RETURN_SCORES: + # NOTE: the returned score is not the same as the reference because we need to adjust as we find new maxes per block. We are not doing that + sd_mask_ptrs = sd_mask_base_ptrs + start_n * stride_sn + tl.store(sd_mask_ptrs, p, mask=qk_mask) + + # -- update output accumulator -- + # alpha is an adjustment factor for acc and li as we loop and find new maxes + # store the diff in maxes to adjust acc and li as we discover new maxes + m_diff = tl.where(m_ij == float("-inf"), + float("-inf"), + m_i - m_ij) + if USE_EXP2: + alpha = tl.math.exp2(m_diff * RCP_LN2) + else: + alpha = tl.math.exp(m_diff) + acc = acc * alpha[:, None] + if not PRE_LOAD_V: + v = tl.load(v_ptrs, mask=v_mask, other=v_mask_other) if PADDED_HEAD else tl.load(v_ptrs) + + # -- update m_i and l_i + l_i = l_i * alpha + l_ij + m_i = m_ij + + if IS_FP8: + scale_p, descale_p = compute_fp8_scaling_factors(p, FP8_MAX) + acc += (tl.dot((p * scale_p).to(v.type.element_ty), v) * descale_p * descale_v) + else: + acc += tl.dot(p.to(v.type.element_ty), v) + + return acc, l_i, m_i @triton.jit -def _attn_fwd_inner(acc, l_i, m_i, q, k_ptrs, v_ptrs, bias_ptrs, stride_kn, stride_vk, stride_bn, stride_sn, start_m, - actual_seqlen_k, actual_seqlen_q, dropout_p, philox_seed, philox_ptrs, sd_mask_ptrs, dropout_mask_ptrs, - block_min, block_max, offs_n_causal, masked_blocks, n_extra_tokens, alibi_slope, +def _attn_fwd_mask(acc, l_i, m_i, + q, k_base_ptrs, v_base_ptrs, bias_base_ptrs, + stride_kn, stride_vk, stride_bn, stride_sn, start_m, + seqlen_k, seqlen_q, + dropout_p, philox_seed, philox_base_ptrs, + sd_mask_base_ptrs, dropout_mask_base_ptrs, + offs_m, offs_n, offs_d, + block_min, block_max, n_extra_tokens, alibi_slope, descale_q, descale_k, descale_v, IS_FP8: tl.constexpr, FP8_MAX: tl.constexpr, IS_CAUSAL: tl.constexpr, BLOCK_M: tl.constexpr, BLOCK_DMODEL: tl.constexpr, BLOCK_N: tl.constexpr, - OFFS_M: tl.constexpr, OFFS_N: tl.constexpr, PRE_LOAD_V: tl.constexpr, MASK_STEPS: tl.constexpr, + PRE_LOAD_V: tl.constexpr, ENABLE_DROPOUT: tl.constexpr, PADDED_HEAD: tl.constexpr, ACTUAL_BLOCK_DMODEL: tl.constexpr, SM_SCALE: tl.constexpr, USE_ALIBI: tl.constexpr, USE_EXP2: tl.constexpr, - RETURN_SCORES: tl.constexpr, ACCUMULATOR_TYPE): + RETURN_SCORES: tl.constexpr, + USE_SLIDING_WINDOW: tl.constexpr, WINDOW_SIZE_LEFT: tl.constexpr, WINDOW_SIZE_RIGHT: tl.constexpr, + ACCUMULATOR_TYPE): if USE_EXP2: RCP_LN2: tl.constexpr = 1.4426950408889634 + + # seqlen diff + seqlen_delta_qk = seqlen_k - seqlen_q # loop over k, v, and update accumulator for start_n in range(block_min, block_max, BLOCK_N): + # get ptrs + k_ptrs = k_base_ptrs + start_n * stride_kn + v_ptrs = v_base_ptrs + start_n * stride_vk + # For padded blocks, we will overrun the tensor size if # we load all BLOCK_N. For others, the blocks are all within range. - if MASK_STEPS: - k_offs_n = start_n + tl.arange(0, BLOCK_N) - else: - k_offs_n = None - k_offs_k = None if not PADDED_HEAD else tl.arange(0, BLOCK_DMODEL) - k = load_fn(k_ptrs, k_offs_k, k_offs_n, ACTUAL_BLOCK_DMODEL, actual_seqlen_k) + kv_offs_n = start_n + tl.arange(0, BLOCK_N) + k_mask = (kv_offs_n[None, :] < seqlen_k) + v_mask = (kv_offs_n[:, None] < seqlen_k) + if PADDED_HEAD: + k_mask = k_mask & (offs_d[:, None] < ACTUAL_BLOCK_DMODEL) + v_mask = v_mask & (offs_d[None, :] < ACTUAL_BLOCK_DMODEL) + + # load k and if preload_v then v + k = tl.load(k_ptrs, mask=k_mask, other = 0.0) if PRE_LOAD_V: - # We can use the same offsets as k, just with dims transposed. - v = load_fn(v_ptrs, k_offs_n, k_offs_k, actual_seqlen_k, ACTUAL_BLOCK_DMODEL) + v = tl.load(v_ptrs, mask=v_mask, other=0.0) + + # setup qk accumlator qk = tl.zeros([BLOCK_M, BLOCK_N], dtype=ACCUMULATOR_TYPE) + # We start from end of seqlen_k so only the first iteration would need # to be checked for padding if it is not a multiple of block_n # TODO: This can be optimized to only be true for the padded block. - if MASK_STEPS: - # If this is the last block / iteration, we want to - # mask if the sequence length is not a multiple of block size - # a solution is to always do BLOCK_M // BLOCK_N + 1 steps if not is_modulo_mn. - # last step might get wasted but that is okay. check if this masking works For - # that case. - if (start_n + BLOCK_N == block_max) and (n_extra_tokens != 0): - boundary_m = tl.full([BLOCK_M], actual_seqlen_k, dtype=tl.int32) - size_n = start_n + OFFS_N[None, :] - mask = size_n < boundary_m[:, None] - qk = tl.where(mask, qk, float("-inf")) - - # compute masks - q_mask = (OFFS_M[:, None] < actual_seqlen_q) - k_mask = ((start_n + tl.arange(0, BLOCK_N))[None, :] < actual_seqlen_k) - p_mask = q_mask & k_mask + # If this is the last block / iteration, we want to + # mask if the sequence length is not a multiple of block size + # a solution is to always do BLOCK_M // BLOCK_N + 1 steps if not is_modulo_mn. + # last step might get wasted but that is okay. check if this masking works For + # that case. + if (n_extra_tokens != 0) and (start_n + BLOCK_N == block_max): + boundary_m = tl.full([BLOCK_M], seqlen_k, dtype=tl.int32) + size_n = start_n + offs_n[None, :] + mask = size_n < boundary_m[:, None] + qk = tl.where(mask, qk, float("-inf")) # -- compute qk ---- if IS_FP8 : @@ -84,27 +207,121 @@ def _attn_fwd_inner(acc, l_i, m_i, q, k_ptrs, v_ptrs, bias_ptrs, stride_kn, stri if USE_ALIBI: # compute the global position of each token within the sequence - global_m_positions = start_m * BLOCK_M + tl.arange(0, BLOCK_M) - global_n_positions = start_n + tl.arange(0, BLOCK_N) - alibi_block = compute_alibi_block(alibi_slope, actual_seqlen_q, actual_seqlen_k, global_m_positions, - global_n_positions) + q_offs_m = start_m * BLOCK_M + tl.arange(0, BLOCK_M) + alibi_block = compute_alibi_block(alibi_slope, seqlen_q, seqlen_k, q_offs_m, + kv_offs_n) qk_scaled += alibi_block - if IS_CAUSAL: - causal_boundary = start_n + offs_n_causal - causal_mask = OFFS_M[:, None] >= causal_boundary[None, :] - qk_scaled = tl.where(causal_mask, qk_scaled, float("-inf")) + if USE_SLIDING_WINDOW: + if IS_CAUSAL: + # ========== CAUSAL SLIDING WINDOW MASKING ========== + # For causal sliding window, we need to apply both constraints: + # 1. Causal: col_idx <= row_idx + (seqlen_k - seqlen_q) + # 2. Sliding window: row_idx - window_left <= col_idx <= row_idx + window_right + + # Get positions + row_idx = offs_m # Query positions + col_idx = kv_offs_n # Key positions + + # Expand for broadcasting + row_idx_expanded = row_idx[:, None] # [BLOCK_M, 1] + col_idx_expanded = col_idx[None, :] # [1, BLOCK_N] + + # Apply causal constraint: can only attend to positions before or at the diagonal + causal_offset = seqlen_k - seqlen_q + causal_mask = col_idx_expanded > (row_idx_expanded + causal_offset) + + # Apply sliding window constraint + if WINDOW_SIZE_LEFT < 0: + # Only right window constraint + window_mask = col_idx_expanded > (row_idx_expanded + causal_offset + WINDOW_SIZE_RIGHT) + else: + # Both left and right window constraints + # Adjust window bounds by causal offset + left_bound = row_idx_expanded + causal_offset - WINDOW_SIZE_LEFT + right_bound = row_idx_expanded + causal_offset + WINDOW_SIZE_RIGHT + + # Can't attend to positions outside the window + window_mask = (col_idx_expanded < left_bound) | (col_idx_expanded > right_bound) + + # Final mask is the union of both constraints (True = cannot attend) + mask = causal_mask | window_mask + + # Apply mask + qk_scaled = tl.where(mask, float("-inf"), qk_scaled) + else: + # ========== NON-CAUSAL SLIDING WINDOW MASKING ========== + # Exactly matching reference construct_local_mask: + # row_idx = query positions, col_idx = key positions + # sk = seqlen_k, sq = seqlen_q + + # Get positions + row_idx = offs_m # Query positions + col_idx = kv_offs_n # Key positions + + # sk and sq from reference (no padding masks in this test) + sk = seqlen_k + sq = seqlen_q + + # Expand for broadcasting + row_idx_expanded = row_idx[:, None] # [BLOCK_M, 1] + col_idx_expanded = col_idx[None, :] # [1, BLOCK_N] + + # Reference logic for mask computation + if WINDOW_SIZE_LEFT < 0: + # Reference: return col_idx > row_idx + sk - sq + window_size[1] + mask = col_idx_expanded > (row_idx_expanded + sk - sq + WINDOW_SIZE_RIGHT) + else: + # Reference: + # sk = torch.full_like(col_idx, seqlen_k) if key_padding_mask is None else sk + # return torch.logical_or( + # col_idx > torch.minimum(row_idx + sk - sq + window_size[1], sk), + # col_idx < row_idx + sk - sq - window_size[0], + # ) + # Create sk tensor with proper shape for broadcasting + # sk represents the key sequence length, which should be compared per column + sk_full = tl.full((1, BLOCK_N), sk, dtype=tl.int32) + + # Compute boundaries + right_bound_val = row_idx_expanded + sk - sq + WINDOW_SIZE_RIGHT + right_bound = tl.minimum(right_bound_val, sk_full) + left_bound = row_idx_expanded + sk - sq - WINDOW_SIZE_LEFT + + # Mask where True = cannot attend (matching reference) + mask = (col_idx_expanded > right_bound) | (col_idx_expanded < left_bound) + + # Apply mask (set to -inf where mask is True) + qk_scaled = tl.where(mask, float("-inf"), qk_scaled) + else: + if IS_CAUSAL: + causal_boundary = start_n + offs_n - seqlen_delta_qk + causal_mask = offs_m[:, None] >= causal_boundary[None, :] + qk_scaled = tl.where(causal_mask, qk_scaled, float("-inf")) - if bias_ptrs is not None: - bias_offs_n = start_n + tl.arange(0, BLOCK_N) if MASK_STEPS else None - bias = load_fn(bias_ptrs, OFFS_M, bias_offs_n, actual_seqlen_q, actual_seqlen_k) + # compute qk mask + qk_mask = (offs_m[:, None] < seqlen_q) & (kv_offs_n[None, :] < seqlen_k) + + # compute bias + if bias_base_ptrs is not None: + bias_ptrs = bias_base_ptrs + start_n * stride_bn + bias = tl.load(bias_ptrs, mask=qk_mask, other=0.0) qk_scaled += bias # get max scores so far m_ij = tl.maximum(m_i, tl.max(qk_scaled, 1)) # scale and subtract max - q_shifted = qk_scaled - m_ij[:, None] + # IMPORTANT: Handle the case where all values are -inf + # When m_ij = -inf and qk_scaled = -inf, subtraction gives NaN + # We need to handle this explicitly + if USE_SLIDING_WINDOW: + # Check if this block has any valid values (m_ij != -inf) + # For rows where everything is -inf, set q_shifted to -inf (not NaN) + q_shifted = tl.where(m_ij[:, None] == float("-inf"), + float("-inf"), + qk_scaled - m_ij[:, None]) + else: + q_shifted = qk_scaled - m_ij[:, None] # Compute scaled QK and softmax probabilities if USE_EXP2: @@ -115,38 +332,44 @@ def _attn_fwd_inner(acc, l_i, m_i, q, k_ptrs, v_ptrs, bias_ptrs, stride_kn, stri # CAVEAT: Must update l_ij before applying dropout l_ij = tl.sum(p, 1) if ENABLE_DROPOUT: + dropout_mask_ptrs = dropout_mask_base_ptrs + start_n * stride_sn + sd_mask_ptrs = sd_mask_base_ptrs + start_n * stride_sn + philox_ptrs = philox_base_ptrs + start_n * stride_sn if tl_DROPOUT_USE_PYTORCH: - dropout_mask = tl.load(dropout_mask_ptrs, mask=p_mask) + dropout_mask = tl.load(dropout_mask_ptrs, mask=qk_mask) else: rng_output = tl.rand(philox_seed, philox_ptrs) # TODO: use tl.randint for better performance dropout_mask = rng_output > dropout_p if tl_DROPOUT_DUMP: - tl.store(dropout_mask_ptrs, dropout_mask, mask=p_mask) + tl.store(dropout_mask_ptrs, dropout_mask, mask=qk_mask) # return scores with negative values for dropped vals sd_mask = tl.where(dropout_mask, p, -p) - tl.store(sd_mask_ptrs, sd_mask, mask=p_mask) + tl.store(sd_mask_ptrs, sd_mask, mask=qk_mask) # apply dropout mask in place p = tl.where(dropout_mask, p, 0.0) elif RETURN_SCORES: # NOTE: the returned score is not the same as the reference because we need to adjust as we find new maxes per block. We are not doing that - tl.store(sd_mask_ptrs, p, mask=p_mask) + sd_mask_ptrs = sd_mask_base_ptrs + start_n * stride_sn + tl.store(sd_mask_ptrs, p, mask=qk_mask) # -- update output accumulator -- # alpha is an adjustment factor for acc and li as we loop and find new maxes # store the diff in maxes to adjust acc and li as we discover new maxes - m_diff = m_i - m_ij + m_diff = tl.where(m_ij == float("-inf"), + float("-inf"), + m_i - m_ij) if USE_EXP2: alpha = tl.math.exp2(m_diff * RCP_LN2) else: alpha = tl.math.exp(m_diff) acc = acc * alpha[:, None] if not PRE_LOAD_V: - v = load_fn(v_ptrs, k_offs_n, k_offs_k, actual_seqlen_k, ACTUAL_BLOCK_DMODEL) + v = tl.load(v_ptrs, mask=v_mask, other=0.0) + # -- update m_i and l_i l_i = l_i * alpha + l_ij - # update m_i and l_i m_i = m_ij if IS_FP8: @@ -154,108 +377,140 @@ def _attn_fwd_inner(acc, l_i, m_i, q, k_ptrs, v_ptrs, bias_ptrs, stride_kn, stri acc += (tl.dot((p * scale_p).to(v.type.element_ty), v) * descale_p * descale_v) else: acc += tl.dot(p.to(v.type.element_ty), v) - - k_ptrs += BLOCK_N * stride_kn - v_ptrs += BLOCK_N * stride_vk - if bias_ptrs is not None: - bias_ptrs += BLOCK_N * stride_bn - if RETURN_SCORES: - sd_mask_ptrs += BLOCK_N * stride_sn - - if ENABLE_DROPOUT: - dropout_mask_ptrs += BLOCK_N * stride_sn - philox_ptrs += BLOCK_N * stride_sn + return acc, l_i, m_i -def get_cdna_autotune_configs(): - return [ - triton.Config({'BLOCK_M': 128, 'BLOCK_N': 128, 'waves_per_eu': 2, 'PRE_LOAD_V': False}, num_stages=1, - num_warps=4), - triton.Config({'BLOCK_M': 128, 'BLOCK_N': 64, 'waves_per_eu': 2, 'PRE_LOAD_V': False}, num_stages=1, - num_warps=4), - triton.Config({'BLOCK_M': 128, 'BLOCK_N': 64, 'waves_per_eu': 3, 'PRE_LOAD_V': False}, num_stages=1, - num_warps=4), - triton.Config({'BLOCK_M': 128, 'BLOCK_N': 64, 'waves_per_eu': 1, 'PRE_LOAD_V': False}, num_stages=1, - num_warps=4), - triton.Config({'BLOCK_M': 128, 'BLOCK_N': 32, 'waves_per_eu': 2, 'PRE_LOAD_V': False}, num_stages=1, - num_warps=4), - triton.Config({'BLOCK_M': 64, 'BLOCK_N': 64, 'waves_per_eu': 1, 'PRE_LOAD_V': False}, num_stages=1, - num_warps=4), - # Fall-back config. - triton.Config({'BLOCK_M': 16, 'BLOCK_N': 16, 'waves_per_eu': 1, 'PRE_LOAD_V': False}, num_stages=1, - num_warps=4), - ], ['IS_CAUSAL', 'dropout_p', 'MAX_SEQLENS_Q', 'MAX_SEQLENS_K', 'ACTUAL_BLOCK_DMODEL', 'IS_VARLEN', 'HQ', 'HK'] - - -def get_rdna_autotune_configs(): - return [ - triton.Config({'BLOCK_M': 32, 'BLOCK_N': 32, 'waves_per_eu': 4, 'PRE_LOAD_V': False}, num_stages=1, - num_warps=2), - triton.Config({'BLOCK_M': 32, 'BLOCK_N': 32, 'waves_per_eu': 2, 'PRE_LOAD_V': False}, num_stages=1, - num_warps=2), - triton.Config({'BLOCK_M': 32, 'BLOCK_N': 16, 'waves_per_eu': 4, 'PRE_LOAD_V': False}, num_stages=1, - num_warps=2), - triton.Config({'BLOCK_M': 32, 'BLOCK_N': 16, 'waves_per_eu': 2, 'PRE_LOAD_V': False}, num_stages=1, - num_warps=2), - triton.Config({'BLOCK_M': 16, 'BLOCK_N': 16, 'waves_per_eu': 4, 'PRE_LOAD_V': False}, num_stages=1, - num_warps=2), - triton.Config({'BLOCK_M': 16, 'BLOCK_N': 16, 'waves_per_eu': 2, 'PRE_LOAD_V': False}, num_stages=1, - num_warps=2), - # Fall-back config. - triton.Config({'BLOCK_M': 16, 'BLOCK_N': 16, 'waves_per_eu': 1, 'PRE_LOAD_V': False}, num_stages=1, - num_warps=2), - ], ['IS_CAUSAL', 'dropout_p', 'MAX_SEQLENS_Q', 'MAX_SEQLENS_K', 'ACTUAL_BLOCK_DMODEL', 'IS_VARLEN', 'HQ', 'HK'] - - -def get_autotune_configs(): - if AUTOTUNE: - if is_rdna(): - return get_rdna_autotune_configs() - elif is_cdna(): - return get_cdna_autotune_configs() +@triton.jit +def compute_masking(seqlen_k, seqlen_q, start_m, + IS_CAUSAL: tl.constexpr, USE_SLIDING_WINDOW: tl.constexpr, + WINDOW_SIZE_LEFT: tl.constexpr, WINDOW_SIZE_RIGHT: tl.constexpr, + BLOCK_M: tl.constexpr, BLOCK_N: tl.constexpr): + """ + Classify K blocks for attention computation with sliding window support. + + Returns: + - n_front_skip_blocks: Blocks completely before the window + - n_front_masked_blocks: Blocks partially overlapping window front + - n_full_blocks: Blocks completely inside the window + - n_back_masked_blocks: Blocks partially overlapping window back + - n_extra_tokens: Padding tokens in last K block + """ + # Example case + # BLOCK_M = 4, BLOCK_N = 4, seqlen_q = 8, seqlen_k = 10 + + # Total K blocks in the key sequence + total_k_blocks = tl.cdiv(seqlen_k, BLOCK_N) + + # check if we will need to do masking due either BLOCK_N being bigger than seqlen_k or seqlen_k not being a factor of BLOCK_N + # n_extra_tokens = 10 % 4 = 2 + # This means the last K block has 2 valid tokens and 2 padding positions + # K blocks visualization: + # Block 0 Block 1 Block 2 (last) + # K0 K1 K2 K3 K4 K5 K6 K7 K8 K9 ?? ?? + # ↑---------↑ ↑---------↑ ↑---↑ ↑---↑ + # full block full block valid pad + if seqlen_k < BLOCK_N: + n_extra_tokens = BLOCK_N - seqlen_k + elif seqlen_k % BLOCK_N: + n_extra_tokens = seqlen_k % BLOCK_N + else: + n_extra_tokens = 0 + + if USE_SLIDING_WINDOW: + # TODO: Optimize by computing which blocks can be fully skipped + # For now, process all blocks with the mask function + if IS_CAUSAL: + return 0, 0, 0, total_k_blocks, n_extra_tokens else: - raise ValueError("Unknown Device Type") + return 0, 0, 0, total_k_blocks, n_extra_tokens else: - arch = get_arch() - if arch == "gfx950": - default_config = triton.Config( - {"BLOCK_M": 128, "BLOCK_N": 128, "waves_per_eu": 2, "PRE_LOAD_V": False}, - num_stages=1, - num_warps=4, - ) - elif arch == "gfx942" and False: # Disabled due shared mem oom in CI when using triton==3.3.0 when using top of tree everything seems fine. - default_config = triton.Config( - {"BLOCK_M": 128, "BLOCK_N": 64, "waves_per_eu": 2, "PRE_LOAD_V": False}, - num_stages=1, - num_warps=4, - ) + if IS_CAUSAL: + # ========== CAUSAL MODE: Classify K Blocks ========== + # Calculate causal boundary for this Q block + # [K0 K1 K2 K3] [K4 K5 K6 K7] [K8 K9 ?? ??] + # Q0-Q3: [ 1 0 0 0] [ 0 0 0 0] [ 0 0 -- --] ← Q0 + # [ 1 1 0 0] [ 0 0 0 0] [ 0 0 -- --] ← Q1 + # [ 1 1 1 0] [ 0 0 0 0] [ 0 0 -- --] ← Q2 + # [ 1 1 1 1] [ 1 1 0 0] [ 0 0 -- --] ← Q3 + # ↑ can see up to K5 + # + # Q4-Q7: [ 1 1 1 1] [ 1 1 1 0] [ 0 0 -- --] ← Q4 + # [ 1 1 1 1] [ 1 1 1 1] [ 0 0 -- --] ← Q5 + # [ 1 1 1 1] [ 1 1 1 1] [ 1 0 -- --] ← Q6 + # [ 1 1 1 1] [ 1 1 1 1] [ 1 1 -- --] ← Q7 + + # ------------------------------------------------------------ + # 1. figure out, in tokens, the right-most K position + # this Q-block may attend to + # ------------------------------------------------------------ + q_start = start_m * BLOCK_M + q_end = tl.minimum((start_m + 1) * BLOCK_M - 1, seqlen_q - 1) + + # causal diagonal offset between the two streams + diag = seqlen_k - seqlen_q # 0 when |Q| == |K| + k_max_token = q_end + diag # last visible K index + + # this Q-block is entirely above the diagonal ⇒ nothing to do + if k_max_token < 0: + return 0, 0, 0, 0, n_extra_tokens + + k_max_token = tl.minimum(k_max_token, seqlen_k - 1) + + # ------------------------------------------------------------ + # 2. translate token indices into K-block indices + # ------------------------------------------------------------ + last_visible_k_block = k_max_token // BLOCK_N + n_visible_k_blocks = tl.minimum(last_visible_k_block + 1, total_k_blocks) + + # ------------------------------------------------------------ + # 3. classify those visible blocks + # – we *never* skip or mask blocks in front, because causal + # attention always starts at K0 + # – the back side can require several masked blocks: + # • intersection of the causal diagonal with K-grid + # (at most ⌈BLOCK_M / BLOCK_N⌉ blocks) + # • plus one extra block if this Q-block stops in the + # middle of a K-block or the last K-block is padded + # ------------------------------------------------------------ + padded_last_k = n_extra_tokens != 0 + is_modulo_mn = (not padded_last_k) & (seqlen_q % BLOCK_M == 0) + + n_back_masked_blocks = BLOCK_M // BLOCK_N + tl.where(is_modulo_mn, 0, 1) + n_back_masked_blocks = tl.minimum(n_back_masked_blocks, n_visible_k_blocks) + + n_front_skip_blocks = 0 # causal never skips the left side + n_front_masked_blocks = 0 # ditto + n_full_blocks = n_visible_k_blocks - n_back_masked_blocks else: - default_config = triton.Config( - {"BLOCK_M": 64, "BLOCK_N": 64, "waves_per_eu": 2, "PRE_LOAD_V": False}, - num_stages=1, - num_warps=4, - ) - - return [ - default_config - ], [ - "IS_CAUSAL", - "dropout_p", - "MAX_SEQLENS_Q", - "MAX_SEQLENS_K", - "ACTUAL_BLOCK_DMODEL", - "IS_VARLEN", - "HQ", - "HK", - ] - - -autotune_configs, autotune_keys = get_autotune_configs() + # ========== NON-CAUSAL MODE ========== + # Without causal mask, all positions can attend to all positions + # Only need to handle the padding in the last block + # [K0 K1 K2 K3] [K4 K5 K6 K7] [K8 K9 ?? ??] + # Q0-Q3: [ 1 1 1 1] [ 1 1 1 1] [ 1 1 -∞ -∞] + # [ 1 1 1 1] [ 1 1 1 1] [ 1 1 -∞ -∞] + # [ 1 1 1 1] [ 1 1 1 1] [ 1 1 -∞ -∞] + # [ 1 1 1 1] [ 1 1 1 1] [ 1 1 -∞ -∞] + # + # Q4-Q7: [ 1 1 1 1] [ 1 1 1 1] [ 1 1 -∞ -∞] + # [ 1 1 1 1] [ 1 1 1 1] [ 1 1 -∞ -∞] + # [ 1 1 1 1] [ 1 1 1 1] [ 1 1 -∞ -∞] + # [ 1 1 1 1] [ 1 1 1 1] [ 1 1 -∞ -∞] + + n_front_skip_blocks = 0 # never skips the left side + n_front_masked_blocks = 0 # ditto + if n_extra_tokens != 0: + n_back_masked_blocks = 1 # Last block needs padding mask + n_full_blocks = total_k_blocks - 1 + else: + n_back_masked_blocks = 0 # All blocks are aligned + n_full_blocks = total_k_blocks + + return n_front_skip_blocks, n_front_masked_blocks, n_full_blocks, n_back_masked_blocks, n_extra_tokens @triton.autotune( - configs=autotune_configs, - key=autotune_keys, + configs=fwd_prefill_autotune_configs, + key=fwd_prefill_autotune_keys, use_cuda_graph=True, ) @triton.jit @@ -267,7 +522,8 @@ def attn_fwd(Q, K, V, bias, Cache_seqlens, Cache_batch_idx, stride_sz, stride_sh, stride_sm, stride_sn, stride_lse_z, stride_lse_h, stride_lse_m, cu_seqlens_q, cu_seqlens_k, dropout_p, philox_seed, philox_offset_base, sd_mask, dropout_mask, alibi_slopes, HQ: tl.constexpr, HK: tl.constexpr, ACTUAL_BLOCK_DMODEL: tl.constexpr, MAX_SEQLENS_Q: tl.constexpr, - MAX_SEQLENS_K: tl.constexpr, IS_VARLEN: tl.constexpr, IS_INFERENCE: tl.constexpr, IS_CAUSAL: tl.constexpr, BLOCK_M: tl.constexpr, + MAX_SEQLENS_K: tl.constexpr, IS_VARLEN: tl.constexpr, IS_CAUSAL: tl.constexpr, + USE_SLIDING_WINDOW: tl.constexpr, WINDOW_SIZE_LEFT: tl.constexpr, WINDOW_SIZE_RIGHT: tl.constexpr, BLOCK_M: tl.constexpr, BLOCK_DMODEL: tl.constexpr, BLOCK_N: tl.constexpr, PRE_LOAD_V: tl.constexpr, USE_BIAS: tl.constexpr, ENABLE_DROPOUT: tl.constexpr, RETURN_SCORES: tl.constexpr, USE_ALIBI: tl.constexpr, USE_EXP2: tl.constexpr, IS_FP8: tl.constexpr, FP8_MAX: tl.constexpr, FP8_OUTPUT: tl.constexpr, FLIP_GRID: tl.constexpr): @@ -276,22 +532,21 @@ def attn_fwd(Q, K, V, bias, Cache_seqlens, Cache_batch_idx, # compute offsets if FLIP_GRID: - #NUM_XCDS: tl.constexpr = 8 off_z = tl.program_id(0) off_h_q = tl.program_id(1) start_m = tl.program_id(2) - - #start_m = (tl.cdiv(MAX_SEQLENS_Q, BLOCK_M) - 1) - start_m - - # Remap heads to the same XCD - #pids_per_xcd = HQ // NUM_XCDS - #xcd_group = off_h_q % NUM_XCDS - #pid_in_xcd = off_h_q // NUM_XCDS - #off_h_q = xcd_group * pids_per_xcd + pid_in_xcd else: start_m = tl.program_id(0) off_h_q = tl.program_id(1) off_z = tl.program_id(2) + # If MQA / GQA, set the K and V head offsets appropriately. + GROUP_SIZE: tl.constexpr = HQ // HK + if GROUP_SIZE != 1: + off_h_k = off_h_q // GROUP_SIZE + else: + off_h_k = off_h_q + # Determine if we need to mask the heads + PADDED_HEAD: tl.constexpr = (ACTUAL_BLOCK_DMODEL != BLOCK_DMODEL) offs_m = start_m * BLOCK_M + tl.arange(0, BLOCK_M) offs_n = tl.arange(0, BLOCK_N) @@ -309,76 +564,57 @@ def attn_fwd(Q, K, V, bias, Cache_seqlens, Cache_batch_idx, cu_seqlens_k_start = tl.load(cu_seqlens_k + off_z) cu_seqlens_k_end = tl.load(cu_seqlens_k + off_z + 1) seqlen_k = cu_seqlens_k_end - cu_seqlens_k_start - elif IS_INFERENCE: - cu_seqlens_q_start = 0 - cu_seqlens_k_start = 0 - seqlen_q = MAX_SEQLENS_Q - seqlen_k = tl.load(Cache_seqlens + off_z) else: cu_seqlens_q_start = 0 cu_seqlens_k_start = 0 seqlen_q = MAX_SEQLENS_Q seqlen_k = MAX_SEQLENS_K - # Now we compute whether we need to exit early due to causal masking. - # This is because for seqlen_q > seqlen_k, M rows of the attn scores - # are completely masked, resulting in 0s written to the output, and - # inf written to LSE. We don't need to do any GEMMs in this case. - # This block of code determines what N is, and if this WG is operating - # on those M rows. - n_blocks = tl.cdiv(seqlen_k, BLOCK_N) - if (IS_CAUSAL): - # If seqlen_q == seqlen_k, the attn scores are a square matrix. - # If seqlen_q != seqlen_k, attn scores are rectangular which means - # the causal mask boundary is bottom right aligned, and ends at either - # the top edge (seqlen_q < seqlen_k) or left edge. - # This captures the decrease in n_blocks if we have a rectangular attn matrix - n_blocks_seqlen = tl.cdiv((start_m + 1) * BLOCK_M + seqlen_k - seqlen_q, BLOCK_N) - # This is what adjusts the block_max for the current WG, only - # if IS_CAUSAL. Otherwise we want to always iterate through all n_blocks - n_blocks = min(n_blocks, n_blocks_seqlen) - # If we have no blocks after adjusting for seqlen deltas, this WG is part of - # the blocks that are all 0. We exit early. - if n_blocks <= 0: - o_offset = Out + off_z * stride_oz + off_h_q * stride_oh + cu_seqlens_q_start * stride_om - o_ptrs = o_offset + offs_m[:, None] * stride_om + offs_d[None, :] * stride_on - acc = tl.zeros([BLOCK_M, BLOCK_DMODEL], dtype=Out.type.element_ty) - o_ptrs_mask = (offs_m[:, None] < seqlen_q) & (offs_d[None, :] < ACTUAL_BLOCK_DMODEL) - # We still need to write 0s to the result - tl.store(o_ptrs, acc, mask=o_ptrs_mask) - # The tensor allocated for L is based on MAX_SEQLENS_Q as that is - # statically known. - l_offset = LSE + off_z * stride_lse_z + off_h_q * stride_lse_h + cu_seqlens_q_start * stride_lse_m - l_ptrs = l_offset + offs_m * stride_lse_m - - l = tl.full([BLOCK_M], value=0.0, dtype=ACCUMULATOR_TYPE) - - # mask_m_offsets = start_m + tl.arange(0, BLOCK_M) - # lse_mask = mask_m_offsets < causal_start_idx - # softmax_lse = tl.where(lse_mask, 0.0, softmax_lse) - l_ptrs_mask = offs_m < MAX_SEQLENS_Q - tl.store(l_ptrs, l, mask=l_ptrs_mask) - # TODO: Should dropout and return encoded softmax be handled here too? - return - - # If MQA / GQA, set the K and V head offsets appropriately. - GROUP_SIZE: tl.constexpr = HQ // HK - if GROUP_SIZE != 1: - off_h_k = off_h_q // GROUP_SIZE + # Load scale factors if IS_FP8. + if IS_FP8: + descale_q = tl.load(Descale_Q + off_z * stride_descale_q_z + off_h_q) + descale_k = tl.load(Descale_K + off_z * stride_descale_k_z + off_h_k) + descale_v = tl.load(Descale_V + off_z * stride_descale_v_z + off_h_k) else: - off_h_k = off_h_q - - n_extra_tokens = 0 - # print("n_extra_tokens:", n_extra_tokens) - # print("seqlen_k:", seqlen_k) - # print("BLOCK_N:", BLOCK_N) - # return - if seqlen_k < BLOCK_N: - n_extra_tokens = BLOCK_N - seqlen_k - elif seqlen_k % BLOCK_N: - n_extra_tokens = seqlen_k % BLOCK_N - PADDED_HEAD: tl.constexpr = (ACTUAL_BLOCK_DMODEL != BLOCK_DMODEL) + descale_q, descale_k, descale_v = 1.0, 1.0, 1.0 + + # figure out masking pattern + n_front_skip_blocks, n_front_masked_blocks, n_full_blocks, n_back_masked_blocks, n_extra_tokens = compute_masking( + seqlen_k, seqlen_q, start_m, IS_CAUSAL, USE_SLIDING_WINDOW, + WINDOW_SIZE_LEFT, WINDOW_SIZE_RIGHT, BLOCK_M, BLOCK_N + ) + + # ============================================================ + # PROGRAM EARLY EXIT (All K Blocks Skipped) + # ============================================================ + total_visible_blocks = n_front_masked_blocks + n_full_blocks + n_back_masked_blocks + if total_visible_blocks == 0: + """ + No K blocks visible - write zeros and exit. + """ + # Write zeros to output + o_offset = Out + off_z * stride_oz + off_h_q * stride_oh + cu_seqlens_q_start * stride_om + o_ptrs = o_offset + offs_m[:, None] * stride_om + offs_d[None, :] * stride_on + o_mask = (offs_m[:, None] < seqlen_q) + if PADDED_HEAD: + o_mask = o_mask & (offs_d[None, :] < ACTUAL_BLOCK_DMODEL) + tl.store(o_ptrs, tl.zeros([BLOCK_M, BLOCK_DMODEL], dtype=Out.type.element_ty), mask=o_mask) + + # Write zeros to LSE + l_ptrs = LSE + off_z * stride_lse_z + off_h_q * stride_lse_h + cu_seqlens_q_start * stride_lse_m + offs_m * stride_lse_m + tl.store(l_ptrs, tl.zeros([BLOCK_M], dtype=tl.float32), mask=offs_m < seqlen_q) + return + + # ============================================================ + # NORMAL PROCESSING (Some K Blocks Visible) + # ============================================================ + """ + This program has visible K blocks to process. + We'll use two calls to handle different block types efficiently. + """ + + # Initialize for processing # Compute pointers for all the tensors used in this kernel. q_offset = Q + off_z * stride_qz + off_h_q * stride_qh + cu_seqlens_q_start * stride_qm q_ptrs = q_offset + offs_m[:, None] * stride_qm + offs_d[None, :] * stride_qk @@ -413,132 +649,206 @@ def attn_fwd(Q, K, V, bias, Cache_seqlens, Cache_batch_idx, else: dropout_mask_ptrs = None philox_ptrs = 0 + # initialize pointer to m and l m_i = tl.full([BLOCK_M], float("-inf"), dtype=ACCUMULATOR_TYPE) l_i = tl.full([BLOCK_M], 1.0, dtype=ACCUMULATOR_TYPE) acc = tl.zeros([BLOCK_M, BLOCK_DMODEL], dtype=ACCUMULATOR_TYPE) + # Q is loaded once at the beginning and shared by all N blocks. q_ptrs_mask = offs_m[:, None] < seqlen_q if PADDED_HEAD: q_ptrs_mask = q_ptrs_mask & (offs_d[None, :] < ACTUAL_BLOCK_DMODEL) q = tl.load(q_ptrs, mask=q_ptrs_mask, other=0.0) - # Load scale factors if IS_FP8. - if IS_FP8: - descale_q = tl.load(Descale_Q + off_z * stride_descale_q_z + off_h_q) - descale_k = tl.load(Descale_K + off_z * stride_descale_k_z + off_h_k) - descale_v = tl.load(Descale_V + off_z * stride_descale_v_z + off_h_k) - else: - descale_q, descale_k, descale_v = 1.0, 1.0, 1.0 - # Here we compute how many full and masked blocks we have. - padded_block_k = n_extra_tokens != 0 - is_modulo_mn = not padded_block_k and (seqlen_q % BLOCK_M == 0) - if IS_CAUSAL: - # There are always at least BLOCK_M // BLOCK_N masked blocks. - # Additionally there might be one more due to dissimilar seqlens. - masked_blocks = BLOCK_M // BLOCK_N + (not is_modulo_mn) - else: - # Padding on Q does not need to be masked in the FA loop. - masked_blocks = padded_block_k - # if IS_CAUSAL, not is_modulo_mn does not always result in an additional block. - # In this case we might exceed n_blocks so pick the min. - masked_blocks = min(masked_blocks, n_blocks) - n_full_blocks = n_blocks - masked_blocks - block_min = 0 - block_max = n_blocks * BLOCK_N - # Compute for full blocks. Here we set causal to false regardless of its actual - # value because there is no masking. Similarly we do not need padding. + # ========== Process MASKED K Blocks in the front ========== + # NOTE: we use USE_SLIDING_WINDOW as guard because the compiler will crash other wise. front masking is only for sliding window so that is fine. + if n_front_masked_blocks > 0 and USE_SLIDING_WINDOW: + block_min = n_front_skip_blocks * BLOCK_N + block_max = (n_front_skip_blocks + n_front_masked_blocks) * BLOCK_N + + acc, l_i, m_i = _attn_fwd_mask( + acc, l_i, m_i, + q, k_ptrs, v_ptrs, bias_ptrs, + stride_kn, stride_vk, stride_bn, stride_sn, + start_m, seqlen_k, seqlen_q, + dropout_p, philox_seed, philox_ptrs, + sd_mask_ptrs, dropout_mask_ptrs, + offs_m, offs_n, offs_d, + block_min, # Start of front masked blocks + block_max, # End of front masked blocks + 0, # n_extra_tokens (0 for front blocks, only relevant for last block) + alibi_slope, + descale_q, descale_k, descale_v, IS_FP8, FP8_MAX, + IS_CAUSAL, + BLOCK_M, BLOCK_DMODEL, BLOCK_N, + PRE_LOAD_V, + ENABLE_DROPOUT, PADDED_HEAD, + ACTUAL_BLOCK_DMODEL, SM_SCALE, + USE_ALIBI=USE_ALIBI, USE_EXP2=USE_EXP2, + RETURN_SCORES=RETURN_SCORES, + USE_SLIDING_WINDOW=USE_SLIDING_WINDOW, + WINDOW_SIZE_LEFT=WINDOW_SIZE_LEFT, + WINDOW_SIZE_RIGHT=WINDOW_SIZE_RIGHT, + ACCUMULATOR_TYPE=ACCUMULATOR_TYPE + ) + + # ========== Process FULL K Blocks (Fast Path) ========== if n_full_blocks > 0: - block_max = (n_blocks - masked_blocks) * BLOCK_N - acc, l_i, m_i = _attn_fwd_inner(acc, l_i, m_i, q, k_ptrs, v_ptrs, bias_ptrs, stride_kn, stride_vk, stride_bn, stride_sn, - start_m, seqlen_k, seqlen_q, dropout_p, philox_seed, philox_ptrs, - sd_mask_ptrs, dropout_mask_ptrs, - # _, _, offs_n_causal, masked_blocks, n_extra_tokens, _ - block_min, block_max, 0, 0, 0, alibi_slope, - descale_q, descale_k, descale_v, IS_FP8, FP8_MAX, - # IS_CAUSAL, .... - False, BLOCK_M, BLOCK_DMODEL, BLOCK_N, offs_m, offs_n, - # _, MASK_STEPS, ... - PRE_LOAD_V, False, ENABLE_DROPOUT, PADDED_HEAD, - ACTUAL_BLOCK_DMODEL, SM_SCALE, USE_ALIBI=USE_ALIBI, USE_EXP2=USE_EXP2, RETURN_SCORES=RETURN_SCORES, ACCUMULATOR_TYPE=ACCUMULATOR_TYPE) - block_min = block_max - block_max = n_blocks * BLOCK_N - - tl.debug_barrier() - # Remaining blocks, if any, are full / not masked. - if (masked_blocks > 0): - if IS_CAUSAL: - offs_n_causal = offs_n + (seqlen_q - seqlen_k) - else: - offs_n_causal = 0 - k_ptrs += n_full_blocks * BLOCK_N * stride_kn - v_ptrs += n_full_blocks * BLOCK_N * stride_vk - if USE_BIAS: - bias_ptrs += n_full_blocks * BLOCK_N * stride_bn - if RETURN_SCORES: - sd_mask_ptrs += n_full_blocks * BLOCK_N * stride_sn - if ENABLE_DROPOUT: - dropout_mask_ptrs += n_full_blocks * BLOCK_N * stride_sn - philox_ptrs += n_full_blocks * BLOCK_N * stride_sn - acc, l_i, m_i = _attn_fwd_inner(acc, l_i, m_i, q, k_ptrs, v_ptrs, bias_ptrs, stride_kn, stride_vk, stride_bn, stride_sn, - start_m, seqlen_k, seqlen_q, dropout_p, philox_seed, philox_ptrs, - sd_mask_ptrs, dropout_mask_ptrs, block_min, block_max, offs_n_causal, masked_blocks, - n_extra_tokens, alibi_slope, descale_q, descale_k, descale_v, IS_FP8, FP8_MAX, - IS_CAUSAL, BLOCK_M, BLOCK_DMODEL, BLOCK_N, offs_m, offs_n, - # _, MASK_STEPS, ... - PRE_LOAD_V, True, ENABLE_DROPOUT, PADDED_HEAD, - ACTUAL_BLOCK_DMODEL, SM_SCALE, USE_ALIBI=USE_ALIBI, USE_EXP2=USE_EXP2, RETURN_SCORES=RETURN_SCORES, ACCUMULATOR_TYPE=ACCUMULATOR_TYPE) - # epilogue + block_min = (n_front_skip_blocks + n_front_masked_blocks) * BLOCK_N + block_max = (n_front_skip_blocks + n_front_masked_blocks + n_full_blocks) * BLOCK_N + + acc, l_i, m_i = _attn_fwd_no_mask( + acc, l_i, m_i, + q, k_ptrs, v_ptrs, bias_ptrs, + stride_kn, stride_vk, stride_bn, stride_sn, + start_m, seqlen_k, seqlen_q, + dropout_p, philox_seed, philox_ptrs, + sd_mask_ptrs, dropout_mask_ptrs, + offs_m, offs_n, offs_d, + block_min, # Start of range: 0 + block_max, # End of range: n_full_blocks * BLOCK_N + alibi_slope, + descale_q, descale_k, descale_v, IS_FP8, FP8_MAX, + BLOCK_M, BLOCK_DMODEL, BLOCK_N, + PRE_LOAD_V, + ENABLE_DROPOUT, PADDED_HEAD, + ACTUAL_BLOCK_DMODEL, SM_SCALE, + USE_ALIBI=USE_ALIBI, USE_EXP2=USE_EXP2, + RETURN_SCORES=RETURN_SCORES, ACCUMULATOR_TYPE=ACCUMULATOR_TYPE + ) + + # ========== Process MASKED K Blocks in the back ========== + if n_back_masked_blocks > 0: + block_min = (n_front_skip_blocks + n_front_masked_blocks + n_full_blocks) * BLOCK_N + block_max = (n_front_skip_blocks + n_front_masked_blocks + n_full_blocks + n_back_masked_blocks) * BLOCK_N + + acc, l_i, m_i = _attn_fwd_mask( + acc, l_i, m_i, + q, k_ptrs, v_ptrs, bias_ptrs, + stride_kn, stride_vk, stride_bn, stride_sn, + start_m, seqlen_k, seqlen_q, + dropout_p, philox_seed, philox_ptrs, + sd_mask_ptrs, dropout_mask_ptrs, + offs_m, offs_n, offs_d, + block_min, # Start of range: n_full_blocks * BLOCK_N + block_max, # End of range: n_visible_k_blocks * BLOCK_N + n_extra_tokens, # Padding tokens in last block + alibi_slope, + descale_q, descale_k, descale_v, IS_FP8, FP8_MAX, + IS_CAUSAL, # Use actual causal flag + BLOCK_M, BLOCK_DMODEL, BLOCK_N, + PRE_LOAD_V, + ENABLE_DROPOUT, PADDED_HEAD, + ACTUAL_BLOCK_DMODEL, SM_SCALE, + USE_ALIBI=USE_ALIBI, USE_EXP2=USE_EXP2, + RETURN_SCORES=RETURN_SCORES, + USE_SLIDING_WINDOW=USE_SLIDING_WINDOW, + WINDOW_SIZE_LEFT=WINDOW_SIZE_LEFT, + WINDOW_SIZE_RIGHT=WINDOW_SIZE_RIGHT, + ACCUMULATOR_TYPE=ACCUMULATOR_TYPE + ) + + # ============================================================ + # EPILOGUE + # ============================================================ # This helps the compiler do Newton Raphson on l_i vs on acc which is much larger. - l_recip = 1 / l_i[:, None] + # Instead of directly computing 1/l_i which can be inf, + # we check for the invalid case first + if USE_SLIDING_WINDOW: + # For rows where m_i is still -inf, no keys were valid + # Set l_i to 1.0 to avoid division by zero (acc is already 0) + invalid_mask = m_i == float("-inf") + l_i_safe = tl.where(invalid_mask, 1.0, l_i) + l_recip = 1 / l_i_safe[:, None] + else: + # Original code path + l_recip = 1 / l_i[:, None] acc = acc * l_recip if ENABLE_DROPOUT: dropout_scale = 1 / (1 - dropout_p) acc = acc * dropout_scale - # If seqlen_q > seqlen_k but the delta is not a multiple of BLOCK_M, - # then we have one block with a row of all NaNs which come from computing - # softmax over a row of all -infs (-inf - inf = NaN). We check for that here - # and store 0s where there are NaNs as these rows should've been zeroed out. - end_m_idx = (start_m + 1) * BLOCK_M - start_m_idx = start_m * BLOCK_M - causal_start_idx = seqlen_q - seqlen_k - if IS_CAUSAL: - if causal_start_idx > start_m_idx and causal_start_idx < end_m_idx: - out_mask_boundary = tl.full((BLOCK_DMODEL, ), causal_start_idx, dtype=tl.int32) - mask_m_offsets = start_m_idx + tl.arange(0, BLOCK_M) - out_ptrs_mask = mask_m_offsets[:, None] >= out_mask_boundary[None, :] - z = 0.0 - acc = tl.where(out_ptrs_mask, acc, z.to(acc.type.element_ty)) - # write back LSE(Log Sum Exponents), the log of the normalization constant - l_offset = LSE + off_z * stride_lse_z + off_h_q * stride_lse_h + cu_seqlens_q_start * stride_lse_m - l_ptrs = l_offset + offs_m * stride_lse_m + # compute log-sum-exp if USE_EXP2: RCP_LN2: tl.constexpr = 1.4426950408889634 LN2: tl.constexpr = 0.6931471824645996 # compute log-sum-exp in base 2 units mi_base2 = m_i * RCP_LN2 - softmax_lse = mi_base2 + tl.math.log2(l_i) + # For invalid rows, log(l_i) would be -inf, but we want LSE to be -inf + # So we handle this case explicitly + if USE_SLIDING_WINDOW: + log_l_i = tl.where(invalid_mask, 0.0, tl.math.log2(l_i)) + softmax_lse = mi_base2 + log_l_i + # Ensure invalid rows have LSE = -inf + softmax_lse = tl.where(invalid_mask, float("-inf"), softmax_lse) + else: + softmax_lse = mi_base2 + tl.math.log2(l_i) # convert back to natural units softmax_lse *= LN2 else: - softmax_lse = m_i + tl.math.log(l_i) + if USE_SLIDING_WINDOW: + log_l_i = tl.where(invalid_mask, 0.0, tl.math.log(l_i)) + softmax_lse = m_i + log_l_i + softmax_lse = tl.where(invalid_mask, float("-inf"), softmax_lse) + else: + softmax_lse = m_i + tl.math.log(l_i) + + # handle masking edge cases + if USE_SLIDING_WINDOW: + if IS_CAUSAL: + pass + else: + pass + else: + if IS_CAUSAL: + # When seqlen_q > seqlen_k, some rows are completely above the causal diagonal + # These rows have all -inf attention scores, resulting in NaN after softmax + # e.g. + # Q length: 6, K length: 4 + # Causal mask (X = can attend, . = cannot): + # K0 K1 K2 K3 + # Q0 . . . . <- All masked, would give NaN + # Q1 . . . . <- All masked, would give NaN + # Q2 X . . . <- First valid row + # Q3 X X . . + # Q4 X X X . + # Q5 X X X X + causal_start_idx = seqlen_q - seqlen_k + start_m_idx = start_m * BLOCK_M + + # Create mask for rows that need zeroing + row_indices = start_m_idx + tl.arange(0, BLOCK_M) + causal_mask = row_indices < causal_start_idx + + # Zero out both acc and LSE for these rows + if causal_start_idx > start_m_idx: + end_m_idx = (start_m + 1) * BLOCK_M + if causal_start_idx < end_m_idx: + # This block contains the boundary - need to mask acc + out_mask_boundary = tl.full((BLOCK_DMODEL, ), causal_start_idx, dtype=tl.int32) + out_ptrs_mask = row_indices[:, None] >= out_mask_boundary[None, :] + z = 0.0 + acc = tl.where(out_ptrs_mask, acc, z.to(acc.type.element_ty)) + + # Zero out LSE for rows above diagonal + softmax_lse = tl.where(causal_mask, 0.0, softmax_lse) - if IS_CAUSAL: - # zero out nans caused by -infs when doing causal - lse_mask = (start_m_idx + tl.arange(0, BLOCK_M)) < causal_start_idx - softmax_lse = tl.where(lse_mask, 0.0, softmax_lse) + # write back LSE(Log Sum Exponents), the log of the normalization constant + l_offset = LSE + off_z * stride_lse_z + off_h_q * stride_lse_h + cu_seqlens_q_start * stride_lse_m + l_ptrs = l_offset + offs_m * stride_lse_m # If seqlen_q not multiple of BLOCK_M, we need to mask out the last few rows. - # This is only true for the last M block. For others, overflow_size will be -ve + # This is only true for the last Q block. For others, overflow_size will be -ve + end_m_idx = (start_m + 1) * BLOCK_M overflow_size = end_m_idx - seqlen_q if overflow_size > 0: boundary = tl.full((BLOCK_M, ), BLOCK_M - overflow_size, dtype=tl.int32) l_ptrs_mask = tl.arange(0, BLOCK_M) < boundary - tl.store(l_ptrs, softmax_lse, mask=l_ptrs_mask) # the log of the normalization constant + tl.store(l_ptrs, softmax_lse, mask=l_ptrs_mask) else: - tl.store(l_ptrs, softmax_lse) # the log of the normalization constant + tl.store(l_ptrs, softmax_lse) # write back O o_offset = Out + off_z * stride_oz + off_h_q * stride_oh + cu_seqlens_q_start * stride_om @@ -556,7 +866,6 @@ def attn_fwd(Q, K, V, bias, Cache_seqlens, Cache_batch_idx, else: tl.store(o_ptrs, acc.to(Out.dtype.element_ty), mask=o_ptrs_mask) - def attention_prefill_forward_triton_impl( q: torch.Tensor, k: torch.Tensor, @@ -565,6 +874,8 @@ def attention_prefill_forward_triton_impl( sm_scale: float, alibi_slopes: Optional[torch.Tensor], causal: bool, + window_size_left: int, + window_size_right: int, bias: Optional[torch.Tensor], layout: Literal["bshd", "bhsd", "thd"], # varlen @@ -613,6 +924,7 @@ def attention_prefill_forward_triton_impl( # check flags IS_VARLEN = layout == "thd" + use_sliding_window = window_size_left != -1 or window_size_right!= -1 use_alibi, (stride_az, stride_ah) = (True, alibi_slopes.stride()) if alibi_slopes is not None else (False, (0, 0)) is_inference = False if cache_seqlens is None else True if is_inference: @@ -629,6 +941,7 @@ def attention_prefill_forward_triton_impl( # shape total_q, nheads_q, head_size = q.shape _, nheads_k, _ = k.shape + assert cu_seqlens_q is not None batch = len(cu_seqlens_q) - 1 # softmax_lse is the log of the normalization constant / sum of expoential score(unnormalzied probablities) @@ -695,7 +1008,7 @@ def attention_prefill_forward_triton_impl( attn_fwd[grid](q, k, v, bias, cache_seqlens, cache_batch_idx, descale_q, descale_k, descale_v, descale_o, stride_descale_q_z, stride_descale_k_z, stride_descale_v_z, stride_descale_o_z, - sm_scale, softmax_lse, o, + sm_scale, softmax_lse, o, stride_qb, stride_qh, stride_qm, stride_qd, stride_kb, stride_kh, stride_kn, stride_kd, stride_vb, stride_vh, stride_vn, stride_vd, @@ -703,13 +1016,15 @@ def attention_prefill_forward_triton_impl( stride_bz, stride_bh, stride_bm, stride_bn, stride_az, stride_ah, stride_sz, stride_sh, stride_sm, stride_sn, - stride_lse_z, stride_lse_h, stride_lse_m, + stride_lse_z, stride_lse_h, stride_lse_m, cu_seqlens_q, cu_seqlens_k, dropout_p=dropout_p, philox_seed=philox_seed, philox_offset_base=philox_offset, sd_mask=sd_mask, dropout_mask=dropout_mask, alibi_slopes=alibi_slopes, HQ=nheads_q, HK=nheads_k, ACTUAL_BLOCK_DMODEL=head_size, MAX_SEQLENS_Q=max_seqlens_q, - MAX_SEQLENS_K=max_seqlens_k, IS_CAUSAL=causal, IS_VARLEN=IS_VARLEN, IS_INFERENCE=is_inference, + MAX_SEQLENS_K=max_seqlens_k, IS_CAUSAL=causal, + USE_SLIDING_WINDOW=use_sliding_window, WINDOW_SIZE_LEFT=window_size_left, WINDOW_SIZE_RIGHT=window_size_right, + IS_VARLEN=IS_VARLEN, BLOCK_DMODEL=padded_d_model, USE_BIAS=False if bias is None else True, - USE_ALIBI=use_alibi, ENABLE_DROPOUT=dropout_p - > 0.0, USE_EXP2=use_exp2, RETURN_SCORES=return_softmax, IS_FP8=IS_FP8, FP8_MAX=FP8_MAX, FP8_OUTPUT=FP8_OUTPUT, FLIP_GRID=FLIP_GRID) + USE_ALIBI=use_alibi, ENABLE_DROPOUT=dropout_p > 0.0, USE_EXP2=use_exp2, RETURN_SCORES=return_softmax, + IS_FP8=IS_FP8, FP8_MAX=FP8_MAX, FP8_OUTPUT=FP8_OUTPUT, FLIP_GRID=FLIP_GRID) - return softmax_lse, sd_mask if return_softmax else None + return softmax_lse, sd_mask if return_softmax else None \ No newline at end of file diff --git a/flash_attn/flash_attn_triton_amd/fwd_ref.py b/flash_attn/flash_attn_triton_amd/fwd_ref.py index 6af99798ae9..2265af12096 100644 --- a/flash_attn/flash_attn_triton_amd/fwd_ref.py +++ b/flash_attn/flash_attn_triton_amd/fwd_ref.py @@ -1,12 +1,16 @@ import torch import math -from typing import Literal, Optional +from typing import Literal, Optional, Union from .utils import compute_alibi_tensor_ref DEBUG = False DEBUG_CORE = False -def attention_forward_core_ref_impl(q, k, v, sm_scale, causal, dropout_p, philox_seed, philox_offset, alibi_slopes, use_exp2): +def attention_forward_core_ref_impl( + q, k, v, sm_scale, causal, window_size_left, window_size_right, + dropout_p, philox_seed, philox_offset, alibi_slopes, use_exp2, + cache_seqlens=None +): if DEBUG_CORE: print() print("attention_forward_core_ref_impl") @@ -15,15 +19,21 @@ def attention_forward_core_ref_impl(q, k, v, sm_scale, causal, dropout_p, philox print("v:", v, v.shape) print("sm_scale:", sm_scale) print("causal:", causal) + print("window_size_left:", window_size_left) + print("window_size_right:", window_size_right) print("dropout_p:", dropout_p) print("philox_seed:", philox_seed) print("philox_offset:", philox_offset) print("use_exp2:", use_exp2) + print("cache_seqlens:", cache_seqlens) # cast to float32 q = q.to(torch.float32) k = k.to(torch.float32) v = v.to(torch.float32) + + # get seqlens + L_q, L_k = q.shape[1], k.shape[1] # Compute attention scores attention_scores = torch.matmul(q, k.transpose(-2, -1)) @@ -37,47 +47,146 @@ def attention_forward_core_ref_impl(q, k, v, sm_scale, causal, dropout_p, philox # Apply ALiBi if slopes are provided if alibi_slopes is not None: - L_q, L_k = q.shape[1], k.shape[1] - if DEBUG_CORE: - print("alibi_slopes:", alibi_slopes, alibi_slopes.shape) - alibi_bias = compute_alibi_tensor_ref(alibi_slopes, L_q, L_k) - if DEBUG_CORE: - print("alibi_bias:", alibi_bias, alibi_bias.shape) - alibi_bias = alibi_bias.reshape(-1, L_q, L_k) - if DEBUG_CORE: - print("alibi_bias_flat:", alibi_bias, alibi_bias.shape) + if cache_seqlens is not None: + # DECODE MODE: Special ALiBi handling + # In decode mode, k has shape [nheads, max_cache_len, head_dim] + # but only cache_seqlens positions are valid + + # The test's attn_bias_from_alibi_slopes uses this formula: + # relative_pos = torch.abs(row_idx + sk - sq - col_idx) + # where sk = actual valid key length, sq = query length + + row_idx = torch.arange(L_q, device=q.device, dtype=torch.float32).unsqueeze(1) + col_idx = torch.arange(L_k, device=q.device, dtype=torch.float32).unsqueeze(0) + + # Compute relative positions + # cache_seqlens is the actual number of valid keys (sk in the test) + # L_q is the query sequence length (sq in the test) + relative_pos = torch.abs(row_idx + cache_seqlens - L_q - col_idx) + + # Apply slopes + if alibi_slopes.dim() == 1: + # Shape: [nheads] -> [nheads, 1, 1] + alibi_slopes_expanded = alibi_slopes.view(-1, 1, 1) + else: + # Already has batch dimension + alibi_slopes_expanded = alibi_slopes + + alibi_bias = -alibi_slopes_expanded * relative_pos + + if DEBUG_CORE: + print(f"Decode ALiBi: cache_seqlens={cache_seqlens}, L_q={L_q}, L_k={L_k}") + print(f"relative_pos shape: {relative_pos.shape}") + print(f"alibi_bias shape: {alibi_bias.shape}") + else: + if DEBUG_CORE: + print("alibi_slopes:", alibi_slopes, alibi_slopes.shape) + alibi_bias = compute_alibi_tensor_ref(alibi_slopes, L_q, L_k) + if DEBUG_CORE: + print("alibi_bias:", alibi_bias, alibi_bias.shape) + alibi_bias = alibi_bias.reshape(-1, L_q, L_k) + if DEBUG_CORE: + print("alibi_bias_flat:", alibi_bias, alibi_bias.shape) + attention_scaled_scores = attention_scaled_scores + alibi_bias if DEBUG_CORE: print("attention_scaled_scores after alibi:", attention_scaled_scores, attention_scaled_scores.shape) - - # Apply causal mask if necessary - if causal: - L_q, L_k = q.shape[1], k.shape[1] - row_idx = torch.arange(L_q, device=q.device).unsqueeze(1) - col_idx = torch.arange(L_k, device=q.device).unsqueeze(0) - col_offset = L_q-L_k - causal_mask = row_idx >= (col_offset + col_idx) + # Apply masks + row_idx = torch.arange(L_q, device=q.device).unsqueeze(1) + col_idx = torch.arange(L_k, device=q.device).unsqueeze(0) + + if cache_seqlens is not None: + # We're in decode mode with a KV cache + # k and v are full allocated size, but only cache_seqlens positions are valid + + # Create a mask for valid cache positions + cache_mask = col_idx < cache_seqlens + + # Use cache_seqlens for offset calculation to match test's construct_local_mask + # which uses key_padding_mask.sum() as the sequence length + col_offset = cache_seqlens - L_q + + if DEBUG_CORE: + print(f"Cache mode: valid_len={cache_seqlens}, L_k={L_k}") + print(f"Using col_offset={col_offset} based on valid cache length") + else: + # Calculate offset for when seqlen_q != seqlen_k + # This offset aligns query positions to key positions + # When L_q < L_k, offset is positive, meaning query i maps to key position (i + offset) + # This is consistent with construct_local_mask in the tests which uses (sk - sq) + col_offset = L_k - L_q + cache_mask = None + + mask_applied = False + if causal and (window_size_left, window_size_right) == (-1, -1): + # Pure causal: ensure query doesn't attend to future keys + # With offset, query i can attend to keys up to position (i + col_offset) + mask = row_idx >= (col_idx - col_offset) + mask_applied = True if DEBUG_CORE: - print("causal_mask:", causal_mask) - # set -inf to places the causal mask is false + print("causal_mask:", mask) + elif (window_size_left, window_size_right) != (-1, -1): + # Handle the case where window sizes exceed sequence length + if window_size_left >= L_k: + window_size_left = -1 # No left limit + if window_size_right >= L_k: + window_size_right = -1 # No right limit + + if causal: + # Causal + sliding window: ensure we don't attend to future + window_size_right = min(window_size_right, 0) if window_size_right != -1 else 0 + + # Create sliding window mask + # Each query at position i attends to keys in [i + offset - left, i + offset + right] + if window_size_left == -1 and window_size_right == -1: + # No window restriction + mask = torch.ones((L_q, L_k), dtype=torch.bool, device=q.device) + else: + mask = torch.ones((L_q, L_k), dtype=torch.bool, device=q.device) + if window_size_left != -1: + # Each query at position i attends to keys from position (i - left) accounting for offset + mask = mask & (col_idx >= (row_idx + col_offset - window_size_left)) + if window_size_right != -1: + # Each query at position i attends to keys up to position (i + right) accounting for offset + mask = mask & (col_idx <= (row_idx + col_offset + window_size_right)) + + # Apply causal constraint + if causal: + causal_mask = row_idx >= (col_idx - col_offset) + mask = mask & causal_mask + + mask_applied = True + if DEBUG_CORE: + print(f"sliding_window_mask (left={window_size_left}, right={window_size_right}):", mask) + + # Apply cache mask if needed + if cache_mask is not None: + if mask_applied: + mask = mask & cache_mask + else: + mask = cache_mask + mask_applied = True + + # Apply the mask if created + if mask_applied: attention_scaled_scores = attention_scaled_scores.masked_fill( - torch.logical_not(causal_mask.unsqueeze(0)), float('-inf') + torch.logical_not(mask.unsqueeze(0)), float('-inf') ) if DEBUG_CORE: - print("attention_scaled_scores after causal:", attention_scaled_scores, attention_scaled_scores.shape) + print("attention_scaled_scores after masking:", attention_scaled_scores, attention_scaled_scores.shape) # Compute max for numerical stability max_scores = torch.max(attention_scaled_scores, dim=-1, keepdim=True)[0] if DEBUG_CORE: print("max_scores:", max_scores, max_scores.shape) - if causal: + if mask_applied: # Replace -inf in max_scores with zeros to avoid NaN in subtraction max_scores = torch.where( torch.isinf(max_scores), torch.zeros_like(max_scores), max_scores ) - if DEBUG: - print("max_scores if causal:", max_scores, max_scores.shape) + if DEBUG_CORE: + print("max_scores after mask handling:", max_scores, max_scores.shape) # Shift scores attention_shifted_scaled_scores = attention_scaled_scores - max_scores @@ -98,7 +207,7 @@ def attention_forward_core_ref_impl(q, k, v, sm_scale, causal, dropout_p, philox sum_exp_scores = torch.sum(exp_scores, dim=-1, keepdim=True) if DEBUG_CORE: print("sum_exp_scores:", sum_exp_scores, sum_exp_scores.shape) - if causal: + if mask_applied: # if sum of exp scores is 0.0 it means scores where -inf, we cannot compute softmax and softmax_lse. Setting to 1 deals with -inf case cleanly sum_exp_scores = torch.where( sum_exp_scores == 0, @@ -158,7 +267,7 @@ def attention_forward_core_ref_impl(q, k, v, sm_scale, causal, dropout_p, philox return o, softmax_lse, sd_mask -def attention_vanilla_forward_pytorch_ref_impl(q, k, v, sm_scale, causal, layout, dropout_p, philox_seed, philox_offset, alibi_slopes, use_exp2): +def attention_vanilla_forward_pytorch_ref_impl(q, k, v, sm_scale, causal, window_size_left, window_size_right, layout, dropout_p, philox_seed, philox_offset, alibi_slopes, use_exp2): """Compute reference output and softmax_lse using PyTorch's built-in function""" # Ensure the layout is 'bhsd' @@ -194,7 +303,7 @@ def attention_vanilla_forward_pytorch_ref_impl(q, k, v, sm_scale, causal, layout # Call the core attention function o, softmax_lse, sd_mask = attention_forward_core_ref_impl( - q, k, v, sm_scale, causal, dropout_p, philox_seed, philox_offset, alibi_slopes, use_exp2 + q, k, v, sm_scale, causal, window_size_left, window_size_right, dropout_p, philox_seed, philox_offset, alibi_slopes, use_exp2 ) if group_size != 1: @@ -224,6 +333,8 @@ def attention_varlen_forward_pytorch_ref_impl( v, sm_scale, causal, + window_size_left, + window_size_right, layout, cu_seqlens_q, cu_seqlens_k, @@ -302,7 +413,7 @@ def attention_varlen_forward_pytorch_ref_impl( alibi_slopes_i = None # Call the core attention function for this sequence - o_i, softmax_lse_i, sd_mask_i = attention_forward_core_ref_impl(q_i, k_i, v_i, sm_scale, causal, dropout_p, philox_seed, philox_offset, alibi_slopes_i, use_exp2) + o_i, softmax_lse_i, sd_mask_i = attention_forward_core_ref_impl(q_i, k_i, v_i, sm_scale, causal, window_size_left, window_size_right, dropout_p, philox_seed, philox_offset, alibi_slopes_i, use_exp2) # Reshape outputs back to original dimensions if group_size != 1: @@ -328,8 +439,6 @@ def attention_varlen_forward_pytorch_ref_impl( return o, softmax_lse, sd_mask - - def attention_forward_pytorch_ref_impl( q: torch.Tensor, k: torch.Tensor, @@ -338,6 +447,8 @@ def attention_forward_pytorch_ref_impl( sm_scale: float, alibi_slopes: Optional[torch.Tensor], causal: bool, + window_size_left: int, + window_size_right: int, layout: Literal["bshd", "bhsd", "thd"], cu_seqlens_q: torch.Tensor, cu_seqlens_k: torch.Tensor, @@ -356,6 +467,8 @@ def attention_forward_pytorch_ref_impl( v.clone(), sm_scale, causal, + window_size_left, + window_size_right, layout, cu_seqlens_q, cu_seqlens_k, @@ -374,6 +487,8 @@ def attention_forward_pytorch_ref_impl( v.clone(), sm_scale, causal, + window_size_left, + window_size_right, layout, dropout_p, philox_seed, @@ -385,3 +500,152 @@ def attention_forward_pytorch_ref_impl( out.copy_(o_ref.to(out.dtype)) return softmax_lse_ref, sd_mask_ref + +def attention_decode_forward_ref_impl( + q: torch.Tensor, + k_cache: torch.Tensor, + v_cache: torch.Tensor, + k_new: Optional[torch.Tensor], + v_new: Optional[torch.Tensor], + out: torch.Tensor, + sm_scale: float, + causal: bool, + window_size_left: int, + window_size_right: int, + alibi_slopes: Optional[torch.Tensor], + layout: Literal["bshd"], + cache_seqlens: Optional[torch.Tensor], + cache_batch_idx: Optional[torch.Tensor], +): + """Compute reference output for decode attention using PyTorch's built-in functions""" + + # get batch size before any layout conversion + batch_size = q.shape[0] + + # handle cache_batch_idx + if cache_batch_idx is not None: + # remap batch indices for cache access + batch_indices = cache_batch_idx + else: + batch_indices = torch.arange(batch_size, device=q.device) + + # copy new keys and values into cache if provided (before any layout conversion) + if k_new is not None and v_new is not None: + _, seq_len_new, _, _ = k_new.shape # shape is [batch, seq_len, nheads, head_dim] for bshd layout + + for b in range(batch_size): + cache_idx = batch_indices[b].item() if torch.is_tensor(batch_indices) else batch_indices + + # determine where to place new k/v in cache + if cache_seqlens is not None: + if torch.is_tensor(cache_seqlens): + start_pos = cache_seqlens[b].item() + else: + start_pos = cache_seqlens + else: + # if no cache_seqlens, assume we're filling from the beginning + start_pos = 0 + + end_pos = start_pos + seq_len_new + + # copy new keys and values into cache (both are in bshd layout) + k_cache[cache_idx, start_pos:end_pos, :, :] = k_new[b, :, :, :] + v_cache[cache_idx, start_pos:end_pos, :, :] = v_new[b, :, :, :] + + # ensure the layout is 'bhsd' + if layout == "bshd": + q = q.transpose(1, 2).contiguous() + k_cache = k_cache.transpose(1, 2).contiguous() + v_cache = v_cache.transpose(1, 2).contiguous() + elif layout != "bhsd": + raise ValueError(f"Unknown layout {layout}") + + # prepare tensors + batch_size_q, nheads_q, seq_len_q, head_dim = q.shape + batch_size_cache, nheads_k, max_cache_len, head_dim_k = k_cache.shape + _, nheads_v, _, head_dim_v = v_cache.shape + + # validate dimensions + assert head_dim == head_dim_k == head_dim_v, f"Head dimensions must match: {head_dim}, {head_dim_k}, {head_dim_v}" + + # handle MQA/GQA + group_size = nheads_q // nheads_k + if nheads_q % nheads_k != 0: + raise ValueError("nheads_q must be divisible by nheads_k") + + # handle cache_batch_idx + if cache_batch_idx is not None: + # remap batch indices for cache access + batch_indices = cache_batch_idx + else: + batch_indices = torch.arange(batch_size, device=q.device) + + # prepare outputs + o = torch.zeros_like(q) + softmax_lse = torch.zeros((batch_size, nheads_q, seq_len_q), dtype=torch.float32, device=q.device) + + # process each batch element + for b in range(batch_size): + cache_idx = batch_indices[b].item() if torch.is_tensor(batch_indices) else batch_indices + + # determine valid cache length for this batch element + if cache_seqlens is not None: + if torch.is_tensor(cache_seqlens): + cache_len = cache_seqlens[b].item() + if k_new is not None: + _, seq_len_new, _, _ = k_new.shape + cache_len += seq_len_new + else: + cache_len = cache_seqlens + if k_new is not None: + _, seq_len_new, _, _ = k_new.shape + cache_len += seq_len_new + else: + cache_len = max_cache_len + + # CHANGE: Extract the full cache, not just valid portion + # This matches what the test does - it uses full k_cache_rep/v_cache_rep + k_b = k_cache[cache_idx, :, :, :] # [nheads_k, max_cache_len, head_dim] + v_b = v_cache[cache_idx, :, :, :] # [nheads_v, max_cache_len, head_dim] + q_b = q[b:b+1, :, :, :] # [1, nheads_q, seq_len_q, head_dim] + + # handle MQA/GQA by expanding k and v + if group_size != 1: + # expand k and v to match q's number of heads + k_b = k_b.unsqueeze(1).expand(-1, group_size, -1, -1) + k_b = k_b.reshape(nheads_q, max_cache_len, head_dim) + + v_b = v_b.unsqueeze(1).expand(-1, group_size, -1, -1) + v_b = v_b.reshape(nheads_q, max_cache_len, head_dim) + + # reshape for attention_forward_core_ref_impl + q_b = q_b.reshape(nheads_q, seq_len_q, head_dim) + + # handle alibi slopes for this batch + alibi_slopes_b = None + if alibi_slopes is not None: + if alibi_slopes.dim() == 2: + alibi_slopes_b = alibi_slopes[b] + else: + alibi_slopes_b = alibi_slopes + + # call core attention function with cache information + o_b, softmax_lse_b, _ = attention_forward_core_ref_impl( + q_b, k_b, v_b, sm_scale, causal, window_size_left, window_size_right, + dropout_p=0.0, philox_seed=None, philox_offset=None, + alibi_slopes=alibi_slopes_b, use_exp2=True, + cache_seqlens=cache_len, # Pass valid cache length + ) + + # store outputs + o[b, :, :, :] = o_b.reshape(nheads_q, seq_len_q, head_dim) + softmax_lse[b, :, :] = softmax_lse_b.reshape(nheads_q, seq_len_q) + + # restore original layout if necessary + if layout == "bshd": + o = o.transpose(1, 2) + + # copy output to the provided tensor + out.copy_(o.to(out.dtype)) + + return softmax_lse \ No newline at end of file diff --git a/flash_attn/flash_attn_triton_amd/interface_fa.py b/flash_attn/flash_attn_triton_amd/interface_fa.py index 9c10b7436c2..4e9bd006aee 100644 --- a/flash_attn/flash_attn_triton_amd/interface_fa.py +++ b/flash_attn/flash_attn_triton_amd/interface_fa.py @@ -5,7 +5,7 @@ from .bwd_prefill_fused_atomics import attention_prefill_backward_triton_fused_atomics_impl from .bwd_prefill_fused_no_atomics import attention_prefill_backward_triton_split_fused_no_atomics_impl from .fwd_decode import attention_decode_forward_triton_impl -from .fwd_ref import attention_forward_pytorch_ref_impl +from .fwd_ref import attention_forward_pytorch_ref_impl, attention_decode_forward_ref_impl from .bwd_ref import attention_backward_pytorch_ref_impl from .utils import DEBUG, USE_REF, MetaData, is_fp8 from einops import rearrange, repeat @@ -101,6 +101,8 @@ def fwd(q: torch.Tensor, metadata.sm_scale, metadata.alibi_slopes, metadata.causal, + window_size_left, + window_size_right, metadata.layout, metadata.cu_seqlens_q, metadata.cu_seqlens_k, @@ -123,6 +125,8 @@ def fwd(q: torch.Tensor, metadata.sm_scale, metadata.alibi_slopes, metadata.causal, + window_size_left, + window_size_right, None, metadata.layout, metadata.cu_seqlens_q, @@ -253,6 +257,8 @@ def bwd( softmax_scale, alibi_slopes, causal, + window_size_left, + window_size_right, "bshd", None, None, @@ -472,6 +478,8 @@ def varlen_fwd( metadata.sm_scale, metadata.alibi_slopes, metadata.causal, + window_size_left, + window_size_right, metadata.layout, metadata.cu_seqlens_q, metadata.cu_seqlens_k, @@ -494,6 +502,8 @@ def varlen_fwd( metadata.sm_scale, metadata.alibi_slopes, metadata.causal, + window_size_left, + window_size_right, None, metadata.layout, metadata.cu_seqlens_q, @@ -626,6 +636,8 @@ def varlen_bwd( softmax_scale, alibi_slopes, causal, + window_size_left, + window_size_right, "thd", cu_seqlens_q, cu_seqlens_k, @@ -801,8 +813,21 @@ def fwd_kvcache( metadata.layout = "bshd" metadata.max_seqlens_q = q.shape[1] metadata.max_seqlens_k = k_cache.shape[1] - metadata.cache_seqlens = cache_seqlens metadata.cache_batch_idx = cache_batch_idx + if isinstance(cache_seqlens, int): + metadata.cache_seqlens = torch.tensor(cache_seqlens, device=q.device) + else: + metadata.cache_seqlens = cache_seqlens + + # window_size can be a tensor sometimes + if isinstance(window_size_left, torch.Tensor): + metadata.window_size_left = int(window_size_left.item()) + else: + metadata.window_size_left = window_size_left + if isinstance(window_size_right, torch.Tensor): + metadata.window_size_right = int(window_size_right.item()) + else: + metadata.window_size_right = window_size_right k_new = k v_new = v @@ -829,7 +854,7 @@ def fwd_kvcache( # Rotary Embedding Implementation if apply_rotary: - if metadata.causal: # NOTE: when support is added. Add `or metadata.local` + if metadata.causal or (window_size_left != -1 or window_size_right !=-1): # NOTE: when support is added. Add `or metadata.local` q_ro = apply_rotary_emb( q, metadata.rotary_cos, @@ -860,8 +885,29 @@ def fwd_kvcache( q, k_new = q_ro.to(q.dtype), k_ro.to(q.dtype) # launch kernel - DECODE_KERNEL= True # os.environ.get('DECODE_KERNEL', '0').lower() in ('1', 'true', 'yes') - if DECODE_KERNEL: + if USE_REF: + if DEBUG: + print("Using reference implementation") + softmax_lse_ref = attention_decode_forward_ref_impl( + q, + k_cache, + v_cache, + k_new, + v_new, + out, + metadata.sm_scale, + metadata.causal, + metadata.window_size_left, + metadata.window_size_right, + metadata.alibi_slopes, + metadata.layout, + metadata.cache_seqlens, + metadata.cache_batch_idx, + ) + softmax_lse=softmax_lse_ref + else: + if DEBUG: + print("Using Triton implementation") softmax_lse_triton = attention_decode_forward_triton_impl( q, k_cache, @@ -871,38 +917,14 @@ def fwd_kvcache( out, metadata.sm_scale, metadata.causal, + metadata.window_size_left, + metadata.window_size_right, metadata.alibi_slopes, metadata.layout, metadata.cache_seqlens, metadata.cache_batch_idx, ) - else: - softmax_lse_triton, sd_mask_triton = attention_prefill_forward_triton_impl( - q, - k_cache, - v_cache, - out, - metadata.sm_scale, - metadata.alibi_slopes, - metadata.causal, - None, - metadata.layout, - metadata.cu_seqlens_q, - metadata.cu_seqlens_k, - metadata.max_seqlens_q, - metadata.max_seqlens_k, - metadata.cache_seqlens, - metadata.cache_batch_idx, - metadata.dropout_p, - metadata.philox_seed, - metadata.philox_offset, - metadata.return_softmax, - USE_EXP2, - None, - None, - None, - None) - softmax_lse = softmax_lse_triton + softmax_lse = softmax_lse_triton if DEBUG: print("out:", out, out.shape) diff --git a/flash_attn/flash_attn_triton_amd/test.py b/flash_attn/flash_attn_triton_amd/test.py index f634103ca69..9cffda127a6 100644 --- a/flash_attn/flash_attn_triton_amd/test.py +++ b/flash_attn/flash_attn_triton_amd/test.py @@ -43,355 +43,6 @@ # ATOL_fp8, RTOL_fp8 = 2e-2, 2e-2 # fp8 EQUAL_NAN = True -@pytest.mark.parametrize( - "BATCH, HQ, HK, N_CTX_Q, N_CTX_K, D_HEAD", - [ - (1, 1, 1, 1, 1, 1), - (1, 1, 1, 2, 4, 16), - (1, 2, 2, 2, 4, 16), - (1, 4, 1, 2, 4, 16), - (1, 4, 2, 2, 4, 16), - (1, 1, 1, 4, 2, 16), - (1, 1, 1, 4, 4, 16), - (1, 2, 2, 4, 4, 16), - (2, 1, 1, 4, 4, 16), - (2, 2, 2, 4, 4, 16), - (1, 1, 1, 128, 64, 16), - (2, 2, 2, 2, 128, 1), - (2, 3, 3, 2, 128, 16), - (3, 2, 2, 256, 512, 16), - (3, 3, 3, 128, 128, 64), - (2, 4, 4, 1024, 1024, 64), - (4, 6, 6, 108, 256, 224), - (4, 8, 8, 2048, 2048, 128), - (4, 16, 16, 4096, 4096, 64), - (2, 4, 4, 8192, 8192, 32), - # fa configs - (4, 6, 1, 113, 203, 256), - (4, 6, 1, 128, 217, 256), - (4, 6, 2, 113, 211, 128), - (4, 6, 2, 108, 256, 128), - (4, 6, 1, 256, 512, 64), - (4, 6, 1, 512, 256, 64), - (4, 6, 2, 1024, 1024, 32), - (4, 6, 2, 1023, 1024, 32), - (4, 6, 6, 1024, 1023, 32), - (4, 6, 6, 2048, 2048, 32), - ], -) -@pytest.mark.parametrize('causal', [True, False]) -@pytest.mark.parametrize('dropout_p', [0.0]) -@pytest.mark.parametrize('alibi_slopes', [None]) -@pytest.mark.parametrize('layout', ["bshd", "thd"]) -@pytest.mark.parametrize('dtype', [torch.float16]) -@pytest.mark.parametrize('use_exp2', [True, False]) # works when use_exp2 is false -@pytest.mark.parametrize('DEBUG_INPUT', [False]) # NOTE: debug input can overflow when the tensors are large. Just use to figure out issues -@pytest.mark.skip() -def test_op_prefill_fwd_impl(BATCH, HQ, HK, N_CTX_Q, N_CTX_K, D_HEAD, causal, dropout_p, alibi_slopes, layout, dtype, use_exp2, DEBUG_INPUT): - torch.manual_seed(42) - device = "cuda" - - q, k, v, do, metadata = input_helper(BATCH, HQ, HK, N_CTX_Q, N_CTX_K, D_HEAD, causal, dropout_p, dtype, layout=layout, device=device) - - if DEBUG: - if HQ // HK != 1: - print("MQA/GQA") - else: - print("MHA") - - # update metadata - if causal: - metadata.need_causal(True) - - # NOTE: the returned score is not the same as the reference because we need to adjust as we find new maxes per block. We are not doing that - metadata.need_dropout(dropout_p, True) - - - # call Triton's forward implementation directly - q_triton = q.clone() - k_triton = k.clone() - v_triton = v.clone() - o_triton = torch.zeros_like(q).contiguous() if DEBUG_INPUT else torch.empty_like(q) - softmax_lse_triton, sd_mask_triton = attention_prefill_forward_triton_impl( - q_triton, - k_triton, - v_triton, - o_triton, - metadata.sm_scale, - metadata.alibi_slopes, - metadata.causal, - metadata.bias, - metadata.layout, - metadata.cu_seqlens_q, - metadata.cu_seqlens_k, - metadata.max_seqlens_q, - metadata.max_seqlens_k, - metadata.cache_seqlens, - metadata.cache_batch_idx, - metadata.dropout_p, - metadata.philox_seed, - metadata.philox_offset, - metadata.return_softmax, - use_exp2, - None, - None, - None, - None) - - # ref forward - q_ref = q.clone() - k_ref = k.clone() - v_ref = v.clone() - o_ref = torch.zeros_like(q).contiguous() if DEBUG_INPUT else torch.empty_like(q) - softmax_lse_ref, sd_mask_ref = attention_forward_pytorch_ref_impl( - q_ref, - k_ref, - v_ref, - o_ref, - metadata.sm_scale, - metadata.alibi_slopes, - causal, - layout, - metadata.cu_seqlens_q, - metadata.cu_seqlens_k, - metadata.max_seqlens_q, - metadata.max_seqlens_k, - metadata.dropout_p, - metadata.philox_seed, - metadata.philox_offset, - use_exp2 - ) - - if DEBUG: - print() - print("Compare Triton Impl with refernce Pytorch Impl") - - # this can be set to true manually or when using dropout - if metadata.return_softmax: - if DEBUG: - print("sd_mask_triton:", sd_mask_triton, sd_mask_triton.shape) - print("sd_mask_ref:", sd_mask_ref, sd_mask_ref.shape) - torch.testing.assert_close(sd_mask_triton.to(sd_mask_ref.dtype), sd_mask_ref, atol=ATOL, rtol=RTOL) - - if DEBUG: - print("softmax_lse_triton:", softmax_lse_triton, softmax_lse_triton.shape) - print("softmax_lse_ref:", softmax_lse_ref, softmax_lse_ref.shape) - torch.testing.assert_close(softmax_lse_triton, softmax_lse_ref, atol=ATOL, rtol=RTOL) - - if DEBUG: - print("output_triton:", o_triton, o_triton.shape) - print("output_ref:", o_ref, o_ref.shape) - torch.testing.assert_close(o_triton, o_ref, atol=ATOL, rtol=RTOL) - -@pytest.mark.parametrize( - "BATCH, HQ, HK, N_CTX_Q, N_CTX_K, D_HEAD", [ - (1, 1, 1, 1, 1, 1), - (1, 1, 1, 4, 4, 4), - (2, 1, 1, 4, 4, 16), - (1, 2, 2, 4, 4, 16), - (1, 4, 1, 2, 4, 16), - (1, 8, 1, 2, 4, 16), - (1, 16, 1, 2, 4, 16), - (1, 32, 1, 2, 4, 16), - (1, 64, 1, 2, 4, 16), - (1, 4, 2, 2, 4, 16), - (2, 2, 2, 4, 4, 16), - (1, 1, 1, 4, 4, 16), - (2, 1, 1, 4, 4 , 16), - (4, 6, 6, 8, 8 , 16), - (1, 1, 1, 4, 4, 32), - (1, 1, 1, 16, 16, 16), - (1, 1, 1, 32, 32, 16), - (1, 1, 1, 64, 64, 16), - (1, 1, 1, 64, 64, 16), - (1, 1, 1, 64, 128, 16), - (1, 1, 1, 64, 64, 32), - (1, 1, 1, 64, 128, 32), - (1, 1, 1, 128, 128, 64), - (1, 1, 1, 128, 256, 45), - (1, 1, 1, 113, 203, 192), - (1, 1, 1, 256, 256, 64), - (1, 1, 1, 256, 512, 16), - (1, 1, 1, 512, 512, 64), - (1, 1, 1, 1024, 1024, 64), - # fa configs - (2, 2, 2, 128, 128, 65), - (2, 2, 2, 128, 128, 224), - (4, 6, 6, 108, 256, 224), - (1, 1, 1, 256, 512, 16), - # old tests that work - (4, 48, 6, 1024, 1024, 64), - (4, 48, 12, 2048, 1024, 64), - (4, 48, 24, 1024, 1024, 64), - (4, 48, 48, 1024, 1024, 64), - (4, 48, 48, 1024, 1024, 73), - (4, 48, 48, 2048, 2048, 64), - (1, 24, 24, 4096, 4096, 64), - (1, 16, 16, 1024, 1024, 64), - (1, 16, 16, 1024, 1024, 128), - # testcase new - # seqlen q == k - (1, 1, 1, 2, 2, 2), # small enough to debug - (1, 1, 1, 128, 128, 32), # only one block - (1, 1, 1, 127, 127, 32), # only one block but with masking - (1, 1, 1, 129, 129, 1), # two blocks with 2nd block small enough to debug - (1, 1, 1, 350, 350, 1), # two blocks with 2nd block small enough to debug - (1, 1, 1, 350, 350, 68), # generic masking on q, k and head - (4, 1, 1, 512, 512, 128), # batch > 1 - (4, 8, 2, 512, 512, 128), # GQA - (4, 8, 2, 512, 512, 68), # non-power-of-2 head_dim - (4, 8, 2, 500, 500, 68), # comprehensive case for seqlen q == k - # seqlen q > k - (1, 1, 1, 64, 32, 8), # seqlen_q > seqlen_k - (1, 1, 1, 192, 128, 32), # seqlen_q > seqlen_k - (4, 8, 2, 1024, 512, 68), # seqlen_q < seqlen_k - (1, 1, 1, 729, 516, 68), # seqlen_q > seqlen_k - (16, 16, 4, 2753, 1528, 68), # a comprehensive seqlen_q > seqlen_k - # seqlen q < k - (1, 1, 1, 32, 64, 8), # seqlen_q > seqlen_k - (1, 1, 1, 128, 192, 32), # seqlen_q < seqlen_k - (4, 8, 2, 512, 1024, 68), # seqlen_q < seqlen_k - (1, 1, 1, 200, 413, 1), # seqlen_q < seqlen_k - (1, 1, 1, 782, 1546, 1), # seqlen_q < seqlen_k - (16, 16, 4, 1528, 2753, 68), # a comprehensive seqlen_q < seqlen_k -]) -@pytest.mark.parametrize('causal', [True, False]) -@pytest.mark.parametrize('dropout_p', [0.0]) -@pytest.mark.parametrize('alibi_slopes', [None]) -@pytest.mark.parametrize('layout', ["bshd", "thd"]) -@pytest.mark.parametrize('dtype', [torch.float16]) -@pytest.mark.parametrize('use_exp2', [False]) # FIXME: using exp2 causes issue when used with causal -@pytest.mark.parametrize('DEBUG_INPUT', [False]) # debug output causes nans on larger tensors -@pytest.mark.skip() -def test_op_prefill_bwd_impl(BATCH, HQ, HK, N_CTX_Q, N_CTX_K, D_HEAD, causal, dropout_p, alibi_slopes, layout, dtype, use_exp2, DEBUG_INPUT): - torch.manual_seed(20) - device="cuda" - - # gen inputs - q, k, v, do, metadata = input_helper(BATCH, HQ, HK, N_CTX_Q, N_CTX_K, D_HEAD, causal, dropout_p, dtype, layout=layout, device=device) - - # NOTE: the returned score is not the same as the reference because we need to adjust as we find new maxes per block. We are not doing that - metadata.need_dropout(dropout_p, True) - - # =============================================== Reference ============================================================== - # fwd - q_ref = q.clone() - k_ref = k.clone() - v_ref = v.clone() - output_ref = torch.zeros_like(q).contiguous() if DEBUG_INPUT else torch.empty_like(q) - softmax_lse_ref, sd_mask_ref = attention_forward_pytorch_ref_impl( - q_ref, - k_ref, - v_ref, - output_ref, - metadata.sm_scale, - metadata.alibi_slopes, - causal, - layout, - metadata.cu_seqlens_q, - metadata.cu_seqlens_k, - metadata.max_seqlens_q, - metadata.max_seqlens_k, - metadata.dropout_p, - metadata.philox_seed, - metadata.philox_offset, - use_exp2 - ) - - # bwd - do_ref = do.clone() - dq_ref = torch.zeros_like(q).contiguous() if DEBUG_INPUT else torch.empty_like(q) - dk_ref = torch.zeros_like(k).contiguous() if DEBUG_INPUT else torch.empty_like(k) - dv_ref = torch.zeros_like(v).contiguous() if DEBUG_INPUT else torch.empty_like(v) - delta_ref = attention_backward_pytorch_ref_impl( - do_ref, - q_ref, - k_ref, - v_ref, - output_ref, - softmax_lse_ref, - dq_ref, - dk_ref, - dv_ref, - metadata.sm_scale, - metadata.alibi_slopes, - causal, - layout, - metadata.cu_seqlens_q, - metadata.cu_seqlens_k, - metadata.max_seqlens_q, - metadata.max_seqlens_k, - metadata.dropout_p, - metadata.philox_seed, - metadata.philox_offset, - use_exp2 - ) - - # =============================================== Triton ============================================================== - do_triton = do.clone() - q_triton = q.clone() - k_triton = k.clone() - v_triton = v.clone() - o_triton = output_ref.clone().contiguous() - softmax_lse_triton = softmax_lse_ref.clone().contiguous() - dq_triton = torch.zeros_like(q_triton, dtype=q.dtype) # NOTE: the kernel does inplace accumlation on dq so dq has to be zeros - dk_triton = torch.zeros_like(k_triton, dtype=k.dtype) if DEBUG_INPUT else torch.empty_like(k_triton, dtype=k.dtype) - dv_triton = torch.zeros_like(v_triton, dtype=v.dtype) if DEBUG_INPUT else torch.empty_like(v_triton, dtype=v.dtype) - delta_triton = attention_prefill_backward_triton_split_fused_no_atomics_impl( - do_triton, - q_triton, - k_triton, - v_triton, - o_triton, - softmax_lse_triton, - dq_triton, - dk_triton, - dv_triton, - metadata.sm_scale, - alibi_slopes, - causal, - layout, - metadata.cu_seqlens_q, - metadata.cu_seqlens_k, - metadata.max_seqlens_q, - metadata.max_seqlens_k, - metadata.dropout_p, - metadata.philox_seed, - metadata.philox_offset, - use_exp2, - None, - None, - None, - None, - None, - None, - None, - None, - ) - - # =============================================== Check ============================================================== - if DEBUG: - print() - if DEBUG: - print("delta_triton:", delta_triton, delta_triton.shape) - print("delta_ref:", delta_ref, delta_ref.shape) - torch.testing.assert_close(delta_triton, delta_ref, atol=ATOL, rtol=RTOL, equal_nan=EQUAL_NAN) - - if DEBUG: - print("dv_triton:", dv_triton, dv_triton.shape) - print("dv_ref:", dv_ref, dv_ref.shape) - torch.testing.assert_close(dv_triton, dv_ref, atol=ATOL, rtol=RTOL, equal_nan=EQUAL_NAN) - - if DEBUG: - print("dk_triton:", dk_triton, dk_triton.shape) - print("dk_ref:", dk_ref, dk_ref.shape) - torch.testing.assert_close(dk_triton, dk_ref, atol=ATOL, rtol=RTOL, equal_nan=EQUAL_NAN) - - if DEBUG: - print("dq_triton:", dq_triton, dq_triton.shape) - print("dq_ref:", dq_ref, dq_ref.shape) - torch.testing.assert_close(dq_triton, dq_ref, atol=ATOL, rtol=RTOL, equal_nan=EQUAL_NAN) - def fp8_assert_close(tensor_a, tensor_b, atol=ATOL_fp8, rtol=RTOL_fp8, max_diff_percentage=0.5): """Assert tensors are close with tolerance for small percentage of elements""" # standard comparison @@ -523,7 +174,7 @@ def test_fp8(Z, HQ, HK, N_CTX_Q, N_CTX_K, D_HEAD, causal, dropout_p, layout, pac # test apis if packing == 'qkv': # generate inputs - qkv, do, metadata = input_helper(Z, HQ, HK, N_CTX_Q, N_CTX_K, D_HEAD, causal, dropout_p, ref_dtype, layout, packing=packing, device=device, DEBUG_INPUT=DEBUG_INPUT) + qkv, do, metadata = input_helper(Z, HQ, HK, N_CTX_Q, N_CTX_K, D_HEAD, causal, dropout_p, ref_dtype, layout, packing=packing, device=device) # ---------------------------------------------------------------- # --- FP8 --- @@ -631,7 +282,7 @@ def test_fp8(Z, HQ, HK, N_CTX_Q, N_CTX_K, D_HEAD, causal, dropout_p, layout, pac elif packing is None: # generate inputs - q, k, v, do, metadata = input_helper(Z, HQ, HK, N_CTX_Q, N_CTX_K, D_HEAD, causal, dropout_p, ref_dtype, layout, device=device, DEBUG_INPUT=DEBUG_INPUT) + q, k, v, do, metadata = input_helper(Z, HQ, HK, N_CTX_Q, N_CTX_K, D_HEAD, causal, dropout_p, ref_dtype, layout, device=device) # ---------------------------------------------------------------- # --- FP8 --- @@ -951,6 +602,7 @@ def clear_compile_cache(): # (16, 48, 16, 256, 512, 64), # MQA test (HQ > HK) ], ) +@pytest.mark.skip() # this works or not depending on the torch and triton version compatibility. Keep the test but use it in docker containers where things are in sync def test_torch_compile(BATCH, HQ, HK, N_CTX_Q, N_CTX_K, D_HEAD): print(f"\n\nTesting with BATCH={BATCH}, HQ={HQ}, HK={HK}, N_CTX_Q={N_CTX_Q}, N_CTX_K={N_CTX_K}, D_HEAD={D_HEAD}") diff --git a/flash_attn/flash_attn_triton_amd/utils.py b/flash_attn/flash_attn_triton_amd/utils.py index 1795e0d1366..96d4f662567 100644 --- a/flash_attn/flash_attn_triton_amd/utils.py +++ b/flash_attn/flash_attn_triton_amd/utils.py @@ -6,7 +6,8 @@ import functools import triton import triton.language as tl -from typing import Literal, Optional, Union +import numpy as np +from typing import Literal, Optional # ------------------------------- # Gloabl Variables @@ -42,7 +43,7 @@ class MetaData(): num_contexts = 0 varlen: bool = False layout: Optional[Literal["bshd", "bhsd", "thd"]] = None - cache_seqlens: Optional[Union[(int, torch.Tensor)]] = None + cache_seqlens: Optional[torch.Tensor] = None cache_batch_idx = None packing: Optional[bool] = None return_softmax: bool = False @@ -54,6 +55,8 @@ class MetaData(): rotary_cos: Optional[torch.Tensor] = None rotary_interleaved: bool = False rotary_conjunction: bool = False + window_size_left: int = -1 + window_size_right: int = -1 def __repr__(self) -> str: @@ -73,6 +76,8 @@ def __repr__(self) -> str: f" cache_batch_idx={self.cache_batch_idx},\n" f" dropout_p={self.dropout_p},\n" f" return_softmax={self.return_softmax}\n" + f" window_size_left={self.window_size_left},\n" + f" window_size_right={self.window_size_right},\n" f")") def __init__(self, sm_scale=1.0): @@ -166,7 +171,7 @@ def generate_varlen_tensor( equal_seqlens: bool = False, device: str = "cuda", dtype: torch.dtype = torch.float16, - DEBUG_INPUT: bool = False + mode: Literal["random", "ones", "incremental", "identity"] = "random" ): if DEBUG: print("total_seqlen", total_seqlen) @@ -200,8 +205,8 @@ def generate_varlen_tensor( cu_seqlens = torch.cat([torch.tensor([0], dtype=torch.int32, device=device), seqlens.cumsum(dim=0)]).to(torch.int32).to(device=device) max_seqlen = torch.max(seqlens).to(torch.int32).item() - # create varlen tensor - if DEBUG_INPUT: + # create varlen tensor based on mode + if mode == "incremental": x = torch.zeros(total_seqlen, num_heads, head_size, dtype=dtype, device=device) for i in range(batch_size): start = cu_seqlens[i].item() @@ -213,8 +218,23 @@ def generate_varlen_tensor( .view(length, 1, 1) .expand(length, num_heads, head_size) ) - else: + elif mode == "identity": + x = torch.zeros(total_seqlen, num_heads, head_size, dtype=dtype, device=device) + # for each batch, create identity pattern within that batch's sequence + for i in range(batch_size): + start = cu_seqlens[i].item() + end = cu_seqlens[i+1].item() + length = end - start + + # create identity pattern for positions within this batch + for pos in range(min(length, head_size)): + x[start + pos, :, pos] = 1.0 + elif mode == "random": x = torch.randn((total_seqlen, num_heads, head_size), dtype=dtype, device=device) + elif mode == "ones": + x = torch.ones((total_seqlen, num_heads, head_size), dtype=dtype, device=device) + else: + raise ValueError(f"Unkown mode {mode}") if is_fp8_dtype: # cast to fp8 @@ -225,19 +245,28 @@ def generate_varlen_tensor( x.requires_grad_() return x, cu_seqlens, max_seqlen -def generate_bshd_tensor(BATCH, SEQ_LEN, NUM_HEADS, D_HEAD, dtype: torch.dtype = torch.float16, device="cuda", DEBUG_INPUT=False): +def generate_bshd_tensor(BATCH, SEQ_LEN, NUM_HEADS, D_HEAD, dtype: torch.dtype = torch.float16, device="cuda", mode: Literal["random", "ones", "incremental", "identity"] = "random"): # save fp8 type is_fp8_dtype = is_dtype_fp8(dtype) if is_fp8_dtype: og_fp8_dtype = dtype dtype = torch.float32 - # gen tensor + # gen tensor based on mode tensor_shape = (BATCH, SEQ_LEN, NUM_HEADS, D_HEAD) - if DEBUG_INPUT: + if mode == "incremental": x = torch.arange(SEQ_LEN, dtype=dtype, device=device).view(1, SEQ_LEN, 1, 1).expand(*tensor_shape).contiguous() - else: + elif mode == "identity": + x = torch.zeros(tensor_shape, dtype=dtype, device=device) + # create identity pattern: position i has value 1 at dimension i + for i in range(min(SEQ_LEN, D_HEAD)): + x[:, i, :, i] = 1.0 + elif mode == "random": x = torch.randn(tensor_shape, dtype=dtype, device=device) + elif mode == "ones": + x = torch.ones(tensor_shape, dtype=dtype, device=device) + else: + raise ValueError(f"Unkown mode {mode}") if is_fp8_dtype: # cast to fp8 @@ -248,26 +277,31 @@ def generate_bshd_tensor(BATCH, SEQ_LEN, NUM_HEADS, D_HEAD, dtype: torch.dtype = x.requires_grad_() return x -def generate_bhsd_tensor(BATCH, NUM_HEADS, SEQ_LEN, D_HEAD, dtype: torch.dtype = torch.float16, device="cuda", DEBUG_INPUT=False): +def generate_bhsd_tensor(BATCH, NUM_HEADS, SEQ_LEN, D_HEAD, dtype: torch.dtype = torch.float16, device="cuda", mode: Literal["random", "ones", "incremental", "identity"] = "random"): # save fp8 type is_fp8_dtype = is_dtype_fp8(dtype) if is_fp8_dtype: og_fp8_dtype = dtype dtype = torch.float32 - # gen tensor + # gen tensor based on mode tensor_shape = (BATCH, NUM_HEADS, SEQ_LEN, D_HEAD) - if DEBUG_INPUT: + if mode == "incremental": x = torch.arange(SEQ_LEN, dtype=dtype, device=device).view(1, 1, SEQ_LEN, 1).expand(*tensor_shape).contiguous() - else: + elif mode == "identity": + x = torch.zeros(tensor_shape, dtype=dtype, device=device) + # create identity pattern: position i has value 1 at dimension i + for i in range(min(SEQ_LEN, D_HEAD)): + x[:, :, i, i] = 1.0 + elif mode == "random": x = torch.randn(tensor_shape, dtype=dtype, device=device) + elif mode == "ones": + x = torch.ones(tensor_shape, dtype=dtype, device=device) + else: + raise ValueError(f"Unkown mode {mode}") - if is_fp8_dtype: - # cast to fp8 - x, descale_x = cast_to_fp8(x, og_fp8_dtype, "bhsd") # FIXME: I don't the casting fn supports this atm - x.requires_grad_() - return x, descale_x + raise ValueError("fp8 not supported for bhsd yet") else: x.requires_grad_() return x @@ -513,8 +547,7 @@ def input_helper( dtype: torch.dtype, layout: Literal["bshd", "bhsd", "thd"], packing: Optional[Literal["kv", "qkv"]] = None, - device: Literal["cpu", "cuda"] = "cuda", - DEBUG_INPUT: bool = False, + device: Literal["cpu", "cuda"] = "cuda" ): torch.manual_seed(20) is_fp8_dtype = is_dtype_fp8(dtype) @@ -529,23 +562,23 @@ def input_helper( if packing is None: # gen tensors if is_fp8_dtype: - q, cu_seqlens_q, max_seqlen_q, descale_q = generate_varlen_tensor(TOTAL_SEQLENS_Q, HQ, D_HEAD, batch_size=BATCH, dtype=dtype, device=device, equal_seqlens=equal_seqlens, DEBUG_INPUT=DEBUG_INPUT) - k, cu_seqlens_k, max_seqlen_k, descale_k = generate_varlen_tensor(TOTAL_SEQLENS_K, HK, D_HEAD, batch_size=BATCH, dtype=dtype, device=device, equal_seqlens=equal_seqlens, DEBUG_INPUT=DEBUG_INPUT) - v, _, _, descale_v = generate_varlen_tensor(TOTAL_SEQLENS_K, HK, D_HEAD, batch_size=BATCH, dtype=dtype, device=device, equal_seqlens=equal_seqlens, DEBUG_INPUT=DEBUG_INPUT) + q, cu_seqlens_q, max_seqlen_q, descale_q = generate_varlen_tensor(TOTAL_SEQLENS_Q, HQ, D_HEAD, batch_size=BATCH, dtype=dtype, device=device, equal_seqlens=equal_seqlens) + k, cu_seqlens_k, max_seqlen_k, descale_k = generate_varlen_tensor(TOTAL_SEQLENS_K, HK, D_HEAD, batch_size=BATCH, dtype=dtype, device=device, equal_seqlens=equal_seqlens) + v, _, _, descale_v = generate_varlen_tensor(TOTAL_SEQLENS_K, HK, D_HEAD, batch_size=BATCH, dtype=dtype, device=device, equal_seqlens=equal_seqlens) do, _, _, descale_do = generate_varlen_tensor(TOTAL_SEQLENS_Q, HQ, D_HEAD, batch_size=BATCH, dtype=dtype, device=device, equal_seqlens=equal_seqlens) else: - q, cu_seqlens_q, max_seqlen_q = generate_varlen_tensor(TOTAL_SEQLENS_Q, HQ, D_HEAD, batch_size=BATCH, dtype=dtype, device=device, equal_seqlens=equal_seqlens, DEBUG_INPUT=DEBUG_INPUT) - k, cu_seqlens_k, max_seqlen_k = generate_varlen_tensor(TOTAL_SEQLENS_K, HK, D_HEAD, batch_size=BATCH, dtype=dtype, device=device, equal_seqlens=equal_seqlens, DEBUG_INPUT=DEBUG_INPUT) - v, _, _ = generate_varlen_tensor(TOTAL_SEQLENS_K, HK, D_HEAD, batch_size=BATCH, dtype=dtype, device=device, equal_seqlens=equal_seqlens, DEBUG_INPUT=DEBUG_INPUT) - do = torch.ones_like(q) if DEBUG_INPUT else torch.randn_like(q) + q, cu_seqlens_q, max_seqlen_q = generate_varlen_tensor(TOTAL_SEQLENS_Q, HQ, D_HEAD, batch_size=BATCH, dtype=dtype, device=device, equal_seqlens=equal_seqlens) + k, cu_seqlens_k, max_seqlen_k = generate_varlen_tensor(TOTAL_SEQLENS_K, HK, D_HEAD, batch_size=BATCH, dtype=dtype, device=device, equal_seqlens=equal_seqlens) + v, _, _ = generate_varlen_tensor(TOTAL_SEQLENS_K, HK, D_HEAD, batch_size=BATCH, dtype=dtype, device=device, equal_seqlens=equal_seqlens) + do, _, _ = generate_varlen_tensor(TOTAL_SEQLENS_Q, HQ, D_HEAD, batch_size=BATCH, dtype=dtype, device=device, equal_seqlens=equal_seqlens) elif packing == "kv": # gen tensors with kv packing if is_fp8_dtype: raise ValueError("FP8 not supported for KV packing yet") else: - q, cu_seqlens_q, max_seqlen_q = generate_varlen_tensor(TOTAL_SEQLENS_Q, HQ, D_HEAD, batch_size=BATCH, dtype=dtype, device=device, equal_seqlens=equal_seqlens, DEBUG_INPUT=DEBUG_INPUT) - kv, cu_seqlens_k, max_seqlen_k = generate_varlen_kv_packed(TOTAL_SEQLENS_K, HK, D_HEAD, batch_size=BATCH, dtype=dtype, device=device, equal_seqlens=equal_seqlens, DEBUG_INPUT=DEBUG_INPUT) - do = torch.ones_like(q) if DEBUG_INPUT else torch.randn_like(q) + q, cu_seqlens_q, max_seqlen_q = generate_varlen_tensor(TOTAL_SEQLENS_Q, HQ, D_HEAD, batch_size=BATCH, dtype=dtype, device=device, equal_seqlens=equal_seqlens) + kv, cu_seqlens_k, max_seqlen_k = generate_varlen_kv_packed(TOTAL_SEQLENS_K, HK, D_HEAD, batch_size=BATCH, dtype=dtype, device=device, equal_seqlens=equal_seqlens) + do, _, _ = generate_varlen_tensor(TOTAL_SEQLENS_Q, HQ, D_HEAD, batch_size=BATCH, dtype=dtype, device=device, equal_seqlens=equal_seqlens) elif packing == "qkv": # qkv packing - requires same sequence length for q and k assert N_CTX_Q == N_CTX_K, "For QKV packing, Q and K must have same sequence length" @@ -554,17 +587,13 @@ def input_helper( if is_fp8_dtype: raise ValueError("FP8 not supported for QKV packing yet") else: - qkv, cu_seqlens_q, max_seqlen_q = generate_varlen_qkv_packed(TOTAL_SEQLENS_Q, HQ, D_HEAD, batch_size=BATCH, dtype=dtype, device=device, equal_seqlens=equal_seqlens, DEBUG_INPUT=DEBUG_INPUT) + qkv, cu_seqlens_q, max_seqlen_q = generate_varlen_qkv_packed(TOTAL_SEQLENS_Q, HQ, D_HEAD, batch_size=BATCH, dtype=dtype, device=device, equal_seqlens=equal_seqlens) cu_seqlens_k = cu_seqlens_q max_seqlen_k = max_seqlen_q - # create dummy do for qkv case - do = torch.ones((TOTAL_SEQLENS_Q, HQ, D_HEAD), dtype=dtype, device=device) if DEBUG_INPUT else torch.randn((TOTAL_SEQLENS_Q, HQ, D_HEAD), dtype=dtype, device=device) + do, _, _ = generate_varlen_tensor(TOTAL_SEQLENS_Q, HQ, D_HEAD, batch_size=BATCH, dtype=dtype, device=device, equal_seqlens=equal_seqlens) # setup metadata - if DEBUG_INPUT: - sm_scale = 1 - else: - sm_scale = D_HEAD**-0.5 + sm_scale = D_HEAD**-0.5 metadata = MetaData(sm_scale=sm_scale) metadata.set_varlen_params(cu_seqlens_q, cu_seqlens_k, max_seqlen_q, max_seqlen_k) metadata.need_causal(CAUSAL) @@ -576,39 +605,38 @@ def input_helper( # gen tensors if layout == "bshd": if is_fp8_dtype: - q, descale_q = generate_bshd_tensor(BATCH, N_CTX_Q, HQ, D_HEAD, dtype=dtype, device=device, DEBUG_INPUT=DEBUG_INPUT) - k, descale_k = generate_bshd_tensor(BATCH, N_CTX_K, HK, D_HEAD, dtype=dtype, device=device, DEBUG_INPUT=DEBUG_INPUT) - v, descale_v = generate_bshd_tensor(BATCH, N_CTX_K, HK, D_HEAD, dtype=dtype, device=device, DEBUG_INPUT=DEBUG_INPUT) + q, descale_q = generate_bshd_tensor(BATCH, N_CTX_Q, HQ, D_HEAD, dtype=dtype, device=device) + k, descale_k = generate_bshd_tensor(BATCH, N_CTX_K, HK, D_HEAD, dtype=dtype, device=device) + v, descale_v = generate_bshd_tensor(BATCH, N_CTX_K, HK, D_HEAD, dtype=dtype, device=device) do, descale_do = generate_bshd_tensor(BATCH, N_CTX_Q, HQ, D_HEAD, dtype=dtype, device=device) else: - q = generate_bshd_tensor(BATCH, N_CTX_Q, HQ, D_HEAD, dtype=dtype, device=device, DEBUG_INPUT=DEBUG_INPUT) - k = generate_bshd_tensor(BATCH, N_CTX_K, HK, D_HEAD, dtype=dtype, device=device, DEBUG_INPUT=DEBUG_INPUT) - v = generate_bshd_tensor(BATCH, N_CTX_K, HK, D_HEAD, dtype=dtype, device=device, DEBUG_INPUT=DEBUG_INPUT) - do = torch.ones_like(q) if DEBUG_INPUT else torch.randn_like(q) + q = generate_bshd_tensor(BATCH, N_CTX_Q, HQ, D_HEAD, dtype=dtype, device=device) + k = generate_bshd_tensor(BATCH, N_CTX_K, HK, D_HEAD, dtype=dtype, device=device) + v = generate_bshd_tensor(BATCH, N_CTX_K, HK, D_HEAD, dtype=dtype, device=device) + do = generate_bshd_tensor(BATCH, N_CTX_Q, HQ, D_HEAD, dtype=dtype, device=device) elif layout == "bhsd": - if is_fp8_dtype: - q, descale_q = generate_bhsd_tensor(BATCH, HQ, N_CTX_Q, D_HEAD, dtype=dtype, device=device, DEBUG_INPUT=DEBUG_INPUT) - k, descale_k = generate_bhsd_tensor(BATCH, HK, N_CTX_K, D_HEAD, dtype=dtype, device=device, DEBUG_INPUT=DEBUG_INPUT) - v, descale_v = generate_bhsd_tensor(BATCH, HK, N_CTX_K, D_HEAD, dtype=dtype, device=device, DEBUG_INPUT=DEBUG_INPUT) + q, descale_q = generate_bhsd_tensor(BATCH, HQ, N_CTX_Q, D_HEAD, dtype=dtype, device=device) + k, descale_k = generate_bhsd_tensor(BATCH, HK, N_CTX_K, D_HEAD, dtype=dtype, device=device) + v, descale_v = generate_bhsd_tensor(BATCH, HK, N_CTX_K, D_HEAD, dtype=dtype, device=device) do, descale_do = generate_bhsd_tensor(BATCH, HQ, N_CTX_Q, D_HEAD, dtype=dtype, device=device) - else: - q = generate_bhsd_tensor(BATCH, HQ, N_CTX_Q, D_HEAD, dtype=dtype, device=device, DEBUG_INPUT=DEBUG_INPUT) - k = generate_bhsd_tensor(BATCH, HK, N_CTX_K, D_HEAD, dtype=dtype, device=device, DEBUG_INPUT=DEBUG_INPUT) - v = generate_bhsd_tensor(BATCH, HK, N_CTX_K, D_HEAD, dtype=dtype, device=device, DEBUG_INPUT=DEBUG_INPUT) - do = torch.ones_like(q) if DEBUG_INPUT else torch.randn_like(q) + else: + q = generate_bhsd_tensor(BATCH, HQ, N_CTX_Q, D_HEAD, dtype=dtype, device=device) + k = generate_bhsd_tensor(BATCH, HK, N_CTX_K, D_HEAD, dtype=dtype, device=device) + v = generate_bhsd_tensor(BATCH, HK, N_CTX_K, D_HEAD, dtype=dtype, device=device) + do = generate_bhsd_tensor(BATCH, HQ, N_CTX_Q, D_HEAD, dtype=dtype, device=device) elif packing == "kv": # gen tensors with kv packing if is_fp8_dtype: raise ValueError("FP8 not supported for KV packing yet") else: if layout == "bshd": - q = generate_bshd_tensor(BATCH, N_CTX_Q, HQ, D_HEAD, dtype=dtype, device=device, DEBUG_INPUT=DEBUG_INPUT) - kv = generate_bshd_kv_packed(BATCH, N_CTX_K, HK, D_HEAD, dtype=dtype, device=device, DEBUG_INPUT=DEBUG_INPUT) - do = torch.ones_like(q) if DEBUG_INPUT else torch.randn_like(q) + q = generate_bshd_tensor(BATCH, N_CTX_Q, HQ, D_HEAD, dtype=dtype, device=device) + kv = generate_bshd_kv_packed(BATCH, N_CTX_K, HK, D_HEAD, dtype=dtype, device=device) + do = generate_bshd_tensor(BATCH, N_CTX_Q, HQ, D_HEAD, dtype=dtype, device=device) elif layout == "bhsd": - q = generate_bhsd_tensor(BATCH, HQ, N_CTX_Q, D_HEAD, dtype=dtype, device=device, DEBUG_INPUT=DEBUG_INPUT) - kv = generate_bhsd_kv_packed(BATCH, HK, N_CTX_K, D_HEAD, dtype=dtype, device=device, DEBUG_INPUT=DEBUG_INPUT) - do = torch.ones_like(q) if DEBUG_INPUT else torch.randn_like(q) + q = generate_bhsd_tensor(BATCH, HQ, N_CTX_Q, D_HEAD, dtype=dtype, device=device) + kv = generate_bhsd_kv_packed(BATCH, HK, N_CTX_K, D_HEAD, dtype=dtype, device=device) + do = generate_bhsd_tensor(BATCH, HQ, N_CTX_Q, D_HEAD, dtype=dtype, device=device) elif packing == "qkv": # qkv packing - requires same sequence length for q and k assert N_CTX_Q == N_CTX_K, "For QKV packing, Q and K must have same sequence length" @@ -618,17 +646,14 @@ def input_helper( raise ValueError("FP8 not supported for QKV packing yet") else: if layout == "bshd": - qkv = generate_bshd_qkv_packed(BATCH, N_CTX_Q, HQ, D_HEAD, dtype=dtype, device=device, DEBUG_INPUT=DEBUG_INPUT) - do = torch.ones((BATCH, N_CTX_Q, HQ, D_HEAD), dtype=dtype, device=device) if DEBUG_INPUT else torch.randn((BATCH, N_CTX_Q, HQ, D_HEAD), dtype=dtype, device=device) + qkv = generate_bshd_qkv_packed(BATCH, N_CTX_Q, HQ, D_HEAD, dtype=dtype, device=device) + do = generate_bshd_tensor(BATCH, N_CTX_Q, HQ, D_HEAD, dtype=dtype, device=device) elif layout == "bhsd": - qkv = generate_bhsd_qkv_packed(BATCH, HQ, N_CTX_Q, D_HEAD, dtype=dtype, device=device, DEBUG_INPUT=DEBUG_INPUT) - do = torch.ones((BATCH, HQ, N_CTX_Q, D_HEAD), dtype=dtype, device=device) if DEBUG_INPUT else torch.randn((BATCH, HQ, N_CTX_Q, D_HEAD), dtype=dtype, device=device) + qkv = generate_bhsd_qkv_packed(BATCH, HQ, N_CTX_Q, D_HEAD, dtype=dtype, device=device) + do = generate_bhsd_tensor(BATCH, HQ, N_CTX_Q, D_HEAD, dtype=dtype, device=device) # setup metadata - if DEBUG_INPUT: - sm_scale = 1 - else: - sm_scale = D_HEAD**-0.5 + sm_scale = D_HEAD**-0.5 metadata = MetaData(sm_scale=sm_scale) metadata.max_seqlens_q = N_CTX_Q metadata.max_seqlens_k = N_CTX_K @@ -957,6 +982,29 @@ def compute_alibi_tensor_ref(alibi_slopes, seqlen_q, seqlen_k): def round_multiple(x, m): return (x + m - 1) // m * m +def save_tensor_to_csv(tensor, filename, decimal_places=2): + """ + save a 2d tensor to csv file + + args: + tensor: torch tensor of shape [rows, cols] + filename: output csv filename + decimal_places: number of decimal places (default: 2) + """ + # ensure tensor is 2d + if tensor.ndim != 2: + raise ValueError(f"tensor must be 2d, got shape {tensor.shape}") + + # ensure filename ends with .csv + if not filename.endswith('.csv'): + filename = filename + '.csv' + + # save to csv using numpy + np.savetxt(filename, + tensor.detach().cpu().numpy(), + delimiter=',', + fmt=f'%.{decimal_places}f') + # ------------------------------- # Dropouts # ------------------------------- @@ -1016,6 +1064,91 @@ def write_dropout_mask(x, tensor_name = "tensor"): else: writer.writerows(dropout_mask) +# ------------------------------- +# Autotune +# ------------------------------- +def get_fwd_prefill_cdna_autotune_configs(): + return [ + triton.Config({'BLOCK_M': 128, 'BLOCK_N': 128, 'waves_per_eu': 2, 'PRE_LOAD_V': False}, num_stages=1, + num_warps=4), + triton.Config({'BLOCK_M': 128, 'BLOCK_N': 64, 'waves_per_eu': 2, 'PRE_LOAD_V': False}, num_stages=1, + num_warps=4), + triton.Config({'BLOCK_M': 128, 'BLOCK_N': 64, 'waves_per_eu': 3, 'PRE_LOAD_V': False}, num_stages=1, + num_warps=4), + triton.Config({'BLOCK_M': 128, 'BLOCK_N': 64, 'waves_per_eu': 1, 'PRE_LOAD_V': False}, num_stages=1, + num_warps=4), + triton.Config({'BLOCK_M': 128, 'BLOCK_N': 32, 'waves_per_eu': 2, 'PRE_LOAD_V': False}, num_stages=1, + num_warps=4), + triton.Config({'BLOCK_M': 64, 'BLOCK_N': 64, 'waves_per_eu': 1, 'PRE_LOAD_V': False}, num_stages=1, + num_warps=4), + # Fall-back config. + triton.Config({'BLOCK_M': 16, 'BLOCK_N': 16, 'waves_per_eu': 1, 'PRE_LOAD_V': False}, num_stages=1, + num_warps=4), + ], ['IS_CAUSAL', 'dropout_p', 'MAX_SEQLENS_Q', 'MAX_SEQLENS_K', 'ACTUAL_BLOCK_DMODEL', 'IS_VARLEN', 'HQ', 'HK'] + + +def get_fwd_prefill_rdna_autotune_configs(): + return [ + triton.Config({'BLOCK_M': 32, 'BLOCK_N': 32, 'waves_per_eu': 4, 'PRE_LOAD_V': False}, num_stages=1, + num_warps=2), + triton.Config({'BLOCK_M': 32, 'BLOCK_N': 32, 'waves_per_eu': 2, 'PRE_LOAD_V': False}, num_stages=1, + num_warps=2), + triton.Config({'BLOCK_M': 32, 'BLOCK_N': 16, 'waves_per_eu': 4, 'PRE_LOAD_V': False}, num_stages=1, + num_warps=2), + triton.Config({'BLOCK_M': 32, 'BLOCK_N': 16, 'waves_per_eu': 2, 'PRE_LOAD_V': False}, num_stages=1, + num_warps=2), + triton.Config({'BLOCK_M': 16, 'BLOCK_N': 16, 'waves_per_eu': 4, 'PRE_LOAD_V': False}, num_stages=1, + num_warps=2), + triton.Config({'BLOCK_M': 16, 'BLOCK_N': 16, 'waves_per_eu': 2, 'PRE_LOAD_V': False}, num_stages=1, + num_warps=2), + # Fall-back config. + triton.Config({'BLOCK_M': 16, 'BLOCK_N': 16, 'waves_per_eu': 1, 'PRE_LOAD_V': False}, num_stages=1, + num_warps=2), + ], ['IS_CAUSAL', 'dropout_p', 'MAX_SEQLENS_Q', 'MAX_SEQLENS_K', 'ACTUAL_BLOCK_DMODEL', 'IS_VARLEN', 'HQ', 'HK'] + + +def get_fwd_prefill_autotune_configs(): + if AUTOTUNE: + if is_rdna(): + return get_fwd_prefill_rdna_autotune_configs() + elif is_cdna(): + return get_fwd_prefill_cdna_autotune_configs() + else: + raise ValueError("Unknown Device Type") + else: + arch = get_arch() + if arch == "gfx950": + default_config = triton.Config( + {"BLOCK_M": 128, "BLOCK_N": 128, "waves_per_eu": 2, "PRE_LOAD_V": False}, + num_stages=1, + num_warps=4, + ) + elif arch == "gfx942" and False: # Disabled due shared mem oom in CI when using triton==3.3.0 when using top of tree everything seems fine. + default_config = triton.Config( + {"BLOCK_M": 128, "BLOCK_N": 64, "waves_per_eu": 2, "PRE_LOAD_V": False}, + num_stages=1, + num_warps=4, + ) + else: + default_config = triton.Config( + {"BLOCK_M": 64, "BLOCK_N": 64, "waves_per_eu": 2, "PRE_LOAD_V": False}, + num_stages=1, + num_warps=4, + ) + + return [ + default_config + ], [ + "IS_CAUSAL", + "dropout_p", + "MAX_SEQLENS_Q", + "MAX_SEQLENS_K", + "ACTUAL_BLOCK_DMODEL", + "IS_VARLEN", + "HQ", + "HK", + ] + # ------------------------------- # Runtime info # ------------------------------- diff --git a/tests/test_flash_attn_triton_amd.py b/tests/test_flash_attn_triton_amd.py index 6073cb1c35a..adebe088a79 100755 --- a/tests/test_flash_attn_triton_amd.py +++ b/tests/test_flash_attn_triton_amd.py @@ -16,7 +16,7 @@ from flash_attn.bert_padding import pad_input, unpad_input from flash_attn.flash_attn_interface import _get_block_size_n from flash_attn.layers.rotary import apply_rotary_emb -from flash_attn.flash_attn_triton_amd.utils import USE_TRITON_ROCM, is_hip, is_rdna +from flash_attn.flash_attn_triton_amd.utils import USE_TRITON_ROCM, is_hip, is_rdna, generate_bshd_tensor, save_tensor_to_csv MAX_HEADDIM_SM8x = 192 @@ -573,7 +573,7 @@ def get_dropout_fraction( # @pytest.mark.parametrize("deterministic", [False]) @pytest.mark.parametrize("alibi", [False, True]) # @pytest.mark.parametrize("alibi", [False]) -@pytest.mark.parametrize("local", [False]) +@pytest.mark.parametrize("local", [False, True]) # @pytest.mark.parametrize("local", [False]) @pytest.mark.parametrize("causal", [False, True]) # @pytest.mark.parametrize("causal", [False]) @@ -712,6 +712,10 @@ def test_flash_attn_qkvpacked(seqlen, d, dropout_p, causal, local, alibi, determ if not alibi: assert abs(dropout_fraction - dropout_p) <= (0.01 if not local else 0.025) + if local == True: + print("Sliding Window not supported in backward yet") + return + if (d <= MAX_HEADDIM_SM8x or dropout_p == 0) or (is_sm80 or is_sm90): assert (dqkv - dqkv_ref).abs().max().item() <= 2 * (dqkv_pt - dqkv_ref).abs().max().item() @@ -722,7 +726,7 @@ def test_flash_attn_qkvpacked(seqlen, d, dropout_p, causal, local, alibi, determ # @pytest.mark.parametrize("deterministic", [True]) @pytest.mark.parametrize("alibi", [False, True]) # @pytest.mark.parametrize("alibi", [True]) -@pytest.mark.parametrize("local", [False]) +@pytest.mark.parametrize("local", [False, True]) # @pytest.mark.parametrize("local", [True]) @pytest.mark.parametrize("causal", [False, True]) # @pytest.mark.parametrize('causal', [False]) @@ -860,6 +864,10 @@ def test_flash_attn_varlen_qkvpacked( if not alibi: assert abs(dropout_fraction - dropout_p) <= (0.01 if not local else 0.025) + if local == True: + print("Sliding Window not supported in backward yet") + return + if (d <= MAX_HEADDIM_SM8x or dropout_p == 0) or (is_sm80 or is_sm90): assert (dqkv - dqkv_ref).abs().max().item() <= 2 * (dqkv_pt - dqkv_ref).abs().max().item() @@ -874,7 +882,7 @@ def test_flash_attn_varlen_qkvpacked( # @pytest.mark.parametrize("deterministic", [True]) @pytest.mark.parametrize("alibi", [False, True]) # @pytest.mark.parametrize("alibi", [False]) -@pytest.mark.parametrize("local", [False]) +@pytest.mark.parametrize("local", [False, True]) # @pytest.mark.parametrize("local", [False]) @pytest.mark.parametrize("causal", [False, True]) # @pytest.mark.parametrize("causal", [True]) @@ -1133,6 +1141,10 @@ def test_flash_attn_output( if not alibi: assert abs(dropout_fraction - dropout_p) <= (0.01 if not local else 0.025) + if local == True: + print("Sliding Window not supported in backward yet") + return + if (d <= MAX_HEADDIM_SM8x or dropout_p == 0) or (is_sm80 or is_sm90): assert (dq - dq_ref).abs().max().item() <= 3 * (dq_pt - dq_ref).abs().max().item() assert (dk - dk_ref).abs().max().item() <= 3 * (dk_pt - dk_ref).abs().max().item() @@ -1149,7 +1161,7 @@ def test_flash_attn_output( # @pytest.mark.parametrize("deterministic", [True]) @pytest.mark.parametrize("alibi", [False, True]) # @pytest.mark.parametrize("alibi", [True]) -@pytest.mark.parametrize("local", [False]) +@pytest.mark.parametrize("local", [False, True]) # @pytest.mark.parametrize("local", [True]) @pytest.mark.parametrize("causal", [False, True]) # @pytest.mark.parametrize('causal', [True]) @@ -1455,6 +1467,10 @@ def test_flash_attn_varlen_output( if not alibi: assert abs(dropout_fraction - dropout_p) <= (0.01 if not local else 0.04) + if local == True: + print("Sliding Window not supported in backward yet") + return + if (d <= MAX_HEADDIM_SM8x or dropout_p == 0) or (is_sm80 or is_sm90): assert (dq - dq_ref).abs().max().item() <= 3 * (dq_pt - dq_ref).abs().max().item() assert (dk - dk_ref).abs().max().item() <= 3 * (dk_pt - dk_ref).abs().max().item() @@ -1463,7 +1479,7 @@ def test_flash_attn_varlen_output( @pytest.mark.parametrize("dtype", ([torch.float16] if skip_bfloat16 else [torch.float16, torch.bfloat16])) # @pytest.mark.parametrize("dtype", [torch.bfloat16]) -@pytest.mark.parametrize("local", [False]) +@pytest.mark.parametrize("local", [False, True]) # @pytest.mark.parametrize("local", [True]) @pytest.mark.parametrize("d", [32, 40, 59, 64, 80, 96, 111, 128, 160, 192, 224, 256]) # @pytest.mark.parametrize("d", [32, 64, 96, 128, 160, 192, 224, 256]) @@ -1569,6 +1585,10 @@ def test_flash_attn_causal(seqlen_q, seqlen_k, swap_sq_sk, d, local, dtype): # of a Pytorch implementation. assert (out - out_ref).abs().max().item() <= 2 * (out_pt - out_ref).abs().max().item() + 1e-5 + if local == True: + print("Sliding Window not supported in backward yet") + return + assert (dq - dq_ref).abs().max().item() <= 2 * (dq_pt - dq_ref).abs().max().item() + 1e-5 assert (dk - dk_ref).abs().max().item() <= 2 * (dk_pt - dk_ref).abs().max().item() + 1e-5 assert (dv - dv_ref).abs().max().item() <= 2 * (dv_pt - dv_ref).abs().max().item() + 1e-5 @@ -1576,7 +1596,7 @@ def test_flash_attn_causal(seqlen_q, seqlen_k, swap_sq_sk, d, local, dtype): @pytest.mark.parametrize("dtype", ([torch.float16] if skip_bfloat16 else [torch.float16, torch.bfloat16])) # @pytest.mark.parametrize("dtype", [torch.bfloat16]) -@pytest.mark.parametrize("local", [False]) +@pytest.mark.parametrize("local", [False, True]) # @pytest.mark.parametrize("local", [True]) @pytest.mark.parametrize("d", [32, 40, 59, 64, 80, 96, 111, 128, 160, 192, 224, 256]) # @pytest.mark.parametrize("d", [32, 64, 96, 128, 160, 192, 224, 256]) @@ -1737,6 +1757,10 @@ def test_flash_attn_varlen_causal( # of a Pytorch implementation. assert (out - out_ref).abs().max().item() <= 2 * (out_pt - out_ref).abs().max().item() + 1e-5 + if local == True: + print("Sliding Window not supported in backward yet") + return + if test_backward: assert (dq - dq_ref).abs().max().item() <= 2 * (dq_pt - dq_ref).abs().max().item() + 1e-5 assert (dk - dk_ref).abs().max().item() <= 2 * (dk_pt - dk_ref).abs().max().item() + 1e-5 @@ -1749,7 +1773,7 @@ def test_flash_attn_varlen_causal( # @pytest.mark.parametrize("deterministic", [True]) @pytest.mark.parametrize("alibi", [False, True]) # @pytest.mark.parametrize("alibi", [True]) -@pytest.mark.parametrize("local", [False]) +@pytest.mark.parametrize("local", [False, True]) # @pytest.mark.parametrize("local", [False]) @pytest.mark.parametrize("causal", [False, True]) # @pytest.mark.parametrize("causal", [True]) @@ -1867,6 +1891,10 @@ def test_flash_attn_splitkv( # of a Pytorch implementation. assert (out - out_ref).abs().max().item() <= 2 * (out_pt - out_ref).abs().max().item() + 1e-5 + if local == True: + print("Sliding Window not supported in backward yet") + return + mult = 2 if not alibi else 8 assert (dq - dq_ref).abs().max().item() <= mult * (dq_pt - dq_ref).abs().max().item() + 2e-4 assert (dk - dk_ref).abs().max().item() <= mult * (dk_pt - dk_ref).abs().max().item() + 2e-4 @@ -1883,7 +1911,7 @@ def test_flash_attn_splitkv( # @pytest.mark.parametrize("new_kv", [False]) @pytest.mark.parametrize("alibi", [False, True]) # @pytest.mark.parametrize("alibi", [False]) -@pytest.mark.parametrize("local", [False]) +@pytest.mark.parametrize("local", [False, True]) # @pytest.mark.parametrize("local", [False]) @pytest.mark.parametrize("causal", [False, True]) # @pytest.mark.parametrize("causal", [False]) @@ -2404,7 +2432,7 @@ def test_flash_attn_bwd_varlen_overflow(d, causal, dtype): @pytest.mark.parametrize("dtype", ([torch.float16] if skip_bfloat16 else [torch.float16, torch.bfloat16])) # @pytest.mark.parametrize("dtype", [torch.bfloat16]) -@pytest.mark.parametrize("local", [False]) +@pytest.mark.parametrize("local", [False, True]) # @pytest.mark.parametrize("local", [True]) @pytest.mark.parametrize("causal", [False, True]) # @pytest.mark.parametrize("causal", [True]) @@ -2452,6 +2480,10 @@ def test_flash_attn_deterministic(seqlen_q, seqlen_k, swap_sq_sk, d, causal, loc v = torch.randn(batch_size, seqlen_k, nheads, d, device=device, dtype=dtype, requires_grad=True) out = flash_attn_func(q, k, v, 0.0, causal=causal, window_size=window_size, deterministic=True) + if local == True: + print("Sliding Window not supported in backward yet") + return + g = torch.randn_like(out) dq0, dk0, dv0 = torch.autograd.grad(out, (q, k, v), g, retain_graph=True) for _ in range(50): @@ -2463,7 +2495,7 @@ def test_flash_attn_deterministic(seqlen_q, seqlen_k, swap_sq_sk, d, causal, loc @pytest.mark.parametrize("dtype", ([torch.float16] if skip_bfloat16 else [torch.float16, torch.bfloat16])) # @pytest.mark.parametrize("dtype", [torch.bfloat16]) -@pytest.mark.parametrize("local", [False]) +@pytest.mark.parametrize("local", [False, True]) # @pytest.mark.parametrize("local", [True]) @pytest.mark.parametrize("causal", [False, True]) # @pytest.mark.parametrize("causal", [True]) @@ -2540,6 +2572,10 @@ def test_flash_attn_varlen_deterministic(seqlen_q, seqlen_k, swap_sq_sk, d, caus deterministic=True, ) + if local == True: + print("Sliding Window not supported in backward yet") + return + g = torch.randn_like(out) dq0, dk0, dv0 = torch.autograd.grad(out, (q_unpad, k_unpad, v_unpad), g, retain_graph=True) for _ in range(50):