diff --git a/ggml/include/ggml.h b/ggml/include/ggml.h index 35f0c44ec421..d38b3ad9fa44 100644 --- a/ggml/include/ggml.h +++ b/ggml/include/ggml.h @@ -2437,6 +2437,11 @@ extern "C" { struct ggml_tensor * a, struct ggml_tensor * sinks); + // provide explicit K/V indices to allow sparse attention calculation + GGML_API void ggml_flash_attn_ext_add_top_k( + struct ggml_tensor * a, + struct ggml_tensor * top_k); + // TODO: needs to be adapted to ggml_flash_attn_ext GGML_API struct ggml_tensor * ggml_flash_attn_back( struct ggml_context * ctx, diff --git a/ggml/src/ggml-cuda/fattn-common.cuh b/ggml/src/ggml-cuda/fattn-common.cuh index e67cc7fdf784..deeb40534c3c 100644 --- a/ggml/src/ggml-cuda/fattn-common.cuh +++ b/ggml/src/ggml-cuda/fattn-common.cuh @@ -18,6 +18,14 @@ // The macro on the following line shifts it by a factor of 2**3=8, as was needed to fix https://github.com/ggml-org/llama.cpp/issues/18606 . #define FATTN_KQ_MAX_OFFSET (3.0f*0.6931f) +// flash attention performance hints +struct fattn_perf_hints { + // top_k values from DSA sparse attention + int32_t * __restrict__ top_k; + // top_k row length, other dimensions are the same as mask + int32_t n_top_k; +}; + typedef void (* fattn_kernel_t)( const char * __restrict__ Q, const char * __restrict__ K, @@ -39,7 +47,8 @@ typedef void (* fattn_kernel_t)( const int32_t nb11, const int32_t nb12, const int64_t nb13, const int32_t nb21, const int32_t nb22, const int64_t nb23, const int32_t ne31, const int32_t ne32, const int32_t ne33, - const int32_t nb31, const int32_t nb32, const int64_t nb33); + const int32_t nb31, const int32_t nb32, const int64_t nb33, + const fattn_perf_hints perf_hints); 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); @@ -972,7 +981,7 @@ static __global__ void flash_attn_combine_results( 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 int warp_size = WARP_SIZE + const int nbatch_fa, const bool need_f16_K, const bool need_f16_V, const bool stream_k, const bool use_top_k, const int warp_size = WARP_SIZE ) { constexpr int ncols = ncols1 * ncols2; @@ -984,6 +993,7 @@ void launch_fattn( const ggml_tensor * mask = dst->src[3]; const ggml_tensor * sinks = dst->src[4]; + const ggml_tensor * top_k = dst->src[5]; ggml_tensor * KQV = dst; @@ -1091,7 +1101,7 @@ 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 (mask && K->ne[1] % FATTN_KQ_STRIDE == 0 && (Q->ne[1] >= 1024 || Q->ne[3] > 1)) { + if (!use_top_k && 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); @@ -1114,7 +1124,7 @@ void launch_fattn( GGML_ASSERT(max_blocks_per_sm > 0); int parallel_blocks = max_blocks_per_sm; - const int ntiles_KV = (K->ne[1] + nbatch_fa - 1) / nbatch_fa; // Max. number of parallel blocks limited by KV cache length. + const int ntiles_KV = ((use_top_k ? top_k->ne[0] : K->ne[1]) + nbatch_fa - 1) / nbatch_fa; // Max. number of parallel blocks limited by KV cache length. dim3 blocks_num; if (stream_k) { @@ -1205,6 +1215,12 @@ void launch_fattn( // TODO other tensor dimensions after removal of WMMA kernel: const uint3 ne01 = init_fastdiv_values(Q->ne[1]); + struct fattn_perf_hints perf_hints = {nullptr, 0}; + if (top_k) { + perf_hints.top_k = (int32_t*) top_k->data; + perf_hints.n_top_k = top_k->ne[0]; + } + GGML_ASSERT(block_dim.x % warp_size == 0); ggml_cuda_kernel_launch_params launch_params = ggml_cuda_kernel_launch_params(blocks_num, block_dim, nbytes_shared, main_stream); @@ -1221,7 +1237,8 @@ void launch_fattn( K->ne[0], K->ne[1], K->ne[2], K->ne[3], nb11, nb12, nb13, nb21, nb22, nb23, mask ? mask->ne[1] : 0, mask ? mask->ne[2] : 0, mask ? mask->ne[3] : 0, - mask ? mask->nb[1] : 0, mask ? mask->nb[2] : 0, mask ? mask->nb[3] : 0 + mask ? mask->nb[1] : 0, mask ? mask->nb[2] : 0, mask ? mask->nb[3] : 0, + perf_hints ); CUDA_CHECK(cudaGetLastError()); diff --git a/ggml/src/ggml-cuda/fattn-mma-f16.cuh b/ggml/src/ggml-cuda/fattn-mma-f16.cuh index 7f4cfd5511ff..3ff994bae4a2 100644 --- a/ggml/src/ggml-cuda/fattn-mma-f16.cuh +++ b/ggml/src/ggml-cuda/fattn-mma-f16.cuh @@ -349,20 +349,20 @@ static __host__ int ggml_cuda_fattn_mma_get_nstages(const int DKQ, const int DV, return cp_async_available(cc) && ncols2 >= 2 ? ggml_cuda_fattn_mma_get_nstages_target(DKQ, DV, ncols1*ncols2, cc) : 0; } -static constexpr __device__ int ggml_cuda_fattn_mma_get_nstages(const int DKQ, const int DV, const int ncols1, const int ncols2) { +static constexpr __device__ int ggml_cuda_fattn_mma_get_nstages(const int DKQ, const int DV, const int ncols1, const int ncols2, const bool use_top_k) { #ifdef CP_ASYNC_AVAILABLE - return ncols2 >= 2 ? ggml_cuda_fattn_mma_get_nstages_target(DKQ, DV, ncols1*ncols2) : 0; + return ncols2 >= 2 && !use_top_k ? ggml_cuda_fattn_mma_get_nstages_target(DKQ, DV, ncols1*ncols2) : 0; #else - GGML_UNUSED_VARS(DKQ, DV, ncols1, ncols2); + GGML_UNUSED_VARS(DKQ, DV, ncols1, ncols2, use_top_k); return 0; #endif // CP_ASYNC_AVAILABLE } // ------------------------------------------------------------------------------------------------------------------ -template +template static __device__ __forceinline__ void flash_attn_ext_f16_load_tile( - const half2 * const __restrict__ KV, half2 * const __restrict__ tile_KV, const int D2, const int stride_KV, const int i_sup) { + const half2 * const __restrict__ KV, half2 * const __restrict__ tile_KV, const int D2, const int stride_KV, const int i_sup, const int32_t * const __restrict__ top_k) { constexpr int warp_size = ggml_cuda_get_physical_warp_size(); // K/V data is loaded with decreasing granularity for D for better memory bandwidth. // The minimum granularity is 16 bytes. @@ -371,6 +371,7 @@ static __device__ __forceinline__ void flash_attn_ext_f16_load_tile( if constexpr (use_cp_async) { static_assert(warp_size == 32, "bad warp_size"); static_assert(!oob_check, "OOB check not compatible with cp_async"); + static_assert(!use_top_k, "top_k not compatible with cp_async"); constexpr int preload = 64; const unsigned int tile_KV_32 = ggml_cuda_cvta_generic_to_shared(tile_KV); @@ -432,8 +433,13 @@ static __device__ __forceinline__ void flash_attn_ext_f16_load_tile( for (int k0 = k0_start; k0 < k0_stop; k0 += stride_k) { const int k = k0 + (stride_k == warp_size ? threadIdx.x : threadIdx.x % stride_k); - ggml_cuda_memcpy_1<16>(tile_KV + i*stride_tile + k*4, - !oob_check || i < i_sup ? KV + i*stride_KV + k*h2_per_chunk : zero); + if constexpr (use_top_k) { + ggml_cuda_memcpy_1<16>(tile_KV + i*stride_tile + k*4, + !oob_check || i < i_sup ? KV + int64_t(top_k[i])*stride_KV + k*h2_per_chunk : zero); + } else { + ggml_cuda_memcpy_1<16>(tile_KV + i*stride_tile + k*4, + !oob_check || i < i_sup ? KV + i*stride_KV + k*h2_per_chunk : zero); + } } } }; @@ -447,14 +453,16 @@ static __device__ __forceinline__ void flash_attn_ext_f16_load_tile( } } -template +template static __device__ __forceinline__ void flash_attn_ext_f16_load_mask( const half * const __restrict__ mask_h, half * const __restrict__ tile_mask, - const int stride_mask, const int i_sup, const int j0, const uint3 ne01) { + const int stride_mask, const int i_sup, const int j0, const uint3 ne01, + const int32_t * const __restrict__ top_k) { constexpr int warp_size = ggml_cuda_get_physical_warp_size(); if constexpr (use_cp_async) { static_assert(nbatch_fa <= 8*warp_size && nbatch_fa % 8 == 0, "bad nbatch_fa"); static_assert(!oob_check, "OOB check incompatible with cp_async"); + static_assert(!use_top_k, "top_k incompatible with cp_async"); constexpr int preload = nbatch_fa >= 32 ? nbatch_fa * sizeof(half) : 64; constexpr int cols_per_warp = 8*warp_size/nbatch_fa; constexpr int stride_j = nwarps * cols_per_warp; @@ -474,7 +482,7 @@ static __device__ __forceinline__ void flash_attn_ext_f16_load_mask( cp_async_cg_16(tile_mask_32 + j_sram*(nbatch_fa*sizeof(half) + 16) + i*sizeof(half), mask_h + int64_t(j_vram)*stride_mask + i); } - } else if constexpr (oob_check) { + } else if constexpr (oob_check || use_top_k) { #pragma unroll for (int j1 = 0; j1 < ncols1; j1 += nwarps) { const int j_sram = j1 + threadIdx.y; @@ -488,7 +496,11 @@ static __device__ __forceinline__ void flash_attn_ext_f16_load_mask( for (int i0 = 0; i0 < nbatch_fa; i0 += warp_size) { const int i = i0 + threadIdx.x; - tile_mask[j_sram*(nbatch_fa + 8) + i] = i < i_sup ? mask_h[int64_t(j_vram)*stride_mask + i] : half(0.0f); + if constexpr (use_top_k) { + tile_mask[j_sram*(nbatch_fa + 8) + i] = i < i_sup ? mask_h[int64_t(j_vram)*stride_mask + top_k[i]] : half(0.0f); + } else { + tile_mask[j_sram*(nbatch_fa + 8) + i] = i < i_sup ? mask_h[int64_t(j_vram)*stride_mask + i] : half(0.0f); + } } } } else if constexpr (nbatch_fa < 2*warp_size) { @@ -528,13 +540,14 @@ static __device__ __forceinline__ void flash_attn_ext_f16_load_mask( } template static __device__ __forceinline__ void flash_attn_ext_f16_iter( const float2 * const __restrict__ Q_f2, const half2 * const __restrict__ K_h2, const half2 * const __restrict__ V_h2, const half * const __restrict__ mask_h, + const int32_t * const __restrict__ top_k, float2 * const __restrict__ dstk, float2 * const __restrict__ dstk_fixup, const float scale, @@ -566,13 +579,16 @@ static __device__ __forceinline__ void flash_attn_ext_f16_iter( constexpr int nbatch_K2 = ggml_cuda_fattn_mma_get_nbatch_K2(DKQ, DV, ncols); constexpr int nbatch_V2 = ggml_cuda_fattn_mma_get_nbatch_V2(DKQ, DV, ncols); constexpr bool Q_in_reg = ggml_cuda_fattn_mma_get_Q_in_reg (DKQ, DV, ncols); - constexpr int nstages = ggml_cuda_fattn_mma_get_nstages (DKQ, DV, ncols1, ncols2); + constexpr int nstages = ggml_cuda_fattn_mma_get_nstages (DKQ, DV, ncols1, ncols2, use_top_k); constexpr int stride_tile_K = nbatch_K2 + 4; constexpr int stride_tile_V = V_is_K_view ? stride_tile_K : nbatch_V2 + 4; const int k_VKQ_0 = kb0 * nbatch_fa; + + const int32_t * const tile_top_k = use_top_k ? top_k + k_VKQ_0 : nullptr; + #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) @@ -583,18 +599,24 @@ static __device__ __forceinline__ void flash_attn_ext_f16_iter( if constexpr (nstages > 1) { static_assert(!oob_check, "OOB check incompatible with multi-stage pipeline"); + static_assert(!use_top_k, "top_k incompatible with multi-stage pipeline"); static_assert(!V_is_K_view, "K data reuse not implemented multi-stage loading"); static_assert(nbatch_K2 == DKQ/2, "batching not implemented for multi stage loading"); constexpr bool use_cp_async = true; cp_async_wait_all(); __syncthreads(); - flash_attn_ext_f16_load_tile - (V_h2 + int64_t(k_VKQ_0)*stride_V, tile_V, nbatch_V2, stride_V, k_VKQ_sup); + flash_attn_ext_f16_load_tile + (V_h2 + int64_t(k_VKQ_0)*stride_V, tile_V, nbatch_V2, stride_V, k_VKQ_sup, nullptr); } else { constexpr bool use_cp_async = nstages == 1; if (ncols2 > 1 || mask_h) { - flash_attn_ext_f16_load_mask - (mask_h + k_VKQ_0, tile_mask, stride_mask, k_VKQ_sup, jt*ncols1, ne01); + if constexpr (use_top_k) { + flash_attn_ext_f16_load_mask + (mask_h, tile_mask, stride_mask, k_VKQ_sup, jt*ncols1, ne01, tile_top_k); + } else { + flash_attn_ext_f16_load_mask + (mask_h + k_VKQ_0, tile_mask, stride_mask, k_VKQ_sup, jt*ncols1, ne01, nullptr); + } } } @@ -607,8 +629,14 @@ static __device__ __forceinline__ void flash_attn_ext_f16_iter( if constexpr (nstages <= 1) { const int k0_diff = k0_stop - k0_start; constexpr bool use_cp_async = nstages == 1; - flash_attn_ext_f16_load_tile - (K_h2 + int64_t(k_VKQ_0)*stride_K + k0_start, tile_K, k0_diff, stride_K, k_VKQ_sup); + if constexpr (use_top_k) { + flash_attn_ext_f16_load_tile + (K_h2 + k0_start, tile_K, k0_diff, stride_K, k_VKQ_sup, tile_top_k); + } else { + flash_attn_ext_f16_load_tile + (K_h2 + int64_t(k_VKQ_0)*stride_K + k0_start, tile_K, k0_diff, stride_K, k_VKQ_sup, nullptr); + } + if (use_cp_async) { cp_async_wait_all(); } @@ -934,17 +962,18 @@ static __device__ __forceinline__ void flash_attn_ext_f16_iter( if constexpr (nstages > 1) { static_assert(!V_is_K_view, "K data reuse not implemented multi-stage loading"); + static_assert(!use_top_k, "top_k not implemented for multi-stage loading"); // Preload K tile for next iteration: constexpr bool use_cp_async = true; cp_async_wait_all(); __syncthreads(); if (!last_iter) { if (ncols2 > 1 || mask_h) { - flash_attn_ext_f16_load_mask - (mask_h + k_VKQ_0 + nbatch_fa, tile_mask, stride_mask, k_VKQ_sup, jt*ncols1, ne01); + flash_attn_ext_f16_load_mask + (mask_h + k_VKQ_0 + nbatch_fa, tile_mask, stride_mask, k_VKQ_sup, jt*ncols1, ne01, nullptr); } - flash_attn_ext_f16_load_tile - (K_h2 + int64_t(k_VKQ_0 + nbatch_fa)*stride_K, tile_K, nbatch_K2, stride_K, k_VKQ_sup); + flash_attn_ext_f16_load_tile + (K_h2 + int64_t(k_VKQ_0 + nbatch_fa)*stride_K, tile_K, nbatch_K2, stride_K, k_VKQ_sup, nullptr); } } @@ -959,8 +988,13 @@ static __device__ __forceinline__ void flash_attn_ext_f16_iter( const int i0_diff = i0_stop - i0_start; if (!V_is_K_view || i0_stop > 2*nbatch_K2) { constexpr bool use_cp_async = nstages == 1; - flash_attn_ext_f16_load_tile - (V_h2 + int64_t(k_VKQ_0)*stride_V + i0_start/2, tile_V, i0_diff/2, stride_V, k_VKQ_sup); + if constexpr (use_top_k) { + flash_attn_ext_f16_load_tile + (V_h2 + i0_start/2, tile_V, i0_diff/2, stride_V, k_VKQ_sup, tile_top_k); + } else { + flash_attn_ext_f16_load_tile + (V_h2 + int64_t(k_VKQ_0)*stride_V + i0_start/2, tile_V, i0_diff/2, stride_V, k_VKQ_sup, nullptr); + } if (use_cp_async) { cp_async_wait_all(); } @@ -1015,11 +1049,11 @@ static __device__ __forceinline__ void flash_attn_ext_f16_iter( } } #else - GGML_UNUSED_VARS(Q_f2, K_h2, V_h2, mask_h, dstk, dstk_fixup, + GGML_UNUSED_VARS(Q_f2, K_h2, V_h2, mask_h, top_k, 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, kb0); + Q_B, VKQ_C, KQ_max, KQ_rowsum, jt, kb0, k_VKQ_sup); NO_DEVICE_CODE; #endif // defined(VOLTA_MMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE) || defined(AMD_MFMA_AVAILABLE) } @@ -1113,12 +1147,13 @@ template struct mma_tile_sizes { }; #endif // defined(TURING_MMA_AVAILABLE) -template +template static __device__ __forceinline__ void flash_attn_ext_f16_process_tile( const float2 * const __restrict__ Q_f2, const half2 * const __restrict__ K_h2, const half2 * const __restrict__ V_h2, const half * const __restrict__ mask_h, + const int32_t * const __restrict__ top_k, const float * const __restrict__ sinks_f, float2 * const __restrict__ dstk, float2 * const __restrict__ dstk_fixup, @@ -1128,7 +1163,7 @@ static __device__ __forceinline__ void flash_attn_ext_f16_process_tile( const uint3 ne01, const int ne02, const int gqa_ratio, - const int ne11, + const int n_kv_iter, const int stride_Q1, const int stride_Q2, const int stride_K, @@ -1158,7 +1193,7 @@ static __device__ __forceinline__ void flash_attn_ext_f16_process_tile( constexpr int nbatch_V2 = ggml_cuda_fattn_mma_get_nbatch_V2 (DKQ, DV, ncols); constexpr int nbatch_combine = ggml_cuda_fattn_mma_get_nbatch_combine(DKQ, DV, ncols); constexpr bool Q_in_reg = ggml_cuda_fattn_mma_get_Q_in_reg (DKQ, DV, ncols); - constexpr int nstages = ggml_cuda_fattn_mma_get_nstages (DKQ, DV, ncols1, ncols2); + constexpr int nstages = ggml_cuda_fattn_mma_get_nstages (DKQ, DV, ncols1, ncols2, use_top_k); if (cols_per_warp > ncols) { NO_DEVICE_CODE; @@ -1257,37 +1292,38 @@ static __device__ __forceinline__ void flash_attn_ext_f16_process_tile( // Preload mask and K data for first iteration when using cp_async with multiple stages: if constexpr (nstages > 1) { + static_assert(!use_top_k, "top_k not implemented for multi-stage loading"); static_assert(nbatch_K2 == DKQ/2, "batching not implemented for multi-stage pipeline"); constexpr bool use_cp_async = true; constexpr bool oob_check = false; constexpr int k_VKQ_sup = nbatch_fa; if (ncols2 > 1 || mask_h) { - flash_attn_ext_f16_load_mask - (mask_h + kb0*nbatch_fa, tile_mask, stride_mask, k_VKQ_sup, jt*ncols1, ne01); + flash_attn_ext_f16_load_mask + (mask_h + kb0*nbatch_fa, tile_mask, stride_mask, k_VKQ_sup, jt*ncols1, ne01, nullptr); } - flash_attn_ext_f16_load_tile - (K_h2 + int64_t(kb0)*nbatch_fa*stride_K, tile_K, nbatch_K2, stride_K, k_VKQ_sup); + flash_attn_ext_f16_load_tile + (K_h2 + int64_t(kb0)*nbatch_fa*stride_K, tile_K, nbatch_K2, stride_K, k_VKQ_sup, nullptr); } // kb0_start is always < kb0_stop so the last iter can be executed unconditionally. - if constexpr (ncols2 == 1) { + if constexpr (ncols2 == 1 || use_top_k) { constexpr bool oob_check = true; for (; kb0 < kb0_stop-1; ++kb0) { constexpr bool last_iter = false; constexpr int k_VKQ_sup = nbatch_fa; flash_attn_ext_f16_iter - - (Q_f2, K_h2, V_h2, mask_h, dstk, dstk_fixup, scale, slope, logit_softcap, + (Q_f2, K_h2, V_h2, mask_h, top_k, 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); } constexpr bool last_iter = true; - const int k_VKQ_sup = ne11 - kb0*nbatch_fa; + const int k_VKQ_sup = n_kv_iter - kb0*nbatch_fa; flash_attn_ext_f16_iter - - (Q_f2, K_h2, V_h2, mask_h, dstk, dstk_fixup, scale, slope, logit_softcap, + (Q_f2, K_h2, V_h2, mask_h, top_k, 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 { @@ -1296,18 +1332,18 @@ static __device__ __forceinline__ void flash_attn_ext_f16_process_tile( constexpr bool last_iter = false; constexpr int k_VKQ_sup = nbatch_fa; flash_attn_ext_f16_iter - - (Q_f2, K_h2, V_h2, mask_h, dstk, dstk_fixup, scale, slope, logit_softcap, + (Q_f2, K_h2, V_h2, mask_h, top_k, 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); } constexpr bool last_iter = true; constexpr int k_VKQ_sup = nbatch_fa; flash_attn_ext_f16_iter - - (Q_f2, K_h2, V_h2, mask_h, dstk, dstk_fixup, scale, slope, logit_softcap, + (Q_f2, K_h2, V_h2, mask_h, top_k, 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); } @@ -1692,15 +1728,15 @@ static __device__ __forceinline__ void flash_attn_ext_f16_process_tile( } } #else - GGML_UNUSED_VARS(Q_f2, K_h2, V_h2, mask_h, sinks_f, dstk, dstk_fixup, - scale, slope, logit_softcap, ne01, ne02, gqa_ratio, + GGML_UNUSED_VARS(Q_f2, K_h2, V_h2, mask_h, top_k, sinks_f, dstk, dstk_fixup, + scale, slope, logit_softcap, ne01, ne02, gqa_ratio, n_kv_iter, stride_Q1, stride_Q2, stride_K, stride_V, stride_mask, - jt, kb0_start, kb0_stop); + jt, zt_gqa, kb0_start, kb0_stop); NO_DEVICE_CODE; #endif // defined(VOLTA_MMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE) || defined(AMD_MFMA_AVAILABLE) } -template +template __launch_bounds__(ggml_cuda_fattn_mma_get_nthreads(DKQ, DV, ncols1*ncols2), ggml_cuda_fattn_mma_get_occupancy(DKQ, DV, ncols1*ncols2)) static __global__ void flash_attn_ext_f16( const char * Q_ptr, @@ -1723,7 +1759,8 @@ static __global__ void flash_attn_ext_f16( const int32_t nb11, const int32_t nb12, const int64_t nb13, const int32_t nb21, const int32_t nb22, const int64_t nb23, const int32_t ne31, const int32_t ne32, const int32_t ne33, - const int32_t nb31, const int32_t nb32, const int64_t nb33) { + const int32_t nb31, const int32_t nb32, const int64_t nb33, + const fattn_perf_hints perf_hints) { ggml_cuda_pdl_sync(); // TODO optimize placement #if defined(FLASH_ATTN_AVAILABLE) && (defined(VOLTA_MMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE) || defined(AMD_MFMA_AVAILABLE)) const char * GGML_CUDA_RESTRICT Q = Q_ptr; @@ -1778,6 +1815,8 @@ static __global__ void flash_attn_ext_f16( constexpr int nthreads = ggml_cuda_fattn_mma_get_nthreads(DKQ, DV, ncols); constexpr int nwarps = nthreads / warp_size; + const int n_kv_iter = use_top_k ? perf_hints.n_top_k : ne11; + const int gqa_ratio = ne02 / ne12; // With grouped query attention there are > 1 Q matrices per K, V matrix. const int stride_Q1 = nb01 / sizeof(float2); @@ -1787,7 +1826,7 @@ static __global__ void flash_attn_ext_f16( const int stride_V = V_is_K_view ? stride_K : nb21 / sizeof(half2); - const int iter_k = (ne11 + (nbatch_fa - 1)) / nbatch_fa; + const int iter_k = (n_kv_iter + (nbatch_fa - 1)) / nbatch_fa; const int iter_j = (ne01.z + (ncols1 - 1)) / ncols1; const int iter_z_gqa = (gqa_ratio + (ncols2 - 1)) / ncols2; @@ -1821,22 +1860,24 @@ 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 * top_k_s = use_top_k ? perf_hints.top_k + ((sequence % ne33) * ne32*ne31 + jt*ncols1) * perf_hints.n_top_k : nullptr; + const float slope = ncols2 == 1 ? get_alibi_slope(max_bias, zt_Q, n_head_log2, m0, m1) : 1.0f; - if (KV_max) { + if (!use_top_k && KV_max) { kb0_stop = min(kb0_stop, KV_max[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, 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); + flash_attn_ext_f16_process_tile + (Q_f2, K_h2, V_h2, mask_h, top_k_s, sinks_f, dstk, dst_meta, scale, slope, logit_softcap, + ne01, ne02, gqa_ratio, n_kv_iter, 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, 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); + flash_attn_ext_f16_process_tile + (Q_f2, K_h2, V_h2, mask_h, top_k_s, sinks_f, dstk, dst_meta, scale, slope, logit_softcap, + ne01, ne02, gqa_ratio, n_kv_iter, stride_Q1, stride_Q2, stride_K, stride_V, stride_mask, jt, zt_gqa, kb0_start, kb0_stop); } kbc += iter_k; @@ -1867,17 +1908,19 @@ 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 * top_k_s = use_top_k ? perf_hints.top_k + ((sequence % ne33) * ne32*ne31 + jt*ncols1) * perf_hints.n_top_k : nullptr; + const float slope = ncols2 == 1 ? get_alibi_slope(max_bias, zt_Q, n_head_log2, m0, m1) : 1.0f; - if (KV_max) { + if (!use_top_k && KV_max) { kb0_stop = min(kb0_stop, KV_max[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, 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); + flash_attn_ext_f16_process_tile + (Q_f2, K_h2, V_h2, mask_h, top_k_s, sinks_f, dstk, dst_meta, scale, slope, logit_softcap, + ne01, ne02, gqa_ratio, n_kv_iter, 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, max_bias, m0, m1, n_head_log2, logit_softcap, @@ -1887,11 +1930,33 @@ static __global__ void flash_attn_ext_f16( nb11, nb12, nb13, nb21, nb22, nb23, ne31, ne32, ne33, - nb31, nb32, nb33); + nb31, nb32, nb33, + perf_hints); NO_DEVICE_CODE; #endif // defined(FLASH_ATTN_AVAILABLE) && (defined(VOLTA_MMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE) || defined(AMD_MFMA_AVAILABLE)) } +constexpr bool ggml_cuda_flash_attn_ext_mma_f16_may_use_top_k(const int DKQ, const int DV, const int ncols1, const int ncols2) { + return (DKQ == 576 && DV == 512 && ncols1 == 1 && ncols2 == 16) || (DKQ == 512 && DV == 512 && ncols1 == 1 && ncols2 == 8); +} + +extern bool ggml_cuda_flash_attn_ext_mma_f16_shall_use_top_k(ggml_backend_cuda_context & ctx, ggml_tensor * dst); + +#if !defined(GGML_USE_MUSA) +#define FATTN_SET_SHARED_MEMORY_LIMIT(kernel, id, nbytes) \ + do { \ + static bool shared_memory_limit_raised[GGML_CUDA_MAX_DEVICES] = {false}; \ + if (!shared_memory_limit_raised[id]) { \ + CUDA_CHECK(cudaFuncSetAttribute( \ + reinterpret_cast(kernel), \ + cudaFuncAttributeMaxDynamicSharedMemorySize, nbytes)); \ + shared_memory_limit_raised[id] = true; \ + } \ + } while (0) +#else +#define FATTN_SET_SHARED_MEMORY_LIMIT(kernel, id, nbytes) +#endif + template void ggml_cuda_flash_attn_ext_mma_f16_case(ggml_backend_cuda_context & ctx, ggml_tensor * dst) { const ggml_tensor * KQV = dst; @@ -1935,32 +2000,35 @@ void ggml_cuda_flash_attn_ext_mma_f16_case(ggml_backend_cuda_context & ctx, ggml using fattn_kernel_ptr_t = fattn_kernel_t; #endif // defined(GGML_USE_HIP) fattn_kernel_t fattn_kernel; + bool top_k_usage = false; + if (logit_softcap == 0.0f) { constexpr bool use_logit_softcap = false; - fattn_kernel = flash_attn_ext_f16; - -#if !defined(GGML_USE_MUSA) - static bool shared_memory_limit_raised[GGML_CUDA_MAX_DEVICES] = {false}; - if (!shared_memory_limit_raised[id]) { - CUDA_CHECK(cudaFuncSetAttribute(reinterpret_cast(fattn_kernel), cudaFuncAttributeMaxDynamicSharedMemorySize, nbytes_shared_total)); - shared_memory_limit_raised[id] = true; + if constexpr (ggml_cuda_flash_attn_ext_mma_f16_may_use_top_k(DKQ, DV, ncols1, ncols2)) { + if (ggml_cuda_flash_attn_ext_mma_f16_shall_use_top_k(ctx, dst)) { + constexpr bool use_top_k = true; + fattn_kernel = flash_attn_ext_f16; + FATTN_SET_SHARED_MEMORY_LIMIT(fattn_kernel, id, nbytes_shared_total); + top_k_usage = true; + } else { + constexpr bool use_top_k = false; + fattn_kernel = flash_attn_ext_f16; + FATTN_SET_SHARED_MEMORY_LIMIT(fattn_kernel, id, nbytes_shared_total); + } + } else { + constexpr bool use_top_k = false; + fattn_kernel = flash_attn_ext_f16; + FATTN_SET_SHARED_MEMORY_LIMIT(fattn_kernel, id, nbytes_shared_total); } -#endif // !defined(GGML_USE_MUSA) } else { constexpr bool use_logit_softcap = true; - fattn_kernel = flash_attn_ext_f16; - -#if !defined(GGML_USE_MUSA) - static bool shared_memory_limit_raised[GGML_CUDA_MAX_DEVICES] = {false}; - if (!shared_memory_limit_raised[id]) { - CUDA_CHECK(cudaFuncSetAttribute(reinterpret_cast(fattn_kernel), cudaFuncAttributeMaxDynamicSharedMemorySize, nbytes_shared_total)); - shared_memory_limit_raised[id] = true; - } -#endif // !defined(GGML_USE_MUSA) + constexpr bool use_top_k = false; + fattn_kernel = flash_attn_ext_f16; + FATTN_SET_SHARED_MEMORY_LIMIT(fattn_kernel, id, nbytes_shared_total); } launch_fattn - (ctx, dst, fattn_kernel, nwarps, nbytes_shared_total, nbatch_fa, true, true, true, warp_size_host); + (ctx, dst, fattn_kernel, nwarps, nbytes_shared_total, nbatch_fa, true, true, true, top_k_usage, warp_size_host); } diff --git a/ggml/src/ggml-cuda/fattn-tile.cuh b/ggml/src/ggml-cuda/fattn-tile.cuh index d1164b8526d3..31ac7a20df5b 100644 --- a/ggml/src/ggml-cuda/fattn-tile.cuh +++ b/ggml/src/ggml-cuda/fattn-tile.cuh @@ -811,7 +811,8 @@ static __global__ void flash_attn_tile( const int32_t nb11, const int32_t nb12, const int64_t nb13, const int32_t nb21, const int32_t nb22, const int64_t nb23, const int32_t ne31, const int32_t ne32, const int32_t ne33, - const int32_t nb31, const int32_t nb32, const int64_t nb33) { + const int32_t nb31, const int32_t nb32, const int64_t nb33, + const fattn_perf_hints perf_hints) { #ifdef FLASH_ATTN_AVAILABLE const char * GGML_CUDA_RESTRICT Q = Q_ptr; const char * GGML_CUDA_RESTRICT K = K_ptr; @@ -1163,7 +1164,7 @@ static void launch_fattn_tile_switch_ncols1(ggml_backend_cuda_context & ctx, ggm const int nbatch_fa = ggml_cuda_fattn_tile_get_nbatch_fa(DKQ, DV, cols_per_block, cc); fattn_kernel_t fattn_kernel = flash_attn_tile; launch_fattn - (ctx, dst, fattn_kernel, nwarps, nbytes_shared, nbatch_fa, true, true, false, warp_size); + (ctx, dst, fattn_kernel, nwarps, nbytes_shared, nbatch_fa, true, true, false, false, warp_size); return; } } @@ -1179,7 +1180,7 @@ static void launch_fattn_tile_switch_ncols1(ggml_backend_cuda_context & ctx, ggm const int nbatch_fa = ggml_cuda_fattn_tile_get_nbatch_fa(DKQ, DV, cols_per_block, cc); fattn_kernel_t fattn_kernel = flash_attn_tile; launch_fattn - (ctx, dst, fattn_kernel, nwarps, nbytes_shared, nbatch_fa, true, true, false, warp_size); + (ctx, dst, fattn_kernel, nwarps, nbytes_shared, nbatch_fa, true, true, false, false, warp_size); return; } } @@ -1191,7 +1192,7 @@ static void launch_fattn_tile_switch_ncols1(ggml_backend_cuda_context & ctx, ggm const int nbatch_fa = ggml_cuda_fattn_tile_get_nbatch_fa(DKQ, DV, cols_per_block, cc); fattn_kernel_t fattn_kernel = flash_attn_tile; launch_fattn - (ctx, dst, fattn_kernel, nwarps, nbytes_shared, nbatch_fa, true, true, false, warp_size); + (ctx, dst, fattn_kernel, nwarps, nbytes_shared, nbatch_fa, true, true, false, false, warp_size); return; } } @@ -1203,7 +1204,7 @@ static void launch_fattn_tile_switch_ncols1(ggml_backend_cuda_context & ctx, ggm const int nbatch_fa = ggml_cuda_fattn_tile_get_nbatch_fa(DKQ, DV, cols_per_block, cc); fattn_kernel_t fattn_kernel = flash_attn_tile; launch_fattn - (ctx, dst, fattn_kernel, nwarps, nbytes_shared, nbatch_fa, true, true, false, warp_size); + (ctx, dst, fattn_kernel, nwarps, nbytes_shared, nbatch_fa, true, true, false, false, warp_size); return; } } @@ -1215,7 +1216,7 @@ static void launch_fattn_tile_switch_ncols1(ggml_backend_cuda_context & ctx, ggm const int nbatch_fa = ggml_cuda_fattn_tile_get_nbatch_fa(DKQ, DV, cols_per_block, cc); fattn_kernel_t fattn_kernel = flash_attn_tile; launch_fattn - (ctx, dst, fattn_kernel, nwarps, nbytes_shared, nbatch_fa, true, true, false, warp_size); + (ctx, dst, fattn_kernel, nwarps, nbytes_shared, nbatch_fa, true, true, false, false, warp_size); return; } } @@ -1226,7 +1227,7 @@ static void launch_fattn_tile_switch_ncols1(ggml_backend_cuda_context & ctx, ggm const int nbatch_fa = ggml_cuda_fattn_tile_get_nbatch_fa(DKQ, DV, cols_per_block, cc); fattn_kernel_t fattn_kernel = flash_attn_tile; launch_fattn - (ctx, dst, fattn_kernel, nwarps, nbytes_shared, nbatch_fa, true, true, false, warp_size); + (ctx, dst, fattn_kernel, nwarps, nbytes_shared, nbatch_fa, true, true, false, false, warp_size); return; } diff --git a/ggml/src/ggml-cuda/fattn-vec.cuh b/ggml/src/ggml-cuda/fattn-vec.cuh index 69dd93686243..c1b08b55f2f6 100644 --- a/ggml/src/ggml-cuda/fattn-vec.cuh +++ b/ggml/src/ggml-cuda/fattn-vec.cuh @@ -39,7 +39,8 @@ static __global__ void flash_attn_ext_vec( const int32_t nb11, const int32_t nb12, const int64_t nb13, const int32_t nb21, const int32_t nb22, const int64_t nb23, const int32_t ne31, const int32_t ne32, const int32_t ne33, - const int32_t nb31, const int32_t nb32, const int64_t nb33) { + const int32_t nb31, const int32_t nb32, const int64_t nb33, + const fattn_perf_hints perf_hints) { ggml_cuda_pdl_lc(); #ifdef FLASH_ATTN_AVAILABLE const char * GGML_CUDA_RESTRICT Q = Q_ptr; @@ -540,7 +541,7 @@ void ggml_cuda_flash_attn_ext_vec_case_impl(ggml_backend_cuda_context & ctx, ggm const bool need_f16_K = type_K == GGML_TYPE_F16; const bool need_f16_V = type_V == GGML_TYPE_F16; constexpr size_t nbytes_shared = 0; - launch_fattn(ctx, dst, fattn_kernel, nwarps, nbytes_shared, D, need_f16_K, need_f16_V, false); + launch_fattn(ctx, dst, fattn_kernel, nwarps, nbytes_shared, D, need_f16_K, need_f16_V, false, false); } template diff --git a/ggml/src/ggml-cuda/fattn.cu b/ggml/src/ggml-cuda/fattn.cu index ab7a3b297c07..dc6f8fcbe1d5 100644 --- a/ggml/src/ggml-cuda/fattn.cu +++ b/ggml/src/ggml-cuda/fattn.cu @@ -5,11 +5,28 @@ #include "fattn-vec.cuh" #include "fattn.cuh" +bool ggml_cuda_flash_attn_ext_mma_f16_shall_use_top_k(ggml_backend_cuda_context & ctx, ggml_tensor * dst) { + GGML_UNUSED_VARS(ctx); + const ggml_tensor * Q = dst->src[0]; + const ggml_tensor * K = dst->src[1]; + const ggml_tensor * top_k = dst->src[5]; + // TODO better tuning of when to use top_k optimization on various hardware + return ((Q->ne[1] == 1 && K->ne[1] >= 4096) || K->ne[1] >= 8192) && top_k != nullptr; +} + template static void ggml_cuda_flash_attn_ext_mma_f16_switch_ncols1(ggml_backend_cuda_context & ctx, ggml_tensor * dst) { const int cc = ggml_cuda_info().devices[ggml_cuda_get_device()].cc; const ggml_tensor * Q = dst->src[0]; + // for DeepSeek V3.2 sparse top-k attention force Q tiles to always correspond to only 1 token (ncols1 = 1) + if constexpr (ggml_cuda_flash_attn_ext_mma_f16_may_use_top_k(DKQ, DV, 1, ncols2)) { + if (ggml_cuda_flash_attn_ext_mma_f16_shall_use_top_k(ctx, dst)) { + ggml_cuda_flash_attn_ext_mma_f16_case(ctx, dst); + return; + } + } + if constexpr (ncols2 <= 8) { if (turing_mma_available(cc) && Q->ne[1] <= 8/ncols2) { ggml_cuda_flash_attn_ext_mma_f16_case(ctx, dst); diff --git a/ggml/src/ggml.c b/ggml/src/ggml.c index a7d1fe7d94be..74d76d96266d 100644 --- a/ggml/src/ggml.c +++ b/ggml/src/ggml.c @@ -5479,6 +5479,25 @@ void ggml_flash_attn_ext_add_sinks( a->src[4] = sinks; } +void ggml_flash_attn_ext_add_top_k( + struct ggml_tensor * a, + struct ggml_tensor * top_k) { + if (!top_k) { + a->src[5] = NULL; + return; + } + + GGML_ASSERT(a->op == GGML_OP_FLASH_ATTN_EXT); + GGML_ASSERT(a->src[5] == NULL); + GGML_ASSERT(a->src[3] != NULL); + GGML_ASSERT(a->src[3]->ne[1] == top_k->ne[1]); + GGML_ASSERT(a->src[3]->ne[2] == top_k->ne[2]); + GGML_ASSERT(a->src[3]->ne[3] == top_k->ne[3]); + GGML_ASSERT(top_k->type == GGML_TYPE_I32); + + a->src[5] = top_k; +} + // ggml_flash_attn_back struct ggml_tensor * ggml_flash_attn_back( diff --git a/src/llama-graph.cpp b/src/llama-graph.cpp index c8ecb0a2854c..86d2a61dcf07 100644 --- a/src/llama-graph.cpp +++ b/src/llama-graph.cpp @@ -2389,6 +2389,7 @@ ggml_tensor * llm_graph_context::build_attn_mha( ggml_tensor * kq_mask, ggml_tensor * sinks, ggml_tensor * v_mla, + ggml_tensor * top_k, float kq_scale, int il) const { const bool v_trans = v->nb[1] > v->nb[2]; @@ -2426,6 +2427,7 @@ ggml_tensor * llm_graph_context::build_attn_mha( res->add_fused_node({LLM_FUSED_OP_FLASH_ATTN, cur, il}); ggml_flash_attn_ext_add_sinks(cur, sinks); + ggml_flash_attn_ext_add_top_k(cur, top_k); ggml_flash_attn_ext_set_prec (cur, GGML_PREC_F32); if (v_mla) { @@ -2575,7 +2577,7 @@ ggml_tensor * llm_graph_context::build_attn( ggml_tensor * k = k_cur; ggml_tensor * v = v_cur; - ggml_tensor * cur = build_attn_mha(q, k, v, kq_b, kq_mask, sinks, v_mla, kq_scale, il); + ggml_tensor * cur = build_attn_mha(q, k, v, kq_b, kq_mask, sinks, v_mla, nullptr, kq_scale, il); cb(cur, "kqv_out", il); if (wo) { @@ -2674,7 +2676,7 @@ ggml_tensor * llm_graph_context::build_attn( ggml_tensor * k = mctx_cur->get_k(ctx0, il); ggml_tensor * v = mctx_cur->get_v(ctx0, il); - ggml_tensor * cur = build_attn_mha(q, k, v, kq_b, kq_mask, sinks, v_mla, kq_scale, il); + ggml_tensor * cur = build_attn_mha(q, k, v, kq_b, kq_mask, sinks, v_mla, nullptr, kq_scale, il); cb(cur, "kqv_out", il); if (inp->self_v_rot) { @@ -2765,7 +2767,7 @@ ggml_tensor * llm_graph_context::build_attn( ggml_tensor * k = mctx_cur->get_k(ctx0, il); ggml_tensor * v = ggml_view_4d(ctx0, k, v_cur->ne[0], k->ne[1], k->ne[2], k->ne[3], k->nb[1], k->nb[2], k->nb[3], 0); - ggml_tensor * cur = build_attn_mha(q, k, v, kq_b, kq_mask, sinks, v_mla, kq_scale, il); + ggml_tensor * cur = build_attn_mha(q, k, v, kq_b, kq_mask, sinks, v_mla, nullptr, kq_scale, il); cb(cur, "kqv_out", il); if (wo) { @@ -2850,7 +2852,7 @@ ggml_tensor * llm_graph_context::build_attn( ggml_tensor * k = mctx_cur->get_k(ctx0, il); ggml_tensor * v = ggml_view_4d(ctx0, k, v_cur->ne[0], k->ne[1], k->ne[2], k->ne[3], k->nb[1], k->nb[2], k->nb[3], 0); - ggml_tensor * cur = build_attn_mha(q, k, v, kq_b, kq_mask_top_k, sinks, v_mla, kq_scale, il); + ggml_tensor * cur = build_attn_mha(q, k, v, kq_b, kq_mask_top_k, sinks, v_mla, top_k, kq_scale, il); cb(cur, "kqv_out", il); if (wo) { @@ -2929,7 +2931,7 @@ ggml_tensor * llm_graph_context::build_attn( ggml_tensor * k = mctx_cur->get_k(ctx0, il); ggml_tensor * v = mctx_cur->get_v(ctx0, il); - ggml_tensor * cur = build_attn_mha(q, k, v, kq_b, kq_mask, sinks, v_mla, kq_scale, il); + ggml_tensor * cur = build_attn_mha(q, k, v, kq_b, kq_mask, sinks, v_mla, nullptr, kq_scale, il); cb(cur, "kqv_out", il); if (v_rot) { @@ -2992,7 +2994,7 @@ ggml_tensor * llm_graph_context::build_attn( ggml_tensor * k = k_cur; ggml_tensor * v = v_cur; - ggml_tensor * cur = build_attn_mha(q, k, v, kq_b, kq_mask, sinks, v_mla, kq_scale, il); + ggml_tensor * cur = build_attn_mha(q, k, v, kq_b, kq_mask, sinks, v_mla, nullptr, kq_scale, il); cb(cur, "kqv_out", il); if (wo) { diff --git a/src/llama-graph.h b/src/llama-graph.h index 7ed490ce6728..9f193c13788c 100644 --- a/src/llama-graph.h +++ b/src/llama-graph.h @@ -1073,6 +1073,7 @@ struct llm_graph_context { ggml_tensor * kq_mask, ggml_tensor * sinks, // [n_head_q] ggml_tensor * v_mla, // [n_embd_head_v_mla, n_embd_head_v, n_head_v] + ggml_tensor * top_k, float kq_scale, int il) const; diff --git a/src/models/deepseek4.cpp b/src/models/deepseek4.cpp index 5ad6473ce203..631f93f35c7b 100644 --- a/src/models/deepseek4.cpp +++ b/src/models/deepseek4.cpp @@ -698,7 +698,24 @@ ggml_tensor * llama_model_deepseek4::graph::build_csa_lid_attention( ggml_tensor * kq_mask = ggml_concat(ctx0, raw_mask, csa_mask, 0); cb(kq_mask, "csa_lid_kq_mask", il); - ggml_tensor * out = build_attn_mha(q, k_all, k_all, nullptr, kq_mask, sinks, nullptr, kq_scale, il); + ggml_tensor * top_k_all = nullptr; + if (csa_k->ne[2] >= 8192) { + // recover indices of non-masked SWA cells from raw_mask + ggml_tensor * raw_mask_f32 = ggml_cast(ctx0, raw_mask, GGML_TYPE_F32); + ggml_tensor * top_k_swa = ggml_cont(ctx0, ggml_top_k(ctx0, raw_mask_f32, GGML_PAD(hparams.n_swa, 256))); + cb(top_k_swa, "top_k_swa", il); + + // offset values of top_k by raw_mask row length + ggml_tensor * top_k_f32 = ggml_cast(ctx0, top_k, GGML_TYPE_F32); + top_k_f32 = ggml_scale_bias(ctx0, top_k_f32, 1.0f, 1.0f*raw_mask->ne[0]); + ggml_tensor * top_k_offset = ggml_cast(ctx0, top_k_f32, GGML_TYPE_I32); + cb(top_k_offset, "top_k_offset", il); + + // merge top_k_swa and top_k_offset + top_k_all = ggml_concat(ctx0, top_k_swa, top_k_offset, 0); + } + + ggml_tensor * out = build_attn_mha(q, k_all, k_all, nullptr, kq_mask, sinks, nullptr, top_k_all, kq_scale, il); if (k_rot) { out = llama_mul_mat_hadamard(ctx0, out, k_rot); } @@ -753,7 +770,7 @@ ggml_tensor * llama_model_deepseek4::graph::build_hca_attention( ggml_tensor * kq_mask = ggml_concat(ctx0, raw_mask, hca_mask, 0); cb(kq_mask, "hca_kq_mask", il); - ggml_tensor * out = build_attn_mha(q, k_all, k_all, nullptr, kq_mask, sinks, nullptr, kq_scale, il); + ggml_tensor * out = build_attn_mha(q, k_all, k_all, nullptr, kq_mask, sinks, nullptr, nullptr, kq_scale, il); if (k_rot) { out = llama_mul_mat_hadamard(ctx0, out, k_rot); } @@ -789,7 +806,7 @@ ggml_tensor * llama_model_deepseek4::graph::build_raw_attention( ggml_tensor * k = mctx_cur->get_k(ctx0, il); - ggml_tensor * out = build_attn_mha(q, k, k, nullptr, kq_mask, sinks, nullptr, kq_scale, il); + ggml_tensor * out = build_attn_mha(q, k, k, nullptr, kq_mask, sinks, nullptr, nullptr, kq_scale, il); if (k_rot) { out = llama_mul_mat_hadamard(ctx0, out, k_rot); } diff --git a/tests/test-backend-ops.cpp b/tests/test-backend-ops.cpp index e7cd6d0cb668..c28688628117 100644 --- a/tests/test-backend-ops.cpp +++ b/tests/test-backend-ops.cpp @@ -221,6 +221,53 @@ static void init_tensor_tril(ggml_tensor * tensor, float min = -1.0f, float max ggml_backend_tensor_set(tensor, data_f32.data(), 0, ggml_nbytes(tensor)); } +static void init_tensor_top_k(ggml_tensor * tensor, int64_t kv, uint32_t top_k_seed) { + const int64_t n_top_k = tensor->ne[0]; + const int64_t nb = tensor->ne[1]; + const int64_t nr3 = tensor->ne[3]; + + std::vector top_k_row(kv); + for(int i = 0; i < kv; ++i) { + top_k_row[i] = i; + } + std::mt19937 rng(top_k_seed); + std::vector top_k_data(n_top_k * nb * 1 * nr3); + if (n_top_k > 0) { + for (int r = 0; r < nb * 1 * nr3; ++r) { + std::shuffle(std::begin(top_k_row), std::end(top_k_row), rng); + memcpy(&top_k_data[r * n_top_k], top_k_row.data(), n_top_k * sizeof(int)); + } + } + ggml_backend_tensor_set(tensor, top_k_data.data(), 0, ggml_nelements(tensor) * sizeof(int)); +} + +static void init_tensor_kq_mask_top_k(ggml_tensor * tensor, int64_t n_top_k, uint32_t top_k_seed) { + const int64_t kv = tensor->ne[0]; + const int64_t nb = tensor->ne[1]; + const int64_t nr3 = tensor->ne[3]; + + std::vector top_k_row(kv); + for(int i = 0; i < kv; ++i) { + top_k_row[i] = i; + } + std::mt19937 rng(top_k_seed); + + // unmask only top_k indices + std::vector data_f32(ggml_nelements(tensor)); + std::vector data_f16(ggml_nelements(tensor)); + for (int i = 0; i < ggml_nelements(tensor); ++i) { + data_f32[i] = -INFINITY; + } + for (int r = 0; r < nb * 1 * nr3; ++r) { + std::shuffle(std::begin(top_k_row), std::end(top_k_row), rng); + for (int i = 0; i < n_top_k; ++i) { + data_f32[r * kv + top_k_row[i]] = 0; + } + } + ggml_fp32_to_fp16_row(data_f32.data(), data_f16.data(), ggml_nelements(tensor)); + ggml_backend_tensor_set(tensor, data_f16.data(), 0, ggml_nelements(tensor) * sizeof(ggml_fp16_t)); +} + static std::vector tensor_to_float(const ggml_tensor * t) { std::vector tv; tv.reserve(ggml_nelements(t)); @@ -6799,6 +6846,7 @@ struct test_flash_attn_ext : public test_case { const std::array nr23; // repeat in dim 2 and 3, tests for grouped-query attention const int64_t kv; // kv size const int64_t nb; // batch size + const int64_t n_top_k; // top_k row length const bool mask; // use mask const bool sinks; // use sinks @@ -6811,8 +6859,10 @@ struct test_flash_attn_ext : public test_case { const ggml_type type_V; std::array permute; + uint32_t top_k_seed; + std::string vars() override { - return VARS_TO_STR14(hsk, hsv, nh, nr23, kv, nb, mask, sinks, max_bias, logit_softcap, prec, type_K, type_V, permute); + return VARS_TO_STR15(hsk, hsv, nh, nr23, kv, nb, n_top_k, mask, sinks, max_bias, logit_softcap, prec, type_K, type_V, permute); } double max_nmse_err() override { @@ -6826,11 +6876,11 @@ struct test_flash_attn_ext : public test_case { return (2 * nh*nr23[0] * nb * (hsk + hsv) * kv)*nr23[1]; } - 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, + 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, int64_t n_top_k = 0, 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}) - : 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) {} + : hsk(hsk), hsv(hsv), nh(nh), nr23(nr23), kv(kv), nb(nb), n_top_k(n_top_k), mask(mask), sinks(sinks), max_bias(max_bias), logit_softcap(logit_softcap), prec(prec), + type_K(type_K), type_V(type_V), permute(permute), top_k_seed(time(NULL)) {} ggml_tensor * build_graph(ggml_context * ctx) override { const int64_t hsk_padded = GGML_PAD(hsk, ggml_blck_size(type_K)); @@ -6882,6 +6932,12 @@ struct test_flash_attn_ext : public test_case { ggml_set_name(m, "m"); } + ggml_tensor * top_k = nullptr; + if (n_top_k > 0) { + top_k = ggml_new_tensor_4d(ctx, GGML_TYPE_I32, n_top_k, nb, 1, nr23[1]); + ggml_set_name(top_k, "top_k"); + } + ggml_tensor * s = nullptr; if (sinks) { s = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, q->ne[2]); @@ -6890,6 +6946,7 @@ struct test_flash_attn_ext : public test_case { ggml_tensor * out = ggml_flash_attn_ext(ctx, q, k, v, m, 1.0f/sqrtf(hsk), max_bias, logit_softcap); ggml_flash_attn_ext_add_sinks(out, s); + ggml_flash_attn_ext_add_top_k(out, top_k); ggml_flash_attn_ext_set_prec (out, prec); ggml_set_name(out, "out"); @@ -6902,7 +6959,13 @@ struct test_flash_attn_ext : public test_case { // 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") == 0) { - init_tensor_kq_mask(t); + if (n_top_k > 0) { + init_tensor_kq_mask_top_k(t, n_top_k, top_k_seed); + } else { + init_tensor_kq_mask(t); + } + } else if (strcmp(t->name, "top_k") == 0) { + init_tensor_top_k(t, kv, top_k_seed); } else { init_tensor_uniform(t); } @@ -9513,11 +9576,11 @@ static std::vector> make_test_cases_eval() { for (ggml_type type_KV : {GGML_TYPE_F32, GGML_TYPE_F16, GGML_TYPE_BF16, GGML_TYPE_Q8_0, GGML_TYPE_Q5_1, GGML_TYPE_Q5_0, GGML_TYPE_Q4_1, GGML_TYPE_Q4_0, GGML_TYPE_IQ4_NL}) { if (type_KV != GGML_TYPE_F16 && hsk != 64 && hsk != 72) continue; test_cases.emplace_back(new test_flash_attn_ext( - hsk, hsv, nh, {nr2, nr3}, kv, nb, mask, sinks, max_bias, logit_softcap, prec, type_KV, type_KV)); + hsk, hsv, nh, {nr2, nr3}, kv, nb, 0, mask, sinks, max_bias, logit_softcap, prec, type_KV, type_KV)); // run fewer test cases permuted if (mask == true && max_bias == 0.0f && logit_softcap == 0 && kv == 512) { test_cases.emplace_back(new test_flash_attn_ext( - hsk, hsv, nh, {nr2, nr3}, kv, nb, mask, sinks, max_bias, logit_softcap, prec, type_KV, type_KV, {0, 2, 1, 3})); + hsk, hsv, nh, {nr2, nr3}, kv, nb, 0, mask, sinks, max_bias, logit_softcap, prec, type_KV, type_KV, {0, 2, 1, 3})); } } } @@ -9534,15 +9597,25 @@ static std::vector> make_test_cases_eval() { } // mixed quant and Q1_0 test cases - test_cases.emplace_back(new test_flash_attn_ext(64, 64, 4, {1, 1}, 128, 2, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_Q8_0, GGML_TYPE_Q4_0)); - test_cases.emplace_back(new test_flash_attn_ext(64, 64, 4, {1, 1}, 128, 2, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_Q4_0, GGML_TYPE_F16)); - test_cases.emplace_back(new test_flash_attn_ext(72, 72, 4, {1, 1}, 96, 2, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_Q4_0, GGML_TYPE_Q8_0)); - test_cases.emplace_back(new test_flash_attn_ext(64, 64, 4, {1, 1}, 96, 2, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_F16, GGML_TYPE_F32)); - test_cases.emplace_back(new test_flash_attn_ext(128, 128, 4, {1, 1}, 256, 1, false, false, 0, 0, GGML_PREC_F32, GGML_TYPE_F16, GGML_TYPE_Q4_0)); - test_cases.emplace_back(new test_flash_attn_ext(128, 128, 4, {1, 1}, 96, 2, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_Q1_0, GGML_TYPE_Q1_0)); - test_cases.emplace_back(new test_flash_attn_ext(128, 64, 4, {1, 1}, 128, 2, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_Q1_0, GGML_TYPE_Q4_0)); - test_cases.emplace_back(new test_flash_attn_ext(64, 128, 4, {1, 1}, 128, 2, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_Q4_0, GGML_TYPE_Q1_0)); - test_cases.emplace_back(new test_flash_attn_ext(128, 64, 4, {1, 1}, 64, 2, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_Q1_0, GGML_TYPE_F16)); + test_cases.emplace_back(new test_flash_attn_ext(64, 64, 4, {1, 1}, 128, 2, 0, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_Q8_0, GGML_TYPE_Q4_0)); + test_cases.emplace_back(new test_flash_attn_ext(64, 64, 4, {1, 1}, 128, 2, 0, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_Q4_0, GGML_TYPE_F16)); + test_cases.emplace_back(new test_flash_attn_ext(72, 72, 4, {1, 1}, 96, 2, 0, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_Q4_0, GGML_TYPE_Q8_0)); + test_cases.emplace_back(new test_flash_attn_ext(64, 64, 4, {1, 1}, 96, 2, 0, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_F16, GGML_TYPE_F32)); + test_cases.emplace_back(new test_flash_attn_ext(128, 128, 4, {1, 1}, 256, 1, 0, false, false, 0, 0, GGML_PREC_F32, GGML_TYPE_F16, GGML_TYPE_Q4_0)); + test_cases.emplace_back(new test_flash_attn_ext(128, 128, 4, {1, 1}, 96, 2, 0, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_Q1_0, GGML_TYPE_Q1_0)); + test_cases.emplace_back(new test_flash_attn_ext(128, 64, 4, {1, 1}, 128, 2, 0, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_Q1_0, GGML_TYPE_Q4_0)); + test_cases.emplace_back(new test_flash_attn_ext(64, 128, 4, {1, 1}, 128, 2, 0, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_Q4_0, GGML_TYPE_Q1_0)); + test_cases.emplace_back(new test_flash_attn_ext(128, 64, 4, {1, 1}, 64, 2, 0, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_Q1_0, GGML_TYPE_F16)); + + test_cases.emplace_back(new test_flash_attn_ext(512, 512, 1, {128, 1}, 131072, 1, 2048, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_F16, GGML_TYPE_F16)); + test_cases.emplace_back(new test_flash_attn_ext(512, 512, 1, {128, 10}, 131072, 1, 2048, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_F16, GGML_TYPE_F16)); + test_cases.emplace_back(new test_flash_attn_ext(512, 512, 1, {128, 1}, 131072, 2, 2048, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_F16, GGML_TYPE_F16)); + test_cases.emplace_back(new test_flash_attn_ext(512, 512, 1, {128, 1}, 131072, 512, 2048, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_F16, GGML_TYPE_F16)); + + test_cases.emplace_back(new test_flash_attn_ext(576, 512, 1, {128, 1}, 131072, 1, 2048, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_F16, GGML_TYPE_F16)); + test_cases.emplace_back(new test_flash_attn_ext(576, 512, 1, {128, 10}, 131072, 1, 2048, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_F16, GGML_TYPE_F16)); + test_cases.emplace_back(new test_flash_attn_ext(576, 512, 1, {128, 1}, 131072, 2, 2048, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_F16, GGML_TYPE_F16)); + test_cases.emplace_back(new test_flash_attn_ext(576, 512, 1, {128, 1}, 131072, 512, 2048, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_F16, GGML_TYPE_F16)); 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})); @@ -9870,23 +9943,28 @@ static std::vector> make_test_cases_perf() { } // Qwen3-VL-8B https://github.com/ggml-org/llama.cpp/issues/17012 - test_cases.emplace_back(new test_flash_attn_ext(72, 72, 16, {1, 1}, 5776, 5776, false, false, 0, 0, GGML_PREC_F32, GGML_TYPE_F16, GGML_TYPE_F16)); + test_cases.emplace_back(new test_flash_attn_ext(72, 72, 16, {1, 1}, 5776, 5776, 0, false, false, 0, 0, GGML_PREC_F32, GGML_TYPE_F16, GGML_TYPE_F16)); - test_cases.emplace_back(new test_flash_attn_ext(64, 64, 8, {8, 1}, 7680, 1, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_F16, GGML_TYPE_F16)); - test_cases.emplace_back(new test_flash_attn_ext(64, 64, 8, {8, 1}, 7680, 4, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_F16, GGML_TYPE_F16)); - test_cases.emplace_back(new test_flash_attn_ext(64, 64, 8, {8, 1}, 7680, 1, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_Q4_0, GGML_TYPE_Q4_0)); - test_cases.emplace_back(new test_flash_attn_ext(64, 64, 8, {8, 1}, 7680, 512, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_Q4_0, GGML_TYPE_Q4_0)); - test_cases.emplace_back(new test_flash_attn_ext(64, 64, 8, {8, 1}, 7680, 1, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_Q8_0, GGML_TYPE_Q8_0)); - test_cases.emplace_back(new test_flash_attn_ext(64, 64, 8, {8, 1}, 7680, 512, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_Q8_0, GGML_TYPE_Q8_0)); + test_cases.emplace_back(new test_flash_attn_ext(64, 64, 8, {8, 1}, 7680, 1, 0, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_F16, GGML_TYPE_F16)); + test_cases.emplace_back(new test_flash_attn_ext(64, 64, 8, {8, 1}, 7680, 4, 0, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_F16, GGML_TYPE_F16)); + test_cases.emplace_back(new test_flash_attn_ext(64, 64, 8, {8, 1}, 7680, 1, 0, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_Q4_0, GGML_TYPE_Q4_0)); + test_cases.emplace_back(new test_flash_attn_ext(64, 64, 8, {8, 1}, 7680, 512, 0, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_Q4_0, GGML_TYPE_Q4_0)); + test_cases.emplace_back(new test_flash_attn_ext(64, 64, 8, {8, 1}, 7680, 1, 0, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_Q8_0, GGML_TYPE_Q8_0)); + test_cases.emplace_back(new test_flash_attn_ext(64, 64, 8, {8, 1}, 7680, 512, 0, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_Q8_0, GGML_TYPE_Q8_0)); for (int kv : { 4096, 8192, 16384, }) { for (int hs : { 64, 128, }) { for (int nr : { 1, 4, }) { - test_cases.emplace_back(new test_flash_attn_ext(hs, hs, 8, {nr, 1}, kv, 1, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_F16, GGML_TYPE_F16)); + test_cases.emplace_back(new test_flash_attn_ext(hs, hs, 8, {nr, 1}, kv, 1, 0, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_F16, GGML_TYPE_F16)); } } } + for (int kv : { 4096, 8192, 16384, 32768, 65536, 131072 }) { + test_cases.emplace_back(new test_flash_attn_ext(576, 512, 1, {128, 1}, kv, 1, 2048, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_F16, GGML_TYPE_F16)); + test_cases.emplace_back(new test_flash_attn_ext(576, 512, 1, {128, 1}, kv, 1, 0, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_F16, GGML_TYPE_F16)); + } + for (int col : {8192, 16384, 32768, 65536, 131072, 262144, 524288}) { for (int rows : {1, 4, 16}){ test_cases.emplace_back(new test_soft_max(GGML_TYPE_F32, {col, rows, 1, 1}, false, false, GGML_TYPE_F32, {1, 1}, 1.0f, 0.0f));