diff --git a/ggml/src/ggml-cuda/fattn-common.cuh b/ggml/src/ggml-cuda/fattn-common.cuh index 6d1ce52dbc77..5bb52f0f6e4f 100644 --- a/ggml/src/ggml-cuda/fattn-common.cuh +++ b/ggml/src/ggml-cuda/fattn-common.cuh @@ -7,6 +7,7 @@ #include #define FATTN_KQ_STRIDE 256 +#define FATTN_MASK_KV_BLOCKS 16 #define HALF_MAX_HALF __float2half(65504.0f/2) // Use neg. of this instead of -INFINITY to initialize KQ max vals to avoid NaN upon subtraction. #define SOFTMAX_FTZ_THRESHOLD -20.0f // Softmax exp. of values smaller than this are flushed to zero to avoid NaNs. @@ -41,6 +42,9 @@ typedef void (* fattn_kernel_t)( const int32_t ne31, const int32_t ne32, const int32_t ne33, const int32_t nb31, const int32_t nb32, const int64_t nb33); +typedef void (*fattn_mask_blocks_kernel_t)( + const half *, uint32_t *, const int, const int, const int, const int64_t, const int64_t); + typedef float (*vec_dot_KQ_t)( const char * __restrict__ K_c, const void * __restrict__ Q_v, const int * __restrict__ Q_q8 , const void * __restrict__ Q_ds); @@ -664,7 +668,8 @@ constexpr __device__ dequantize_V_t get_dequantize_V() { template __launch_bounds__(FATTN_KQ_STRIDE/2, 1) static __global__ void flash_attn_mask_to_KV_max( - const half2 * mask_ptr, int * KV_max_ptr, const int ne30, const int64_t s31, const int64_t s33) { + const half2 * mask_ptr, int * KV_max_ptr, const int ne30, const int ne01, const int ne33, + const int64_t s31, const int64_t s33) { const half2 * GGML_CUDA_RESTRICT mask = mask_ptr; int * GGML_CUDA_RESTRICT KV_max = KV_max_ptr; @@ -673,7 +678,7 @@ static __global__ void flash_attn_mask_to_KV_max( const int sequence = blockIdx.y; const int jt = blockIdx.x; - mask += sequence*s33 + jt*ncols1*s31; + mask += (sequence % ne33)*s33 + jt*ncols1*s31; __shared__ int buf_iw[WARP_SIZE]; if (tid < WARP_SIZE) { @@ -687,7 +692,7 @@ static __global__ void flash_attn_mask_to_KV_max( int all_inf = 1; #pragma unroll - for (int j = 0; j < ncols1; ++j) { + for (int j = 0; j < min(ncols1, ne01 - jt*ncols1); ++j) { const float2 tmp = __half22float2(mask[j*s31 + KV_max_sj/2 + tid]); all_inf = all_inf && int(isinf(tmp.x)) && int(isinf(tmp.y)); } @@ -715,7 +720,60 @@ static __global__ void flash_attn_mask_to_KV_max( return; } - KV_max[sequence*ne31 + jt] = KV_max_sj; + KV_max[int64_t(sequence)*ne31 + jt] = KV_max_sj; +} + +template +__launch_bounds__(WARP_SIZE * FATTN_MASK_KV_BLOCKS, 1) +static __global__ void flash_attn_mask_to_KV_blocks( + const half * mask, uint32_t * KV_blocks, const int ne00, const int ne01, + const int nbatch_fa, const int64_t s31, const int64_t s33) { + const int lane = threadIdx.x % WARP_SIZE; + const int warp = threadIdx.x / WARP_SIZE; + const int kb = blockIdx.x * FATTN_MASK_KV_BLOCKS + warp; + const int jt = blockIdx.y; + const int sequence = blockIdx.z; + const int64_t k0 = int64_t(kb) * nbatch_fa; + const int j0 = jt * ncols1; + const int j1 = j0 + min(ncols1, ne01 - j0); + const int64_t ib = (int64_t(sequence)*gridDim.y + jt)*gridDim.x + blockIdx.x; + mask += int64_t(sequence)*s33; + ggml_cuda_pdl_sync(); + + int all_masked = 1; + int all_zero = 0; + if (k0 < ne00) { + const float corner = __half2float(mask[int64_t(j1 - 1)*s31 + k0]); + all_masked = corner == -INFINITY; + all_zero = corner == 0.0f; + if (all_masked || all_zero) { + for (int j = j0; j < j1; ++j) { + for (int i = lane; i < nbatch_fa && k0 + i < ne00; i += WARP_SIZE) { + const float value = __half2float(mask[int64_t(j)*s31 + k0 + i]); + all_masked &= value == -INFINITY; + all_zero &= value == 0.0f; + } + } + } + } + all_masked = warp_reduce_all(all_masked); + all_zero = warp_reduce_all(all_zero); + + __shared__ int block_types[FATTN_MASK_KV_BLOCKS]; + if (lane == 0) { + block_types[warp] = all_masked ? 1 : all_zero ? 2 : 0; + } + __syncthreads(); + if (threadIdx.x < WARP_SIZE) { + uint32_t bits = lane < FATTN_MASK_KV_BLOCKS ? uint32_t(block_types[lane]) << (2*lane) : 0; +#pragma unroll + for (int offset = WARP_SIZE/2; offset > 0; offset /= 2) { + bits |= __shfl_xor_sync(0xffffffff, bits, offset, WARP_SIZE); + } + if (lane == 0) { + KV_blocks[ib] = bits; + } + } } void ggml_cuda_flash_attn_ext_compact_mask( @@ -976,9 +1034,10 @@ template void launch_fattn( ggml_backend_cuda_context & ctx, ggml_tensor * dst, fattn_kernel_t fattn_kernel, const int nwarps, const size_t nbytes_shared, const int nbatch_fa, const bool need_f16_K, const bool need_f16_V, const bool stream_k, const bool use_sparse, - const int warp_size = WARP_SIZE + const int warp_size = WARP_SIZE, fattn_mask_blocks_kernel_t mask_blocks_kernel = nullptr ) { constexpr int ncols = ncols1 * ncols2; + const bool use_mask_blocks = mask_blocks_kernel != nullptr; const ggml_tensor * Q = dst->src[0]; const ggml_tensor * K = dst->src[1]; @@ -1010,6 +1069,7 @@ void launch_fattn( ggml_cuda_flash_attn_ext_get_f16_extra_data(KQV, need_f16_K, need_f16_V); ggml_cuda_pool_alloc KV_max(pool); + ggml_cuda_pool_alloc KV_blocks(pool); ggml_cuda_pool_alloc dst_tmp(pool); ggml_cuda_pool_alloc dst_tmp_meta(pool); @@ -1109,23 +1169,61 @@ void launch_fattn( // Optional optimization where the mask is scanned to determine whether part of the calculation can be skipped. // Only worth the overhead if there is at lease one FATTN_KQ_STRIDE x FATTN_KQ_STRIDE square to be skipped or // multiple sequences of possibly different lengths. - if (!use_sparse && mask && K->ne[1] % FATTN_KQ_STRIDE == 0 && (Q->ne[1] >= 1024 || Q->ne[3] > 1)) { + if (!use_sparse && !use_mask_blocks && mask && K->ne[1] % FATTN_KQ_STRIDE == 0 && (Q->ne[1] >= 1024 || Q->ne[3] > 1)) { const int64_t s31 = mask->nb[1] / sizeof(half2); const int64_t s33 = mask->nb[3] / sizeof(half2); const dim3 blocks_num_KV_max(ntiles_x, Q->ne[3], 1); const dim3 block_dim_KV_max(FATTN_KQ_STRIDE/2, 1, 1); - const int ne_KV_max = blocks_num_KV_max.x*blocks_num_KV_max.y; + const size_t ne_KV_max = size_t(blocks_num_KV_max.x)*blocks_num_KV_max.y; const int iter_k = K->ne[1] / FATTN_KQ_STRIDE; KV_max.alloc(ne_KV_max); ggml_cuda_kernel_launch_params launch_params = ggml_cuda_kernel_launch_params(blocks_num_KV_max, block_dim_KV_max, 0, main_stream); ggml_cuda_kernel_launch(flash_attn_mask_to_KV_max, launch_params, - (const half2 *) mask->data, KV_max.ptr, iter_k, s31, s33); + (const half2 *) mask->data, KV_max.ptr, iter_k, Q->ne[1], mask->ne[3], s31, s33); CUDA_CHECK(cudaGetLastError()); } + if (use_mask_blocks) { + // AMD WMMA dispatch requires a dense, single-head mask. + GGML_ASSERT(!use_sparse && mask && mask->ne[2] == 1); + const int nblocks_KV = (K->ne[1] - 1) / (nbatch_fa * FATTN_MASK_KV_BLOCKS) + 1; + const int64_t s31 = mask->nb[1] / sizeof(half); + const int64_t s33 = mask->nb[3] / sizeof(half); + KV_blocks.alloc(size_t(nblocks_KV)*ntiles_x*mask->ne[3]); + + auto launch_blocks = [&](const int jt0, const int ntiles, const int64_t sequence, const int nsequences) { + const int64_t j0 = int64_t(jt0)*ncols1; + const int nq = std::min(Q->ne[1] - j0, int64_t(ntiles)*ncols1); + const dim3 blocks_num_KV(nblocks_KV, ntiles, nsequences); + ggml_cuda_kernel_launch_params launch_params(blocks_num_KV, dim3(WARP_SIZE * FATTN_MASK_KV_BLOCKS, 1, 1), 0, main_stream); + ggml_cuda_kernel_launch(mask_blocks_kernel, launch_params, + (const half *) mask->data + sequence*s33 + j0*s31, + KV_blocks.ptr + (sequence*ntiles_x + jt0)*nblocks_KV, + K->ne[1], nq, nbatch_fa, s31, s33); + CUDA_CHECK(cudaGetLastError()); + }; + + // Bound grid.y and grid.z for HIP runtimes that enforce their limits. + constexpr int max_grid_yz = 65535; + if (ntiles_x <= max_grid_yz && mask->ne[3] <= max_grid_yz) { + launch_blocks(0, ntiles_x, 0, mask->ne[3]); + } else { + for (int64_t sequence = 0; sequence < mask->ne[3];) { + // Query slices need separate sequences to preserve the packed row stride. + const int nsequences = ntiles_x <= max_grid_yz ? std::min(mask->ne[3] - sequence, max_grid_yz) : 1; + for (int jt0 = 0; jt0 < ntiles_x;) { + const int ntiles = std::min(ntiles_x - jt0, max_grid_yz); + launch_blocks(jt0, ntiles, sequence, nsequences); + jt0 += ntiles; + } + sequence += nsequences; + } + } + } + const dim3 block_dim(warp_size, nwarps, 1); int max_blocks_per_sm = 1; // Max. number of active blocks limited by occupancy. CUDA_CHECK(cudaOccupancyMaxActiveBlocksPerMultiprocessor(&max_blocks_per_sm, fattn_kernel, block_dim.x * block_dim.y * block_dim.z, nbytes_shared)); @@ -1242,7 +1340,7 @@ void launch_fattn( V_data, mask ? ((const char *) mask->data) : nullptr, sinks ? ((const char *) sinks->data) : nullptr, - KV_max.ptr, + use_mask_blocks ? (const int *) KV_blocks.ptr : KV_max.ptr, !stream_k && parallel_blocks > 1 ? dst_tmp.ptr : (float *) KQV->data, dst_tmp_meta.ptr, scale, max_bias, m0, m1, n_head_log2, logit_softcap, Q->ne[0], ne01, Q->ne[2], Q->ne[3], Q->nb[1], Q->nb[2], Q->nb[3], diff --git a/ggml/src/ggml-cuda/fattn-mma-f16.cuh b/ggml/src/ggml-cuda/fattn-mma-f16.cuh index 84219ca9259d..cec3e3076a2c 100644 --- a/ggml/src/ggml-cuda/fattn-mma-f16.cuh +++ b/ggml/src/ggml-cuda/fattn-mma-f16.cuh @@ -588,6 +588,7 @@ static __device__ __forceinline__ void flash_attn_ext_f16_iter( const half2 * const __restrict__ V_h2, const half * const __restrict__ mask_h, const int32_t * const __restrict__ indices, + const bool skip_mask, float2 * const __restrict__ dstk, float2 * const __restrict__ dstk_fixup, const float scale, @@ -626,6 +627,7 @@ static __device__ __forceinline__ void flash_attn_ext_f16_iter( constexpr int stride_tile_V = V_is_K_view ? stride_tile_K : ggml_cuda_fattn_mma_get_stride_tile(nbatch_V2, swz); const int k_VKQ_0 = kb0 * nbatch_fa; + const bool use_mask = !skip_mask && (ncols2 > 1 || mask_h); #if defined(TURING_MMA_AVAILABLE) T_C_KQ KQ_C[nbatch_fa/(np*(cols_per_warp == 8 ? T_C_KQ::I : T_C_KQ::J))]; #elif defined(AMD_WMMA_AVAILABLE) || defined(AMD_MFMA_AVAILABLE) @@ -646,7 +648,7 @@ static __device__ __forceinline__ void flash_attn_ext_f16_iter( } else { // the sparse mask values are gathered per element, always load them synchronously constexpr bool use_cp_async = nstages == 1 && !use_sparse; - if (ncols2 > 1 || mask_h) { + if (use_mask) { flash_attn_ext_f16_load_mask (mask_h, tile_mask, stride_mask, k_VKQ_0, k_VKQ_sup, jt*ncols1, ne01, indices); } @@ -746,7 +748,7 @@ static __device__ __forceinline__ void flash_attn_ext_f16_iter( float KQ_rowsum_add[cols_per_thread] = {0.0f}; if constexpr (cols_per_warp == 8) { - if (ncols2 > 1 || mask_h) { + if (use_mask) { #pragma unroll for (int i00 = 0; i00 < nbatch_fa; i00 += np*T_C_KQ::I) { const int i0 = i00 + (threadIdx.y % np)*T_C_KQ::I; @@ -808,7 +810,7 @@ static __device__ __forceinline__ void flash_attn_ext_f16_iter( } } } else { // not Turing mma or T_B_KQ::I > 8 - if (ncols2 > 1 || mask_h) { + if (use_mask) { #pragma unroll for (int i00 = 0; i00 < nbatch_fa; i00 += np*T_C_KQ::J) { const int i0 = i00 + (threadIdx.y % np)*T_C_KQ::J; @@ -994,7 +996,7 @@ static __device__ __forceinline__ void flash_attn_ext_f16_iter( cp_async_wait_all(); __syncthreads(); if (!last_iter) { - if (ncols2 > 1 || mask_h) { + if (use_mask) { flash_attn_ext_f16_load_mask (mask_h, tile_mask, stride_mask, k_VKQ_0 + nbatch_fa, k_VKQ_sup, jt*ncols1, ne01, nullptr); } @@ -1071,7 +1073,7 @@ static __device__ __forceinline__ void flash_attn_ext_f16_iter( } } #else - GGML_UNUSED_VARS(Q_f2, K_h2, V_h2, mask_h, indices, dstk, dstk_fixup, + GGML_UNUSED_VARS(Q_f2, K_h2, V_h2, mask_h, indices, skip_mask, dstk, dstk_fixup, scale, slope, logit_softcap, ne01, ne02, stride_K, stride_V, stride_mask, tile_Q, tile_K, tile_V, tile_mask, @@ -1202,6 +1204,7 @@ static __device__ __forceinline__ void flash_attn_ext_f16_process_tile( const half2 * const __restrict__ V_h2, const half * const __restrict__ mask_h, const int32_t * const __restrict__ indices, + const uint32_t * const __restrict__ KV_blocks, const float * const __restrict__ sinks_f, float2 * const __restrict__ dstk, float2 * const __restrict__ dstk_fixup, @@ -1355,6 +1358,34 @@ static __device__ __forceinline__ void flash_attn_ext_f16_process_tile( (K_h2, tile_K, nbatch_K2, stride_K, kb0*nbatch_fa, k_VKQ_sup, nullptr); } +#if defined(AMD_WMMA_AVAILABLE) + static_assert(nstages == 0 && !use_sparse, "mask runs require synchronous dense attention"); + while (kb0 < kb0_stop) { + const int offset = uint32_t(kb0) % FATTN_MASK_KV_BLOCKS; + const uint32_t bits = uint32_t(__builtin_amdgcn_readfirstlane(int(KV_blocks[uint32_t(kb0) / FATTN_MASK_KV_BLOCKS]))) >> (2*offset); + const uint32_t block_type = bits & 3; + const uint32_t changed = bits ^ (block_type * 0x55555555u); + const int remaining = min(kb0_stop - kb0, FATTN_MASK_KV_BLOCKS - offset); + const int run_stop = kb0 + min(remaining, changed ? (__ffs(changed) - 1)/2 : FATTN_MASK_KV_BLOCKS); + if (block_type == 1) { + kb0 = run_stop; + continue; + } + const bool skip_mask = block_type == 2; + constexpr bool last_iter = false; + constexpr bool oob_check = ncols2 == 1; + for (; kb0 < run_stop; ++kb0) { + const int k_VKQ_sup = oob_check ? min(nbatch_fa, ne11 - kb0*nbatch_fa) : nbatch_fa; + flash_attn_ext_f16_iter + + (Q_f2, K_h2, V_h2, mask_h, indices, skip_mask, dstk, dstk_fixup, scale, slope, logit_softcap, + ne01, ne02, stride_K, stride_V, stride_mask, tile_Q, tile_K, tile_V, tile_mask, Q_B, VKQ_C, + KQ_max, KQ_rowsum, jt, kb0, k_VKQ_sup); + } + } +#else + GGML_UNUSED(KV_blocks); // kb0_start is always < kb0_stop so the last iter can be executed unconditionally. if constexpr (ncols2 == 1 || use_sparse) { constexpr bool oob_check = true; @@ -1364,7 +1395,7 @@ static __device__ __forceinline__ void flash_attn_ext_f16_process_tile( flash_attn_ext_f16_iter - (Q_f2, K_h2, V_h2, mask_h, indices, dstk, dstk_fixup, scale, slope, logit_softcap, + (Q_f2, K_h2, V_h2, mask_h, indices, false, dstk, dstk_fixup, scale, slope, logit_softcap, ne01, ne02, stride_K, stride_V, stride_mask, tile_Q, tile_K, tile_V, tile_mask, Q_B, VKQ_C, KQ_max, KQ_rowsum, jt, kb0, k_VKQ_sup); } @@ -1373,7 +1404,7 @@ static __device__ __forceinline__ void flash_attn_ext_f16_process_tile( flash_attn_ext_f16_iter - (Q_f2, K_h2, V_h2, mask_h, indices, dstk, dstk_fixup, scale, slope, logit_softcap, + (Q_f2, K_h2, V_h2, mask_h, indices, false, dstk, dstk_fixup, scale, slope, logit_softcap, ne01, ne02, stride_K, stride_V, stride_mask, tile_Q, tile_K, tile_V, tile_mask, Q_B, VKQ_C, KQ_max, KQ_rowsum, jt, kb0, k_VKQ_sup); } else { @@ -1384,7 +1415,7 @@ static __device__ __forceinline__ void flash_attn_ext_f16_process_tile( flash_attn_ext_f16_iter - (Q_f2, K_h2, V_h2, mask_h, indices, dstk, dstk_fixup, scale, slope, logit_softcap, + (Q_f2, K_h2, V_h2, mask_h, indices, false, dstk, dstk_fixup, scale, slope, logit_softcap, ne01, ne02, stride_K, stride_V, stride_mask, tile_Q, tile_K, tile_V, tile_mask, Q_B, VKQ_C, KQ_max, KQ_rowsum, jt, kb0, k_VKQ_sup); } @@ -1393,11 +1424,13 @@ static __device__ __forceinline__ void flash_attn_ext_f16_process_tile( flash_attn_ext_f16_iter - (Q_f2, K_h2, V_h2, mask_h, indices, dstk, dstk_fixup, scale, slope, logit_softcap, + (Q_f2, K_h2, V_h2, mask_h, indices, false, dstk, dstk_fixup, scale, slope, logit_softcap, ne01, ne02, stride_K, stride_V, stride_mask, tile_Q, tile_K, tile_V, tile_mask, Q_B, VKQ_C, KQ_max, KQ_rowsum, jt, kb0, k_VKQ_sup); } +#endif + // With multi-stage loading there is no __syncthreads at the end of the iter, // there can be a race condition on shared memory access for combining/writing back results. if constexpr (nstages > 1 && nwarps*cols_per_warp > nbatch_fa) { @@ -1786,7 +1819,7 @@ static __device__ __forceinline__ void flash_attn_ext_f16_process_tile( } } #else - GGML_UNUSED_VARS(Q_f2, K_h2, V_h2, mask_h, indices, sinks_f, dstk, dstk_fixup, + GGML_UNUSED_VARS(Q_f2, K_h2, V_h2, mask_h, indices, KV_blocks, sinks_f, dstk, dstk_fixup, scale, slope, logit_softcap, ne01, ne02, gqa_ratio, stride_Q1, stride_Q2, stride_K, stride_V, stride_mask, jt, kb0_start, kb0_stop); @@ -1810,7 +1843,12 @@ static __global__ void flash_attn_ext_f16( const char * V_ptr, const char * mask_ptr, const char * sinks_ptr, +#if defined(GGML_USE_HIP) || defined(GGML_USE_MUSA) + const int * GGML_CUDA_RESTRICT KV_max_ptr, +#else + // Keep CUDA parameter qualifiers independent of __CUDA_ARCH__. const int * KV_max_ptr, +#endif float * dst_ptr, float2 * dst_meta_ptr, const float scale, @@ -1835,7 +1873,14 @@ static __global__ void flash_attn_ext_f16( const char * GGML_CUDA_RESTRICT sinks = sinks_ptr; // sparse: one index list per (sequence, query tile), the live count of each list follows the lists const int * GGML_CUDA_RESTRICT sparse_indices = use_sparse ? KV_max_ptr : nullptr; - const int * GGML_CUDA_RESTRICT KV_max = KV_max_ptr; +#if defined(AMD_WMMA_AVAILABLE) + static_assert(!use_sparse, "packed mask classes require dense attention"); + // AMD WMMA uses the metadata argument for packed mask classes. + const uint32_t * GGML_CUDA_RESTRICT KV_blocks_ptr = (const uint32_t *) KV_max_ptr; + const int * KV_max = nullptr; +#else + const int * GGML_CUDA_RESTRICT KV_max = KV_max_ptr; +#endif float * GGML_CUDA_RESTRICT dst = dst_ptr; float2 * GGML_CUDA_RESTRICT dst_meta = dst_meta_ptr; @@ -1934,24 +1979,28 @@ static __global__ void flash_attn_ext_f16( const half2 * V_h2 = V_is_K_view ? K_h2 : (const half2 *) (V + nb23*sequence + nb22*z_KV); const float * sinks_f = sinks ? (const float *) sinks + zt_Q : nullptr; const int32_t * indices = use_sparse ? sparse_indices + (int64_t(sequence % ne33)*iter_j + jt)*ne11 : nullptr; + const uint32_t * KV_blocks = nullptr; +#if defined(AMD_WMMA_AVAILABLE) + KV_blocks = KV_blocks_ptr + (int64_t(sequence % ne33)*iter_j + jt)*((iter_k - 1) / FATTN_MASK_KV_BLOCKS + 1); +#endif const float slope = ncols2 == 1 ? get_alibi_slope(max_bias, zt_Q, n_head_log2, m0, m1) : 1.0f; if (use_sparse) { kb0_stop = min(kb0_stop, (KV_max[(sequence % ne33)*iter_j + jt] + nbatch_fa - 1) / nbatch_fa); } else if (KV_max) { - kb0_stop = min(kb0_stop, KV_max[sequence*iter_j + jt] / nbatch_fa); + kb0_stop = min(kb0_stop, KV_max[int64_t(sequence)*iter_j + jt] / nbatch_fa); } constexpr bool is_fixup = false; // All but (potentially) the last iterations write their data to dst rather than the fixup buffer. if (kb0_start == 0) { constexpr bool needs_fixup = false; // CUDA block is working on an entire tile. flash_attn_ext_f16_process_tile - (Q_f2, K_h2, V_h2, mask_h, indices, sinks_f, dstk, dst_meta, scale, slope, logit_softcap, + (Q_f2, K_h2, V_h2, mask_h, indices, KV_blocks, sinks_f, dstk, dst_meta, scale, slope, logit_softcap, ne01, ne02, gqa_ratio, ne11, stride_Q1, stride_Q2, stride_K, stride_V, stride_mask, jt, zt_gqa, kb0_start, kb0_stop); } else { constexpr bool needs_fixup = true; // CUDA block is missing the beginning of a tile. flash_attn_ext_f16_process_tile - (Q_f2, K_h2, V_h2, mask_h, indices, sinks_f, dstk, dst_meta, scale, slope, logit_softcap, + (Q_f2, K_h2, V_h2, mask_h, indices, KV_blocks, sinks_f, dstk, dst_meta, scale, slope, logit_softcap, ne01, ne02, gqa_ratio, ne11, stride_Q1, stride_Q2, stride_K, stride_V, stride_mask, jt, zt_gqa, kb0_start, kb0_stop); } @@ -1983,19 +2032,23 @@ static __global__ void flash_attn_ext_f16( const half2 * V_h2 = V_is_K_view ? K_h2 : (const half2 *) (V + nb23*sequence + nb22*z_KV); const float * sinks_f = sinks ? (const float *) sinks + zt_Q : nullptr; const int32_t * indices = use_sparse ? sparse_indices + (int64_t(sequence % ne33)*iter_j + jt)*ne11 : nullptr; + const uint32_t * KV_blocks = nullptr; +#if defined(AMD_WMMA_AVAILABLE) + KV_blocks = KV_blocks_ptr + (int64_t(sequence % ne33)*iter_j + jt)*((iter_k - 1) / FATTN_MASK_KV_BLOCKS + 1); +#endif const float slope = ncols2 == 1 ? get_alibi_slope(max_bias, zt_Q, n_head_log2, m0, m1) : 1.0f; if (use_sparse) { kb0_stop = min(kb0_stop, (KV_max[(sequence % ne33)*iter_j + jt] + nbatch_fa - 1) / nbatch_fa); } else if (KV_max) { - kb0_stop = min(kb0_stop, KV_max[sequence*iter_j + jt] / nbatch_fa); + kb0_stop = min(kb0_stop, KV_max[int64_t(sequence)*iter_j + jt] / nbatch_fa); } constexpr bool is_fixup = true; // Last index writes its data to fixup buffer to avoid data races with other blocks. constexpr bool needs_fixup = false; flash_attn_ext_f16_process_tile - (Q_f2, K_h2, V_h2, mask_h, indices, sinks_f, dstk, dst_meta, scale, slope, logit_softcap, + (Q_f2, K_h2, V_h2, mask_h, indices, KV_blocks, sinks_f, dstk, dst_meta, scale, slope, logit_softcap, ne01, ne02, gqa_ratio, ne11, stride_Q1, stride_Q2, stride_K, stride_V, stride_mask, jt, zt_gqa, kb0_start, kb0_stop); #else GGML_UNUSED_VARS(Q_ptr, K_ptr, V_ptr, mask_ptr, sinks_ptr, KV_max_ptr, dst_ptr, dst_meta_ptr, scale, @@ -2113,8 +2166,14 @@ void ggml_cuda_flash_attn_ext_mma_f16_case(ggml_backend_cuda_context & ctx, ggml #endif // !defined(GGML_USE_MUSA) } + fattn_mask_blocks_kernel_t mask_blocks_kernel = nullptr; +#if defined(GGML_USE_HIP) + if (amd_wmma_available(cc)) { + mask_blocks_kernel = flash_attn_mask_to_KV_blocks; + } +#endif launch_fattn - (ctx, dst, fattn_kernel, nwarps, nbytes_shared_total, nbatch_fa, true, true, true, use_sparse, warp_size_host); + (ctx, dst, fattn_kernel, nwarps, nbytes_shared_total, nbatch_fa, true, true, true, use_sparse, warp_size_host, mask_blocks_kernel); } diff --git a/tests/test-backend-ops.cpp b/tests/test-backend-ops.cpp index cc68e9ca7bff..aaab32cbe352 100644 --- a/tests/test-backend-ops.cpp +++ b/tests/test-backend-ops.cpp @@ -7754,9 +7754,11 @@ struct test_flash_attn_ext : public test_case { const bool kv_view; // create K/V as views of a larger buffer (like a KV cache) const bool v_is_view_of_k; const int64_t n_kv_max; + const int mask_pattern; + const bool mask_broadcast; std::string vars() override { - return VARS_TO_STR17(hsk, hsv, nh, nr23, kv, nb, mask, sinks, max_bias, logit_softcap, prec, type_K, type_V, permute, kv_view, v_is_view_of_k, n_kv_max); + return VARS_TO_STR17(hsk, hsv, nh, nr23, kv, nb, mask, sinks, max_bias, logit_softcap, prec, type_K, type_V, permute, kv_view, v_is_view_of_k, n_kv_max) + "," + VAR_TO_STR(mask_pattern) + "," + VAR_TO_STR(mask_broadcast); } double max_nmse_err() override { @@ -7773,9 +7775,9 @@ struct test_flash_attn_ext : public test_case { test_flash_attn_ext(int64_t hsk = 128, int64_t hsv = 128, int64_t nh = 32, std::array nr23 = {1, 1}, int64_t kv = 96, int64_t nb = 8, bool mask = true, bool sinks = false, float max_bias = 0.0f, float logit_softcap = 0.0f, ggml_prec prec = GGML_PREC_F32, ggml_type type_K = GGML_TYPE_F16, ggml_type type_V = GGML_TYPE_F16, std::array permute = {0, 1, 2, 3}, - bool kv_view = true, bool v_is_view_of_k = false, int64_t n_kv_max = 0) + bool kv_view = true, bool v_is_view_of_k = false, int64_t n_kv_max = 0, int mask_pattern = 0, bool mask_broadcast = false) : hsk(hsk), hsv(hsv), nh(nh), nr23(nr23), kv(kv), nb(nb), mask(mask), sinks(sinks), max_bias(max_bias), logit_softcap(logit_softcap), prec(prec), - type_K(type_K), type_V(type_V), permute(permute), kv_view(kv_view), v_is_view_of_k(v_is_view_of_k), n_kv_max(n_kv_max) {} + type_K(type_K), type_V(type_V), permute(permute), kv_view(kv_view), v_is_view_of_k(v_is_view_of_k), n_kv_max(n_kv_max), mask_pattern(mask_pattern), mask_broadcast(mask_broadcast) {} ggml_tensor * build_graph(ggml_context * ctx) override { const int64_t hsk_padded = GGML_PAD(hsk, ggml_blck_size(type_K)); @@ -7823,7 +7825,11 @@ struct test_flash_attn_ext : public test_case { ggml_tensor * m = nullptr; if (mask) { - m = ggml_new_tensor_4d(ctx, GGML_TYPE_F16, kv, nb, 1, nr23[1]); + m = ggml_new_tensor_4d(ctx, GGML_TYPE_F16, kv, mask_broadcast ? GGML_PAD(nb, 64) : nb, 1, nr23[1]); + if (mask_broadcast) { + ggml_set_name(m, "m_storage"); + m = ggml_view_4d(ctx, m, kv, nb, 1, 1, m->nb[1], m->nb[2], m->nb[3], 0); + } ggml_set_name(m, "m"); } @@ -7847,8 +7853,50 @@ struct test_flash_attn_ext : public test_case { if (strcmp(t->name, "s") == 0) { // make the sink values more noticeable in order to trigger a test failure when the implementation is wrong init_tensor_uniform(t, -10.0f, 10.0f); + } else if (strcmp(t->name, "m_storage") == 0) { + // Poison the backing storage outside the logical broadcast mask. + std::vector data(ggml_nelements(t), ggml_fp32_to_fp16(-INFINITY)); + ggml_backend_tensor_set(t, data.data(), 0, data.size()*sizeof(ggml_fp16_t)); } else if (strcmp(t->name, "m") == 0) { - if (n_kv_max > 0) { + if (mask_pattern > 0) { + std::vector data(ggml_nelements(t), -INFINITY); + for (int64_t row = 0; row < ggml_nrows(t); ++row) { + if (mask_pattern == 5) { + for (int64_t i = 0; i < t->ne[0]; ++i) { + if (i < t->ne[0]/3 || i >= 2*t->ne[0]/3) { + data[row*t->ne[0] + i] = 0.0f; + } + } + if (row % t->ne[1] == 0) { + const int64_t sequence = row / t->ne[1]; + data[row*t->ne[0] + 17 + (sequence % 2)*(t->ne[0]/4)] = 0.75f; + data[row*t->ne[0] + t->ne[0] - 17] = -0.5f; + } + continue; + } + if (mask_pattern == 4 || (mask_pattern == 3 && row % t->ne[1] == 0)) { + continue; // Fully masked queries with attention sinks. + } + if (mask_pattern == 2) { + data[row*t->ne[0]] = 0.0f; + if (row % t->ne[1] == 0) { + const int64_t sequence = row / t->ne[1]; + const int64_t col = t->ne[0]/2 + 17 - (t->ne[3] > 1 && sequence % 2 == 0)*(t->ne[0]/4); + data[row*t->ne[0] + col] = 0.75f; + } + } else { + const bool tail_only = (mask_pattern == 1 || mask_pattern == 3) && (row / t->ne[1]) % 2 != 0; + for (int64_t i = 0; i < t->ne[0]; ++i) { + if ((!tail_only && i < t->ne[0]/4) || i >= 3*t->ne[0]/4) { + data[row*t->ne[0] + i] = -0.25f; + } + } + } + } + std::vector data_f16(data.size()); + ggml_fp32_to_fp16_row(data.data(), data_f16.data(), data.size()); + ggml_backend_tensor_set(t, data_f16.data(), 0, data_f16.size()*sizeof(ggml_fp16_t)); + } else if (n_kv_max > 0) { init_tensor_kq_mask_sparse(t, n_kv_max); } else { init_tensor_kq_mask(t); @@ -10893,6 +10941,30 @@ static std::vector> make_test_cases_eval() { test_cases.emplace_back(new test_flash_attn_ext(256, 256, 4, {2, 1}, 1024, 32, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_F16, GGML_TYPE_F16)); test_cases.emplace_back(new test_flash_attn_ext(512, 512, 4, {2, 1}, 1024, 4, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_F16, GGML_TYPE_F16)); + // Broadcast one logical mask; backing storage also covers partial query tiles. + for (int64_t hs : {64, 128, 256}) { + for (int64_t nb : {1, 17, 64}) { + test_cases.emplace_back(new test_flash_attn_ext(hs, hs, 2, {8, 2}, 512, nb, true, false, 0, 0, + GGML_PREC_F32, GGML_TYPE_F16, GGML_TYPE_F16, {0, 1, 2, 3}, true, false, 0, 1, true)); + } + } + + // Zero tiles, interior holes, and isolated finite biases cross packed-word boundaries. + for (int64_t hs : {64, 128, 256}) { + test_cases.emplace_back(new test_flash_attn_ext(hs, hs, 2, {8, 2}, 2048, 17, true, false, 0, 0, + GGML_PREC_F32, GGML_TYPE_F16, GGML_TYPE_F16, {0, 1, 2, 3}, true, false, 0, 5)); + } + + // Masked interior tiles, one visible element in an otherwise masked tile, and fully masked queries with sinks. + for (int64_t hs : {64, 128, 256}) { + for (int mask_pattern : {1, 2, 3, 4}) { + test_cases.emplace_back(new test_flash_attn_ext(hs, hs, 2, {8, 2}, 512, 17, true, mask_pattern >= 3, 0, 0, + GGML_PREC_F32, GGML_TYPE_F16, GGML_TYPE_F16, {0, 1, 2, 3}, true, false, 0, mask_pattern)); + } + } + test_cases.emplace_back(new test_flash_attn_ext(256, 256, 4, {6, 1}, 1024, 64, true, false, 0, 0, + GGML_PREC_F32, GGML_TYPE_Q8_0, GGML_TYPE_Q8_0, {0, 1, 2, 3}, true, false, 0, 2)); + test_cases.emplace_back(new test_cross_entropy_loss (GGML_TYPE_F32, { 10, 5, 4, 3})); test_cases.emplace_back(new test_cross_entropy_loss (GGML_TYPE_F32, {30000, 1, 1, 1})); test_cases.emplace_back(new test_cross_entropy_loss_back(GGML_TYPE_F32, { 10, 5, 4, 3}));