From d6cab1248ebb031eaffc491eecbe7b09aaade6a5 Mon Sep 17 00:00:00 2001 From: Gaetan Puleo <12990773+gaetan-puleo@users.noreply.github.com> Date: Thu, 10 Sep 2026 13:01:37 +0200 Subject: [PATCH] ggml: optimize Flash Next inference paths --- ggml/src/ggml-cpu/ops.cpp | 8 +- ggml/src/ggml-cuda/binbcast.cu | 2 +- ggml/src/ggml-cuda/fattn-common.cuh | 118 +++++++++- ggml/src/ggml-cuda/fattn-masked-qk-layout.h | 54 +++++ ggml/src/ggml-cuda/fattn-mma-f16.cuh | 154 +++++++++++-- ggml/src/ggml-cuda/fattn.cu | 4 +- ggml/src/ggml-cuda/ggml-cuda.cu | 41 ++++ ggml/src/ggml-cuda/moe-weighted-reduction.cu | 31 ++- ggml/src/ggml-cuda/unary.cu | 50 +++++ ggml/src/ggml-cuda/unary.cuh | 4 + ggml/src/ggml-vulkan/ggml-vulkan.cpp | 207 ++++++++++++++++- .../vulkan-shaders/count_experts.comp | 11 +- .../vulkan-shaders/flash_attn_base.glsl | 9 + .../vulkan-shaders/sigmoid_mul_f32.comp | 28 +++ .../vulkan-shaders/vulkan-shaders-gen.cpp | 1 + src/llama-context.cpp | 4 + tests/CMakeLists.txt | 8 + tests/test-backend-ops.cpp | 158 +++++++++++++ tests/test-flash-attn-reference.cpp | 55 +++++ tests/test-masked-qk-attention.cpp | 208 ++++++++++++++++++ tests/test-masked-qk-flags.cu | 95 ++++++++ tests/test-query-pieces.cpp | 35 +++ tests/test-vulkan-fa-order.cpp | 110 +++++++++ 23 files changed, 1346 insertions(+), 49 deletions(-) create mode 100644 ggml/src/ggml-cuda/fattn-masked-qk-layout.h create mode 100644 ggml/src/ggml-vulkan/vulkan-shaders/sigmoid_mul_f32.comp create mode 100644 tests/test-flash-attn-reference.cpp create mode 100644 tests/test-masked-qk-attention.cpp create mode 100644 tests/test-masked-qk-flags.cu create mode 100644 tests/test-query-pieces.cpp create mode 100644 tests/test-vulkan-fa-order.cpp diff --git a/ggml/src/ggml-cpu/ops.cpp b/ggml/src/ggml-cpu/ops.cpp index 66e3659f943c..6cc220da6c83 100644 --- a/ggml/src/ggml-cpu/ops.cpp +++ b/ggml/src/ggml-cpu/ops.cpp @@ -8735,6 +8735,8 @@ static void ggml_compute_forward_flash_attn_ext_f16_one_chunk( GGML_ASSERT(( q_to_vec_dot) && "fattn: unsupported K-type"); GGML_ASSERT((v->type == GGML_TYPE_F32 || v_to_float ) && "fattn: unsupported V-type"); + const bool ref_f32_v = params->use_ref && v->type == GGML_TYPE_F16; + int ith = params->ith; for (int ir = ir0; ir < ir1; ++ir) { @@ -8754,7 +8756,7 @@ static void ggml_compute_forward_flash_attn_ext_f16_one_chunk( ggml_fp16_t * VKQ16 = (ggml_fp16_t *) (VKQ32 + 1*DV); // (temporary) FP16 VKQ accumulator ggml_fp16_t * Q_q = (ggml_fp16_t *) (VKQ32 + 2*DV); // (temporary) buffer for Q converted to quantized/FP16 - if (v->type == GGML_TYPE_F16) { + if (v->type == GGML_TYPE_F16 && !ref_f32_v) { memset(VKQ16, 0, DV*sizeof(ggml_fp16_t)); } else { memset(VKQ32, 0, DV*sizeof(float)); @@ -8803,7 +8805,7 @@ static void ggml_compute_forward_flash_attn_ext_f16_one_chunk( const char * v_data = ((const char *) v->data + (ic*nbv1 + iv2*nbv2 + iv3*nbv3)); - if (v->type == GGML_TYPE_F16) { + if (v->type == GGML_TYPE_F16 && !ref_f32_v) { if (s > M) { // s is new maximum, ms < 1.0f, vs == expf(s - s) == 1.0f M = s; @@ -8844,7 +8846,7 @@ static void ggml_compute_forward_flash_attn_ext_f16_one_chunk( S = S*ms + vs; // scale and increment sum with partial sum } - if (v->type == GGML_TYPE_F16) { + if (v->type == GGML_TYPE_F16 && !ref_f32_v) { for (int64_t d = 0; d < DV; ++d) { VKQ32[d] = GGML_CPU_FP16_TO_FP32(VKQ16[d]); } diff --git a/ggml/src/ggml-cuda/binbcast.cu b/ggml/src/ggml-cuda/binbcast.cu index 20dd6682d397..12cee1734f07 100644 --- a/ggml/src/ggml-cuda/binbcast.cu +++ b/ggml/src/ggml-cuda/binbcast.cu @@ -45,7 +45,7 @@ static bool try_binary_contiguous_f32(ggml_backend_cuda_context & ctx, ggml_tens const ggml_tensor * src0 = dst->src[0]; const ggml_tensor * src1 = dst->src[1]; if (src0->type != GGML_TYPE_F32 || src1->type != GGML_TYPE_F32 || dst->type != GGML_TYPE_F32 || - src0->ne[0] != 2048 || + (src0->ne[0] != 2048 && src0->ne[0] != 2560) || !ggml_are_same_shape(src0, src1) || !ggml_are_same_shape(src0, dst) || !ggml_is_contiguous(src0) || !ggml_is_contiguous(src1) || !ggml_is_contiguous(dst)) { return false; diff --git a/ggml/src/ggml-cuda/fattn-common.cuh b/ggml/src/ggml-cuda/fattn-common.cuh index 3edd936d86e7..7e0c5fc3e614 100644 --- a/ggml/src/ggml-cuda/fattn-common.cuh +++ b/ggml/src/ggml-cuda/fattn-common.cuh @@ -972,7 +972,72 @@ static __global__ void flash_attn_combine_results( dst[tid] = VKQ_numerator / VKQ_denominator; } -template +#include "fattn-masked-qk-layout.h" + +#if defined(GGML_USE_HIP) +// KV_max[a], Q finite[a], K finite[2*b], mask empty[a*b]. +static __global__ void masked_qk_q_finite(const char * q, int * flags, size_t nb1, size_t nb2, float scale, + int t, int * full_bound, int n) { + int valid = 1; + for (int i = threadIdx.x; i < 8*24*128; i += blockDim.x) { + const int row = blockIdx.x*8 + i/(24*128); + if (row >= t) { continue; } // The main kernel zero-pads Q, without reading it. + const float2 x = ((const float2 *) (q + row*nb1 + (i/128%24)*nb2))[i%128]; + const half2 h = make_half2(scale, scale) * make_half2(x.x, x.y); + valid &= (__half_as_ushort(__low2half(h)) & 0x7c00) != 0x7c00; + valid &= (__half_as_ushort(__high2half(h)) & 0x7c00) != 0x7c00; + } + const int all_valid = __syncthreads_and(valid); + if (threadIdx.x == 0) { + flags[blockIdx.x] = all_valid; + if (full_bound) { full_bound[blockIdx.x] = n; } + } +} + +static __global__ void masked_qk_k_finite(const char * k, int * flags, size_t nb1, size_t nb2, int b) { + int valid = 1; + for (int i = threadIdx.x; i < 32*256; i += blockDim.x) { + const auto * row = (const unsigned short *) (k + (blockIdx.x%b*32 + i/256)*nb1 + (blockIdx.x/b)*nb2); + valid &= (row[i%256] & 0x7c00) != 0x7c00; + } + const int all_valid = __syncthreads_and(valid); + if (threadIdx.x == 0) { flags[blockIdx.x] = all_valid; } +} + +static __global__ void masked_qk_mask_empty(const unsigned short * mask, int * flags, int t, int n, int b) { + const int tile = blockIdx.x*8 + threadIdx.y; + int empty = 1; + for (int j = 0; j < 8; ++j) { + empty &= mask[((size_t(tile/b)*8 + j)%t)*n + tile%b*32 + threadIdx.x] == 0xfc00; + } + const int all_empty = __all(empty); + if (threadIdx.x == 0) { flags[tile] = all_empty; } +} +#endif + +static constexpr bool fattn_query_pieces_eligible(int64_t q, int64_t kv, int cc) { + return q == 2048 && (kv == 2048 || kv >= 32768) && masked_qk_flag_count(q, kv) != 0 && + cc == GGML_CUDA_CC_OFFSET_AMD + 0x1151; +} +static_assert( fattn_query_pieces_eligible(2048, 66048, GGML_CUDA_CC_OFFSET_AMD + 0x1151)); +static_assert(!fattn_query_pieces_eligible(2047, 66048, GGML_CUDA_CC_OFFSET_AMD + 0x1151)); +static_assert(!fattn_query_pieces_eligible(2049, 66048, GGML_CUDA_CC_OFFSET_AMD + 0x1151)); +static_assert( fattn_query_pieces_eligible(2048, 32768, GGML_CUDA_CC_OFFSET_AMD + 0x1151)); +static_assert(!fattn_query_pieces_eligible(2048, 32512, GGML_CUDA_CC_OFFSET_AMD + 0x1151)); +static_assert(!fattn_query_pieces_eligible(2048, 32769, GGML_CUDA_CC_OFFSET_AMD + 0x1151)); +static_assert( fattn_query_pieces_eligible(2048, 65792, GGML_CUDA_CC_OFFSET_AMD + 0x1151)); +static_assert( fattn_query_pieces_eligible(2048, 66304, GGML_CUDA_CC_OFFSET_AMD + 0x1151)); +static_assert( fattn_query_pieces_eligible(2048, 130048, GGML_CUDA_CC_OFFSET_AMD + 0x1151)); +static_assert( fattn_query_pieces_eligible(2048, 1002240, GGML_CUDA_CC_OFFSET_AMD + 0x1151)); +static_assert( fattn_query_pieces_eligible(2048, 1050624, GGML_CUDA_CC_OFFSET_AMD + 0x1151)); +static_assert(!fattn_query_pieces_eligible(2048, INT64_MAX, GGML_CUDA_CC_OFFSET_AMD + 0x1151)); +static_assert( fattn_query_pieces_eligible(2048, 2048, GGML_CUDA_CC_OFFSET_AMD + 0x1151)); +static_assert(!fattn_query_pieces_eligible(2048, 2304, GGML_CUDA_CC_OFFSET_AMD + 0x1151)); +static_assert(!fattn_query_pieces_eligible(1, 66048, GGML_CUDA_CC_OFFSET_AMD + 0x1151)); +static_assert(!fattn_query_pieces_eligible(2048, 66048, GGML_CUDA_CC_OFFSET_AMD + 0x1150)); +static_assert(!fattn_query_pieces_eligible(2048, 66048, GGML_CUDA_CC_OFFSET_AMD + 0x942)); + +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, @@ -1105,7 +1170,14 @@ 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)) { + const bool scan_KV_max = !use_sparse && mask && K->ne[1] % FATTN_KQ_STRIDE == 0 && + (Q->ne[1] >= 1024 || Q->ne[3] > 1); + if constexpr (masked_qk) { + const size_t count = masked_qk_flag_count(Q->ne[1], K->ne[1]); + GGML_ASSERT(count); + KV_max.alloc(count); + } + if (scan_KV_max) { const int64_t s31 = mask->nb[1] / sizeof(half2); const int64_t s33 = mask->nb[3] / sizeof(half2); @@ -1115,7 +1187,7 @@ void launch_fattn( const int ne_KV_max = 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); + if constexpr (!masked_qk) { 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); @@ -1217,6 +1289,21 @@ void launch_fattn( scale /= logit_softcap; } +#if defined(GGML_USE_HIP) + if constexpr (masked_qk) { + GGML_ASSERT(KV_max.ptr); + const int a = ntiles_x, b = K->ne[1]/32; + ggml_cuda_kernel_launch(masked_qk_q_finite, ggml_cuda_kernel_launch_params(dim3(a), dim3(256), 0, main_stream), + (const char *) Q->data, KV_max.ptr + a, Q->nb[1], Q->nb[2], scale, + int(Q->ne[1]), scan_KV_max ? nullptr : KV_max.ptr, int(K->ne[1])); + ggml_cuda_kernel_launch(masked_qk_k_finite, ggml_cuda_kernel_launch_params(dim3(2*b), dim3(256), 0, main_stream), + K_data, KV_max.ptr + 2*a, nb11, nb12, b); + ggml_cuda_kernel_launch(masked_qk_mask_empty, ggml_cuda_kernel_launch_params(dim3(a*b/8), dim3(32, 8), 0, main_stream), + (const unsigned short *) mask->data, KV_max.ptr + 2*a + 2*b, int(Q->ne[1]), int(K->ne[1]), b); + CUDA_CHECK(cudaGetLastError()); + } +#endif + const uint32_t n_head = Q->ne[2]; const uint32_t n_head_log2 = 1u << uint32_t(floorf(log2f(float(n_head)))); @@ -1228,7 +1315,28 @@ void launch_fattn( 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); + dim3 main_blocks_num = blocks_num; +#if defined(GGML_USE_HIP) + if constexpr (DV == 256 && ncols1 == 8 && ncols2 == 8) { + // Derive the physical grid from the current ranges, not a fixed context length. + // Keep the logical grid for partial allocation and fixup. + if ((masked_qk || K->ne[1] == 2048) && + fattn_query_pieces_eligible(Q->ne[1], K->ne[1], cc) && stream_k && nbatch_fa == 32 && + Q->ne[0] == 256 && Q->ne[2] == 24 && Q->ne[3] == 1 && + K->ne[0] == 256 && K->ne[2] == 2 && K->ne[3] == 1 && + !V_is_K_view && !sinks && max_bias == 0.0f && logit_softcap == 0.0f && scan_KV_max && + block_dim.x == 32 && block_dim.y == 8 && block_dim.z == 1 && nbytes_shared == 51328 && + blocks_num.x > 0 && blocks_num.y == 1 && blocks_num.z == 1) { + main_blocks_num.y = fattn_query_piece_count(ntiles_dst, ntiles_KV, blocks_num.x); + } + } +#endif + + // The masked path requires a singleton mask-head axis. Its unused stride + // can exceed int32 for long contexts; pass zero instead of narrowing it. + const size_t mask_head_stride = mask && !(masked_qk && mask->ne[2] == 1) ? mask->nb[2] : 0; + + ggml_cuda_kernel_launch_params launch_params = ggml_cuda_kernel_launch_params(main_blocks_num, block_dim, nbytes_shared, main_stream); ggml_cuda_kernel_launch(fattn_kernel, launch_params, (const char *) Q->data, K_data, @@ -1242,7 +1350,7 @@ void launch_fattn( K->ne[0], n_kv, 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_head_stride, mask ? mask->nb[3] : 0 ); CUDA_CHECK(cudaGetLastError()); diff --git a/ggml/src/ggml-cuda/fattn-masked-qk-layout.h b/ggml/src/ggml-cuda/fattn-masked-qk-layout.h new file mode 100644 index 000000000000..9e17d156932c --- /dev/null +++ b/ggml/src/ggml-cuda/fattn-masked-qk-layout.h @@ -0,0 +1,54 @@ +#pragma once +#include +#include +#include + +// Zero means unsupported. Bound every int-indexed slice and Stream-K product +// before allocation; all intermediate arithmetic here is unsigned 64-bit. +constexpr size_t masked_qk_flag_count(int64_t t, int64_t n) { + if (t <= 0 || n <= 0 || t > INT_MAX - 7 || n > INT_MAX - 255 || n % 256 != 0) { + return 0; + } + const uint64_t a = (uint64_t(t) + 7)/8, b = uint64_t(n)/32; + const uint64_t count = 2*a + 2*b + a*b; + return 4*a*b <= INT_MAX && 400*a <= INT_MAX && count <= INT_MAX && count <= SIZE_MAX/sizeof(int) + ? size_t(count) : 0; +} + +static_assert(masked_qk_flag_count(2048, 34816) == 281216, "winning layout"); +static_assert(masked_qk_flag_count(9, 256) == 36, "partial tile layout"); +static_assert(masked_qk_flag_count(0, 256) == 0, "empty query"); +static_assert(masked_qk_flag_count(2048, 257) == 0, "key alignment"); +static_assert(masked_qk_flag_count(INT64_MAX, 256) == 0, "query overflow"); +static_assert(masked_qk_flag_count(2048, INT64_MAX) == 0, "key overflow"); +static_assert(masked_qk_flag_count(1048576, 1048576) == 0, "Stream-K overflow"); + +// Subdivide physical work only at output-tile boundaries. Logical Stream-K +// ranges (and hence their partial sums and fixup order) remain unchanged. +constexpr int fattn_query_tiles_per_piece = 4; +constexpr unsigned fattn_query_piece_count(int tiles, int kv_tiles, unsigned logical_blocks) { + const int64_t work = int64_t(tiles)*kv_tiles; + if (tiles <= 0 || kv_tiles <= 0 || logical_blocks == 0 || work > INT_MAX) { + return 1; + } + int longest = 0; + for (unsigned b = 0; b < logical_blocks; ++b) { + const int64_t begin = int64_t(b)*work/logical_blocks; + const int64_t end = (int64_t(b) + 1)*work/logical_blocks; + if (begin < end) { + const int span = int((end - 1)/kv_tiles - begin/kv_tiles + 1); + longest = span > longest ? span : longest; + } + } + return longest ? unsigned((longest + fattn_query_tiles_per_piece - 1)/fattn_query_tiles_per_piece) : 1; +} + +static_assert(fattn_query_piece_count(1024, 66048/32, 20) == 13); +static_assert(fattn_query_piece_count(1024, 130048/32, 20) == 13); +static_assert(fattn_query_piece_count(1024, 1002240/32, 20) == 13); +static_assert(fattn_query_piece_count(1024, 1050624/32, 20) == 13); +static_assert(fattn_query_piece_count(1024, 130048/32, 40) == 7); +static_assert(fattn_query_piece_count(1024, 130048/32, 1024) == 1); +static_assert(fattn_query_piece_count(0, 4064, 20) == 1); +static_assert(fattn_query_piece_count(1024, 4064, 0) == 1); +static_assert(fattn_query_piece_count(INT_MAX, INT_MAX, 20) == 1); diff --git a/ggml/src/ggml-cuda/fattn-mma-f16.cuh b/ggml/src/ggml-cuda/fattn-mma-f16.cuh index 5c230481af02..6a2f006dc14d 100644 --- a/ggml/src/ggml-cuda/fattn-mma-f16.cuh +++ b/ggml/src/ggml-cuda/fattn-mma-f16.cuh @@ -565,7 +565,8 @@ static __device__ __forceinline__ void flash_attn_ext_f16_load_mask( template + typename T_A_KQ, typename T_B_KQ, typename T_C_KQ, typename T_A_VKQ, typename T_B_VKQ, typename T_C_VKQ, + bool masked_qk = false> static __device__ __forceinline__ void flash_attn_ext_f16_iter( const float2 * const __restrict__ Q_f2, const half2 * const __restrict__ K_h2, @@ -592,7 +593,7 @@ static __device__ __forceinline__ void flash_attn_ext_f16_iter( float * const __restrict__ KQ_rowsum, const int jt, const int kb0, - const int k_VKQ_sup) { + const int k_VKQ_sup, const int * empty = nullptr, const int * kfinite = nullptr) { #if defined(VOLTA_MMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE) || defined(AMD_MFMA_AVAILABLE) constexpr int warp_size = ggml_cuda_get_physical_warp_size(); constexpr int ncols = ncols1 * ncols2; @@ -638,6 +639,17 @@ static __device__ __forceinline__ void flash_attn_ext_f16_iter( } } + bool skip_qk = false; + if constexpr (masked_qk) { + static_assert(DKQ == 256 && DV == 256 && ncols1 == 8 && ncols2 == 8 && !V_is_K_view && !use_logit_softcap); + static_assert(nstages == 0 && nbatch_fa == 32 && nbatch_K2 == 128 && !Q_in_reg); + // Finite half operands bound each FP32 dot partial below 1.1e12. Adding -inf then matches zero + -inf. + skip_qk = empty && empty[kb0] && kfinite[kb0]; + } + if (skip_qk) { + __syncthreads(); // Publish mask, as in the dense K load phase. + __syncthreads(); // Keep the K/V shared-memory phase boundary. + } else { // For MLA K and V have the same data. // Therefore, iterate over K in reverse and later re-use the data if possible. #pragma unroll @@ -712,6 +724,8 @@ static __device__ __forceinline__ void flash_attn_ext_f16_iter( } } + } + if (use_logit_softcap) { constexpr int stride = cols_per_warp == 8 ? np*T_C_KQ::I : np*T_C_KQ::J; static_assert(nbatch_fa % stride == 0, "bad loop size"); @@ -1154,7 +1168,8 @@ 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, @@ -1179,7 +1194,7 @@ static __device__ __forceinline__ void flash_attn_ext_f16_process_tile( const int jt, const int zt_gqa, const int kb0_start, - const int kb0_stop) { + const int kb0_stop, const int * empty = nullptr, const int * kfinite = nullptr) { #if defined(VOLTA_MMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE) || defined(AMD_MFMA_AVAILABLE) //In this kernel Q, K, V are matrices while i, j, k are matrix indices. @@ -1322,19 +1337,19 @@ static __device__ __forceinline__ void flash_attn_ext_f16_process_tile( constexpr int k_VKQ_sup = nbatch_fa; flash_attn_ext_f16_iter + T_A_KQ, T_B_KQ, T_C_KQ, T_A_VKQ, T_B_VKQ, T_C_VKQ, masked_qk> (Q_f2, K_h2, V_h2, mask_h, indices, 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); + KQ_max, KQ_rowsum, jt, kb0, k_VKQ_sup, empty, kfinite); } constexpr bool last_iter = true; const int k_VKQ_sup = ne11 - kb0*nbatch_fa; flash_attn_ext_f16_iter + T_A_KQ, T_B_KQ, T_C_KQ, T_A_VKQ, T_B_VKQ, T_C_VKQ, masked_qk> (Q_f2, K_h2, V_h2, mask_h, indices, 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); + KQ_max, KQ_rowsum, jt, kb0, k_VKQ_sup, empty, kfinite); } else { constexpr bool oob_check = false; for (; kb0 < kb0_stop-1; ++kb0) { @@ -1342,19 +1357,19 @@ static __device__ __forceinline__ void flash_attn_ext_f16_process_tile( constexpr int k_VKQ_sup = nbatch_fa; flash_attn_ext_f16_iter + T_A_KQ, T_B_KQ, T_C_KQ, T_A_VKQ, T_B_VKQ, T_C_VKQ, masked_qk> (Q_f2, K_h2, V_h2, mask_h, indices, 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); + KQ_max, KQ_rowsum, jt, kb0, k_VKQ_sup, empty, kfinite); } constexpr bool last_iter = true; constexpr int k_VKQ_sup = nbatch_fa; flash_attn_ext_f16_iter + T_A_KQ, T_B_KQ, T_C_KQ, T_A_VKQ, T_B_VKQ, T_C_VKQ, masked_qk> (Q_f2, K_h2, V_h2, mask_h, indices, 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); + KQ_max, KQ_rowsum, jt, kb0, k_VKQ_sup, empty, kfinite); } // With multi-stage loading there is no __syncthreads at the end of the iter, @@ -1763,7 +1778,8 @@ static constexpr __host__ __device__ bool ggml_cuda_flash_attn_ext_mma_f16_may_u (DKQ == 576 && DV == 512 && ncols1 == 1 && ncols2 == 16); } -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, @@ -1789,6 +1805,13 @@ static __global__ void flash_attn_ext_f16( const int32_t nb31, const int32_t nb32, const int64_t nb33) { 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)) +#if !defined(RDNA3_5) + // The host runtime check does not exclude masked-QK instantiation on CDNA. + if constexpr (masked_qk) { + NO_DEVICE_CODE; + return; + } else { +#endif const char * GGML_CUDA_RESTRICT Q = Q_ptr; const char * GGML_CUDA_RESTRICT K = K_ptr; const char * GGML_CUDA_RESTRICT V = V_ptr; @@ -1828,7 +1851,7 @@ static __global__ void flash_attn_ext_f16( #endif // __CUDA_ARCH__ == GGML_CUDA_CC_TURING #if defined(AMD_WMMA_AVAILABLE) - if (ncols1*ncols2 < 16 || ncols2 == 1 || DKQ > 256 || (DKQ > 128 && DKQ != 256)) { + if (ncols1*ncols2 < 16 || DKQ > 256 || (DKQ > 128 && DKQ != 256)) { NO_DEVICE_CODE; return; } @@ -1862,7 +1885,27 @@ static __global__ void flash_attn_ext_f16( // kbc == k block continuous, current index in continuous ijk space. int kbc = int64_t(blockIdx.x + 0)*(iter_k*iter_j*iter_z_gqa*ne12*ne03) / gridDim.x; - const int kbc_stop = int64_t(blockIdx.x + 1)*(iter_k*iter_j*iter_z_gqa*ne12*ne03) / gridDim.x; + int kbc_stop = int64_t(blockIdx.x + 1)*(iter_k*iter_j*iter_z_gqa*ne12*ne03) / gridDim.x; + +#if defined(GGML_USE_HIP) && defined(__gfx1151__) + // Physical pieces retain the original X-range and its numerical KV seams. + if constexpr (DKQ == 256 && DV == 256 && ncols1 == 8 && ncols2 == 8 && + !use_logit_softcap && !V_is_K_view) { + if (gridDim.y > 1 && ne01.z == 2048 && ((masked_qk && ne11 >= 32768) || ne11 == 2048) && ne02 == 24 && ne12 == 2 && + ne03 == 1 && gridDim.z == 1) { + const int original_begin = kbc; + const int original_end = kbc_stop; + const int first_tile = original_begin / iter_k + fattn_query_tiles_per_piece*blockIdx.y; + kbc = max(original_begin, first_tile*iter_k); + kbc_stop = min(original_end, (first_tile + fattn_query_tiles_per_piece)*iter_k); + if (kbc >= kbc_stop) { + return; + } + // Interior edges are tile-aligned: the existing while consumes them. + // Only original KV seams reach needs_fixup or the final is_fixup call. + } + } +#endif // If the seams of 2 CUDA blocks fall within an output tile their results need to be combined. // For this we need to track both the block that starts the tile (needs_fixup) and the block that finishes the tile (is_fixup). @@ -1899,14 +1942,18 @@ static __global__ void flash_attn_ext_f16( 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 + flash_attn_ext_f16_process_tile (Q_f2, K_h2, V_h2, mask_h, indices, 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); + ne01, ne02, gqa_ratio, ne11, stride_Q1, stride_Q2, stride_K, stride_V, stride_mask, jt, zt_gqa, kb0_start, kb0_stop, + masked_qk && KV_max[iter_j + jt] ? KV_max + 2*iter_j + 2*iter_k + jt*iter_k : nullptr, + masked_qk ? KV_max + 2*iter_j + z_KV*iter_k : nullptr); } else { constexpr bool needs_fixup = true; // CUDA block is missing the beginning of a tile. - flash_attn_ext_f16_process_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, - ne01, ne02, gqa_ratio, ne11, stride_Q1, stride_Q2, stride_K, stride_V, stride_mask, jt, zt_gqa, kb0_start, kb0_stop); + ne01, ne02, gqa_ratio, ne11, stride_Q1, stride_Q2, stride_K, stride_V, stride_mask, jt, zt_gqa, kb0_start, kb0_stop, + masked_qk && KV_max[iter_j + jt] ? KV_max + 2*iter_j + 2*iter_k + jt*iter_k : nullptr, + masked_qk ? KV_max + 2*iter_j + z_KV*iter_k : nullptr); } kbc += iter_k; @@ -1946,9 +1993,14 @@ static __global__ void flash_attn_ext_f16( 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 + flash_attn_ext_f16_process_tile (Q_f2, K_h2, V_h2, mask_h, indices, 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); + ne01, ne02, gqa_ratio, ne11, stride_Q1, stride_Q2, stride_K, stride_V, stride_mask, jt, zt_gqa, kb0_start, kb0_stop, + masked_qk && KV_max[iter_j + jt] ? KV_max + 2*iter_j + 2*iter_k + jt*iter_k : nullptr, + masked_qk ? KV_max + 2*iter_j + z_KV*iter_k : nullptr); +#if !defined(RDNA3_5) + } +#endif #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, @@ -1965,6 +2017,17 @@ static __global__ void flash_attn_ext_f16( bool ggml_cuda_flash_attn_ext_mma_f16_shall_use_sparse(ggml_backend_cuda_context & ctx, ggml_tensor * dst); +// Conservative cutoff: paired PP is neutral at N=14080 and gains at N=34048/66048. +// The crossover between these measurements has not been measured. +static constexpr bool masked_qk_work_eligible(int64_t n_kv) { + return n_kv >= 32768; +} +static_assert(!masked_qk_work_eligible(32767)); +static_assert( masked_qk_work_eligible(32768)); +static_assert( masked_qk_work_eligible(32769)); +static_assert(!masked_qk_work_eligible(32512)); +static_assert( masked_qk_work_eligible(33024)); + template void ggml_cuda_flash_attn_ext_mma_f16_case(ggml_backend_cuda_context & ctx, ggml_tensor * dst) { const ggml_tensor * KQV = dst; @@ -2064,6 +2127,55 @@ void ggml_cuda_flash_attn_ext_mma_f16_case(ggml_backend_cuda_context & ctx, ggml #endif // !defined(GGML_USE_MUSA) } +#if defined(GGML_USE_HIP) + if constexpr (DKQ == 256 && DV == 256 && ncols1 == 8 && ncols2 == 8) { + const ggml_tensor * Q = dst->src[0], * K = dst->src[1], * V = dst->src[2], * mask = dst->src[3]; + float max_bias; + memcpy(&max_bias, (const float *) dst->op_params + 1, sizeof(float)); + const bool aliased = V->data == K->data || (V->view_src && (V->view_src == K || + (V->view_src == K->view_src && V->view_offs == K->view_offs))); + bool aligned = true; + for (const ggml_tensor * t : {Q, K, V}) { + const size_t alignment = t == Q ? sizeof(float2) : sizeof(half2); + aligned &= uintptr_t(t->data) % alignment == 0; + for (int i = 1; i < 4; ++i) { aligned &= t->nb[i] % alignment == 0 && t->nb[i] <= INT32_MAX; } + // Device address expressions multiply int head/row indices by int strides. + uint64_t span = 256*ggml_element_size(t); + for (int i = 1; i < 4; ++i) { + if (t->ne[i] <= 0 || t->nb[i] > INT32_MAX || uint64_t(t->ne[i]) > INT32_MAX) { + aligned = false; + } else { + span += uint64_t(t->ne[i] - 1)*t->nb[i]; + } + } + aligned &= span <= INT32_MAX; + } + const size_t flag_count = masked_qk_flag_count(Q->ne[1], K->ne[1]); + if (masked_qk_work_eligible(K->ne[1]) && GGML_CUDA_CC_IS_RDNA3_5(cc) && nstages == 0 && nbatch_fa == 32 && nbatch_K2 == 128 && + nbatch_V2 == 64 && nbytes_shared_total == 51328 && flag_count && + nthreads == 256 && warp_size_host == 32 && !Q_in_reg && !aliased && aligned && + Q->type == GGML_TYPE_F32 && K->type == GGML_TYPE_F16 && V->type == GGML_TYPE_F16 && + Q->ne[0] == 256 && Q->ne[2] == 24 && Q->ne[3] == 1 && + K->ne[0] == 256 && K->ne[2] == 2 && K->ne[3] == 1 && + ggml_are_same_shape(K, V) && mask && mask->type == GGML_TYPE_F16 && ggml_is_contiguous(mask) && + mask->ne[0] == K->ne[1] && mask->ne[1] >= Q->ne[1] && mask->ne[1] <= INT32_MAX && + (Q->ne[1] < 1024 || mask->ne[1] >= 8*((Q->ne[1] + 7)/8)) && + mask->nb[1] <= INT32_MAX && + uintptr_t(mask->data) % sizeof(half2) == 0 && mask->ne[2] == 1 && mask->ne[3] == 1 && + !dst->src[4] && max_bias == 0.0f && logit_softcap == 0.0f) { + fattn_kernel_t candidate = flash_attn_ext_f16; + CUDA_CHECK(cudaFuncSetAttribute(reinterpret_cast(candidate), cudaFuncAttributeMaxDynamicSharedMemorySize, nbytes_shared_total)); + int stock_occupancy, candidate_occupancy; + CUDA_CHECK(cudaOccupancyMaxActiveBlocksPerMultiprocessor(&stock_occupancy, fattn_kernel, nthreads, nbytes_shared_total)); + CUDA_CHECK(cudaOccupancyMaxActiveBlocksPerMultiprocessor(&candidate_occupancy, candidate, nthreads, nbytes_shared_total)); + if (stock_occupancy > 0 && candidate_occupancy == stock_occupancy) { + launch_fattn + (ctx, dst, candidate, nwarps, nbytes_shared_total, nbatch_fa, true, true, true, false, warp_size_host); + return; + } + } + } +#endif launch_fattn (ctx, dst, fattn_kernel, nwarps, nbytes_shared_total, nbatch_fa, true, true, true, use_sparse, warp_size_host); } diff --git a/ggml/src/ggml-cuda/fattn.cu b/ggml/src/ggml-cuda/fattn.cu index cafe1eeaf887..2f93dc40b543 100644 --- a/ggml/src/ggml-cuda/fattn.cu +++ b/ggml/src/ggml-cuda/fattn.cu @@ -669,7 +669,9 @@ static best_fattn_kernel ggml_cuda_get_best_fattn_kernel(const int device, const // On RDNA3.5 the D=256 tile kernel is FMA-bound for prefill batches; the WMMA kernel with 64 columns is ~30-45% faster. // For smaller batches (ncols 16/32 configs) the tile kernel is still faster. - if (GGML_CUDA_CC_IS_RDNA3_5(cc) && gqa_opt_applies && Q->ne[0] == 256 && V->ne[0] == 256 && Q->ne[1] * gqa_ratio_eff > 32) { + // Without GQA reuse, queries alone can fill the tile; keep ALiBi on its existing fallback. + if (GGML_CUDA_CC_IS_RDNA3_5(cc) && max_bias == 0.0f && Q->ne[0] == 256 && V->ne[0] == 256 && + Q->ne[1] * (gqa_opt_applies ? gqa_ratio_eff : 1) > 32) { return BEST_FATTN_KERNEL_MMA_F16; } diff --git a/ggml/src/ggml-cuda/ggml-cuda.cu b/ggml/src/ggml-cuda/ggml-cuda.cu index 08164ce00192..b8d9e524f520 100644 --- a/ggml/src/ggml-cuda/ggml-cuda.cu +++ b/ggml/src/ggml-cuda/ggml-cuda.cu @@ -3784,6 +3784,47 @@ static int ggml_cuda_try_fuse(ggml_backend_cuda_context * cuda_ctx, ggml_cgraph ggml_tensor * node = cgraph->nodes[i]; +#if defined(__HIP_PLATFORM_AMD__) + // Do not skip a fork/join or change stream assignment in the execution loop. + if (node->op == GGML_OP_UNARY && ggml_get_unary_op(node) == GGML_UNARY_OP_RELU && + GGML_CUDA_CC_IS_RDNA3_5(ggml_cuda_info().devices[cuda_ctx->device].cc) && + cuda_ctx->curr_stream_no == 0 && cuda_ctx->stream_context().concurrent_events.empty() && + ggml_can_fuse_subgraph(cgraph, i, { GGML_OP_UNARY, GGML_OP_PERMUTE, GGML_OP_CONT, GGML_OP_SUM_ROWS }, { i + 3 })) { + const ggml_tensor * src = node->src[0]; + const ggml_tensor * perm = cgraph->nodes[i + 1]; + const ggml_tensor * cont = cgraph->nodes[i + 2]; + ggml_tensor * sum = cgraph->nodes[i + 3]; + bool ok = src && perm->src[0] == node && cont->src[0] == perm && sum->src[0] == cont && + node->view_src == nullptr && cont->view_src == nullptr && perm->view_src == node && perm->view_offs == 0; + for (const ggml_tensor * t : { src, static_cast(node), perm, cont, static_cast(sum) }) { + ok = ok && t && t->type == GGML_TYPE_F32 && t->data && t->buffer && + t->buffer->buft == ggml_backend_cuda_buffer_type(cuda_ctx->device); + } + for (const ggml_tensor * t : { static_cast(node), perm, cont, static_cast(sum) }) { + for (int j = 1; j < GGML_MAX_SRC; ++j) { + ok = ok && t->src[j] == nullptr; + } + } + if (ok) { + ok = src->ne[1] == 4 && ggml_are_same_shape(src, node) && ggml_is_contiguous(src) && + ggml_is_contiguous(node) && ggml_is_contiguous(cont) && ggml_is_contiguous(sum); + const int axes[4] = { 1, 0, 2, 3 }; + for (int d = 0; d < 4; ++d) { + ok = ok && ggml_get_op_params_i32(perm, d) == axes[d] && + perm->ne[d] == node->ne[axes[d]] && perm->nb[d] == node->nb[axes[d]] && + cont->ne[d] == perm->ne[d] && sum->ne[d] == (d == 0 ? 1 : cont->ne[d]); + } + const int nsm = ggml_cuda_info().devices[cuda_ctx->device].nsm; + // sumrows.cu uses 512 threads below this boundary; only reproduce its 32-thread branch. + ok = ok && nsm > 0 && ggml_nrows(cont) / nsm >= 2 && ggml_nelements(cont) <= INT_MAX; + } + if (ok) { + ggml_cuda_op_relu_sum4(*cuda_ctx, src, sum); + return 3; + } + } +#endif + // RDNA3.5 decode: consecutive single-column MUL_MATs that read the same activation vector (plain Q8_0/F32 // matvecs, or a [mul_mat, mul_mat, glu] gate/up pair) are launched as one grouped kernel. The segments are // mutually independent, so the launch position of the first one is valid for all of them. diff --git a/ggml/src/ggml-cuda/moe-weighted-reduction.cu b/ggml/src/ggml-cuda/moe-weighted-reduction.cu index 11ec58497f1e..9fc373461300 100644 --- a/ggml/src/ggml-cuda/moe-weighted-reduction.cu +++ b/ggml/src/ggml-cuda/moe-weighted-reduction.cu @@ -1,5 +1,6 @@ #include "moe-weighted-reduction.cuh" +template static __global__ void moe_weighted_reduction_f32(const float * __restrict__ experts, const float * __restrict__ expert_scale, const float * __restrict__ weights, @@ -7,21 +8,31 @@ static __global__ void moe_weighted_reduction_f32(const float * __restrict__ exp const int64_t n_embd, const int n_expert_used) { const int64_t token = blockIdx.x; - const int64_t col = (int64_t) blockIdx.y * blockDim.x + threadIdx.x; + const int64_t col = ((int64_t) blockIdx.y * blockDim.x + threadIdx.x) * values_per_thread; if (col >= n_embd) { return; } const uint64_t first_row = (uint64_t) token * n_expert_used; const float first_scale = expert_scale != nullptr ? expert_scale[first_row] : 1.0f; - float sum = (experts[first_row * n_embd + col] * first_scale) * weights[first_row]; + float sum[values_per_thread]; +#pragma unroll + for (int value = 0; value < values_per_thread; ++value) { + sum[value] = (experts[first_row * n_embd + col + value] * first_scale) * weights[first_row]; + } for (int expert = 1; expert < n_expert_used; ++expert) { const uint64_t row = first_row + expert; const float scale = expert_scale != nullptr ? expert_scale[row] : 1.0f; - sum += (experts[row * n_embd + col] * scale) * weights[row]; +#pragma unroll + for (int value = 0; value < values_per_thread; ++value) { + sum[value] += (experts[row * n_embd + col + value] * scale) * weights[row]; + } + } +#pragma unroll + for (int value = 0; value < values_per_thread; ++value) { + dst[token * n_embd + col + value] = sum[value]; } - dst[token * n_embd + col] = sum; } static void launch_moe_weighted_reduction(const float * experts, @@ -33,9 +44,15 @@ static void launch_moe_weighted_reduction(const float * experts, int n_expert_used, cudaStream_t stream) { constexpr int threads = 256; - const dim3 blocks(n_tokens, (n_embd + threads - 1) / threads, 1); - moe_weighted_reduction_f32 - <<>>(experts, expert_scale, weights, dst, n_embd, n_expert_used); + if (n_embd % 4 == 0) { + const dim3 blocks(n_tokens, (n_embd + 4 * threads - 1) / (4 * threads), 1); + moe_weighted_reduction_f32<4> + <<>>(experts, expert_scale, weights, dst, n_embd, n_expert_used); + } else { + const dim3 blocks(n_tokens, (n_embd + threads - 1) / threads, 1); + moe_weighted_reduction_f32<1> + <<>>(experts, expert_scale, weights, dst, n_embd, n_expert_used); + } } void ggml_cuda_op_moe_weighted_reduction(ggml_backend_cuda_context & ctx, diff --git a/ggml/src/ggml-cuda/unary.cu b/ggml/src/ggml-cuda/unary.cu index 1508bb857a89..4401169f090e 100644 --- a/ggml/src/ggml-cuda/unary.cu +++ b/ggml/src/ggml-cuda/unary.cu @@ -45,6 +45,56 @@ static __device__ __forceinline__ float op_relu(float x) { return fmaxf(x, 0); } +#if defined(__HIP_PLATFORM_AMD__) +static __device__ __forceinline__ float relu_sum4_add_rn(const float a, const float b) { + float result; + asm volatile("v_add_f32_e32 %0, %1, %2" : "=v"(result) : "v"(a), "v"(b)); + return result; +} + +static __global__ void relu_sum4_f32(const float * src, float * dst, const int64_t nb, const int64_t nrows) { + const int64_t row = int64_t(blockIdx.x) * blockDim.x + threadIdx.x; + if (row >= nrows) { + return; + } + const int64_t base = (row / nb) * (4 * nb) + row % nb; + float h[4]; +#pragma unroll + for (int head = 0; head < 4; ++head) { + // reduce_rows.cuh: sum_temp[0], then sum over eight sum_temp slots. + float u = relu_sum4_add_rn(0.0f, op_relu(src[base + head * nb])); + u = relu_sum4_add_rn(0.0f, u); +#pragma unroll + for (int j = 0; j < 7; ++j) { + u = relu_sum4_add_rn(u, 0.0f); + } + // The 32-thread branch has zero partners at xor offsets 16, 8, 4. +#pragma unroll + for (int j = 0; j < 3; ++j) { + u = relu_sum4_add_rn(u, 0.0f); + } + h[head] = u; + } + dst[row] = relu_sum4_add_rn(relu_sum4_add_rn(h[0], h[2]), relu_sum4_add_rn(h[1], h[3])); +} + +void ggml_cuda_op_relu_sum4(ggml_backend_cuda_context & ctx, const ggml_tensor * src, ggml_tensor * dst) { + const int64_t nrows = ggml_nelements(dst); + const uintptr_t in = reinterpret_cast(src->data); + const uintptr_t out = reinterpret_cast(dst->data); + const bool overlap = in <= out ? out - in < ggml_nbytes(src) : in - out < ggml_nbytes(dst); + // CONT can end the input lifetime before SUM_ROWS reuses its allocation. + ggml_cuda_pool_alloc staged(ctx.pool()); + float * result = overlap ? staged.alloc(nrows) : static_cast(dst->data); + constexpr int block_size = 256; // Initial launch size, not hardware-tuned. + const ggml_cuda_kernel_launch_params params((nrows + block_size - 1) / block_size, block_size, 0, ctx.stream()); + ggml_cuda_kernel_launch(relu_sum4_f32, params, static_cast(src->data), result, src->ne[0], nrows); + if (overlap) { + CUDA_CHECK(cudaMemcpyAsync(dst->data, result, ggml_nbytes(dst), cudaMemcpyDeviceToDevice, ctx.stream())); + } +} +#endif + static __device__ __forceinline__ float op_sigmoid(float x) { return 1.0f / (1.0f + expf(-x)); } diff --git a/ggml/src/ggml-cuda/unary.cuh b/ggml/src/ggml-cuda/unary.cuh index f2a543cffbe2..c19875f9eaef 100644 --- a/ggml/src/ggml-cuda/unary.cuh +++ b/ggml/src/ggml-cuda/unary.cuh @@ -19,6 +19,10 @@ #define CUDA_GLU_BLOCK_SIZE 256 #define CUDA_XIELU_BLOCK_SIZE 256 +#if defined(__HIP_PLATFORM_AMD__) +void ggml_cuda_op_relu_sum4(ggml_backend_cuda_context & ctx, const ggml_tensor * src, ggml_tensor * dst); +#endif + void ggml_cuda_op_abs(ggml_backend_cuda_context & ctx, ggml_tensor * dst); void ggml_cuda_op_sgn(ggml_backend_cuda_context & ctx, ggml_tensor * dst); diff --git a/ggml/src/ggml-vulkan/ggml-vulkan.cpp b/ggml/src/ggml-vulkan/ggml-vulkan.cpp index 353edb87c466..f335bc52a487 100644 --- a/ggml/src/ggml-vulkan/ggml-vulkan.cpp +++ b/ggml/src/ggml-vulkan/ggml-vulkan.cpp @@ -940,6 +940,10 @@ struct vk_device_struct { vk::DescriptorSetLayout dsl; vk_matmul_pipeline pipeline_matmul_f32 {}; + vk_pipeline pipeline_prefill_q8_0; + vk_pipeline pipeline_prefill_q8_0_aligned; + vk_pipeline pipeline_prefill_q8_0_short_k; + vk_pipeline pipeline_prefill_q8_0_short_k_aligned; vk_matmul_pipeline pipeline_matmul_f32_f16 {}; vk_matmul_pipeline pipeline_matmul_bf16 {}; vk_matmul_pipeline2 pipeline_matmul_f16; @@ -1072,6 +1076,7 @@ struct vk_device_struct { vk_pipeline pipeline_geglu[2]; vk_pipeline pipeline_reglu[2]; vk_pipeline pipeline_swiglu[2]; + vk_pipeline pipeline_sigmoid_mul_f32; vk_pipeline pipeline_swiglu_oai[2]; vk_pipeline pipeline_swiglu_clamp[2]; vk_pipeline pipeline_geglu_erf[2]; @@ -1303,6 +1308,14 @@ struct vk_buffer_struct { } }; +struct vk_sigmoid_mul_push_constants { + uint32_t n; + uint32_t a_offset; + uint32_t b_offset; + uint32_t d_offset; +}; +static_assert(sizeof(vk_sigmoid_mul_push_constants) == 16, "sigmoid_mul push constant ABI"); + struct vk_subbuffer { vk_buffer buffer; uint64_t offset; @@ -5107,6 +5120,26 @@ static void ggml_vk_load_shaders(vk_device& device, vk_pipeline requested) { #endif // defined(VK_NV_cooperative_matrix2) && defined(GGML_VULKAN_COOPMAT2_GLSLC_SUPPORT) #if defined(VK_KHR_cooperative_matrix) && defined(GGML_VULKAN_COOPMAT_GLSLC_SUPPORT) if (device->coopmat_support) { + static const bool prefill_enabled = [] { + const char * env = getenv("GGML_VK_FLASHNEXT_PREFILL"); + return env && strcmp(env, "1") == 0; + }(); + if (prefill_enabled && ggml_vk_dense_f16b_mode() == 1 && + device->vendor_id == VK_VENDOR_ID_AMD && device->uma && + device->coopmat_acc_f16_support && device->coopmat_support_16x16x16_f16acc && + device->subgroup_size_control && device->subgroup_require_full_support && + device->subgroup_min_size <= 32 && device->subgroup_max_size >= 64 && + device->properties.limits.maxComputeSharedMemorySize >= 65536 && + device->properties.limits.maxComputeWorkGroupInvocations >= 256 && + device->properties.limits.maxComputeWorkGroupSize[0] >= 256) { + const std::vector standard = {256, 128, 128, 32, 128, 32, 2, 16, 16, 16, 64}; + const std::vector short_k = {256, 64, 64, 32, 32, 16, 2, 16, 16, 16, 32}; + ggml_vk_create_pipeline(device, device->pipeline_prefill_q8_0, "prefill_q8_0", matmul_q8_0_f16_f16acc_cm1_len, matmul_q8_0_f16_f16acc_cm1_data, "main", 3, sizeof(vk_mat_mat_push_constants), {128, 128, 1}, ggml_vk_mul_mm_spec(standard, false), 1, false, true, 64); + ggml_vk_create_pipeline(device, device->pipeline_prefill_q8_0_aligned, "prefill_q8_0_aligned", matmul_q8_0_f16_f16acc_cm1_len, matmul_q8_0_f16_f16acc_cm1_data, "main", 3, sizeof(vk_mat_mat_push_constants), {128, 128, 1}, ggml_vk_mul_mm_spec(standard, true), 128, false, true, 64); + ggml_vk_create_pipeline(device, device->pipeline_prefill_q8_0_short_k, "prefill_q8_0_short_k", matmul_q8_0_f16_f16acc_cm1_len, matmul_q8_0_f16_f16acc_cm1_data, "main", 3, sizeof(vk_mat_mat_push_constants), {64, 64, 1}, ggml_vk_mul_mm_spec(short_k, false), 1, false, true, 32); + ggml_vk_create_pipeline(device, device->pipeline_prefill_q8_0_short_k_aligned, "prefill_q8_0_short_k_aligned", matmul_q8_0_f16_f16acc_cm1_len, matmul_q8_0_f16_f16acc_cm1_data, "main", 3, sizeof(vk_mat_mat_push_constants), {64, 64, 1}, ggml_vk_mul_mm_spec(short_k, true), 128, false, true, 32); + } + // Deterministic subgroup sizing for the dense coopmat pipelines. Two parts: // // (1) The required subgroup size is the tile's own WARP, not the driver's choice. The cm1 @@ -6403,6 +6436,14 @@ static void ggml_vk_load_shaders(vk_device& device, vk_pipeline requested) { CREATE_UNARY(neg) CREATE_UNARY(tanh) CREATE_UNARY(sigmoid) + const char * sigmoid_mul_env = getenv("GGML_VK_FUSE_SIGMOID_MUL"); + if (sigmoid_mul_env && strcmp(sigmoid_mul_env, "1") == 0 && + device->properties.limits.maxComputeWorkGroupInvocations >= 256 && + device->properties.limits.maxComputeWorkGroupSize[0] >= 256) { + ggml_vk_create_pipeline(device, device->pipeline_sigmoid_mul_f32, "sigmoid_mul_f32", + sigmoid_mul_f32_len, sigmoid_mul_f32_data, "main", 3, + sizeof(vk_sigmoid_mul_push_constants), {1024, 1, 1}, {}, 1); + } CREATE_UNARY(hardsigmoid) CREATE_UNARY(hardswish) CREATE_UNARY(abs) @@ -10186,6 +10227,23 @@ static void ggml_vk_mul_mat_q_f16(ggml_backend_vk_context * ctx, vk_context& sub vk_pipeline pipeline = ggml_vk_guess_matmul_pipeline(ctx, mmp, ne01, ne11, aligned, qx_needs_dequant ? f16_type : src0->type, effective_src1_type); + if (ctx->device->pipeline_prefill_q8_0 && dense_f16b && qy_needs_dequant && !qx_needs_dequant && + src0->type == GGML_TYPE_Q8_0 && src1->type == GGML_TYPE_F32 && dst->type == GGML_TYPE_F32 && + (ggml_prec) dst->op_params[0] == GGML_PREC_DEFAULT && ctx->num_additional_fused_ops == 0 && + ggml_is_contiguous(src0) && ggml_is_contiguous(src1) && ggml_is_contiguous(dst) && + ne02 == 1 && ne03 == 1 && ne12 == 1 && ne13 == 1 && dst->ne[2] == 1 && dst->ne[3] == 1 && + ne00 == ne10 && ne11 == 2048 && dst->ne[0] == (int64_t) ne01 && dst->ne[1] == 2048 && + ((ne10 == 2560 && (ne01 == 6144 || ne01 == 10240 || ne01 == 12288 || + ne01 == 2560 || ne01 == 512 || ne01 == 640)) || + (ne10 == 640 && ne01 == 2560) || (ne10 == 6144 && ne01 == 2560) || + (ne10 == 10240 && ne01 == 320) || (ne10 == 320 && ne01 == 10240))) { + // Retain base K alignment even for short K: K=320 uses the unaligned shader. + const bool prefill_aligned = ne10 % 128 == 0; + pipeline = ne10 == 320 ? + (prefill_aligned ? ctx->device->pipeline_prefill_q8_0_short_k_aligned : ctx->device->pipeline_prefill_q8_0_short_k) : + (prefill_aligned ? ctx->device->pipeline_prefill_q8_0_aligned : ctx->device->pipeline_prefill_q8_0); + } + if (ggml_nbytes(src0) > ctx->device->properties.limits.maxStorageBufferRange) { pipeline = ggml_vk_get_64b_indexing_pipeline(ctx, pipeline); } @@ -10998,6 +11056,28 @@ static bool ggml_vk_can_use_fwht(const ggml_backend_vk_context * ctx, const ggml return true; } +static void ggml_vk_sigmoid_mul(ggml_backend_vk_context * ctx, vk_context& subctx, + const ggml_tensor * a, const ggml_tensor * b, const ggml_tensor * dst) { + vk_pipeline pipeline = ctx->device->pipeline_sigmoid_mul_f32; + ggml_pipeline_request_descriptor_sets(ctx, pipeline, 1); + if (!subctx) { + return; + } + const vk_subbuffer a_buf = ggml_vk_tensor_subbuffer(ctx, a, true); + const vk_subbuffer b_buf = ggml_vk_tensor_subbuffer(ctx, b, true); + const vk_subbuffer d_buf = ggml_vk_tensor_subbuffer(ctx, dst, true); + const uint64_t bytes = ggml_nbytes(dst); + // The expanded descriptor range includes the residual, also for UMA host mappings. + const vk_sigmoid_mul_push_constants pc = { + (uint32_t) ggml_nelements(dst), + (uint32_t) ((a_buf.size - bytes) / sizeof(float)), + (uint32_t) ((b_buf.size - bytes) / sizeof(float)), + (uint32_t) ((d_buf.size - bytes) / sizeof(float)), + }; + ggml_vk_sync_buffers(ctx, subctx); + ggml_vk_dispatch_pipeline(ctx, subctx, pipeline, {a_buf, b_buf, d_buf}, pc, {pc.n, 1, 1}); +} + static void ggml_vk_fwht(ggml_backend_vk_context * ctx, vk_context& subctx, const ggml_tensor * src, ggml_tensor * dst) { const int idx = ggml_vk_fwht_pipeline_idx(src->ne[0]); vk_pipeline pipeline = ctx->device->pipeline_fwht_f32[idx]; @@ -11123,7 +11203,7 @@ static void ggml_vk_mul_mat_id_q_f16(ggml_backend_vk_context * ctx, vk_context& // n_as counts, n_as offsets, one total, then one packed row id per (expert, token). // Hoisting requires 16-bit indices for the packing and a table that fits one binding. const uint64_t hoisted_row_id_words = 2 * n_as + 1 + nei0 * nei1; - const bool hoist_row_ids = n_as <= 256 && nei0 <= 0xffff && nei1 <= 0xffff && + const bool hoist_row_ids = n_as <= 512 && nei0 <= 0xffff && nei1 <= 0xffff && hoisted_row_id_words * sizeof(uint32_t) <= ctx->device->properties.limits.maxStorageBufferRange; @@ -12867,6 +12947,17 @@ static void ggml_vk_flash_attn(ggml_backend_vk_context * ctx, vk_context& subctx mask != nullptr, use_mask_opt, logit_softcap != 0, k_type_eff, v_type_eff, fa_compact.dynamic_kv); + // Interleave complete query/head tiles for KV reuse on Strix Halo (measured at 64K). + // Keep the original KV traversal and only enable when split-K is not needed. + if (ctx->device->vendor_id == VK_VENDOR_ID_AMD && ctx->device->properties.deviceID == 0x1586 && + tuning_params.path == FA_COOPMAT1 && f32acc && N == 2048 && KV >= 32768 && + HSK == 256 && HSV == 256 && k_type_eff == GGML_TYPE_F16 && v_type_eff == GGML_TYPE_F16 && + neq2 == 24 && nek2 == 2 && nev2 == 2 && neq3 == 1 && nek3 == 1 && nev3 == 1 && + gqa_ratio == 1 && CEIL_DIV(N, tuning_params.block_rows)*workgroups_y >= + 2*std::max(ctx->device->shader_core_count, 16u)) { + fa_pipeline_state.flags |= 32; + } + vk_pipeline pipeline = nullptr; { @@ -18014,8 +18105,15 @@ static bool ggml_vk_build_graph(ggml_backend_vk_context * ctx, ggml_cgraph * cgr break; } + if (ctx->num_additional_fused_ops == 1 && ggml_get_unary_op(node) == GGML_UNARY_OP_SIGMOID) { + const ggml_tensor * mul = cgraph->nodes[node_idx + 1]; + const ggml_tensor * other = mul->src[0] == node ? mul->src[1] : mul->src[0]; + ggml_vk_sigmoid_mul(ctx, compute_ctx, node->src[0], other, mul); + break; + } + // Fused silu(x)*y: run it as a swiglu split, writing straight to the MUL's destination. - if (ctx->num_additional_fused_ops == 1) { + if (ctx->num_additional_fused_ops == 1 && ggml_get_unary_op(node) == GGML_UNARY_OP_SILU) { ggml_tensor * mul = cgraph->nodes[node_idx + 1]; ggml_tensor * other = (mul->src[0] == node) ? mul->src[1] : mul->src[0]; @@ -19069,13 +19167,107 @@ static bool ggml_vk_can_fuse(const ggml_backend_vk_context * ctx, const struct g // path, so the silu result makes a full round trip through memory. That is the same shape // swiglu-split already computes in one pass, so route the pair to the existing GLU pipeline. if (ops.size() == 2 && ops.begin()[0] == GGML_OP_UNARY && ops.begin()[1] == GGML_OP_MUL) { + const ggml_tensor * unary = cgraph->nodes[node_idx]; + const ggml_tensor * mul = cgraph->nodes[node_idx + 1]; + if (ggml_get_unary_op(unary) == GGML_UNARY_OP_SIGMOID) { + static const char * sigmoid_env = getenv("GGML_VK_FUSE_SIGMOID_MUL"); + if (!sigmoid_env || strcmp(sigmoid_env, "1") != 0 || !ctx->device->pipeline_sigmoid_mul_f32 || + (unary->flags & (GGML_TENSOR_FLAG_INPUT | GGML_TENSOR_FLAG_OUTPUT)) || + (mul->src[0] == unary) == (mul->src[1] == unary)) { + return false; + } + const ggml_tensor * other = mul->src[0] == unary ? mul->src[1] : mul->src[0]; + if (!other || (other->op == GGML_OP_UNARY && ggml_get_unary_op(other) == GGML_UNARY_OP_SIGMOID)) { + return false; + } + for (const ggml_tensor * t = unary->view_src; t; t = t->view_src) { + if (t->flags & (GGML_TENSOR_FLAG_INPUT | GGML_TENSOR_FLAG_OUTPUT)) { + return false; + } + } + const ggml_tensor * tensors[] = {unary->src[0], other, mul, unary}; + const auto & limits = ctx->device->properties.limits; + const int64_t n = ggml_nelements(mul); + // Bound CEIL_DIV's uint addition and every launched index, including tail lanes. + if (n <= 0 || uint64_t(n) > uint64_t(UINT32_MAX) - 1023 || + (uint64_t(n) + 1023) / 1024 > limits.maxComputeWorkGroupCount[0] || + limits.maxComputeWorkGroupCount[1] == 0 || limits.maxComputeWorkGroupCount[2] == 0 || + limits.maxComputeWorkGroupInvocations < 256 || limits.maxComputeWorkGroupSize[0] < 256) { + return false; + } + const uint64_t bytes = uint64_t(n) * sizeof(float); + vk_buffer buffers[4]; + uint64_t offsets[4]; + for (int j = 0; j < 4; ++j) { + const ggml_tensor * t = tensors[j]; + if (!t || t->type != GGML_TYPE_F32 || !ggml_is_contiguous(t) || + !ggml_are_same_shape(t, mul) || !t->data || !t->buffer || !t->buffer->context || + !ggml_backend_buffer_is_vk(t->buffer)) { + return false; + } + vk_buffer buffer = nullptr; + size_t offset = 0; + if (ctx->device->uma) { + ggml_vk_host_get(ctx->device, t->data, buffer, offset); + } + if (!buffer) { + buffer = ((ggml_backend_vk_buffer_context *) t->buffer->context)->dev_buffer; + const uint64_t base = vk_tensor_offset(t); + if (t->view_offs > SIZE_MAX - base) { + return false; + } + offset = base + t->view_offs; + } + const uint64_t residual = offset & (limits.minStorageBufferOffsetAlignment - 1); + if (!buffer || offset % sizeof(float) != 0 || residual % sizeof(float) != 0 || + offset > buffer->size || bytes > buffer->size - offset || + residual > limits.maxStorageBufferRange || bytes > limits.maxStorageBufferRange - residual || + residual / sizeof(float) + uint64_t(n) - 1 > UINT32_MAX) { + return false; + } + const vk_subbuffer sb = ggml_vk_tensor_subbuffer(ctx, t, true); + if (sb.buffer != buffer || sb.offset != offset - residual || sb.size != bytes + residual) { + return false; + } + buffers[j] = buffer; + offsets[j] = offset; + } + // An in-place sigmoid must not change B before the unfused MUL reads it. + if (buffers[1] == buffers[3] && + (offsets[1] < offsets[3] ? offsets[3] - offsets[1] : offsets[1] - offsets[3]) < bytes) { + return false; + } + // Check the actual read operands even for OP_NONE; the generic alias gate skips them. + for (int j : {0, 1, 3}) { + if (buffers[j] != buffers[2] || + (offsets[j] < offsets[2] ? offsets[2] - offsets[j] : offsets[j] - offsets[2]) >= bytes) { + continue; + } + if (offsets[j] != offsets[2]) { + return false; + } + for (const ggml_tensor * t = tensors[j]; t; t = t->view_src) { + if (t->flags & (GGML_TENSOR_FLAG_INPUT | GGML_TENSOR_FLAG_OUTPUT)) { + return false; + } + } + // Exact aliases are allowed only when the overwritten temporary is dead. + for (int k = node_idx + 2; k < cgraph->n_nodes; ++k) { + for (const ggml_tensor * src : cgraph->nodes[k]->src) { + for (const ggml_tensor * t = src; t && t != mul; t = t->view_src) { + if (t == tensors[j]) { + return false; + } + } + } + } + } + return true; + } static const char * env = getenv("GGML_VK_FUSE_UNARY_MUL"); if (!(env && atoi(env) != 0)) { return false; } - const ggml_tensor * unary = cgraph->nodes[node_idx]; - const ggml_tensor * mul = cgraph->nodes[node_idx + 1]; - if (ggml_get_unary_op(unary) != GGML_UNARY_OP_SILU) { return false; } @@ -19879,7 +20071,9 @@ static ggml_status ggml_backend_vk_graph_compute(ggml_backend_t backend, ggml_cg op_srcs_fused_elementwise[1] = true; } else if (ggml_vk_can_fuse(ctx, cgraph, i, { GGML_OP_UNARY, GGML_OP_MUL })) { ctx->num_additional_fused_ops = 1; - fusion_string = "SILU_MUL"; + fusion_string = ggml_get_unary_op(cgraph->nodes[i]) == GGML_UNARY_OP_SIGMOID ? "SIGMOID_MUL" : "SILU_MUL"; + op_srcs_fused_elementwise[0] = true; + op_srcs_fused_elementwise[1] = true; } else if (ggml_can_fuse_subgraph(cgraph, i, { GGML_OP_RMS_NORM, GGML_OP_MUL, GGML_OP_ROPE, GGML_OP_VIEW, GGML_OP_SET_ROWS }, { i + 4 }) && ggml_check_edges(cgraph, i, rms_norm_mul_rope_view_set_rows_edges) && ggml_vk_can_fuse_rms_norm_mul_rope(ctx, cgraph, i) && @@ -20052,6 +20246,7 @@ static ggml_status ggml_backend_vk_graph_compute(ggml_backend_t backend, ggml_cg } } if (need_disable) { + fusion_string = nullptr; ctx->num_additional_fused_ops = 0; ctx->fused_ops_write_mask = 1; ctx->fused_topk_moe_mode = TOPK_MOE_COUNT; diff --git a/ggml/src/ggml-vulkan/vulkan-shaders/count_experts.comp b/ggml/src/ggml-vulkan/vulkan-shaders/count_experts.comp index ef659959d950..3e55fd39030a 100644 --- a/ggml/src/ggml-vulkan/vulkan-shaders/count_experts.comp +++ b/ggml/src/ggml-vulkan/vulkan-shaders/count_experts.comp @@ -24,15 +24,16 @@ layout (push_constant) uniform parameter } p; #define BLOCK_SIZE 256 +#define MAX_EXPERTS 512 layout(local_size_x = BLOCK_SIZE, local_size_y = 1, local_size_z = 1) in; layout (binding = 0) readonly buffer A {uint data_a[];}; layout (binding = 1) writeonly buffer D {uint data_d[];}; -shared uint vals[BLOCK_SIZE]; -shared uint offsets[BLOCK_SIZE]; -shared uint cursors[BLOCK_SIZE]; +shared uint vals[MAX_EXPERTS]; +shared uint offsets[MAX_EXPERTS]; +shared uint cursors[MAX_EXPERTS]; // data_d layout when p.hoist_row_ids is set: // [0, n_experts) per-expert row count @@ -46,8 +47,8 @@ void main() { const uint tid = gl_LocalInvocationID.x; if (p.hoist_row_ids != 0) { - if (tid < p.n_experts) { - vals[tid] = 0; + for (uint expert = tid; expert < p.n_experts; expert += BLOCK_SIZE) { + vals[expert] = 0; } barrier(); diff --git a/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_base.glsl b/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_base.glsl index b562c5d78749..026af98e095a 100644 --- a/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_base.glsl +++ b/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_base.glsl @@ -29,6 +29,7 @@ const bool OLD_AMD_WINDOWS = (Flags & 8) != 0; // workgroup counts derive from neq1/neq2/neq3 and never from KV, so no indirect dispatch is // needed: only this loop bound changes. Folds away for every other pipeline. const bool DYNAMIC_KV = (Flags & 16) != 0; +const bool QUERY_HEAD_INTERLEAVE = (Flags & 32) != 0; // Round up head sizes to a multiple of 16, for coopmat1/coopmat2 paths const uint32_t HSK_pad = (HSK + 15) & ~15; @@ -196,6 +197,14 @@ void init_indices() iq2 = gl_WorkGroupID.y * gqa_ratio; iq3 = gl_WorkGroupID.z; + // Host enables this only for non-GQA, unsplit prefill. All original tile + // coordinates occur once; each retains the same KV loop and arithmetic. + if (QUERY_HEAD_INTERLEAVE) { + const uint32_t tile = gl_WorkGroupID.x + gl_NumWorkGroups.x * gl_WorkGroupID.y; + i = tile / p.neq2; + iq2 = tile % p.neq2; + } + // broadcast factors rk2 = p.neq2/p.nek2; rk3 = p.neq3/p.nek3; diff --git a/ggml/src/ggml-vulkan/vulkan-shaders/sigmoid_mul_f32.comp b/ggml/src/ggml-vulkan/vulkan-shaders/sigmoid_mul_f32.comp new file mode 100644 index 000000000000..43414309dc03 --- /dev/null +++ b/ggml/src/ggml-vulkan/vulkan-shaders/sigmoid_mul_f32.comp @@ -0,0 +1,28 @@ +#version 450 + +layout(local_size_x = 256, local_size_y = 1, local_size_z = 1) in; + +layout(binding = 0, std430) readonly buffer A { float data_a[]; }; +layout(binding = 1, std430) readonly buffer B { float data_b[]; }; +layout(binding = 2, std430) writeonly buffer D { float data_d[]; }; + +layout(push_constant) uniform PushConstants { + uint n; + uint a_offset; + uint b_offset; + uint d_offset; +} pc; + +void main() { + const uint base = gl_WorkGroupID.x * 1024 + gl_LocalInvocationID.x; + for (uint lane = 0; lane < 4; ++lane) { + const uint i = base + lane * 256; + if (i < pc.n) { + const float a = data_a[pc.a_offset + i]; + const float b = data_b[pc.b_offset + i]; + precise float sigmoid = 1.0 / (1.0 + exp(-a)); + precise float result = sigmoid * b; + data_d[pc.d_offset + i] = result; + } + } +} diff --git a/ggml/src/ggml-vulkan/vulkan-shaders/vulkan-shaders-gen.cpp b/ggml/src/ggml-vulkan/vulkan-shaders/vulkan-shaders-gen.cpp index 295e17b7b088..1552d9e701cc 100644 --- a/ggml/src/ggml-vulkan/vulkan-shaders/vulkan-shaders-gen.cpp +++ b/ggml/src/ggml-vulkan/vulkan-shaders/vulkan-shaders-gen.cpp @@ -970,6 +970,7 @@ void process_shaders() { string_to_spv("gelu_quick_f32", "unary.comp", {{"A_TYPE", "float"}, {"D_TYPE", "float"}, {"OP", "op_gelu_quick"}}); string_to_spv("silu_f16", "unary.comp", {{"A_TYPE", "float16_t"}, {"D_TYPE", "float16_t"}, {"OP", "op_silu"}}); string_to_spv("silu_f32", "unary.comp", {{"A_TYPE", "float"}, {"D_TYPE", "float"}, {"OP", "op_silu"}}); + string_to_spv("sigmoid_mul_f32", "sigmoid_mul_f32.comp", {}); string_to_spv("relu_f16", "unary.comp", {{"A_TYPE", "float16_t"}, {"D_TYPE", "float16_t"}, {"OP", "op_relu"}}); string_to_spv("relu_f32", "unary.comp", {{"A_TYPE", "float"}, {"D_TYPE", "float"}, {"OP", "op_relu"}}); string_to_spv("sqr_f16", "unary.comp", {{"A_TYPE", "float16_t"}, {"D_TYPE", "float16_t"}, {"OP", "op_sqr"}}); diff --git a/src/llama-context.cpp b/src/llama-context.cpp index e202c68400f6..8245972b1e97 100644 --- a/src/llama-context.cpp +++ b/src/llama-context.cpp @@ -1407,6 +1407,10 @@ llm_graph_result * llama_context::process_ubatch(const llama_ubatch & ubatch, ll { const auto t_inputs_us = ggml_time_us(); + // Integrated backends can read host inputs directly, even without pipeline parallelism. + // Finish their previous reads before overwriting masks or recurrent-state indices. + ggml_backend_sched_synchronize(sched.get()); + // FIXME this call causes a crash if any model inputs were not used in the graph and were therefore not allocated res->set_inputs(&ubatch); diff --git a/tests/CMakeLists.txt b/tests/CMakeLists.txt index f52433a4a2b6..158db7d51343 100644 --- a/tests/CMakeLists.txt +++ b/tests/CMakeLists.txt @@ -1,5 +1,12 @@ llama_add_compile_flags() +if (GGML_VULKAN) + add_executable(test-vulkan-fa-order test-vulkan-fa-order.cpp) + target_link_libraries(test-vulkan-fa-order PRIVATE ggml) + add_test(NAME test-vulkan-fa-order COMMAND test-vulkan-fa-order) + set_tests_properties(test-vulkan-fa-order PROPERTIES LABELS "vulkan" TIMEOUT 120 SKIP_RETURN_CODE 77) +endif() + function(llama_build source) set(TEST_SOURCES ${source} ${ARGN}) @@ -298,6 +305,7 @@ if (NOT LLAMA_SANITIZE_ADDRESS AND NOT GGML_SCHED_NO_REALLOC) llama_build_and_test(test-opt.cpp) endif() llama_build_and_test(test-backend-ops.cpp) +llama_build_and_test(test-flash-attn-reference.cpp) # the tensor API kernels come from a separate metallib - check they produce correct results # ref: https://github.com/ggml-org/llama.cpp/issues/27473 diff --git a/tests/test-backend-ops.cpp b/tests/test-backend-ops.cpp index 964c51728667..8442c8d0d33a 100644 --- a/tests/test-backend-ops.cpp +++ b/tests/test-backend-ops.cpp @@ -33,6 +33,7 @@ #include #include #include +#include #include #include #include @@ -7380,6 +7381,107 @@ struct test_sum_rows : public test_case { } }; +// Compound correctness fixture. Do not register in make_test_cases_perf(): it repeats only the last node. +struct test_relu_sum4 : public test_case { + const std::array ne; + const ggml_type type; + const std::string variant; + ggml_tensor * input = nullptr; + std::vector verify; + + test_relu_sum4(std::array ne, std::string variant = "packed", ggml_type type = GGML_TYPE_F32) + : ne(ne), type(type), variant(std::move(variant)) {} + + std::string op_desc(ggml_tensor *) override { return "RELU_SUM4"; } + std::string vars() override { return VARS_TO_STR3(ne, type, variant); } + bool run_whole_graph() override { return true; } + std::vector fusion_test_nodes() override { return verify; } + + ggml_tensor * build_graph(ggml_context * ctx) override { + verify.clear(); + input = ggml_new_tensor_4d(ctx, type, ne[0], ne[1], ne[2], ne[3]); + ggml_set_name(input, "scores"); + ggml_tensor * src = input; + if (variant == "strided") { + src = ggml_view_4d(ctx, input, ne[0] / 2, ne[1], ne[2], ne[3], input->nb[1], input->nb[2], input->nb[3], 0); + } else if (variant == "reshape") { + src = ggml_reshape_4d(ctx, ggml_reshape_2d(ctx, input, ne[0], ne[1] * ne[2] * ne[3]), ne[0], ne[1], ne[2], ne[3]); + } + ggml_tensor * relu = ggml_relu(ctx, src); + ggml_tensor * perm = variant == "wrong_permute" ? ggml_permute(ctx, relu, 0, 1, 2, 3) : ggml_permute(ctx, relu, 1, 0, 2, 3); + if (variant == "perm_strided") { + perm->ne[1] /= 2; + perm->nb[1] *= 2; + } + ggml_tensor * cont = ggml_cont(ctx, perm); + // Stock HIP SUM_ROWS requires F32; retain the F16 unary chain with an explicit conversion. + ggml_tensor * sum = ggml_sum_rows(ctx, type == GGML_TYPE_F32 ? cont : ggml_cast(ctx, cont, GGML_TYPE_F32)); + if (variant == "alias" || variant == "alias_offset") { + // Model allocation reuse explicitly; the ordinary context allocator does not reuse lifetimes. + sum->view_src = input; + sum->view_offs = variant == "alias_offset" ? ne[0] * sizeof(float) : 0; + } + ggml_tensor * out = ggml_reshape_3d(ctx, sum, sum->ne[1], sum->ne[2], sum->ne[3]); + verify.push_back(out); + ggml_tensor * intermediate = variant == "output_perm" || variant == "consumer_view" ? perm : + variant == "output_cont" || variant == "consumer_cont" ? cont : relu; + if (variant.compare(0, 7, "output_") == 0) { + ggml_set_output(intermediate); + verify.push_back(intermediate); + } else if (variant.compare(0, 9, "consumer_") == 0) { + ggml_build_forward_expand(gf, out); + if (variant == "consumer_view") { + intermediate = ggml_view_4d(ctx, intermediate, intermediate->ne[0], intermediate->ne[1], + intermediate->ne[2], intermediate->ne[3], intermediate->nb[1], intermediate->nb[2], intermediate->nb[3], 0); + ggml_set_output(intermediate); + } + ggml_tensor * external = ggml_cont(ctx, intermediate); + ggml_build_forward_expand(gf, external); + verify.push_back(external); + } + return out; + } + + void initialize_tensors(ggml_context * ctx) override { + for (ggml_tensor * t = ggml_get_first_tensor(ctx); t; t = ggml_get_next_tensor(ctx, t)) { + if (t->op == GGML_OP_NONE && !t->view_src) { + init_tensor_uniform(t, -4.0f, 4.0f); + } + } + if (variant != "exceptional") { + return; + } + const float tiny = std::numeric_limits::denorm_min(); + const float inf = std::numeric_limits::infinity(); + const float nan = std::numeric_limits::quiet_NaN(); + const float cases[][4] = { + { 0.0f, -0.0f, 0.0f, -0.0f }, { tiny, -tiny, tiny, 0.0f }, + { nan, -inf, 1.0f, 2.0f }, { inf, 1.0f, -inf, nan }, + { FLT_MAX * 0.75f, 0.0f, FLT_MAX * 0.75f, 0.0f }, + { 16777216.0f, 1.0f, 1.0f, 0.0f }, { 16777216.0f, 1.0f, 0.0f, 1.0f }, + }; + std::vector values(ggml_nelements(input)); + for (int64_t row = 0; row < ne[0] * ne[2] * ne[3]; ++row) { + for (int h = 0; h < 4; ++h) { + values[(row / ne[0]) * (4 * ne[0]) + h * ne[0] + row % ne[0]] = cases[row % 7][h]; + } + } + ggml_backend_tensor_set(input, values.data(), 0, ggml_nbytes(input)); + } + + double err(const float * a, const float * b, size_t n) override { + // The callback checks infinity signs. Exclude matching infinities from the finite CPU tolerance. + double result = 0.0; + for (size_t j = 0; j < n; ++j) { + if (std::isfinite(a[j]) && std::isfinite(b[j])) { + result = std::max(result, std::abs(double(a[j]) - b[j]) / std::max(1.0, std::abs(double(b[j])))); + } + } + return result; + } + double max_nmse_err() override { return 1e-6; } // Not the bitwise stock-GPU acceptance gate. +}; + // GGML_OP_MEAN struct test_mean : public test_case { const ggml_type type; @@ -9959,6 +10061,16 @@ static std::vector> make_test_cases_eval() { //add_test_bin_bcast(type, {3, 3, 2560, 1280}, {2, 1, 1, 1}); } + // Contiguous binary dispatch widths, neighboring fallbacks, and Flash Next prefill/decode shapes. + for (int64_t width : {2047, 2048, 2049, 2559, 2560, 2561}) { + for (int64_t tokens : {1, 2048}) { + add_test_bin_bcast(GGML_TYPE_F32, {width, tokens, 1, 1}, {1, 1, 1, 1}); + } + } + add_test_bin_bcast(GGML_TYPE_F32, {2560, 3, 2, 1}, {1, 1, 1, 1}, true); + add_test_bin_bcast(GGML_TYPE_F32, {2560, 3, 2, 1}, {1, 2, 1, 1}); + add_test_bin_bcast(GGML_TYPE_F16, {2560, 3, 2, 1}, {1, 1, 1, 1}); + // single inplace tests, especially important for WebGPU backend since kernels for inplace vs. not are different test_cases.emplace_back(new test_bin_bcast(ggml_add_inplace, GGML_TYPE_F32, {16, 5, 4, 3}, {1, 1, 1, 1}, 16)); test_cases.emplace_back(new test_bin_bcast(ggml_mul_inplace, GGML_TYPE_F32, {16, 5, 4, 3}, {1, 1, 1, 1}, 16)); @@ -10941,6 +11053,25 @@ static std::vector> make_test_cases_eval() { test_cases.emplace_back(new test_sum(GGML_TYPE_F32, { 33, 256, 1, 1 })); test_cases.emplace_back(new test_sum(GGML_TYPE_F32, { 33, 256, 1, 1 }, { 1, 0, 2, 3 })); // sum dst not-contiguous test_cases.emplace_back(new test_sum_rows()); + test_cases.emplace_back(new test_relu_sum4({ 8704, 4, 2048, 1 })); + test_cases.emplace_back(new test_relu_sum4({ 257, 4, 3, 1 })); // Partial 256-thread block. + test_cases.emplace_back(new test_relu_sum4({ 129, 4, 3, 2 })); + for (const char * variant : { "reshape", "exceptional", "alias", "alias_offset", "wrong_permute", "strided", "perm_strided", + "output_relu", "output_perm", "output_cont", "consumer_relu", "consumer_cont", "consumer_view" }) { + test_cases.emplace_back(new test_relu_sum4({ 257, 4, 3, 2 }, variant)); + } + for (int64_t heads : { 3, 5 }) { + test_cases.emplace_back(new test_relu_sum4({ 257, heads, 3, 2 })); + } + test_cases.emplace_back(new test_relu_sum4({ 257, 4, 3, 2 }, "packed", GGML_TYPE_F16)); + // Supply the selected GPU's nsm for semantic boundary coverage without a backend-specific test dependency. + if (const char * value = getenv("GGML_TEST_RELU_SUM4_NSM")) { + const int nsm = atoi(value); + GGML_ASSERT(nsm > 0 && nsm <= 1024); + for (int delta : { -1, 0, 1 }) { + test_cases.emplace_back(new test_relu_sum4({ 2 * nsm + delta, 4, 1, 1 })); + } + } test_cases.emplace_back(new test_sum_rows(GGML_TYPE_F32, { 11, 5, 6, 3 }, true, false)); test_cases.emplace_back(new test_sum_rows(GGML_TYPE_F32, { 11, 5, 6, 3 }, false, true)); test_cases.emplace_back(new test_sum_rows(GGML_TYPE_F32, { 11, 5, 6, 3 }, true, true)); @@ -11151,6 +11282,32 @@ static std::vector> make_test_cases_eval() { 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_Q2_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_Q2_0, GGML_TYPE_F16)); + // D256 GQA12 dense fallback at the 256-key clamp boundary, with native and cache strides. + for (int64_t kv : {255, 256, 257}) { + for (const std::array permute : {std::array{0, 1, 2, 3}, std::array{0, 2, 1, 3}}) { + test_cases.emplace_back(new test_flash_attn_ext(256, 256, 2, {12, 1}, kv, 64, true, false, 0, 0, + GGML_PREC_F32, GGML_TYPE_F16, GGML_TYPE_F16, permute)); + } + } + + // D256 WMMA without GQA reuse: query-width crossover, partial keys, masks and cache strides. + for (int64_t nb : {31, 32, 33}) { + for (bool mask : {false, true}) { + test_cases.emplace_back(new test_flash_attn_ext(256, 256, 2, {12, 1}, 2049, nb, mask, false, 0, 0, + GGML_PREC_F32, GGML_TYPE_F16, GGML_TYPE_F16, {0, 2, 1, 3})); + } + } + + // Narrow masked-QK dispatch: production cache layout and neighboring query/key lengths. + for (int64_t nb : {2047, 2048, 2049}) { + test_cases.emplace_back(new test_flash_attn_ext(256, 256, 2, {12, 1}, 34816, nb, true, false, 0, 0, + GGML_PREC_F32, GGML_TYPE_F16, GGML_TYPE_F16, {0, 2, 1, 3})); + } + for (int64_t kv : {34815, 34817}) { + test_cases.emplace_back(new test_flash_attn_ext(256, 256, 2, {12, 1}, kv, 2048, true, false, 0, 0, + GGML_PREC_F32, GGML_TYPE_F16, GGML_TYPE_F16, {0, 2, 1, 3})); + } + // q8_0 KV cases: decode and prompt batches, KV pad, permuted KV, feature flags, and long context test_cases.emplace_back(new test_flash_attn_ext(256, 256, 2, {16, 1}, 113, 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(256, 256, 2, {16, 1}, 1024, 1, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_Q8_0, GGML_TYPE_Q8_0)); @@ -11319,6 +11476,7 @@ static std::vector> make_test_cases_eval() { test_cases.emplace_back(new test_moe_weighted_reduction(63, 12, 33, true, true, true)); test_cases.emplace_back(new test_moe_weighted_reduction(2048, 15, 40, false, true)); test_cases.emplace_back(new test_moe_weighted_reduction(2048, 16, 32, false, true)); + test_cases.emplace_back(new test_moe_weighted_reduction(2560, 10, 2048)); test_cases.emplace_back(new test_gated_delta_net(GGML_TYPE_F32, 32, 128, 1, 1)); test_cases.emplace_back(new test_gated_delta_net(GGML_TYPE_F32, 32, 16, 1, 1)); diff --git a/tests/test-flash-attn-reference.cpp b/tests/test-flash-attn-reference.cpp new file mode 100644 index 000000000000..754f767807c7 --- /dev/null +++ b/tests/test-flash-attn-reference.cpp @@ -0,0 +1,55 @@ +#include "ggml.h" +#include "ggml-cpu.h" + +#include +#include +#include + +// Uniform logits have an analytical mean; alternating exact logits also exercise +// online max rescaling against an independent double-precision two-pass oracle. +int main() { + bool ok = true; + for (const int nk : {2049, 34815, 34816, 34817}) { + for (const bool varying : {false, true}) { + constexpr int d = 64; + ggml_context * ctx = ggml_init({32*1024*1024, nullptr, false}); + ggml_tensor * q = ggml_new_tensor_4d(ctx, GGML_TYPE_F32, d, 1, 1, 1); + ggml_tensor * k = ggml_new_tensor_4d(ctx, GGML_TYPE_F16, d, nk, 1, 1); + ggml_tensor * v = ggml_new_tensor_4d(ctx, GGML_TYPE_F16, d, nk, 1, 1); + ggml_set_zero(q); + ggml_set_zero(k); + static_cast(q->data)[0] = 1.0f; + double denominator = 0.0; + double numerator[d] = {}; + for (int i = 0; i < nk; ++i) { + const float score = varying ? float(i % 9 - 4)*0.5f : 0.0f; + static_cast(k->data)[i*d] = ggml_fp32_to_fp16(score); + const double weight = std::exp(double(score) - (varying ? 2.0 : 0.0)); + denominator += weight; + for (int j = 0; j < d; ++j) { + const float value = j % 2 == 0 ? 1.0f : float((i + j) % 17 - 8)*0.125f; + static_cast(v->data)[i*d + j] = ggml_fp32_to_fp16(value); + numerator[j] += weight*value; + } + } + ggml_tensor * out = ggml_flash_attn_ext(ctx, q, k, v, nullptr, 1.0f, 0.0f, 0.0f); + ggml_cgraph * graph = ggml_new_graph(ctx); + ggml_build_forward_expand(graph, out); + ggml_cplan plan = ggml_graph_plan(graph, 1, nullptr); + plan.use_ref = true; + std::vector work(plan.work_size); + plan.work_data = work.data(); + ok &= ggml_graph_compute(graph, &plan) == GGML_STATUS_SUCCESS; + double max_error = 0.0; + for (int j = 0; j < d; ++j) { + const double actual = static_cast(out->data)[j]; + const double error = std::abs(actual - numerator[j]/denominator); + ok &= std::isfinite(actual) && error < 2e-5; + max_error = std::fmax(max_error, error); + } + std::printf("CPU reference nk=%d varying=%d max_abs=%.9g %s\n", nk, varying, max_error, ok ? "OK" : "FAIL"); + ggml_free(ctx); + } + } + return ok ? 0 : 1; +} diff --git a/tests/test-masked-qk-attention.cpp b/tests/test-masked-qk-attention.cpp new file mode 100644 index 000000000000..3705d8d1f066 --- /dev/null +++ b/tests/test-masked-qk-attention.cpp @@ -0,0 +1,208 @@ +#include "ggml.h" +#include "ggml-alloc.h" +#include "ggml-backend.h" +#include +#include +#include +#include +#include +#include +#include +#include +#include + +static float next_value(uint32_t & s) { + s^=s<<13; s^=s>>17; s^=s<<5; + return (int32_t(s&2047)-1024)*(1.0f/1024.0f); +} +static uint64_t hash(const void * p,size_t n) { + uint64_t h=14695981039346656037ull; + for (size_t i=0;i=1024?(nq+7)/8*8:nq; + const int64_t past0=getenv("ATTN_DEPTH")?atoll(getenv("ATTN_DEPTH")):std::max(int64_t(0),nk-nq); + const int64_t hq=mode=="heads" ? 16 : 24,hk=2; + const bool no_mask=mode=="no-mask",padded=mode=="strided"; + const char * backenddir=getenv("ATTN_BACKEND_DIR"); + if (!backenddir) { return 2; } + auto ctx=ggml_init({ggml_tensor_overhead()*32+ggml_graph_overhead(),nullptr,true}); + if (!backend || !ctx) { return 3; } + const bool qwide=mode=="q-strided",native=mode=="native"; + auto qb=ggml_new_tensor_4d(ctx,GGML_TYPE_F32,d,native?nq:hq*(qwide?2:1),native?hq:nq,1); + auto q=qwide?ggml_view_4d(ctx,qb,d,hq,nq,1,qb->nb[1],qb->nb[2],qb->nb[3],0):qb; + if (!native) { q=ggml_permute(ctx,q,0,2,1,3); } + auto kb=ggml_new_tensor_4d(ctx,GGML_TYPE_F16,d,hk*(padded?2:1),nk,1); + auto vb=ggml_new_tensor_4d(ctx,GGML_TYPE_F16,d,hk*(padded?2:1),nk,1); + auto k=padded ? ggml_view_4d(ctx,kb,d,hk,nk,1,kb->nb[1],kb->nb[2],kb->nb[3],0) : kb; + auto v=padded ? ggml_view_4d(ctx,vb,d,hk,nk,1,vb->nb[1],vb->nb[2],vb->nb[3],0) : vb; + k=ggml_permute(ctx,k,0,2,1,3); v=ggml_permute(ctx,v,0,2,1,3); + auto m=ggml_new_tensor_4d(ctx,GGML_TYPE_F16,nk,nm,1,1); + auto out=ggml_flash_attn_ext(ctx,q,k,v,no_mask?nullptr:m,mode=="qoverflow"?2.0f:0.0625f,0,0); + ggml_flash_attn_ext_set_prec(out,GGML_PREC_F32); + ggml_set_name(out,"continuation-attention"); + auto graph=ggml_new_graph(ctx); ggml_build_forward_expand(graph,out); + auto buffer=ggml_backend_alloc_ctx_tensors(ctx,backend); + if (!buffer || !ggml_backend_supports_op(backend,out)) { return 3; } + uint32_t seed=0x13579bdf; + std::vector original_q(ggml_nelements(qb)); + std::vector original_k(ggml_nelements(kb)),original_v(ggml_nelements(vb)); + for (auto & x:original_q) { x=next_value(seed); } + for (auto & x:original_k) { x=ggml_fp32_to_fp16(next_value(seed)); } + for (auto & x:original_v) { x=ggml_fp32_to_fp16(next_value(seed)); } + const int steps=synthetic_replay?5:real_replay?3:1; + const int only=getenv("ATTN_STEP")?atoi(getenv("ATTN_STEP")):-1; + printf("CONFIG mode=%s nk=%lld hq=%lld padded=%d graph_default=1 q=%p k=%p v=%p mask=%p graph=%p backend=%s\n", + mode.c_str(),(long long)nk,(long long)hq,padded,q->data,k->data,v->data,m->data,(void *)graph,backenddir); + for (int step=0;step=0 && step!=only) { continue; } + auto qdata=original_q; + auto kdata=original_k,vdata=original_v; + std::vector mask(size_t(nk)*nm,mode=="pad-finite"?0:0xfc00); + for (int j=0;jpast-2048); + mask[size_t(j)*nk+i]=visible?0:0xfc00; + } + } + if (mode=="qoverflow") { qdata[23*256]=65504.0f; } + if (mode=="qnan") { qdata[23*256]=NAN; } + if (mode=="kinf") { kdata[32*(padded?1024:512)+256]=0x7c00; } + if (mode=="knan") { kdata[32*(padded?1024:512)+256]=0x7e01; } + if (mode=="vinf") { vdata[32*512]=0x7c00; } + if (mode=="vnan") { vdata[32*512]=0x7e01; } + if (mode=="minf" || mode=="mnan") { mask[32]=mode=="minf"?0x7c00:0x7e01; } + if (maskfile!="-") { + std::string path=maskfile; + if (real_replay && step==1) { + const auto pos=path.rfind("mask-1.f16"); + if (pos==std::string::npos) { return 2; } + path.replace(pos,10,"mask-12.f16"); + } + load(path,mask.data(),mask.size()*2); + } + if (synthetic_replay && step==1) { + qdata[23*256]=1.0e10f; + kdata[32*512+256]=0x7c00; + vdata[32*512]=0x7e01; + for (int j=0;j expected(ggml_nelements(out)),output(expected.size()); + const bool check=reference!="-"; + if (check) { + load(reference+"-step"+std::to_string(step)+".f32",expected.data(),expected.size()*4); + std::vector input_reference(strlen(inputs)); + load(reference+"-step"+std::to_string(step)+".inputs",input_reference.data(),input_reference.size()); + if (memcmp(inputs,input_reference.data(),input_reference.size())) { return 8; } + } + // In replay mode, verify each execution, including capture and subsequent replay. + for (int repeat=0;repeat<(replay?3:1);++repeat) { + compute(); + ggml_backend_tensor_get(out,output.data(),0,output.size()*4); + if (check && memcmp(output.data(),expected.data(),output.size()*4)) { + save(label+".FIRST_FAILURE.f32",output.data(),output.size()*4); + fprintf(stderr,"BIT_FAIL step=%d repeat=%d\n",step,repeat); return 7; + } + } + save(label+".f32",output.data(),output.size()*4); + size_t nonfinite=0; for (float x:output) { nonfinite+=!std::isfinite(x); } + const bool exceptional=mode=="qoverflow" || mode=="qnan" || mode=="kinf" || mode=="knan" || + mode=="vinf" || mode=="vnan" || mode=="minf" || mode=="mnan" || (synthetic_replay && step==1); + if (nonfinite && !exceptional) { return 6; } + printf("RAW_PASS step=%d elements=%zu nonfinite=%zu checked=%d %s",step,output.size(),nonfinite,check,inputs); + if (perf) { + if (!check || replay) { return 2; } + compute(); // One untimed warmup after the correctness invocation. + int64_t samples[6]; + for (int i=0;i<6;++i) { + ggml_backend_synchronize(backend); + const auto start=std::chrono::steady_clock::now(); compute(); + samples[i]=std::chrono::duration_cast(std::chrono::steady_clock::now()-start).count(); + } + ggml_backend_tensor_get(out,output.data(),0,output.size()*4); + save(label+".after.f32",output.data(),output.size()*4); + if (memcmp(output.data(),expected.data(),output.size()*4)) { fprintf(stderr,"BIT_FAIL after timing\n"); return 7; } + printf("TIMING_BIT_PASS step=%d warmup=1 samples=6 scope=whole_graph_sync\n",step); + for (int i=0;i<6;++i) { printf("sample_ns[%d]=%lld\n",i,(long long)samples[i]); } + } + ggml_backend_tensor_get(qb,qdata.data(),0,qdata.size()*4); + ggml_backend_tensor_get(kb,kdata.data(),0,kdata.size()*2); + ggml_backend_tensor_get(vb,vdata.data(),0,vdata.size()*2); + ggml_backend_tensor_get(m,mask.data(),0,mask.size()*2); + const uint64_t after[]={hash(qdata.data(),qdata.size()*4),hash(kdata.data(),kdata.size()*2), + hash(vdata.data(),vdata.size()*2),hash(mask.data(),mask.size()*2)}; + if (memcmp(hashes,after,sizeof(hashes))) { return 8; } + printf("INPUT_AFTER_PASS step=%d\n",step); + } + std::ifstream maps("/proc/self/maps"); std::string line; + while (std::getline(maps,line)) { if (line.find("libggml")!=std::string::npos) { printf("MAP %s\n",line.c_str()); } } + ggml_backend_buffer_free(buffer); ggml_free(ctx); + return 0; +} + +int main(int argc,char ** argv) { + if (argc!=5 || !getenv("ATTN_BACKEND_DIR")) { return 2; } + ggml_backend_load_all_from_path(getenv("ATTN_BACKEND_DIR")); + auto dev=ggml_backend_dev_by_name("ROCm0"); + if (!dev || ggml_backend_dev_type(dev)!=GGML_BACKEND_DEVICE_TYPE_IGPU) { return 3; } + auto backend=ggml_backend_dev_init(dev,nullptr); + if (!backend) { return 3; } + int rc=0; + if (strcmp(argv[1],"resize-replay")) { + rc=run_case(argc,argv,backend); + } else { + // Reuse the GPU backend/graph cache across changing Q and KV extents. + const int shapes[][2]={{63,256},{1025,1280},{2048,14080},{9,256},{2048,34816}}; + for (int i=0;i<5 && !rc;++i) { + const std::string t=std::to_string(shapes[i][0]),n=std::to_string(shapes[i][1]); + setenv("ATTN_T",t.c_str(),1); setenv("ATTN_N",n.c_str(),1); setenv("ATTN_DEPTH","0",1); + std::string mode="replay-synthetic",prefix=std::string(argv[2])+"-shape"+std::to_string(i); + std::string reference=strcmp(argv[4],"-")?std::string(argv[4])+"-shape"+std::to_string(i):"-"; + char * args[]={argv[0],mode.data(),prefix.data(),argv[3],reference.data()}; + rc=run_case(5,args,backend); + } + } + ggml_backend_free(backend); + return rc; +} diff --git a/tests/test-masked-qk-flags.cu b/tests/test-masked-qk-flags.cu new file mode 100644 index 000000000000..85ecea9a2fc3 --- /dev/null +++ b/tests/test-masked-qk-flags.cu @@ -0,0 +1,95 @@ +#include "ggml-cuda/fattn-common.cuh" +#include +#include +#include + +int main(int argc, char ** argv) { + if (argc != 4) { return 2; } + const int t=atoi(argv[1]), n=atoi(argv[2]), wide=atoi(argv[3]); + const int a=(t+7)/8, b=n/32, m=t>=1024?8*a:t; + const size_t count=masked_qk_flag_count(t,n); + if (!count) { return 2; } + const size_t qs=24*256*(wide?2:1), ks=2*256*(wide?2:1); + // No padded Q rows allocated; last-head exceptions include adjacent half fields. + std::vector q(size_t(t)*qs,0.5f); + std::vector k(size_t(n)*ks,0x3800), mask(size_t(m)*n,0xfc00); + for (int j=0;j=512?n-256:n) && (i<3 || (i<=n-t+j && i>n-t+j-64))) { mask[size_t(j)*n+i]=0; } + } + } + float * dq; ggml_fp16_t * dk,* dm; int * df; + CUDA_CHECK(cudaMalloc(&dq,q.size()*4)); CUDA_CHECK(cudaMalloc(&dk,k.size()*2)); + CUDA_CHECK(cudaMalloc(&dm,mask.size()*2)); CUDA_CHECK(cudaMalloc(&df,count*4)); + int previous_last_bound=-1; + for (int step=0;step<5;++step) { + const float scale=step==1?2.0f:0.0625f; + q[(t-1)*qs+23*256+255]=step==1?65504.0f:step==2?NAN:0.5f; + k[(n-1)*ks+256+255]=step==2?0x7c00:step==3?0x7e01:0x3800; + mask[32]=step==1?0x7c00:step==2?0x7e01:0xfc00; + if (n>=512) { + for (int j=0;jt) { std::fill(mask.begin()+size_t(t)*n,mask.end(),step==4?0:0xfc00); } + CUDA_CHECK(cudaMemcpy(dq,q.data(),q.size()*4,cudaMemcpyHostToDevice)); + CUDA_CHECK(cudaMemcpy(dk,k.data(),k.size()*2,cudaMemcpyHostToDevice)); + CUDA_CHECK(cudaMemcpy(dm,mask.data(),mask.size()*2,cudaMemcpyHostToDevice)); + CUDA_CHECK(cudaMemset(df,0xa5,count*4)); + if (t>=1024) { + flash_attn_mask_to_KV_max<8><<>>((const half2 *)dm,df,n/256,n/2,size_t(m)*n/2); + } + masked_qk_q_finite<<>>((const char *)dq,df+a,qs*4,256*4,scale,t,t<1024?df:nullptr,n); + masked_qk_k_finite<<>>((const char *)dk,df+2*a,ks*2,256*2,b); + masked_qk_mask_empty<<>>((const unsigned short *)dm,df+2*a+2*b,t,n,b); + CUDA_CHECK(cudaGetLastError()); + std::vector actual(count), expected(count,1); + CUDA_CHECK(cudaMemcpy(actual.data(),df,count*4,cudaMemcpyDeviceToHost)); + for (int jt=0;jt=1024) { + for (int i=0;it && n>=512 && step==4) { + if (previous_last_bound>n-256 || actual[a-1]!=n) { return 6; } + printf("PHYSICAL_PADDING_PASS T=%d N=%d before=%d after=%d\n",t,n,previous_last_bound,actual[a-1]); + } + previous_last_bound=actual[a-1]; + printf("FLAG_PASS T=%d N=%d wide=%d step=%d all_ints=%zu qgrid=%d kgrid=%d emptygrid=%d\n",t,n,wide,step,count,a,2*b,a*b/8); + } + CUDA_CHECK(cudaFree(dq)); CUDA_CHECK(cudaFree(dk)); CUDA_CHECK(cudaFree(dm)); CUDA_CHECK(cudaFree(df)); +} diff --git a/tests/test-query-pieces.cpp b/tests/test-query-pieces.cpp new file mode 100644 index 000000000000..528f5d09bca8 --- /dev/null +++ b/tests/test-query-pieces.cpp @@ -0,0 +1,35 @@ +#include "fattn-masked-qk-layout.h" +#include +#include +#include + +int main() { + size_t ranges = 0; + for (int n = 32768; n <= 1050624; n += 256) { + const int kv_tiles = n/32; + const int tiles = 1024; + const int64_t work = int64_t(tiles)*kv_tiles; + for (unsigned blocks : {1u, 3u, 7u, 20u, 40u, 64u, 1024u}) { + const unsigned pieces = fattn_query_piece_count(tiles, kv_tiles, blocks); + assert(pieces > 0); + for (unsigned b = 0; b < blocks; ++b) { + const int64_t begin = int64_t(b)*work/blocks; + const int64_t end = (int64_t(b) + 1)*work/blocks; + int64_t cursor = begin; + for (unsigned p = 0; p < pieces; ++p) { + const int64_t tile = begin/kv_tiles + int64_t(p)*fattn_query_tiles_per_piece; + const int64_t lo = std::max(begin, tile*kv_tiles); + const int64_t hi = std::min(end, (tile + fattn_query_tiles_per_piece)*kv_tiles); + if (lo >= hi) { continue; } + assert(lo == cursor); + assert(lo == begin || lo % kv_tiles == 0); + assert(hi == end || hi % kv_tiles == 0); + cursor = hi; + } + assert(cursor == end); + ++ranges; + } + } + } + std::printf("PASS: %zu logical ranges partitioned exactly, KV lengths 32768..1050624\n", ranges); +} diff --git a/tests/test-vulkan-fa-order.cpp b/tests/test-vulkan-fa-order.cpp new file mode 100644 index 000000000000..e3ee03356610 --- /dev/null +++ b/tests/test-vulkan-fa-order.cpp @@ -0,0 +1,110 @@ +#include "ggml.h" +#include "ggml-alloc.h" +#include "ggml-backend.h" + +#include +#include +#include +#include +#include +#include +#include + +static float random_value(uint32_t & state) { + state ^= state << 13; + state ^= state >> 17; + state ^= state << 5; + return (int(state & 2047) - 1024)/1024.0f; +} + +// Smaller query views select the unchanged Vulkan path. Query rows are +// independent, and both paths use the same CM1 tile and unsplit KV traversal. +static bool check(ggml_backend_t backend, int64_t nk, int64_t nq) { + constexpr int64_t d = 256, hq = 24, hk = 2, chunk = 512; + const size_t context_size = 128*ggml_tensor_overhead() + 6*ggml_graph_overhead(); + ggml_context * ctx = ggml_init({context_size, nullptr, true}); + GGML_ASSERT(ctx); + auto qb = ggml_new_tensor_4d(ctx, GGML_TYPE_F32, d, hq, nq, 1); + auto kb = ggml_new_tensor_4d(ctx, GGML_TYPE_F16, d, hk, nk, 1); + auto vb = ggml_new_tensor_4d(ctx, GGML_TYPE_F16, d, hk, nk, 1); + auto q = ggml_permute(ctx, qb, 0, 2, 1, 3); + auto k = ggml_permute(ctx, kb, 0, 2, 1, 3); + auto v = ggml_permute(ctx, vb, 0, 2, 1, 3); + auto mask = ggml_new_tensor_4d(ctx, GGML_TYPE_F16, nk, nq, 1, 1); + auto full = ggml_flash_attn_ext(ctx, q, k, v, mask, 0.0625f, 0, 0); + ggml_flash_attn_ext_set_prec(full, GGML_PREC_F32); + auto graph = ggml_new_graph(ctx); + ggml_build_forward_expand(graph, full); + std::vector parts; + std::vector graphs; + for (int64_t start = 0; start < nq; start += chunk) { + const int64_t rows = std::min(chunk, nq - start); + auto qs = ggml_view_4d(ctx, q, d, rows, hq, 1, q->nb[1], q->nb[2], q->nb[3], start*q->nb[1]); + auto ms = ggml_view_4d(ctx, mask, nk, rows, 1, 1, mask->nb[1], mask->nb[2], mask->nb[3], start*mask->nb[1]); + auto part = ggml_flash_attn_ext(ctx, qs, k, v, ms, 0.0625f, 0, 0); + ggml_flash_attn_ext_set_prec(part, GGML_PREC_F32); + auto part_graph = ggml_new_graph(ctx); + ggml_build_forward_expand(part_graph, part); + parts.push_back(part); + graphs.push_back(part_graph); + } + auto buffer = ggml_backend_alloc_ctx_tensors(ctx, backend); + GGML_ASSERT(buffer && ggml_backend_supports_op(backend, full)); + uint32_t seed = 0x13579bdf; + std::vector qdata(ggml_nelements(qb)); + std::vector kdata(ggml_nelements(kb)), vdata(ggml_nelements(vb)); + for (auto & x : qdata) { x = random_value(seed); } + for (auto & x : kdata) { x = ggml_fp32_to_fp16(random_value(seed)); } + for (auto & x : vdata) { x = ggml_fp32_to_fp16(random_value(seed)); } + std::vector mdata(size_t(nk)*nq, 0xfc00); + for (int64_t row = 0; row < nq; ++row) { + std::fill_n(mdata.data() + row*nk, nk - nq + row + 1, ggml_fp16_t(0)); + } + ggml_backend_tensor_set(qb, qdata.data(), 0, qdata.size()*sizeof(float)); + ggml_backend_tensor_set(kb, kdata.data(), 0, kdata.size()*sizeof(ggml_fp16_t)); + ggml_backend_tensor_set(vb, vdata.data(), 0, vdata.size()*sizeof(ggml_fp16_t)); + ggml_backend_tensor_set(mask, mdata.data(), 0, mdata.size()*sizeof(ggml_fp16_t)); + std::vector reference(ggml_nelements(full)), actual(reference.size()); + size_t offset = 0; + for (size_t i = 0; i < parts.size(); ++i) { + ggml_backend_tensor_memset(parts[i], 0xa5, 0, ggml_nbytes(parts[i])); + GGML_ASSERT(ggml_backend_graph_compute(backend, graphs[i]) == GGML_STATUS_SUCCESS); + ggml_backend_synchronize(backend); + ggml_backend_tensor_get(parts[i], reference.data() + offset, 0, ggml_nbytes(parts[i])); + offset += ggml_nelements(parts[i]); + } + GGML_ASSERT(offset == reference.size()); + ggml_backend_tensor_memset(full, 0x7f, 0, ggml_nbytes(full)); + GGML_ASSERT(ggml_backend_graph_compute(backend, graph) == GGML_STATUS_SUCCESS); + ggml_backend_synchronize(backend); + ggml_backend_tensor_get(full, actual.data(), 0, ggml_nbytes(full)); + bool valid = memcmp(reference.data(), actual.data(), ggml_nbytes(full)) == 0; + for (float x : actual) { valid &= std::isfinite(x); } + printf("%s N=%lld Q=%lld words=%zu reference_query_chunk=512\n", + valid ? "PASS" : "FAIL", (long long) nk, (long long) nq, actual.size()); + ggml_backend_buffer_free(buffer); + ggml_free(ctx); + return valid; +} + +int main(int argc, char ** argv) { + ggml_backend_load_all(); + auto dev = ggml_backend_dev_by_name("Vulkan0"); + if (!dev || std::string(ggml_backend_dev_description(dev)).find("RADV STRIX_HALO") == std::string::npos) { + puts("SKIP: test targets RADV STRIX_HALO"); + return 77; + } + auto backend = ggml_backend_dev_init(dev, nullptr); + GGML_ASSERT(backend); + bool valid = true; + if (argc == 2) { + const int64_t n = std::strtoll(argv[1], nullptr, 10); + GGML_ASSERT(n >= 2048 && n <= 130048); + valid = check(backend, n, 2048); + } else { + for (int64_t n : {32512, 32768, 33024}) { valid &= check(backend, n, 2048); } + valid &= check(backend, 32768, 2047); + } + ggml_backend_free(backend); + return valid ? 0 : 1; +}