diff --git a/ggml/src/ggml-cuda/fattn-mma-f16.cuh b/ggml/src/ggml-cuda/fattn-mma-f16.cuh index a95dc4dc..2a558bce 100644 --- a/ggml/src/ggml-cuda/fattn-mma-f16.cuh +++ b/ggml/src/ggml-cuda/fattn-mma-f16.cuh @@ -480,32 +480,83 @@ static __device__ __forceinline__ void flash_attn_ext_f16_store_q4_0_block( } // Dequantize one q8_0 block (32 weights) given its scale and the 32 int8 bytes as 8 words, store 16 half2. +// Each byte pair becomes the half2 1024 + (q + 128) by placing it under the exponent byte 0x64, minus 1152 that is +// exactly q. d*q is exact in float (11 x 8 significant bits), so the correctly rounded half product equals the +// rounding of d*float(q) that dequantize_q8_0 -> half does: the tile is bit-identical, at a fraction of the cost. static __device__ __forceinline__ void flash_attn_ext_f16_store_q8_0_block( - half2 * const __restrict__ dst, const float d, const uint32_t * const __restrict__ q) { + half2 * const __restrict__ dst, const half d, const uint32_t * const __restrict__ q) { + const half2 d2 = __half2half2(d); + const half2 bias = make_half2(1152.0f, 1152.0f); #pragma unroll for (int m = 0; m < 8; ++m) { - const uint32_t v = q[m]; - dst[2*m ] = __floats2half2_rn(d*float(int(int8_t(v & 0xFF))), d*float(int(int8_t((v >> 8) & 0xFF)))); - dst[2*m + 1] = __floats2half2_rn(d*float(int(int8_t((v >> 16) & 0xFF))), d*float(int(int8_t((v >> 24) & 0xFF)))); + const uint32_t u = q[m] ^ 0x80808080u; // int8 -> uint8 offset by 128 + const uint32_t lo = __byte_perm(u, 0x64646464u, 0x4140); + const uint32_t hi = __byte_perm(u, 0x64646464u, 0x4342); + dst[2*m ] = __hmul2(__hsub2(*((const half2 *) &lo), bias), d2); + dst[2*m + 1] = __hmul2(__hsub2(*((const half2 *) &hi), bias), d2); + } +} + +// A q4_0 / q8_0 K or V row is read as block pairs (64 weights): a q4_0 pair is 36 bytes and a q8_0 pair 68 bytes, +// both 4-byte aligned at even block offsets, so the pair is read as aligned 32-bit words and the odd block's bytes +// are realigned with byte permutes. +template +struct flash_attn_ext_f16_q_pair { + static constexpr int bs = 32; // weights per block, both types + static constexpr int ts = type_KV == GGML_TYPE_Q4_0 ? 18 : 34; // bytes per block + static constexpr int nwords = 2*ts/4; // 32-bit words per block pair: 9 or 17 + static_assert((2*ts) % 4 == 0, "block pair is not word aligned"); +}; + +template +static __device__ __forceinline__ void flash_attn_ext_f16_load_q_pair( + const char * const __restrict__ src, uint32_t * const __restrict__ w) { +#pragma unroll + for (int k = 0; k < flash_attn_ext_f16_q_pair::nwords; ++k) { + w[k] = __ldg((const uint32_t *) src + k); + } +} + +// Dequantize a block pair from its words into 32 half2 (64 weights). +template +static __device__ __forceinline__ void flash_attn_ext_f16_dequantize_q_pair( + const uint32_t * const __restrict__ w, half2 * const __restrict__ dst) { + static_assert(type_KV == GGML_TYPE_Q4_0 || type_KV == GGML_TYPE_Q8_0, "unsupported quantized K/V type"); + constexpr int bs = flash_attn_ext_f16_q_pair::bs; + if constexpr (type_KV == GGML_TYPE_Q4_0) { + // block 0: d = bytes 0-1, qs = bytes 2-17. block 1: d = bytes 18-19, qs = bytes 20-35 (= w5..w8). + const float d0 = __half2float(__ushort_as_half((unsigned short) (w[0] & 0xFFFF))); + const float d1 = __half2float(__ushort_as_half((unsigned short) (w[4] >> 16))); + flash_attn_ext_f16_store_q4_0_block(dst, d0, + __byte_perm(w[0], w[1], 0x5432), __byte_perm(w[1], w[2], 0x5432), + __byte_perm(w[2], w[3], 0x5432), __byte_perm(w[3], w[4], 0x5432)); + flash_attn_ext_f16_store_q4_0_block(dst + bs/2, d1, w[5], w[6], w[7], w[8]); + } else { + // block 0: d = bytes 0-1, qs = bytes 2-33. block 1: d = bytes 34-35, qs = bytes 36-67 (= w9..w16). + const half d0 = __ushort_as_half((unsigned short) (w[0] & 0xFFFF)); + const half d1 = __ushort_as_half((unsigned short) (w[8] >> 16)); + uint32_t q0[8]; +#pragma unroll + for (int k = 0; k < 8; ++k) { + q0[k] = __byte_perm(w[k], w[k + 1], 0x5432); + } + flash_attn_ext_f16_store_q8_0_block(dst, d0, q0); + flash_attn_ext_f16_store_q8_0_block(dst + bs/2, d1, w + 9); } } // Loads a [nbatch_fa x ne] tile of a q4_0 / q8_0 K or V cache and dequantizes it into the F16 shared tile, in the // same layout the F16 loader produces. KV points at byte 0 of tile row 0, e0 / ne are the element offset and count -// within a row (multiples of 64), stride_KV is the row stride in bytes. Each thread handles two adjacent blocks -// (64 weights): a q4_0 pair is 36 bytes and a q8_0 pair 68 bytes, both 4-byte aligned at even block offsets, so the -// pair is read as aligned 32-bit words and the odd block's bytes are realigned with byte permutes. +// within a row (multiples of 64), stride_KV is the row stride in bytes. Each thread handles one block pair at a time. template static __device__ __forceinline__ void flash_attn_ext_f16_load_tile_q( const char * const __restrict__ KV, half2 * const __restrict__ tile_KV, const int e0, const int ne, const int stride_KV, const int i_sup) { - static_assert(type_KV == GGML_TYPE_Q4_0 || type_KV == GGML_TYPE_Q8_0, "unsupported quantized K/V type"); constexpr int warp_size = ggml_cuda_get_physical_warp_size(); constexpr int nthreads = nwarps*warp_size; - constexpr int bs = 32; // weights per block, both types - constexpr int ts = type_KV == GGML_TYPE_Q4_0 ? 18 : 34; // bytes per block - constexpr int nwords = 2*ts/4; // 32-bit words per block pair: 9 or 17 - static_assert((2*ts) % 4 == 0, "block pair is not word aligned"); + constexpr int bs = flash_attn_ext_f16_q_pair::bs; + constexpr int ts = flash_attn_ext_f16_q_pair::ts; + constexpr int nwords = flash_attn_ext_f16_q_pair::nwords; const int pairs_per_row = ne / (2*bs); const int ntotal = nbatch_fa * pairs_per_row; @@ -523,32 +574,89 @@ static __device__ __forceinline__ void flash_attn_ext_f16_load_tile_q( continue; } - const uint32_t * src = (const uint32_t *) (KV + int64_t(i)*stride_KV + ((e0 + p*2*bs)/bs)*ts); uint32_t w[nwords]; -#pragma unroll - for (int k = 0; k < nwords; ++k) { - w[k] = __ldg(src + k); + flash_attn_ext_f16_load_q_pair(KV + int64_t(i)*stride_KV + ((e0 + p*2*bs)/bs)*ts, w); + flash_attn_ext_f16_dequantize_q_pair(w, dst); + } +} + +// Prompt-sized tiles over a quantized cache are software pipelined: the raw rows of the next K or V tile and the next +// mask tile are copied into shared memory with cp.async while the tensor cores work on the current tile, then +// dequantized shared -> shared. Needs whole 16-byte aligned rows (head size 256: 272 bytes of q8_0, 144 of q4_0) +// and no OOB rows (ncols2 > 1 implies a padded cache); the staging buffers fit in the Q tile's shared memory. +static constexpr __host__ __device__ int ggml_cuda_fattn_mma_q_row_bytes(const ggml_type type_KV, const int D) { + return D/32 * (type_KV == GGML_TYPE_Q4_0 ? 18 : 34); +} + +static constexpr __host__ __device__ bool ggml_cuda_fattn_mma_kv_prefetch_shape( + const ggml_type type_KV, const int DKQ, const int DV, const int ncols1, const int ncols2) { + return ggml_cuda_fattn_mma_kv_native(type_KV) && ncols1*ncols2 >= 64 && ncols2 > 1 && + ggml_cuda_fattn_mma_q_row_bytes(type_KV, DKQ) % 16 == 0 && ggml_cuda_fattn_mma_q_row_bytes(type_KV, DV) % 16 == 0; +} + +static __host__ bool ggml_cuda_fattn_mma_kv_prefetch( + const ggml_type type_KV, const int DKQ, const int DV, const int ncols1, const int ncols2, const int cc) { + return cp_async_available(cc) && ggml_cuda_fattn_mma_kv_prefetch_shape(type_KV, DKQ, DV, ncols1, ncols2); +} + +static constexpr __device__ bool ggml_cuda_fattn_mma_kv_prefetch( + const ggml_type type_KV, const int DKQ, const int DV, const int ncols1, const int ncols2) { +#ifdef CP_ASYNC_AVAILABLE + return ggml_cuda_fattn_mma_kv_prefetch_shape(type_KV, DKQ, DV, ncols1, ncols2); +#else + GGML_UNUSED_VARS(type_KV, DKQ, DV, ncols1, ncols2); + return false; +#endif // CP_ASYNC_AVAILABLE +} + +// Starts copying the raw rows of a [nbatch_fa x row_bytes] quantized K/V tile into shared memory, row after row. +template +static __device__ __forceinline__ void flash_attn_ext_f16_load_raw_q( + const char * const __restrict__ KV, char * const __restrict__ raw, const int stride_KV) { + constexpr int warp_size = ggml_cuda_get_physical_warp_size(); + constexpr int nthreads = nwarps*warp_size; + constexpr int cpr = row_bytes/16; // 16-byte chunks per row + constexpr int nchunks = nbatch_fa*cpr; + static_assert(row_bytes % 16 == 0, "rows must be whole 16-byte chunks"); + + const unsigned int raw_32 = ggml_cuda_cvta_generic_to_shared(raw); + // Not unrolled: the 64-bit addresses of all chunks would spill the Q and VKQ registers at head size 256. +#pragma unroll 1 + for (int c0 = 0; c0 < nchunks; c0 += nthreads) { + const int c = c0 + threadIdx.y*warp_size + threadIdx.x; + if (nchunks % nthreads == 0 || c < nchunks) { + const int i = c / cpr; + const int k = c % cpr; + cp_async_cg_16<0>(raw_32 + i*row_bytes + 16*k, KV + int64_t(i)*stride_KV + 16*k); } + } +} - if constexpr (type_KV == GGML_TYPE_Q4_0) { - // block 0: d = bytes 0-1, qs = bytes 2-17. block 1: d = bytes 18-19, qs = bytes 20-35 (= w5..w8). - const float d0 = __half2float(__ushort_as_half((unsigned short) (w[0] & 0xFFFF))); - const float d1 = __half2float(__ushort_as_half((unsigned short) (w[4] >> 16))); - flash_attn_ext_f16_store_q4_0_block(dst, d0, - __byte_perm(w[0], w[1], 0x5432), __byte_perm(w[1], w[2], 0x5432), - __byte_perm(w[2], w[3], 0x5432), __byte_perm(w[3], w[4], 0x5432)); - flash_attn_ext_f16_store_q4_0_block(dst + bs/2, d1, w[5], w[6], w[7], w[8]); - } else { - // block 0: d = bytes 0-1, qs = bytes 2-33. block 1: d = bytes 34-35, qs = bytes 36-67 (= w9..w16). - const float d0 = __half2float(__ushort_as_half((unsigned short) (w[0] & 0xFFFF))); - const float d1 = __half2float(__ushort_as_half((unsigned short) (w[8] >> 16))); - uint32_t q0[8]; -#pragma unroll - for (int k = 0; k < 8; ++k) { - q0[k] = __byte_perm(w[k], w[k + 1], 0x5432); +// Dequantizes a raw [nbatch_fa x D] tile from shared memory into the F16 tile, block pair by block pair. +template +static __device__ __forceinline__ void flash_attn_ext_f16_dequantize_raw_q( + const char * const __restrict__ raw, half2 * const __restrict__ tile_KV) { + constexpr int warp_size = ggml_cuda_get_physical_warp_size(); + constexpr int nthreads = nwarps*warp_size; + constexpr int row_bytes = ggml_cuda_fattn_mma_q_row_bytes(type_KV, D); + constexpr int ts = flash_attn_ext_f16_q_pair::ts; + constexpr int nwords = flash_attn_ext_f16_q_pair::nwords; + constexpr int ppr = D/64; // block pairs per row + constexpr int ntotal = nbatch_fa*ppr; + +#pragma unroll + for (int t0 = 0; t0 < ntotal; t0 += nthreads) { + const int t = t0 + threadIdx.y*warp_size + threadIdx.x; + if (ntotal % nthreads == 0 || t < ntotal) { + const int i = t / ppr; + const int p = t % ppr; + const uint32_t * src = (const uint32_t *) (raw + i*row_bytes + p*(2*ts)); + uint32_t w[nwords]; +#pragma unroll + for (int k = 0; k < nwords; ++k) { + w[k] = src[k]; } - flash_attn_ext_f16_store_q8_0_block(dst, d0, q0); - flash_attn_ext_f16_store_q8_0_block(dst + bs/2, d1, w + 9); + flash_attn_ext_f16_dequantize_q_pair(w, tile_KV + i*stride_tile + p*32); } } } @@ -683,6 +791,15 @@ static __device__ __forceinline__ void flash_attn_ext_f16_iter( static_assert(type_KV == GGML_TYPE_F16 || ((2*nbatch_K2) % 64 == 0 && (2*nbatch_V2) % 64 == 0 && !V_is_K_view), "quantized K/V loading needs whole block pairs per pass"); + constexpr bool kv_prefetch = ggml_cuda_fattn_mma_kv_prefetch(type_KV, DKQ, DV, ncols1, ncols2); + static_assert(!kv_prefetch || (nstages == 0 && nbatch_K2 == DKQ/2 && nbatch_V2 == DV/2 && !oob_check), + "K/V prefetch needs whole rows per pass and no OOB rows"); + constexpr int row_bytes_K = ggml_cuda_fattn_mma_q_row_bytes(type_KV, DKQ); + constexpr int row_bytes_V = ggml_cuda_fattn_mma_q_row_bytes(type_KV, DV); + // Raw staging tiles behind the mask tile, see flash_attn_ext_f16_process_tile. + char * raw_K = (char *) (tile_mask + ncols1*(nbatch_fa + 8)); + char * raw_V = raw_K + nbatch_fa*row_bytes_K; + const int k_VKQ_0 = kb0 * nbatch_fa; #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))]; @@ -701,6 +818,13 @@ static __device__ __forceinline__ void flash_attn_ext_f16_iter( __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); + } else if constexpr (kv_prefetch) { + // Raw K and the mask of this tile were started by the previous iteration (or before the first one); + // the barrier also retires every warp's VKQ reads of tile_V before tile_K (the same memory) is written. + cp_async_wait_all(); + __syncthreads(); + flash_attn_ext_f16_load_raw_q + ((const char *) V_h2 + int64_t(k_VKQ_0)*stride_V, raw_V, stride_V); } else { constexpr bool use_cp_async = nstages == 1; if (ncols2 > 1 || mask_h) { @@ -724,6 +848,8 @@ static __device__ __forceinline__ void flash_attn_ext_f16_iter( if (use_cp_async) { cp_async_wait_all(); } + } else if constexpr (kv_prefetch) { + flash_attn_ext_f16_dequantize_raw_q(raw_K, tile_K); } else { flash_attn_ext_f16_load_tile_q ((const char *) K_h2 + int64_t(k_VKQ_0)*stride_K, tile_K, 2*k0_start, 2*k0_diff, stride_K, k_VKQ_sup); @@ -783,7 +909,7 @@ static __device__ __forceinline__ void flash_attn_ext_f16_iter( } } - if constexpr (nstages <= 1) { + if constexpr (nstages <= 1 && !kv_prefetch) { __syncthreads(); // Only needed if tile_K == tile_V. } } @@ -1081,6 +1207,17 @@ static __device__ __forceinline__ void flash_attn_ext_f16_iter( if (use_cp_async) { cp_async_wait_all(); } + } else if constexpr (kv_prefetch) { + // Every warp is past KQ and the softmax, so raw K and the mask tile can take the next tile. + cp_async_wait_all(); + __syncthreads(); + if constexpr (!last_iter) { + 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_raw_q + ((const char *) K_h2 + int64_t(k_VKQ_0 + nbatch_fa)*stride_K, raw_K, stride_K); + } + flash_attn_ext_f16_dequantize_raw_q(raw_V, tile_V); } else { flash_attn_ext_f16_load_tile_q ((const char *) V_h2 + int64_t(k_VKQ_0)*stride_V, tile_V, i0_start, i0_diff, stride_V, k_VKQ_sup); @@ -1131,7 +1268,8 @@ static __device__ __forceinline__ void flash_attn_ext_f16_iter( } #endif // defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE) || defined(AMD_MFMA_AVAILABLE) - if constexpr (nstages <= 1) { + // With prefetch the barrier at the start of the next iteration (or after the loop) protects tile_V. + if constexpr (nstages <= 1 && !(kv_prefetch && !last_iter)) { __syncthreads(); // Only needed if tile_K == tile_V. } } @@ -1376,6 +1514,15 @@ static __device__ __forceinline__ void flash_attn_ext_f16_process_tile( int kb0 = kb0_start; + // Quantized prompt tiles: start the raw K and mask tiles of the first iteration, see flash_attn_ext_f16_iter. + if constexpr (ggml_cuda_fattn_mma_kv_prefetch(type_KV, DKQ, DV, ncols1, ncols2)) { + char * raw_K = (char *) (tile_mask + ncols1*(nbatch_fa + 8)); + flash_attn_ext_f16_load_mask + (mask_h + kb0*nbatch_fa, tile_mask, stride_mask, nbatch_fa, jt*ncols1, ne01); + flash_attn_ext_f16_load_raw_q + ((const char *) K_h2 + int64_t(kb0)*nbatch_fa*stride_K, raw_K, stride_K); + } + // Preload mask and K data for first iteration when using cp_async with multiple stages: if constexpr (nstages > 1) { static_assert(nbatch_K2 == DKQ/2, "batching not implemented for multi-stage pipeline"); @@ -2051,7 +2198,11 @@ void ggml_cuda_flash_attn_ext_mma_f16_case(ggml_backend_cuda_context & ctx, ggml const size_t nbytes_shared_mask = ncols1 * (nbatch_fa/2 + 4) * sizeof(half2); const size_t nbytes_shared_combine = nwarps*cols_per_warp * (nbatch_combine + 4) * sizeof(half2); - const size_t nbytes_shared_KV = nstages <= 1 ? nbytes_shared_KV_1stage : nbytes_shared_KV_2stage; + // Raw staging tiles of the quantized prompt-tile pipeline, behind the mask tile. + const size_t nbytes_shared_raw = ggml_cuda_fattn_mma_kv_prefetch(type_KV, DKQ, DV, ncols1, ncols2, cc) ? + nbatch_fa * (ggml_cuda_fattn_mma_q_row_bytes(type_KV, DKQ) + ggml_cuda_fattn_mma_q_row_bytes(type_KV, DV)) : 0; + + const size_t nbytes_shared_KV = (nstages <= 1 ? nbytes_shared_KV_1stage : nbytes_shared_KV_2stage) + nbytes_shared_raw; const size_t nbytes_shared_total = std::max(nbytes_shared_combine, Q_in_reg ? std::max(nbytes_shared_Q, nbytes_shared_KV + nbytes_shared_mask) : @@ -2140,9 +2291,11 @@ static inline bool ggml_cuda_fattn_mma_kv_native_supported(const ggml_tensor * d if (!(Q->ne[0] == 128 || Q->ne[0] == 256) || V->ne[0] != Q->ne[0]) { return false; } - // Rows must be 32-bit addressable for the block-pair loads. + // Rows must be 32-bit addressable for the block-pair loads, and whole 16-byte chunks for the cp.async copies of + // the prompt tiles at head size 256 (see ggml_cuda_fattn_mma_kv_prefetch). + const size_t align = Q->ne[0] == 256 ? 16 : 4; for (const ggml_tensor * t : {K, V}) { - if (t->nb[1] % 4 != 0 || t->nb[2] % 4 != 0 || t->nb[3] % 4 != 0) { + if (t->nb[1] % align != 0 || t->nb[2] % align != 0 || t->nb[3] % align != 0 || (uintptr_t) t->data % align != 0) { return false; } } diff --git a/ggml/src/ggml-cuda/fattn.cu b/ggml/src/ggml-cuda/fattn.cu index 62e8e9b8..39655868 100644 --- a/ggml/src/ggml-cuda/fattn.cu +++ b/ggml/src/ggml-cuda/fattn.cu @@ -42,6 +42,21 @@ static void ggml_cuda_flash_attn_ext_mma_f16_switch_ncols1(ggml_backend_cuda_con ggml_cuda_flash_attn_ext_mma_f16_case(ctx, dst); } +// Whether launch_fattn gives every output tile of ncols1 x ncols2 columns to one block, i.e. does not use stream-k +// (mirrors its rule, with the occupancy the MMA config targets). +static bool ggml_cuda_fattn_mma_whole_tiles(const ggml_tensor * dst, const int DKQ, const int DV, const int ncols1, const int ncols2, const int cc) { + if (cc >= GGML_CUDA_CC_ADA_LOVELACE || amd_wmma_available(cc)) { + return false; + } + const ggml_tensor * Q = dst->src[0]; + const ggml_tensor * K = dst->src[1]; + const int64_t gqa_ratio = Q->ne[2] / K->ne[2]; + const int64_t ntiles = (Q->ne[1] + ncols1 - 1)/ncols1 * ((gqa_ratio + ncols2 - 1)/ncols2) * K->ne[2] * Q->ne[3]; + const int64_t max_blocks = int64_t(ggml_cuda_fattn_mma_get_occupancy(DKQ, DV, ncols1*ncols2, cc)) * ggml_cuda_info().devices[ggml_cuda_get_device()].nsm; + const int64_t nwaves = (ntiles + max_blocks - 1) / max_blocks; + return 100 * ntiles / (max_blocks*nwaves) >= 75; +} + template static void ggml_cuda_flash_attn_ext_mma_f16_switch_ncols2(ggml_backend_cuda_context & ctx, ggml_tensor * dst) { const int cc = ggml_cuda_info().devices[ggml_cuda_get_device()].cc; @@ -97,6 +112,24 @@ static void ggml_cuda_flash_attn_ext_mma_f16_switch_ncols2(ggml_backend_cuda_con } } + // Prompt batches over a quantized cache: tile the Q heads of a KV head without padding columns (24 over 4 KV heads + // as 3 pairs, not 6 of 8), so no tensor-core work is wasted and every K/V tile is streamed fewer times. Each + // column's arithmetic does not depend on ncols2, so the result is the same as long as neither tiling splits a + // column's KV range across blocks (stream-k): small batches and decode/verify keep the 8-wide tile. + if constexpr (ggml_cuda_fattn_mma_kv_native(type_KV)) { + if (use_gqa_opt && Q->ne[1] > 16 && gqa_ratio > 4 && gqa_ratio % 8 != 0 && gqa_ratio % 2 == 0) { + if (gqa_ratio % 4 == 0) { + if (ggml_cuda_fattn_mma_whole_tiles(dst, DKQ, DV, 16, 4, cc) && ggml_cuda_fattn_mma_whole_tiles(dst, DKQ, DV, 8, 8, cc)) { + ggml_cuda_flash_attn_ext_mma_f16_switch_ncols1(ctx, dst); + return; + } + } else if (ggml_cuda_fattn_mma_whole_tiles(dst, DKQ, DV, 32, 2, cc) && ggml_cuda_fattn_mma_whole_tiles(dst, DKQ, DV, 8, 8, cc)) { + ggml_cuda_flash_attn_ext_mma_f16_switch_ncols1(ctx, dst); + return; + } + } + } + if (use_gqa_opt && gqa_ratio > 4) { ggml_cuda_flash_attn_ext_mma_f16_switch_ncols1(ctx, dst); return; diff --git a/tests/test-backend-ops.cpp b/tests/test-backend-ops.cpp index c93cacf5..57ead864 100644 --- a/tests/test-backend-ops.cpp +++ b/tests/test-backend-ops.cpp @@ -10171,6 +10171,15 @@ static std::vector> make_test_cases_eval() { } } + // prompt-shaped cases with the same shape: the quantized cache is tiled two query heads per column group + for (ggml_type type_KV : { GGML_TYPE_F16, GGML_TYPE_Q8_0, GGML_TYPE_Q4_0, }) { + for (int nb : { 9, 33, 64, 300, 1024, }) { + for (int kv : { 512, 4096, }) { + test_cases.emplace_back(new test_flash_attn_ext(256, 256, 4, {6, 1}, kv, nb, true, false, 0, 0, GGML_PREC_F32, type_KV, type_KV)); + } + } + } + // prefill-shaped cases with long KV (nb >= 32, kv >= 1024): covers the // XMX/GEMM-accelerated SYCL FA path which only activates for these shapes. for (int kv : { 1024, 2048, }) { @@ -10746,6 +10755,24 @@ static std::vector> make_test_cases_perf() { test_cases.emplace_back(new test_flash_attn_ext(256, 256, 2, {16, 1}, 10000, 512, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_F16, GGML_TYPE_F16)); test_cases.emplace_back(new test_flash_attn_ext(256, 256, 2, {16, 1}, 20000, 512, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_F16, GGML_TYPE_F16)); + // prefill at depth with the Qwen 3.x full-attention shape (24 query heads over 4 KV heads) + for (int kv : { 17408, 33792, 66560, }) { + for (int nb : { 512, 1024, }) { + for (ggml_type type_KV : { GGML_TYPE_F16, GGML_TYPE_Q8_0, }) { + test_cases.emplace_back(new test_flash_attn_ext(256, 256, 4, {6, 1}, kv, nb, true, false, 0, 0, GGML_PREC_F32, type_KV, type_KV)); + } + } + } + + // prompt-shaped cases with the same shape: the quantized cache is tiled two query heads per column group + for (ggml_type type_KV : { GGML_TYPE_F16, GGML_TYPE_Q8_0, GGML_TYPE_Q4_0, }) { + for (int nb : { 9, 33, 64, 300, 1024, }) { + for (int kv : { 512, 4096, }) { + test_cases.emplace_back(new test_flash_attn_ext(256, 256, 4, {6, 1}, kv, nb, true, false, 0, 0, GGML_PREC_F32, type_KV, type_KV)); + } + } + } + for (int kv : { 4096, 8192, 16384, }) { for (int hs : { 64, 128, }) { for (int nr : { 1, 4, }) {