From 2bb034b4660520c9954c93c2afe257b38e8c1f82 Mon Sep 17 00:00:00 2001 From: Michael Date: Wed, 4 Jun 2025 10:43:51 -0500 Subject: [PATCH 01/10] Compress SWA work test case set up debug inputs add fwd ref one mask ref fwd first pass save ref doesnot work for bigger seqlens save new version some causal cases failing found bad cases working new attn new atten works new attn_fwd works reorg n_extra_tokens use seqlen_delta_qk ref fwd works add sliding window to bwd ref test kvcache decode ref work with everything except sliding window add debug code for 12 failing sliding window cases for decode attention_decode_forward_ref_impl mostly works except for alibi fix alibi in attention_decode_forward_ref_impl ref works with normal, varlen & kvcache move stuff around figure out masking old attn inner two inner functions remove load_fn do Lk - Lq like ref unify IS_CAUSAL code in epilogue clean up add args rm inference stuff simplify compute_masking simpler compute mask stub out returning front masking variables remove pointer pass compute ptrs inloop compute block min and max window stub inside inner mask loop trying to use attn_fwd_mask causes issues fix compiler bug when front masking gen specifc types add sliding window and debug statements use identity for v add more taste cases add comments save use k_max_token for clarity disable debug configs basic NON-CAUSAL SLIDING WINDOW non causal sliding window works on the all the shapes non sliding window working in fwd clean up fused bwd seperate old fwd_prefill move configs to utils.py --- flash_attn/flash_attn_triton_amd/bwd_ref.py | 111 +- .../flash_attn_triton_amd/fwd_decode.py | 13 +- .../flash_attn_triton_amd/fwd_prefill.py | 950 ++++++++++++------ .../flash_attn_triton_amd/fwd_prefill_old.py | 473 +++++++++ flash_attn/flash_attn_triton_amd/fwd_ref.py | 328 +++++- .../flash_attn_triton_amd/interface_fa.py | 76 +- flash_attn/flash_attn_triton_amd/utils.py | 268 +++-- tests/test_flash_attn_triton_amd.py | 148 ++- 8 files changed, 1848 insertions(+), 519 deletions(-) create mode 100644 flash_attn/flash_attn_triton_amd/fwd_prefill_old.py diff --git a/flash_attn/flash_attn_triton_amd/bwd_ref.py b/flash_attn/flash_attn_triton_amd/bwd_ref.py index 8bdccb1d329..9569551d87a 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 (causal and/or sliding window) + 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, @@ -300,6 +355,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 +423,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 +480,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 +502,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 +525,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 +541,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..fe969027536 100644 --- a/flash_attn/flash_attn_triton_amd/fwd_decode.py +++ b/flash_attn/flash_attn_triton_amd/fwd_decode.py @@ -571,10 +571,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,7 +588,7 @@ 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 SPLIT_K = None NUM_QUANT_GROUPS = 1 @@ -602,11 +604,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 diff --git a/flash_attn/flash_attn_triton_amd/fwd_prefill.py b/flash_attn/flash_attn_triton_amd/fwd_prefill.py index e33982bb6a7..49f15d521b9 100644 --- a/flash_attn/flash_attn_triton_amd/fwd_prefill.py +++ b/flash_attn/flash_attn_triton_amd/fwd_prefill.py @@ -1,79 +1,206 @@ +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 + + +from flash_attn.flash_attn_triton_amd.fwd_prefill_old import attn_fwd_old DEBUG = False + # NOTE: triton fails to import tl.constexprs so create them here for the file 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 +211,87 @@ 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: + pass + 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 +302,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 +347,145 @@ 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: impl sliding window + if IS_CAUSAL: + return 0, 0, 0, 0, 0 else: - raise ValueError("Unknown Device Type") + # ========== NON-CAUSAL SLIDING WINDOW ========== + # For now, we'll use a simple approach: + # Process all blocks with masking since we have a sliding window + # This is less efficient than skipping blocks, but ensures correctness + + # TODO: Optimize by computing which blocks can be fully skipped + # For now, process all blocks with the mask function + 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 +497,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 +507,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 +539,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 +624,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) - 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) + # 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) + + # 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 +841,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 +849,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 +899,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: @@ -693,7 +980,9 @@ def attention_prefill_forward_triton_impl( else: stride_bz, stride_bh, stride_bm, stride_bn = (0, 0, 0, 0) - attn_fwd[grid](q, k, v, bias, cache_seqlens, cache_batch_idx, + USE_OLD = os.environ.get('USE_OLD', '0').lower() in ('1', 'true', 'yes') + if USE_OLD: + attn_fwd_old[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, stride_qb, stride_qh, stride_qm, stride_qd, @@ -711,5 +1000,16 @@ def attention_prefill_forward_triton_impl( 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) + else: + 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, *q_strides, *k_strides, *v_strides, *o_strides, + *bias_strides, stride_az, stride_ah, *scores_strides, 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, 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) return softmax_lse, sd_mask if return_softmax else None diff --git a/flash_attn/flash_attn_triton_amd/fwd_prefill_old.py b/flash_attn/flash_attn_triton_amd/fwd_prefill_old.py new file mode 100644 index 00000000000..cab91a604e1 --- /dev/null +++ b/flash_attn/flash_attn_triton_amd/fwd_prefill_old.py @@ -0,0 +1,473 @@ +import os +import torch +import triton +import triton.language as tl +from typing import Literal, Optional, Union +from .utils import DEBUG, DROPOUT_USE_PYTORCH, DROPOUT_DUMP, AUTOTUNE, compute_alibi_block, compute_fp8_scaling_factors, get_arch, get_shapes_from_layout, get_strides_from_layout, is_cdna, is_fp8, is_rdna, create_dropout_mask + +# NOTE: triton fails to import tl.constexprs so create them here for the file +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. +@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 + + + +@triton.jit +def _attn_fwd_inner_old(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, + 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, + 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): + # 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) + 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) + 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 + + # -- 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 + 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) + 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 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) + 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] + + # 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: + if tl_DROPOUT_USE_PYTORCH: + dropout_mask = tl.load(dropout_mask_ptrs, mask=p_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) + + # 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) + + # 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) + + # -- 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 + 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) + # -- 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: + 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) + + 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 + +# @triton.autotune( +# configs=autotune_configs, +# key=autotune_keys, +# use_cuda_graph=True, +# ) +@triton.jit +def attn_fwd_old(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: tl.constexpr, LSE, Out, stride_qz, stride_qh, stride_qm, stride_qk, + stride_kz, stride_kh, stride_kn, stride_kk, stride_vz, stride_vh, stride_vk, stride_vn, + stride_oz, stride_oh, stride_om, stride_on, 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, 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, + 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): + # set params + ACCUMULATOR_TYPE = tl.float32 + + # 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) + + offs_m = start_m * BLOCK_M + tl.arange(0, BLOCK_M) + offs_n = tl.arange(0, BLOCK_N) + offs_d = tl.arange(0, BLOCK_DMODEL) + + # handle seqlen + if IS_VARLEN: + cu_seqlens_q_start = tl.load(cu_seqlens_q + off_z) + cu_seqlens_q_end = tl.load(cu_seqlens_q + off_z + 1) + seqlen_q = cu_seqlens_q_end - cu_seqlens_q_start + + # we have a one-size-fits-all grid in id(0). Some seqlens might be too small for all start_m so for those we return early. + if start_m * BLOCK_M > seqlen_q: + return + 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 + 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) + + # 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 + k_offset = K + off_z * stride_kz + off_h_k * stride_kh + cu_seqlens_k_start * stride_kn + k_ptrs = k_offset + offs_d[:, None] * stride_kk + offs_n[None, :] * stride_kn + v_offset = V + off_z * stride_vz + off_h_k * stride_vh + cu_seqlens_k_start * stride_vk + v_ptrs = v_offset + offs_n[:, None] * stride_vk + offs_d[None, :] * stride_vn + if USE_BIAS: + # Note: this might get large enough to overflow on some configs + bias_offset = off_h_q * stride_bh + bias_ptrs = bias + bias_offset + offs_m[:, None] * stride_bm + offs_n[None, :] * stride_bn + else: + bias_ptrs = None + + if USE_ALIBI: + a_offset = off_z * stride_az + off_h_q * stride_ah + alibi_slope = tl.load(alibi_slopes + a_offset) + else: + alibi_slope = None + + if RETURN_SCORES: + sd_mask_offset = sd_mask + off_z * stride_sz + off_h_q * stride_sh #+ cu_seqlens_q_start * stride_sm + sd_mask_ptrs = sd_mask_offset + offs_m[:, None] * stride_sm + offs_n[None, :] * stride_sn + else: + sd_mask_ptrs = None + + if ENABLE_DROPOUT: + dropout_mask_offset = dropout_mask + off_z * stride_sz + off_h_q * stride_sh #+ cu_seqlens_q_start * stride_sm + dropout_mask_ptrs = dropout_mask_offset + offs_m[:, None] * stride_sm + offs_n[None, :] * stride_sn + batch_philox_offset = philox_offset_base + off_z * stride_sz + off_h_q * stride_sh #+ cu_seqlens_q_start * stride_sm + philox_ptrs = batch_philox_offset + offs_m[:, None] * stride_sm + offs_n[None, :] * stride_sn + 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. + if n_full_blocks > 0: + block_max = (n_blocks - masked_blocks) * BLOCK_N + acc, l_i, m_i = _attn_fwd_inner_old(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_old(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 + # This helps the compiler do Newton Raphson on l_i vs on acc which is much larger. + 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 + 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) + # convert back to natural units + softmax_lse *= LN2 + else: + softmax_lse = m_i + tl.math.log(l_i) + + 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) + + # 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 + 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 + else: + tl.store(l_ptrs, softmax_lse) # the log of the normalization constant + + # write back O + 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_ptrs_mask = tl.full([BLOCK_M, BLOCK_DMODEL], 1, dtype=tl.int1) + if overflow_size > 0: + o_ptrs_mask = o_ptrs_mask & (offs_m[:, None] < seqlen_q) + if PADDED_HEAD: + o_ptrs_mask = o_ptrs_mask & (offs_d[None, :] < ACTUAL_BLOCK_DMODEL) + + if FP8_OUTPUT: + scale_acc, descale_acc = compute_fp8_scaling_factors(acc, FP8_MAX) + tl.store(Descale_O + off_z * stride_descale_o_z + off_h_q, descale_acc) + tl.store(o_ptrs, (acc * scale_acc).to(Out.type.element_ty), mask=o_ptrs_mask) + else: + tl.store(o_ptrs, acc.to(Out.dtype.element_ty), mask=o_ptrs_mask) 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..4e49d3a73d0 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,11 @@ 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 k_new = k v_new = v @@ -829,7 +844,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 +875,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, + window_size_left, + 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 +907,14 @@ def fwd_kvcache( out, metadata.sm_scale, metadata.causal, + window_size_left, + 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/utils.py b/flash_attn/flash_attn_triton_amd/utils.py index 1795e0d1366..a00e0fb856d 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 @@ -165,8 +166,8 @@ def generate_varlen_tensor( batch_size: Optional[int] = None, equal_seqlens: bool = False, device: str = "cuda", - dtype: torch.dtype = torch.float16, - DEBUG_INPUT: bool = False + dtype: torch.dtype = torch.float32, + mode: Literal["random", "ones", "incremental", "identity"] = "random" ): if DEBUG: print("total_seqlen", total_seqlen) @@ -200,8 +201,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 +214,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 +241,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, 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 +273,35 @@ 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, 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") + # 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 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,14 @@ 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_qkv_packed(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 +606,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 +647,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_qkv_packed(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_qkv_packed(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 +983,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 +1065,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..bb5d252fd0d 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, generate_bshd_tensor, is_hip, is_rdna, 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]) @@ -722,7 +722,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]) @@ -874,9 +874,9 @@ 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", [True]) # @pytest.mark.parametrize("local", [False]) -@pytest.mark.parametrize("causal", [False, True]) +@pytest.mark.parametrize("causal", [False]) # @pytest.mark.parametrize("causal", [True]) @pytest.mark.parametrize("d", [32, 40, 59, 64, 96, 111, 128, 160, 192, 224, 256]) # @pytest.mark.parametrize("d", [32, 64, 96, 128, 160, 192, 224, 256]) @@ -887,6 +887,13 @@ def test_flash_attn_varlen_qkvpacked( @pytest.mark.parametrize( "seqlen_q,seqlen_k", [ + # debug + # (32, 32), + # (32, 64), + # (64, 32), + # (64, 64), + # (128, 128), + # og (113, 203), (128, 217), (113, 211), @@ -906,6 +913,14 @@ def test_flash_attn_varlen_qkvpacked( def test_flash_attn_output( seqlen_q, seqlen_k, d, dropout_p, causal, local, alibi, deterministic, mha_type, dtype, kvpacked, softcap ): + DEBUG = False + BATCH_AND_HEAD_ONE = False + SAVE_TO_CSV = False + TEST_BACKWARD = False + if DEBUG: + print() + print("Debugging") + if USE_TRITON_ROCM: if causal: if seqlen_q ==1024 and seqlen_k==1024 and d==160: @@ -920,12 +935,23 @@ def test_flash_attn_output( device = "cuda" # set seed torch.random.manual_seed(0) - batch_size = 4 - nheads = 6 if softcap == 0.0 else 4 # softcap reference impl takes more memory - nheads_k = nheads if mha_type == "mha" else (1 if mha_type == "mqa" else 2) + if BATCH_AND_HEAD_ONE: + batch_size = 1 + nheads = 1 + nheads_k = 1 + else: + batch_size = 4 + nheads = 6 if softcap == 0.0 else 4 # softcap reference impl takes more memory + nheads_k = nheads if mha_type == "mha" else (1 if mha_type == "mqa" else 2) assert nheads % nheads_k == 0 - window_size = (-1, -1) if not local else torch.randint(0, seqlen_k, (2,)) - q = torch.randn(batch_size, seqlen_q, nheads, d, device=device, dtype=dtype, requires_grad=True) + if DEBUG: + window_size = (-1, -1) if not local else (seqlen_k//4, (seqlen_k*3)//4) + else: + window_size = (-1, -1) if not local else torch.randint(0, seqlen_k, (2,)) + if DEBUG: + q = generate_bshd_tensor(batch_size, seqlen_q, nheads, d, dtype=dtype, device=device, mode="ones") + else: + q = torch.randn(batch_size, seqlen_q, nheads, d, device=device, dtype=dtype, requires_grad=True) if softcap > 0: # Ensure the values of qk are at least within softcap range. q = q * softcap @@ -934,18 +960,28 @@ def test_flash_attn_output( batch_size, seqlen_k, 2, nheads_k, d, device=device, dtype=dtype, requires_grad=True ) else: - k = torch.randn( - batch_size, seqlen_k, nheads_k, d, device=device, dtype=dtype, requires_grad=True - ) - v = torch.randn( - batch_size, seqlen_k, nheads_k, d, device=device, dtype=dtype, requires_grad=True - ) + if DEBUG: + k = generate_bshd_tensor(batch_size, seqlen_k, nheads_k, d, dtype=dtype, device=device, mode="ones") + v = generate_bshd_tensor(batch_size, seqlen_k, nheads_k, d, dtype=dtype, device=device, mode="identity") + else: + k = torch.randn( + batch_size, seqlen_k, nheads_k, d, device=device, dtype=dtype, requires_grad=True + ) + v = torch.randn( + batch_size, seqlen_k, nheads_k, d, device=device, dtype=dtype, requires_grad=True + ) if alibi: alibi_slopes = torch.rand(batch_size, nheads, device=device, dtype=torch.float32) * 0.3 attn_bias = attn_bias_from_alibi_slopes(alibi_slopes, seqlen_q, seqlen_k, causal=causal) else: alibi_slopes, attn_bias = None, None + if DEBUG: + print("q:", q, q.shape) + print("k:", k, k.shape) + print("v:", v, v.shape) + print("window_size:", window_size) + if kvpacked: out, lse, S_dmask = flash_attn_kvpacked_func( q, @@ -1075,7 +1111,10 @@ def test_flash_attn_output( print(f"Attention max diff: {(attn - attn_ref).abs().max().item()}") print(f"Attention Pytorch max diff: {(attn_pt - attn_ref).abs().max().item()}") - g = torch.randn_like(out) + if DEBUG: + g = torch.ones_like(out) + else: + g = torch.randn_like(out) do_o = (g.float() * out.float()).sum(-1) if (d <= MAX_HEADDIM_SM8x or dropout_p == 0) or (is_sm80 or is_sm90): if kvpacked: @@ -1123,9 +1162,24 @@ def test_flash_attn_output( print(f"dK Pytorch mean diff: {(dk_pt - dk_ref).abs().mean().item()}") print(f"dV Pytorch mean diff: {(dv_pt - dv_ref).abs().mean().item()}") + if DEBUG: + print("window_size:", window_size) + print("out:", out, out.shape) + print("out_ref:", out_ref, out_ref.shape) + if SAVE_TO_CSV: + for batch_idx in range(batch_size): + for head_idx in range(nheads): + save_tensor_to_csv(out[batch_idx, :, head_idx, :], + f"out_b{batch_idx}_h{head_idx}.csv") + save_tensor_to_csv(out_ref[batch_idx, :, head_idx, :], + f"out_ref_b{batch_idx}_h{head_idx}.csv") + + # Check that FlashAttention's numerical error is at most twice the numerical error # of a Pytorch implementation. assert (out - out_ref).abs().max().item() <= 2 * (out_pt - out_ref).abs().max().item() + if not TEST_BACKWARD: + return if dropout_p > 0.0: # assert (attn - attn_ref).abs().max().item() <= 2 * (attn_pt - attn_ref).abs().max().item() @@ -1133,6 +1187,14 @@ def test_flash_attn_output( if not alibi: assert abs(dropout_fraction - dropout_p) <= (0.01 if not local else 0.025) + + if DEBUG: + print("dq:", dq) + print("dq_ref:", dq_ref) + print("dk:", dk) + print("dk_ref:", dk_ref) + print("dv:", dv) + print("dv_ref:", dv_ref) 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 +1211,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]) @@ -1463,7 +1525,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]) @@ -1576,7 +1638,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]) @@ -1749,7 +1811,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]) @@ -1883,7 +1945,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]) @@ -1940,6 +2002,7 @@ def test_flash_attn_kvcache( num_splits, dtype, ): + DEBUG = False if seqlen_q > seqlen_k and new_kv: pytest.skip() if not new_kv and rotary_fraction > 0.0: @@ -1951,24 +2014,40 @@ def test_flash_attn_kvcache( device = "cuda" # set seed torch.random.manual_seed(0) - batch_size = 2 - batch_size_cache = batch_size if not has_batch_idx else batch_size * 2 - nheads = 6 + if DEBUG: + batch_size = 1 + batch_size_cache = batch_size if not has_batch_idx else batch_size * 2 + nheads = 1 + else: + batch_size = 2 + batch_size_cache = batch_size if not has_batch_idx else batch_size * 2 + nheads = 6 # rotary_dim must be a multiple of 16, and must be <= d rotary_dim = math.floor(int(rotary_fraction * d) / 16) * 16 nheads_k = nheads if mha_type == "mha" else (1 if mha_type == "mqa" else 3) assert nheads % nheads_k == 0 window_size = (-1, -1) if not local else torch.randint(0, seqlen_k, (2,)) - q = torch.randn(batch_size, seqlen_q, nheads, d, device=device, dtype=dtype) + if DEBUG: + q = torch.randn(batch_size, seqlen_q, nheads, d, device=device, dtype=dtype) + else: + q = torch.randn(batch_size, seqlen_q, nheads, d, device=device, dtype=dtype) seqlen_new = seqlen_q if seqlen_new_eq_seqlen_q else torch.randint(1, seqlen_q + 1, (1,)).item() if new_kv: - k = torch.randn(batch_size, seqlen_new, nheads_k, d, device=device, dtype=dtype) - v = torch.randn(batch_size, seqlen_new, nheads_k, d, device=device, dtype=dtype) + if DEBUG: + k = torch.randn(batch_size, seqlen_new, nheads_k, d, device=device, dtype=dtype) + v = torch.randn(batch_size, seqlen_new, nheads_k, d, device=device, dtype=dtype) + else: + k = torch.randn(batch_size, seqlen_new, nheads_k, d, device=device, dtype=dtype) + v = torch.randn(batch_size, seqlen_new, nheads_k, d, device=device, dtype=dtype) else: k, v = None, None if paged_kv_block_size is None: - k_cache = torch.randn(batch_size_cache, seqlen_k, nheads_k, d, device=device, dtype=dtype) - v_cache = torch.randn(batch_size_cache, seqlen_k, nheads_k, d, device=device, dtype=dtype) + if DEBUG: + k_cache = torch.randn(batch_size_cache, seqlen_k, nheads_k, d, device=device, dtype=dtype) + v_cache = torch.randn(batch_size_cache, seqlen_k, nheads_k, d, device=device, dtype=dtype) + else: + k_cache = torch.randn(batch_size_cache, seqlen_k, nheads_k, d, device=device, dtype=dtype) + v_cache = torch.randn(batch_size_cache, seqlen_k, nheads_k, d, device=device, dtype=dtype) block_table = None else: ( @@ -2155,6 +2234,11 @@ def test_flash_attn_kvcache( assert torch.allclose(k_cache_select, k_cache_ref, rtol=1e-3, atol=1e-3) assert torch.equal(v_cache_select, v_cache_ref) mult = 3 if not alibi else 5 + if DEBUG: + print("out:", out, out.shape) + print("out_ref:", out_ref, out_ref.shape) + # save_tensor_to_csv(out, "out.csv") + # save_tensor_to_csv(out_ref, "out_ref.csv") assert (out - out_ref).abs().max().item() <= mult * (out_pt - out_ref).abs().max().item() + 1e-5 @@ -2404,7 +2488,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]) @@ -2463,7 +2547,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]) From 6d185c49c263da63b38a4466640aeed19f1adabb Mon Sep 17 00:00:00 2001 From: Michael Date: Wed, 9 Jul 2025 10:04:17 -0500 Subject: [PATCH 02/10] fix bwd ref bug --- flash_attn/flash_attn_triton_amd/bwd_ref.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/flash_attn/flash_attn_triton_amd/bwd_ref.py b/flash_attn/flash_attn_triton_amd/bwd_ref.py index 9569551d87a..1f384cab068 100644 --- a/flash_attn/flash_attn_triton_amd/bwd_ref.py +++ b/flash_attn/flash_attn_triton_amd/bwd_ref.py @@ -325,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) @@ -333,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 From f5e8dbdb88b8c2836fbf86167821ad44bb1e5fb7 Mon Sep 17 00:00:00 2001 From: Michael Date: Wed, 9 Jul 2025 11:48:15 -0500 Subject: [PATCH 03/10] skip local cases so that fa output --- .../flash_attn_triton_amd/fwd_prefill.py | 25 +++++-- tests/test_flash_attn_triton_amd.py | 69 +++++++------------ 2 files changed, 42 insertions(+), 52 deletions(-) diff --git a/flash_attn/flash_attn_triton_amd/fwd_prefill.py b/flash_attn/flash_attn_triton_amd/fwd_prefill.py index 49f15d521b9..2bbcc6ff909 100644 --- a/flash_attn/flash_attn_triton_amd/fwd_prefill.py +++ b/flash_attn/flash_attn_triton_amd/fwd_prefill.py @@ -916,6 +916,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) @@ -998,18 +999,28 @@ def attention_prefill_forward_triton_impl( 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, 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) else: 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, *q_strides, *k_strides, *v_strides, *o_strides, - *bias_strides, stride_az, stride_ah, *scores_strides, stride_lse_z, stride_lse_h, stride_lse_m, cu_seqlens_q, cu_seqlens_k, + 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, + stride_ob, stride_oh, stride_om, stride_od, + 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, + 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, USE_SLIDING_WINDOW=use_sliding_window, WINDOW_SIZE_LEFT=window_size_left, WINDOW_SIZE_RIGHT=window_size_right, IS_VARLEN=is_varlen, + 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 diff --git a/tests/test_flash_attn_triton_amd.py b/tests/test_flash_attn_triton_amd.py index bb5d252fd0d..d95d45df749 100755 --- a/tests/test_flash_attn_triton_amd.py +++ b/tests/test_flash_attn_triton_amd.py @@ -874,9 +874,9 @@ 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", [True]) +@pytest.mark.parametrize("local", [False, True]) # @pytest.mark.parametrize("local", [False]) -@pytest.mark.parametrize("causal", [False]) +@pytest.mark.parametrize("causal", [False, True]) # @pytest.mark.parametrize("causal", [True]) @pytest.mark.parametrize("d", [32, 40, 59, 64, 96, 111, 128, 160, 192, 224, 256]) # @pytest.mark.parametrize("d", [32, 64, 96, 128, 160, 192, 224, 256]) @@ -914,9 +914,6 @@ def test_flash_attn_output( seqlen_q, seqlen_k, d, dropout_p, causal, local, alibi, deterministic, mha_type, dtype, kvpacked, softcap ): DEBUG = False - BATCH_AND_HEAD_ONE = False - SAVE_TO_CSV = False - TEST_BACKWARD = False if DEBUG: print() print("Debugging") @@ -935,7 +932,7 @@ def test_flash_attn_output( device = "cuda" # set seed torch.random.manual_seed(0) - if BATCH_AND_HEAD_ONE: + if DEBUG: batch_size = 1 nheads = 1 nheads_k = 1 @@ -1166,20 +1163,21 @@ def test_flash_attn_output( print("window_size:", window_size) print("out:", out, out.shape) print("out_ref:", out_ref, out_ref.shape) - if SAVE_TO_CSV: - for batch_idx in range(batch_size): - for head_idx in range(nheads): - save_tensor_to_csv(out[batch_idx, :, head_idx, :], - f"out_b{batch_idx}_h{head_idx}.csv") - save_tensor_to_csv(out_ref[batch_idx, :, head_idx, :], - f"out_ref_b{batch_idx}_h{head_idx}.csv") + # if True: + # for batch_idx in range(batch_size): + # for head_idx in range(nheads): + # save_tensor_to_csv(out[batch_idx, :, head_idx, :], + # f"out_b{batch_idx}_h{head_idx}.csv") + # save_tensor_to_csv(out_ref[batch_idx, :, head_idx, :], + # f"out_ref_b{batch_idx}_h{head_idx}.csv") + + if local == True and causal == True: + pytest.skip("Sliding Window and Causal not supported") # Check that FlashAttention's numerical error is at most twice the numerical error # of a Pytorch implementation. assert (out - out_ref).abs().max().item() <= 2 * (out_pt - out_ref).abs().max().item() - if not TEST_BACKWARD: - return if dropout_p > 0.0: # assert (attn - attn_ref).abs().max().item() <= 2 * (attn_pt - attn_ref).abs().max().item() @@ -1187,6 +1185,9 @@ 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 DEBUG: print("dq:", dq) @@ -2002,7 +2003,6 @@ def test_flash_attn_kvcache( num_splits, dtype, ): - DEBUG = False if seqlen_q > seqlen_k and new_kv: pytest.skip() if not new_kv and rotary_fraction > 0.0: @@ -2014,40 +2014,24 @@ def test_flash_attn_kvcache( device = "cuda" # set seed torch.random.manual_seed(0) - if DEBUG: - batch_size = 1 - batch_size_cache = batch_size if not has_batch_idx else batch_size * 2 - nheads = 1 - else: - batch_size = 2 - batch_size_cache = batch_size if not has_batch_idx else batch_size * 2 - nheads = 6 + batch_size = 2 + batch_size_cache = batch_size if not has_batch_idx else batch_size * 2 + nheads = 6 # rotary_dim must be a multiple of 16, and must be <= d rotary_dim = math.floor(int(rotary_fraction * d) / 16) * 16 nheads_k = nheads if mha_type == "mha" else (1 if mha_type == "mqa" else 3) assert nheads % nheads_k == 0 window_size = (-1, -1) if not local else torch.randint(0, seqlen_k, (2,)) - if DEBUG: - q = torch.randn(batch_size, seqlen_q, nheads, d, device=device, dtype=dtype) - else: - q = torch.randn(batch_size, seqlen_q, nheads, d, device=device, dtype=dtype) + q = torch.randn(batch_size, seqlen_q, nheads, d, device=device, dtype=dtype) seqlen_new = seqlen_q if seqlen_new_eq_seqlen_q else torch.randint(1, seqlen_q + 1, (1,)).item() if new_kv: - if DEBUG: - k = torch.randn(batch_size, seqlen_new, nheads_k, d, device=device, dtype=dtype) - v = torch.randn(batch_size, seqlen_new, nheads_k, d, device=device, dtype=dtype) - else: - k = torch.randn(batch_size, seqlen_new, nheads_k, d, device=device, dtype=dtype) - v = torch.randn(batch_size, seqlen_new, nheads_k, d, device=device, dtype=dtype) + k = torch.randn(batch_size, seqlen_new, nheads_k, d, device=device, dtype=dtype) + v = torch.randn(batch_size, seqlen_new, nheads_k, d, device=device, dtype=dtype) else: k, v = None, None if paged_kv_block_size is None: - if DEBUG: - k_cache = torch.randn(batch_size_cache, seqlen_k, nheads_k, d, device=device, dtype=dtype) - v_cache = torch.randn(batch_size_cache, seqlen_k, nheads_k, d, device=device, dtype=dtype) - else: - k_cache = torch.randn(batch_size_cache, seqlen_k, nheads_k, d, device=device, dtype=dtype) - v_cache = torch.randn(batch_size_cache, seqlen_k, nheads_k, d, device=device, dtype=dtype) + k_cache = torch.randn(batch_size_cache, seqlen_k, nheads_k, d, device=device, dtype=dtype) + v_cache = torch.randn(batch_size_cache, seqlen_k, nheads_k, d, device=device, dtype=dtype) block_table = None else: ( @@ -2234,11 +2218,6 @@ def test_flash_attn_kvcache( assert torch.allclose(k_cache_select, k_cache_ref, rtol=1e-3, atol=1e-3) assert torch.equal(v_cache_select, v_cache_ref) mult = 3 if not alibi else 5 - if DEBUG: - print("out:", out, out.shape) - print("out_ref:", out_ref, out_ref.shape) - # save_tensor_to_csv(out, "out.csv") - # save_tensor_to_csv(out_ref, "out_ref.csv") assert (out - out_ref).abs().max().item() <= mult * (out_pt - out_ref).abs().max().item() + 1e-5 From e0300769581e1f0f6427d4d6bde6a91363919e66 Mon Sep 17 00:00:00 2001 From: Michael Date: Wed, 9 Jul 2025 12:23:41 -0500 Subject: [PATCH 04/10] no sliding window causal green --- .../flash_attn_triton_amd/fwd_prefill.py | 51 +++++++++++++++---- tests/test_flash_attn_triton_amd.py | 40 --------------- 2 files changed, 40 insertions(+), 51 deletions(-) diff --git a/flash_attn/flash_attn_triton_amd/fwd_prefill.py b/flash_attn/flash_attn_triton_amd/fwd_prefill.py index 2bbcc6ff909..0fe312616d1 100644 --- a/flash_attn/flash_attn_triton_amd/fwd_prefill.py +++ b/flash_attn/flash_attn_triton_amd/fwd_prefill.py @@ -218,7 +218,41 @@ def _attn_fwd_mask(acc, l_i, m_i, if USE_SLIDING_WINDOW: if IS_CAUSAL: - pass + # ========== 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: @@ -387,17 +421,12 @@ def compute_masking(seqlen_k, seqlen_q, start_m, else: n_extra_tokens = 0 - if USE_SLIDING_WINDOW: # TODO: impl sliding window + 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, 0, 0 + return 0, 0, 0, total_k_blocks, n_extra_tokens else: - # ========== NON-CAUSAL SLIDING WINDOW ========== - # For now, we'll use a simple approach: - # Process all blocks with masking since we have a sliding window - # This is less efficient than skipping blocks, but ensures correctness - - # TODO: Optimize by computing which blocks can be fully skipped - # For now, process all blocks with the mask function return 0, 0, 0, total_k_blocks, n_extra_tokens else: if IS_CAUSAL: @@ -1023,4 +1052,4 @@ def attention_prefill_forward_triton_impl( 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/tests/test_flash_attn_triton_amd.py b/tests/test_flash_attn_triton_amd.py index d95d45df749..f0dc4c97496 100755 --- a/tests/test_flash_attn_triton_amd.py +++ b/tests/test_flash_attn_triton_amd.py @@ -887,13 +887,6 @@ def test_flash_attn_varlen_qkvpacked( @pytest.mark.parametrize( "seqlen_q,seqlen_k", [ - # debug - # (32, 32), - # (32, 64), - # (64, 32), - # (64, 64), - # (128, 128), - # og (113, 203), (128, 217), (113, 211), @@ -914,10 +907,6 @@ def test_flash_attn_output( seqlen_q, seqlen_k, d, dropout_p, causal, local, alibi, deterministic, mha_type, dtype, kvpacked, softcap ): DEBUG = False - if DEBUG: - print() - print("Debugging") - if USE_TRITON_ROCM: if causal: if seqlen_q ==1024 and seqlen_k==1024 and d==160: @@ -973,12 +962,6 @@ def test_flash_attn_output( else: alibi_slopes, attn_bias = None, None - if DEBUG: - print("q:", q, q.shape) - print("k:", k, k.shape) - print("v:", v, v.shape) - print("window_size:", window_size) - if kvpacked: out, lse, S_dmask = flash_attn_kvpacked_func( q, @@ -1159,22 +1142,6 @@ def test_flash_attn_output( print(f"dK Pytorch mean diff: {(dk_pt - dk_ref).abs().mean().item()}") print(f"dV Pytorch mean diff: {(dv_pt - dv_ref).abs().mean().item()}") - if DEBUG: - print("window_size:", window_size) - print("out:", out, out.shape) - print("out_ref:", out_ref, out_ref.shape) - # if True: - # for batch_idx in range(batch_size): - # for head_idx in range(nheads): - # save_tensor_to_csv(out[batch_idx, :, head_idx, :], - # f"out_b{batch_idx}_h{head_idx}.csv") - # save_tensor_to_csv(out_ref[batch_idx, :, head_idx, :], - # f"out_ref_b{batch_idx}_h{head_idx}.csv") - - if local == True and causal == True: - pytest.skip("Sliding Window and Causal not supported") - - # Check that FlashAttention's numerical error is at most twice the numerical error # of a Pytorch implementation. assert (out - out_ref).abs().max().item() <= 2 * (out_pt - out_ref).abs().max().item() @@ -1189,13 +1156,6 @@ def test_flash_attn_output( print("Sliding Window not supported in backward yet") return - if DEBUG: - print("dq:", dq) - print("dq_ref:", dq_ref) - print("dk:", dk) - print("dk_ref:", dk_ref) - print("dv:", dv) - print("dv_ref:", dv_ref) 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() From 81703b59b0d22f87806938ce9d5e3aeafcf5c56f Mon Sep 17 00:00:00 2001 From: Michael Date: Wed, 9 Jul 2025 12:31:04 -0500 Subject: [PATCH 05/10] add backward test skip for sliding window --- tests/test_flash_attn_triton_amd.py | 75 +++++++++++++++++------------ 1 file changed, 44 insertions(+), 31 deletions(-) diff --git a/tests/test_flash_attn_triton_amd.py b/tests/test_flash_attn_triton_amd.py index f0dc4c97496..0a81dad72f1 100755 --- a/tests/test_flash_attn_triton_amd.py +++ b/tests/test_flash_attn_triton_amd.py @@ -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() @@ -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() @@ -906,7 +914,6 @@ def test_flash_attn_varlen_qkvpacked( def test_flash_attn_output( seqlen_q, seqlen_k, d, dropout_p, causal, local, alibi, deterministic, mha_type, dtype, kvpacked, softcap ): - DEBUG = False if USE_TRITON_ROCM: if causal: if seqlen_q ==1024 and seqlen_k==1024 and d==160: @@ -921,23 +928,12 @@ def test_flash_attn_output( device = "cuda" # set seed torch.random.manual_seed(0) - if DEBUG: - batch_size = 1 - nheads = 1 - nheads_k = 1 - else: - batch_size = 4 - nheads = 6 if softcap == 0.0 else 4 # softcap reference impl takes more memory - nheads_k = nheads if mha_type == "mha" else (1 if mha_type == "mqa" else 2) + batch_size = 4 + nheads = 6 if softcap == 0.0 else 4 # softcap reference impl takes more memory + nheads_k = nheads if mha_type == "mha" else (1 if mha_type == "mqa" else 2) assert nheads % nheads_k == 0 - if DEBUG: - window_size = (-1, -1) if not local else (seqlen_k//4, (seqlen_k*3)//4) - else: - window_size = (-1, -1) if not local else torch.randint(0, seqlen_k, (2,)) - if DEBUG: - q = generate_bshd_tensor(batch_size, seqlen_q, nheads, d, dtype=dtype, device=device, mode="ones") - else: - q = torch.randn(batch_size, seqlen_q, nheads, d, device=device, dtype=dtype, requires_grad=True) + window_size = (-1, -1) if not local else torch.randint(0, seqlen_k, (2,)) + q = torch.randn(batch_size, seqlen_q, nheads, d, device=device, dtype=dtype, requires_grad=True) if softcap > 0: # Ensure the values of qk are at least within softcap range. q = q * softcap @@ -946,16 +942,12 @@ def test_flash_attn_output( batch_size, seqlen_k, 2, nheads_k, d, device=device, dtype=dtype, requires_grad=True ) else: - if DEBUG: - k = generate_bshd_tensor(batch_size, seqlen_k, nheads_k, d, dtype=dtype, device=device, mode="ones") - v = generate_bshd_tensor(batch_size, seqlen_k, nheads_k, d, dtype=dtype, device=device, mode="identity") - else: - k = torch.randn( - batch_size, seqlen_k, nheads_k, d, device=device, dtype=dtype, requires_grad=True - ) - v = torch.randn( - batch_size, seqlen_k, nheads_k, d, device=device, dtype=dtype, requires_grad=True - ) + k = torch.randn( + batch_size, seqlen_k, nheads_k, d, device=device, dtype=dtype, requires_grad=True + ) + v = torch.randn( + batch_size, seqlen_k, nheads_k, d, device=device, dtype=dtype, requires_grad=True + ) if alibi: alibi_slopes = torch.rand(batch_size, nheads, device=device, dtype=torch.float32) * 0.3 attn_bias = attn_bias_from_alibi_slopes(alibi_slopes, seqlen_q, seqlen_k, causal=causal) @@ -1091,10 +1083,7 @@ def test_flash_attn_output( print(f"Attention max diff: {(attn - attn_ref).abs().max().item()}") print(f"Attention Pytorch max diff: {(attn_pt - attn_ref).abs().max().item()}") - if DEBUG: - g = torch.ones_like(out) - else: - g = torch.randn_like(out) + g = torch.randn_like(out) do_o = (g.float() * out.float()).sum(-1) if (d <= MAX_HEADDIM_SM8x or dropout_p == 0) or (is_sm80 or is_sm90): if kvpacked: @@ -1478,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() @@ -1592,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 @@ -1760,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 @@ -1890,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 @@ -2475,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): @@ -2563,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): From bc74ab1db9f42a37c2ac65f151226c0a5877a46f Mon Sep 17 00:00:00 2001 From: Michael Date: Wed, 9 Jul 2025 14:33:49 -0500 Subject: [PATCH 06/10] clean reduce in fwd_kvcache. no is_CASUAL branching --- .../flash_attn_triton_amd/fwd_decode.py | 36 +++++++++---------- .../flash_attn_triton_amd/interface_fa.py | 18 +++++++--- flash_attn/flash_attn_triton_amd/utils.py | 4 +++ tests/test_flash_attn_triton_amd.py | 2 +- 4 files changed, 35 insertions(+), 25 deletions(-) diff --git a/flash_attn/flash_attn_triton_amd/fwd_decode.py b/flash_attn/flash_attn_triton_amd/fwd_decode.py index fe969027536..56c05cb3fa3 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) @@ -387,7 +390,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 +428,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 +448,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 @@ -590,6 +583,7 @@ def attention_decode_forward_triton_impl( is_new_kv = True if k_new is not None and v_new is not None else False 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 @@ -653,7 +647,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) @@ -745,6 +739,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, ) @@ -806,8 +803,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/interface_fa.py b/flash_attn/flash_attn_triton_amd/interface_fa.py index 4e49d3a73d0..4e9bd006aee 100644 --- a/flash_attn/flash_attn_triton_amd/interface_fa.py +++ b/flash_attn/flash_attn_triton_amd/interface_fa.py @@ -819,6 +819,16 @@ def fwd_kvcache( 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 @@ -887,8 +897,8 @@ def fwd_kvcache( out, metadata.sm_scale, metadata.causal, - window_size_left, - window_size_right, + metadata.window_size_left, + metadata.window_size_right, metadata.alibi_slopes, metadata.layout, metadata.cache_seqlens, @@ -907,8 +917,8 @@ def fwd_kvcache( out, metadata.sm_scale, metadata.causal, - window_size_left, - window_size_right, + metadata.window_size_left, + metadata.window_size_right, metadata.alibi_slopes, metadata.layout, metadata.cache_seqlens, diff --git a/flash_attn/flash_attn_triton_amd/utils.py b/flash_attn/flash_attn_triton_amd/utils.py index a00e0fb856d..fdc9e953cda 100644 --- a/flash_attn/flash_attn_triton_amd/utils.py +++ b/flash_attn/flash_attn_triton_amd/utils.py @@ -55,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: @@ -74,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): diff --git a/tests/test_flash_attn_triton_amd.py b/tests/test_flash_attn_triton_amd.py index 0a81dad72f1..25d7230cd66 100755 --- a/tests/test_flash_attn_triton_amd.py +++ b/tests/test_flash_attn_triton_amd.py @@ -1911,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, True]) +@pytest.mark.parametrize("local", [False]) # @pytest.mark.parametrize("local", [False]) @pytest.mark.parametrize("causal", [False, True]) # @pytest.mark.parametrize("causal", [False]) From f508ee9fcbb90a55a9225e9299b57d5a400080ae Mon Sep 17 00:00:00 2001 From: Michael Date: Wed, 9 Jul 2025 20:58:54 -0500 Subject: [PATCH 07/10] add kvcache masking --- .../flash_attn_triton_amd/fwd_decode.py | 49 +++++++++++++++---- 1 file changed, 39 insertions(+), 10 deletions(-) diff --git a/flash_attn/flash_attn_triton_amd/fwd_decode.py b/flash_attn/flash_attn_triton_amd/fwd_decode.py index 56c05cb3fa3..76af2fd21c8 100644 --- a/flash_attn/flash_attn_triton_amd/fwd_decode.py +++ b/flash_attn/flash_attn_triton_amd/fwd_decode.py @@ -300,17 +300,46 @@ 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) + + # 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")) + # 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? From 5d56989474114baab46c06d275a7b62c323e6831 Mon Sep 17 00:00:00 2001 From: Michael Date: Wed, 9 Jul 2025 22:06:13 -0500 Subject: [PATCH 08/10] kvcache working --- .../flash_attn_triton_amd/fwd_decode.py | 22 +++++++++---------- tests/test_flash_attn_triton_amd.py | 4 ++-- 2 files changed, 12 insertions(+), 14 deletions(-) diff --git a/flash_attn/flash_attn_triton_amd/fwd_decode.py b/flash_attn/flash_attn_triton_amd/fwd_decode.py index 76af2fd21c8..d3f7f9c32b9 100644 --- a/flash_attn/flash_attn_triton_amd/fwd_decode.py +++ b/flash_attn/flash_attn_triton_amd/fwd_decode.py @@ -346,18 +346,16 @@ def _fwd_kernel_splitK( 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 -- diff --git a/tests/test_flash_attn_triton_amd.py b/tests/test_flash_attn_triton_amd.py index 25d7230cd66..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, generate_bshd_tensor, is_hip, is_rdna, save_tensor_to_csv +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 @@ -1911,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]) From a6b190a344bb796595f3d2d0aa74854ae68c13a8 Mon Sep 17 00:00:00 2001 From: Michael Date: Thu, 10 Jul 2025 09:56:00 -0500 Subject: [PATCH 09/10] fix some bugs in test.py --- flash_attn/flash_attn_triton_amd/test.py | 4 ++-- flash_attn/flash_attn_triton_amd/utils.py | 15 +++++++-------- 2 files changed, 9 insertions(+), 10 deletions(-) diff --git a/flash_attn/flash_attn_triton_amd/test.py b/flash_attn/flash_attn_triton_amd/test.py index f634103ca69..08004fc80e7 100644 --- a/flash_attn/flash_attn_triton_amd/test.py +++ b/flash_attn/flash_attn_triton_amd/test.py @@ -523,7 +523,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 +631,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 --- diff --git a/flash_attn/flash_attn_triton_amd/utils.py b/flash_attn/flash_attn_triton_amd/utils.py index fdc9e953cda..fc8b98d2656 100644 --- a/flash_attn/flash_attn_triton_amd/utils.py +++ b/flash_attn/flash_attn_triton_amd/utils.py @@ -170,7 +170,7 @@ def generate_varlen_tensor( batch_size: Optional[int] = None, equal_seqlens: bool = False, device: str = "cuda", - dtype: torch.dtype = torch.float32, + dtype: torch.dtype = torch.float16, mode: Literal["random", "ones", "incremental", "identity"] = "random" ): if DEBUG: @@ -245,7 +245,7 @@ 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, device="cuda", mode: Literal["random", "ones", "incremental", "identity"] = "random"): +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: @@ -277,7 +277,7 @@ def generate_bshd_tensor(BATCH, SEQ_LEN, NUM_HEADS, D_HEAD, dtype, device="cuda" x.requires_grad_() return x -def generate_bhsd_tensor(BATCH, NUM_HEADS, SEQ_LEN, D_HEAD, dtype, device="cuda", mode: Literal["random", "ones", "incremental", "identity"] = "random"): +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: @@ -582,7 +582,7 @@ def input_helper( 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) 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) + 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" @@ -594,8 +594,7 @@ def input_helper( 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 = generate_varlen_qkv_packed(TOTAL_SEQLENS_Q, HQ, 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) # setup metadata sm_scale = D_HEAD**-0.5 @@ -652,10 +651,10 @@ def input_helper( else: if layout == "bshd": qkv = generate_bshd_qkv_packed(BATCH, N_CTX_Q, HQ, D_HEAD, dtype=dtype, device=device) - do = 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) - do = 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 sm_scale = D_HEAD**-0.5 From d1e68b8725fd398fc5f02b680d888f9d36243909 Mon Sep 17 00:00:00 2001 From: Michael Date: Thu, 10 Jul 2025 11:02:52 -0500 Subject: [PATCH 10/10] clean up --- flash_attn/flash_attn_triton_amd/bwd_ref.py | 2 +- .../flash_attn_triton_amd/fwd_prefill.py | 37 +- .../flash_attn_triton_amd/fwd_prefill_old.py | 473 ------------------ flash_attn/flash_attn_triton_amd/test.py | 350 +------------ flash_attn/flash_attn_triton_amd/utils.py | 4 - 5 files changed, 8 insertions(+), 858 deletions(-) delete mode 100644 flash_attn/flash_attn_triton_amd/fwd_prefill_old.py diff --git a/flash_attn/flash_attn_triton_amd/bwd_ref.py b/flash_attn/flash_attn_triton_amd/bwd_ref.py index 1f384cab068..cb1637157a0 100644 --- a/flash_attn/flash_attn_triton_amd/bwd_ref.py +++ b/flash_attn/flash_attn_triton_amd/bwd_ref.py @@ -58,7 +58,7 @@ def attention_backward_core_ref_impl( if DEBUG_CORE: print("attention_scaled_scores after alibi:", attention_scaled_scores, attention_scaled_scores.shape) - # Apply masks (causal and/or sliding window) + # 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) diff --git a/flash_attn/flash_attn_triton_amd/fwd_prefill.py b/flash_attn/flash_attn_triton_amd/fwd_prefill.py index 0fe312616d1..59fe8bfaf4e 100644 --- a/flash_attn/flash_attn_triton_amd/fwd_prefill.py +++ b/flash_attn/flash_attn_triton_amd/fwd_prefill.py @@ -5,12 +5,8 @@ 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, get_fwd_prefill_autotune_configs - -from flash_attn.flash_attn_triton_amd.fwd_prefill_old import attn_fwd_old - DEBUG = False - # NOTE: triton fails to import tl.constexprs so create them here for the file tl_DROPOUT_USE_PYTORCH: tl.constexpr = triton.language.constexpr(DROPOUT_USE_PYTORCH) tl_DROPOUT_DUMP: tl.constexpr = triton.language.constexpr(DROPOUT_DUMP) @@ -1010,11 +1006,9 @@ def attention_prefill_forward_triton_impl( else: stride_bz, stride_bh, stride_bm, stride_bn = (0, 0, 0, 0) - USE_OLD = os.environ.get('USE_OLD', '0').lower() in ('1', 'true', 'yes') - if USE_OLD: - attn_fwd_old[grid](q, k, v, bias, cache_seqlens, cache_batch_idx, + 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, @@ -1022,34 +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) - else: - 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, - stride_qb, stride_qh, stride_qm, stride_qd, - stride_kb, stride_kh, stride_kn, stride_kd, - stride_vb, stride_vh, stride_vn, stride_vd, - stride_ob, stride_oh, stride_om, stride_od, - 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, - 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, - 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) 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_prefill_old.py b/flash_attn/flash_attn_triton_amd/fwd_prefill_old.py deleted file mode 100644 index cab91a604e1..00000000000 --- a/flash_attn/flash_attn_triton_amd/fwd_prefill_old.py +++ /dev/null @@ -1,473 +0,0 @@ -import os -import torch -import triton -import triton.language as tl -from typing import Literal, Optional, Union -from .utils import DEBUG, DROPOUT_USE_PYTORCH, DROPOUT_DUMP, AUTOTUNE, compute_alibi_block, compute_fp8_scaling_factors, get_arch, get_shapes_from_layout, get_strides_from_layout, is_cdna, is_fp8, is_rdna, create_dropout_mask - -# NOTE: triton fails to import tl.constexprs so create them here for the file -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. -@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 - - - -@triton.jit -def _attn_fwd_inner_old(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, - 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, - 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): - # 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) - 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) - 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 - - # -- 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 - 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) - 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 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) - 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] - - # 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: - if tl_DROPOUT_USE_PYTORCH: - dropout_mask = tl.load(dropout_mask_ptrs, mask=p_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) - - # 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) - - # 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) - - # -- 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 - 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) - # -- 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: - 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) - - 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 - -# @triton.autotune( -# configs=autotune_configs, -# key=autotune_keys, -# use_cuda_graph=True, -# ) -@triton.jit -def attn_fwd_old(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: tl.constexpr, LSE, Out, stride_qz, stride_qh, stride_qm, stride_qk, - stride_kz, stride_kh, stride_kn, stride_kk, stride_vz, stride_vh, stride_vk, stride_vn, - stride_oz, stride_oh, stride_om, stride_on, 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, 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, - 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): - # set params - ACCUMULATOR_TYPE = tl.float32 - - # 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) - - offs_m = start_m * BLOCK_M + tl.arange(0, BLOCK_M) - offs_n = tl.arange(0, BLOCK_N) - offs_d = tl.arange(0, BLOCK_DMODEL) - - # handle seqlen - if IS_VARLEN: - cu_seqlens_q_start = tl.load(cu_seqlens_q + off_z) - cu_seqlens_q_end = tl.load(cu_seqlens_q + off_z + 1) - seqlen_q = cu_seqlens_q_end - cu_seqlens_q_start - - # we have a one-size-fits-all grid in id(0). Some seqlens might be too small for all start_m so for those we return early. - if start_m * BLOCK_M > seqlen_q: - return - 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 - 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) - - # 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 - k_offset = K + off_z * stride_kz + off_h_k * stride_kh + cu_seqlens_k_start * stride_kn - k_ptrs = k_offset + offs_d[:, None] * stride_kk + offs_n[None, :] * stride_kn - v_offset = V + off_z * stride_vz + off_h_k * stride_vh + cu_seqlens_k_start * stride_vk - v_ptrs = v_offset + offs_n[:, None] * stride_vk + offs_d[None, :] * stride_vn - if USE_BIAS: - # Note: this might get large enough to overflow on some configs - bias_offset = off_h_q * stride_bh - bias_ptrs = bias + bias_offset + offs_m[:, None] * stride_bm + offs_n[None, :] * stride_bn - else: - bias_ptrs = None - - if USE_ALIBI: - a_offset = off_z * stride_az + off_h_q * stride_ah - alibi_slope = tl.load(alibi_slopes + a_offset) - else: - alibi_slope = None - - if RETURN_SCORES: - sd_mask_offset = sd_mask + off_z * stride_sz + off_h_q * stride_sh #+ cu_seqlens_q_start * stride_sm - sd_mask_ptrs = sd_mask_offset + offs_m[:, None] * stride_sm + offs_n[None, :] * stride_sn - else: - sd_mask_ptrs = None - - if ENABLE_DROPOUT: - dropout_mask_offset = dropout_mask + off_z * stride_sz + off_h_q * stride_sh #+ cu_seqlens_q_start * stride_sm - dropout_mask_ptrs = dropout_mask_offset + offs_m[:, None] * stride_sm + offs_n[None, :] * stride_sn - batch_philox_offset = philox_offset_base + off_z * stride_sz + off_h_q * stride_sh #+ cu_seqlens_q_start * stride_sm - philox_ptrs = batch_philox_offset + offs_m[:, None] * stride_sm + offs_n[None, :] * stride_sn - 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. - if n_full_blocks > 0: - block_max = (n_blocks - masked_blocks) * BLOCK_N - acc, l_i, m_i = _attn_fwd_inner_old(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_old(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 - # This helps the compiler do Newton Raphson on l_i vs on acc which is much larger. - 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 - 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) - # convert back to natural units - softmax_lse *= LN2 - else: - softmax_lse = m_i + tl.math.log(l_i) - - 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) - - # 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 - 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 - else: - tl.store(l_ptrs, softmax_lse) # the log of the normalization constant - - # write back O - 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_ptrs_mask = tl.full([BLOCK_M, BLOCK_DMODEL], 1, dtype=tl.int1) - if overflow_size > 0: - o_ptrs_mask = o_ptrs_mask & (offs_m[:, None] < seqlen_q) - if PADDED_HEAD: - o_ptrs_mask = o_ptrs_mask & (offs_d[None, :] < ACTUAL_BLOCK_DMODEL) - - if FP8_OUTPUT: - scale_acc, descale_acc = compute_fp8_scaling_factors(acc, FP8_MAX) - tl.store(Descale_O + off_z * stride_descale_o_z + off_h_q, descale_acc) - tl.store(o_ptrs, (acc * scale_acc).to(Out.type.element_ty), mask=o_ptrs_mask) - else: - tl.store(o_ptrs, acc.to(Out.dtype.element_ty), mask=o_ptrs_mask) diff --git a/flash_attn/flash_attn_triton_amd/test.py b/flash_attn/flash_attn_triton_amd/test.py index 08004fc80e7..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 @@ -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 fc8b98d2656..96d4f662567 100644 --- a/flash_attn/flash_attn_triton_amd/utils.py +++ b/flash_attn/flash_attn_triton_amd/utils.py @@ -301,11 +301,7 @@ def generate_bhsd_tensor(BATCH, NUM_HEADS, SEQ_LEN, D_HEAD, dtype: torch.dtype = raise ValueError(f"Unkown mode {mode}") if is_fp8_dtype: - # cast to fp8 raise ValueError("fp8 not supported for bhsd yet") - # 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 else: x.requires_grad_() return x