diff --git a/ggml/src/ggml-cuda/fattn-mma-f16.cuh b/ggml/src/ggml-cuda/fattn-mma-f16.cuh index 7f4cfd55..a95dc4dc 100644 --- a/ggml/src/ggml-cuda/fattn-mma-f16.cuh +++ b/ggml/src/ggml-cuda/fattn-mma-f16.cuh @@ -358,6 +358,22 @@ static constexpr __device__ int ggml_cuda_fattn_mma_get_nstages(const int DKQ, c #endif // CP_ASYNC_AVAILABLE } +// K/V read straight from a quantized cache. The tile loader dequantizes into the F16 shared tile, so there is no +// F16 copy of the whole cache in the compute buffer (n_ctx*D*n_head_kv*2 bytes for K and again for V, reserved +// for the worst case at graph allocation) and no full-cache conversion pass per micro-batch. cp.async cannot +// dequantize, so these variants always run the single-stage register loader. +static constexpr __host__ __device__ bool ggml_cuda_fattn_mma_kv_native(const ggml_type type_KV) { + return type_KV == GGML_TYPE_Q4_0 || type_KV == GGML_TYPE_Q8_0; +} + +static __host__ int ggml_cuda_fattn_mma_get_nstages(const int DKQ, const int DV, const int ncols1, const int ncols2, const int cc, const ggml_type type_KV) { + return type_KV == GGML_TYPE_F16 ? ggml_cuda_fattn_mma_get_nstages(DKQ, DV, ncols1, ncols2, cc) : 0; +} + +static constexpr __device__ int ggml_cuda_fattn_mma_get_nstages(const int DKQ, const int DV, const int ncols1, const int ncols2, const ggml_type type_KV) { + return type_KV == GGML_TYPE_F16 ? ggml_cuda_fattn_mma_get_nstages(DKQ, DV, ncols1, ncols2) : 0; +} + // ------------------------------------------------------------------------------------------------------------------ template @@ -447,6 +463,96 @@ static __device__ __forceinline__ void flash_attn_ext_f16_load_tile( } } +// Dequantize one q4_0 block (32 weights) given its scale and the 16 nibble bytes as 4 words, store 16 half2. +// Same arithmetic as dequantize_q4_0 -> half, so the tile is bit-identical to a to_fp16 pass over the cache. +static __device__ __forceinline__ void flash_attn_ext_f16_store_q4_0_block( + half2 * const __restrict__ dst, const float d, const uint32_t q0, const uint32_t q1, const uint32_t q2, const uint32_t q3) { + const uint32_t q[4] = {q0, q1, q2, q3}; +#pragma unroll + for (int m = 0; m < 4; ++m) { + const uint32_t v = q[m]; + // low nibbles are weights 4m..4m+3, high nibbles are weights 16+4m..16+4m+3 + dst[ 2*m ] = __floats2half2_rn(d*(float(int((v ) & 0xF)) - 8.0f), d*(float(int((v >> 8) & 0xF)) - 8.0f)); + dst[ 2*m + 1] = __floats2half2_rn(d*(float(int((v >> 16) & 0xF)) - 8.0f), d*(float(int((v >> 24) & 0xF)) - 8.0f)); + dst[8 + 2*m ] = __floats2half2_rn(d*(float(int((v >> 4) & 0xF)) - 8.0f), d*(float(int((v >> 12) & 0xF)) - 8.0f)); + dst[8 + 2*m + 1] = __floats2half2_rn(d*(float(int((v >> 20) & 0xF)) - 8.0f), d*(float(int((v >> 28) & 0xF)) - 8.0f)); + } +} + +// Dequantize one q8_0 block (32 weights) given its scale and the 32 int8 bytes as 8 words, store 16 half2. +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) { +#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)))); + } +} + +// 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. +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"); + + const int pairs_per_row = ne / (2*bs); + const int ntotal = nbatch_fa * pairs_per_row; + + for (int t = threadIdx.y*warp_size + threadIdx.x; t < ntotal; t += nthreads) { + const int i = t / pairs_per_row; + const int p = t - i*pairs_per_row; + half2 * dst = tile_KV + i*stride_tile + p*bs; // 2 blocks = 64 weights = 32 half2 per pair + + if (oob_check && i >= i_sup) { +#pragma unroll + for (int k = 0; k < bs; ++k) { + dst[k] = make_half2(0.0f, 0.0f); + } + 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); + } + + 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); + } + 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); + } + } +} + template static __device__ __forceinline__ void flash_attn_ext_f16_load_mask( const half * const __restrict__ mask_h, half * const __restrict__ tile_mask, @@ -529,7 +635,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, + ggml_type type_KV> static __device__ __forceinline__ void flash_attn_ext_f16_iter( const float2 * const __restrict__ Q_f2, const half2 * const __restrict__ K_h2, @@ -566,12 +673,16 @@ static __device__ __forceinline__ void flash_attn_ext_f16_iter( constexpr int nbatch_K2 = ggml_cuda_fattn_mma_get_nbatch_K2(DKQ, DV, ncols); constexpr int nbatch_V2 = ggml_cuda_fattn_mma_get_nbatch_V2(DKQ, DV, ncols); constexpr bool Q_in_reg = ggml_cuda_fattn_mma_get_Q_in_reg (DKQ, DV, ncols); - constexpr int nstages = ggml_cuda_fattn_mma_get_nstages (DKQ, DV, ncols1, ncols2); + constexpr int nstages = ggml_cuda_fattn_mma_get_nstages (DKQ, DV, ncols1, ncols2, type_KV); constexpr int stride_tile_K = nbatch_K2 + 4; constexpr int stride_tile_V = V_is_K_view ? stride_tile_K : nbatch_V2 + 4; + // For quantized K/V, stride_K / stride_V are byte strides and the loader dequantizes block pairs (64 weights). + 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"); + 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))]; @@ -606,11 +717,16 @@ static __device__ __forceinline__ void flash_attn_ext_f16_iter( if constexpr (nstages <= 1) { const int k0_diff = k0_stop - k0_start; - constexpr bool use_cp_async = nstages == 1; - flash_attn_ext_f16_load_tile - (K_h2 + int64_t(k_VKQ_0)*stride_K + k0_start, tile_K, k0_diff, stride_K, k_VKQ_sup); - if (use_cp_async) { - cp_async_wait_all(); + if constexpr (type_KV == GGML_TYPE_F16) { + constexpr bool use_cp_async = nstages == 1; + flash_attn_ext_f16_load_tile + (K_h2 + int64_t(k_VKQ_0)*stride_K + k0_start, tile_K, k0_diff, stride_K, k_VKQ_sup); + if (use_cp_async) { + cp_async_wait_all(); + } + } 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); } __syncthreads(); } @@ -958,11 +1074,16 @@ static __device__ __forceinline__ void flash_attn_ext_f16_iter( if constexpr (nstages <= 1) { const int i0_diff = i0_stop - i0_start; if (!V_is_K_view || i0_stop > 2*nbatch_K2) { - constexpr bool use_cp_async = nstages == 1; - flash_attn_ext_f16_load_tile - (V_h2 + int64_t(k_VKQ_0)*stride_V + i0_start/2, tile_V, i0_diff/2, stride_V, k_VKQ_sup); - if (use_cp_async) { - cp_async_wait_all(); + if constexpr (type_KV == GGML_TYPE_F16) { + constexpr bool use_cp_async = nstages == 1; + flash_attn_ext_f16_load_tile + (V_h2 + int64_t(k_VKQ_0)*stride_V + i0_start/2, tile_V, i0_diff/2, stride_V, k_VKQ_sup); + if (use_cp_async) { + cp_async_wait_all(); + } + } 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); } __syncthreads(); } @@ -1113,7 +1234,7 @@ 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, @@ -1158,7 +1279,7 @@ static __device__ __forceinline__ void flash_attn_ext_f16_process_tile( constexpr int nbatch_V2 = ggml_cuda_fattn_mma_get_nbatch_V2 (DKQ, DV, ncols); constexpr int nbatch_combine = ggml_cuda_fattn_mma_get_nbatch_combine(DKQ, DV, ncols); constexpr bool Q_in_reg = ggml_cuda_fattn_mma_get_Q_in_reg (DKQ, DV, ncols); - constexpr int nstages = ggml_cuda_fattn_mma_get_nstages (DKQ, DV, ncols1, ncols2); + constexpr int nstages = ggml_cuda_fattn_mma_get_nstages (DKQ, DV, ncols1, ncols2, type_KV); if (cols_per_warp > ncols) { NO_DEVICE_CODE; @@ -1277,7 +1398,7 @@ 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, type_KV> (Q_f2, K_h2, V_h2, mask_h, 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); @@ -1286,7 +1407,7 @@ static __device__ __forceinline__ void flash_attn_ext_f16_process_tile( 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, type_KV> (Q_f2, K_h2, V_h2, mask_h, 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); @@ -1297,7 +1418,7 @@ 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, type_KV> (Q_f2, K_h2, V_h2, mask_h, 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); @@ -1306,7 +1427,7 @@ 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, type_KV> (Q_f2, K_h2, V_h2, mask_h, 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); @@ -1700,7 +1821,7 @@ static __device__ __forceinline__ void flash_attn_ext_f16_process_tile( #endif // defined(VOLTA_MMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE) || defined(AMD_MFMA_AVAILABLE) } -template +template __launch_bounds__(ggml_cuda_fattn_mma_get_nthreads(DKQ, DV, ncols1*ncols2), ggml_cuda_fattn_mma_get_occupancy(DKQ, DV, ncols1*ncols2)) static __global__ void flash_attn_ext_f16( const char * Q_ptr, @@ -1782,10 +1903,11 @@ static __global__ void flash_attn_ext_f16( const int stride_Q1 = nb01 / sizeof(float2); const int stride_Q2 = nb02 / sizeof(float2); - const int stride_K = nb11 / sizeof(half2); + // Quantized K/V: byte strides, the tile loader dequantizes block pairs straight out of the cache. + const int stride_K = type_KV == GGML_TYPE_F16 ? nb11 / sizeof(half2) : nb11; const int stride_mask = nb31 / sizeof(half); - const int stride_V = V_is_K_view ? stride_K : nb21 / sizeof(half2); + const int stride_V = V_is_K_view ? stride_K : (type_KV == GGML_TYPE_F16 ? nb21 / sizeof(half2) : nb21); const int iter_k = (ne11 + (nbatch_fa - 1)) / nbatch_fa; const int iter_j = (ne01.z + (ncols1 - 1)) / ncols1; @@ -1829,12 +1951,12 @@ 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, sinks_f, dstk, dst_meta, scale, slope, logit_softcap, ne01, ne02, gqa_ratio, ne11, stride_Q1, stride_Q2, stride_K, stride_V, stride_mask, jt, zt_gqa, kb0_start, kb0_stop); } else { constexpr bool needs_fixup = true; // CUDA block is missing the beginning of a tile. - flash_attn_ext_f16_process_tile + flash_attn_ext_f16_process_tile (Q_f2, K_h2, V_h2, mask_h, sinks_f, dstk, dst_meta, scale, slope, logit_softcap, ne01, ne02, gqa_ratio, ne11, stride_Q1, stride_Q2, stride_K, stride_V, stride_mask, jt, zt_gqa, kb0_start, kb0_stop); } @@ -1875,7 +1997,7 @@ 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, sinks_f, dstk, dst_meta, scale, slope, logit_softcap, ne01, ne02, gqa_ratio, ne11, stride_Q1, stride_Q2, stride_K, stride_V, stride_mask, jt, zt_gqa, kb0_start, kb0_stop); #else @@ -1892,7 +2014,7 @@ static __global__ void flash_attn_ext_f16( #endif // defined(FLASH_ATTN_AVAILABLE) && (defined(VOLTA_MMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE) || defined(AMD_MFMA_AVAILABLE)) } -template +template void ggml_cuda_flash_attn_ext_mma_f16_case(ggml_backend_cuda_context & ctx, ggml_tensor * dst) { const ggml_tensor * KQV = dst; const int id = ggml_cuda_get_device(); @@ -1900,13 +2022,22 @@ void ggml_cuda_flash_attn_ext_mma_f16_case(ggml_backend_cuda_context & ctx, ggml constexpr int ncols = ncols1 * ncols2; + if constexpr (type_KV != GGML_TYPE_F16) { + // The kernel reads the quantized cache in place: same K/V type, whole block pairs per row, word-aligned rows. + const ggml_tensor * K = dst->src[1]; + const ggml_tensor * V = dst->src[2]; + GGML_ASSERT(K->type == type_KV && V->type == type_KV); + GGML_ASSERT(K->nb[1] % 4 == 0 && K->nb[2] % 4 == 0 && K->nb[3] % 4 == 0); + GGML_ASSERT(V->nb[1] % 4 == 0 && V->nb[2] % 4 == 0 && V->nb[3] % 4 == 0); + } + const int nthreads = ggml_cuda_fattn_mma_get_nthreads (DKQ, DV, ncols, cc); const int nbatch_fa = ggml_cuda_fattn_mma_get_nbatch_fa (DKQ, DV, ncols, cc); const int nbatch_K2 = ggml_cuda_fattn_mma_get_nbatch_K2 (DKQ, DV, ncols, cc); const int nbatch_V2 = ggml_cuda_fattn_mma_get_nbatch_V2 (DKQ, DV, ncols, cc); const int nbatch_combine = ggml_cuda_fattn_mma_get_nbatch_combine(DKQ, DV, ncols, cc); const bool Q_in_reg = ggml_cuda_fattn_mma_get_Q_in_reg (DKQ, DV, ncols, cc); - const int nstages = ggml_cuda_fattn_mma_get_nstages (DKQ, DV, ncols1, ncols2, cc); + const int nstages = ggml_cuda_fattn_mma_get_nstages (DKQ, DV, ncols1, ncols2, cc, type_KV); const int cols_per_warp = std::min(ncols, get_cols_per_warp(cc)); const int warp_size_host = ggml_cuda_info().devices[ctx.device].warp_size; @@ -1937,7 +2068,7 @@ void ggml_cuda_flash_attn_ext_mma_f16_case(ggml_backend_cuda_context & ctx, ggml fattn_kernel_t fattn_kernel; if (logit_softcap == 0.0f) { constexpr bool use_logit_softcap = false; - fattn_kernel = flash_attn_ext_f16; + fattn_kernel = flash_attn_ext_f16; #if !defined(GGML_USE_MUSA) static bool shared_memory_limit_raised[GGML_CUDA_MAX_DEVICES] = {false}; @@ -1948,7 +2079,7 @@ void ggml_cuda_flash_attn_ext_mma_f16_case(ggml_backend_cuda_context & ctx, ggml #endif // !defined(GGML_USE_MUSA) } else { constexpr bool use_logit_softcap = true; - fattn_kernel = flash_attn_ext_f16; + fattn_kernel = flash_attn_ext_f16; #if !defined(GGML_USE_MUSA) static bool shared_memory_limit_raised[GGML_CUDA_MAX_DEVICES] = {false}; @@ -1959,8 +2090,9 @@ void ggml_cuda_flash_attn_ext_mma_f16_case(ggml_backend_cuda_context & ctx, ggml #endif // !defined(GGML_USE_MUSA) } + constexpr bool need_f16_KV = type_KV == GGML_TYPE_F16; // quantized variants read the cache in place launch_fattn - (ctx, dst, fattn_kernel, nwarps, nbytes_shared_total, nbatch_fa, true, true, true, warp_size_host); + (ctx, dst, fattn_kernel, nwarps, nbytes_shared_total, nbatch_fa, need_f16_KV, need_f16_KV, true, warp_size_host); } @@ -1968,6 +2100,55 @@ void ggml_cuda_flash_attn_ext_mma_f16_case(ggml_backend_cuda_context & ctx, ggml template void ggml_cuda_flash_attn_ext_mma_f16_case \ (ggml_backend_cuda_context & ctx, ggml_tensor * dst) \ +// In-place quantized K/V variants, instantiated for the head sizes and GQA tilings a quantized cache is served with. +#define DECL_FATTN_MMA_F16_CASE_KV(DKQ, DV, ncols1, ncols2, type_KV) \ + template void ggml_cuda_flash_attn_ext_mma_f16_case \ + (ggml_backend_cuda_context & ctx, ggml_tensor * dst) \ + +#define DECL_FATTN_MMA_F16_CASE_KV_ALL_NCOLS2(DKQ, DV, ncols, type_KV) \ + extern DECL_FATTN_MMA_F16_CASE_KV(DKQ, DV, (ncols)/ 1, 1, type_KV); \ + extern DECL_FATTN_MMA_F16_CASE_KV(DKQ, DV, (ncols)/ 2, 2, type_KV); \ + extern DECL_FATTN_MMA_F16_CASE_KV(DKQ, DV, (ncols)/ 4, 4, type_KV); \ + extern DECL_FATTN_MMA_F16_CASE_KV(DKQ, DV, (ncols)/ 8, 8, type_KV); \ + +DECL_FATTN_MMA_F16_CASE_KV_ALL_NCOLS2(128, 128, 8, GGML_TYPE_Q4_0) +DECL_FATTN_MMA_F16_CASE_KV_ALL_NCOLS2(128, 128, 16, GGML_TYPE_Q4_0) +DECL_FATTN_MMA_F16_CASE_KV_ALL_NCOLS2(128, 128, 32, GGML_TYPE_Q4_0) +DECL_FATTN_MMA_F16_CASE_KV_ALL_NCOLS2(128, 128, 64, GGML_TYPE_Q4_0) +DECL_FATTN_MMA_F16_CASE_KV_ALL_NCOLS2(256, 256, 8, GGML_TYPE_Q4_0) +DECL_FATTN_MMA_F16_CASE_KV_ALL_NCOLS2(256, 256, 16, GGML_TYPE_Q4_0) +DECL_FATTN_MMA_F16_CASE_KV_ALL_NCOLS2(256, 256, 32, GGML_TYPE_Q4_0) +DECL_FATTN_MMA_F16_CASE_KV_ALL_NCOLS2(256, 256, 64, GGML_TYPE_Q4_0) + +DECL_FATTN_MMA_F16_CASE_KV_ALL_NCOLS2(128, 128, 8, GGML_TYPE_Q8_0) +DECL_FATTN_MMA_F16_CASE_KV_ALL_NCOLS2(128, 128, 16, GGML_TYPE_Q8_0) +DECL_FATTN_MMA_F16_CASE_KV_ALL_NCOLS2(128, 128, 32, GGML_TYPE_Q8_0) +DECL_FATTN_MMA_F16_CASE_KV_ALL_NCOLS2(128, 128, 64, GGML_TYPE_Q8_0) +DECL_FATTN_MMA_F16_CASE_KV_ALL_NCOLS2(256, 256, 8, GGML_TYPE_Q8_0) +DECL_FATTN_MMA_F16_CASE_KV_ALL_NCOLS2(256, 256, 16, GGML_TYPE_Q8_0) +DECL_FATTN_MMA_F16_CASE_KV_ALL_NCOLS2(256, 256, 32, GGML_TYPE_Q8_0) +DECL_FATTN_MMA_F16_CASE_KV_ALL_NCOLS2(256, 256, 64, GGML_TYPE_Q8_0) + +// Head sizes with in-place quantized K/V kernels (matches the instantiations above). +static inline bool ggml_cuda_fattn_mma_kv_native_supported(const ggml_tensor * dst) { + const ggml_tensor * Q = dst->src[0]; + const ggml_tensor * K = dst->src[1]; + const ggml_tensor * V = dst->src[2]; + if (K->type != V->type || !ggml_cuda_fattn_mma_kv_native(K->type)) { + return false; + } + 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. + for (const ggml_tensor * t : {K, V}) { + if (t->nb[1] % 4 != 0 || t->nb[2] % 4 != 0 || t->nb[3] % 4 != 0) { + return false; + } + } + return true; +} + #define DECL_FATTN_MMA_F16_CASE_ALL_NCOLS2(DKQ, DV, ncols) \ extern DECL_FATTN_MMA_F16_CASE(DKQ, DV, (ncols)/ 1, 1); \ extern DECL_FATTN_MMA_F16_CASE(DKQ, DV, (ncols)/ 2, 2); \ diff --git a/ggml/src/ggml-cuda/fattn.cu b/ggml/src/ggml-cuda/fattn.cu index 58620e74..62e8e9b8 100644 --- a/ggml/src/ggml-cuda/fattn.cu +++ b/ggml/src/ggml-cuda/fattn.cu @@ -5,44 +5,44 @@ #include "fattn-vec.cuh" #include "fattn.cuh" -template +template static void ggml_cuda_flash_attn_ext_mma_f16_switch_ncols1(ggml_backend_cuda_context & ctx, ggml_tensor * dst) { const int cc = ggml_cuda_info().devices[ggml_cuda_get_device()].cc; const ggml_tensor * Q = dst->src[0]; - // batch-invariant mode: one tile width for 1 to 4 queries, so a token verified in a speculative batch + // batch-invariant mode: one tile width for 1 to 8 queries, so a token verified in a speculative batch // attends with the same arithmetic as a token decoded alone if constexpr (ncols2 <= 8) { - if (ggml_cuda_batch_invariant() && turing_mma_available(cc) && Q->ne[1] <= 4 && Q->ne[1]*ncols2 <= 32) { - ggml_cuda_flash_attn_ext_mma_f16_case(ctx, dst); + if (ggml_cuda_batch_invariant() && turing_mma_available(cc) && Q->ne[1] <= 8 && Q->ne[1]*ncols2 <= 32) { + ggml_cuda_flash_attn_ext_mma_f16_case(ctx, dst); return; } } if constexpr (ncols2 <= 8) { if (turing_mma_available(cc) && Q->ne[1] <= 8/ncols2) { - ggml_cuda_flash_attn_ext_mma_f16_case(ctx, dst); + ggml_cuda_flash_attn_ext_mma_f16_case(ctx, dst); return; } } if constexpr (ncols2 <= 16) { if (Q->ne[1] <= 16/ncols2) { - ggml_cuda_flash_attn_ext_mma_f16_case(ctx, dst); + ggml_cuda_flash_attn_ext_mma_f16_case(ctx, dst); return; } } if (Q->ne[1] <= 32/ncols2 || (GGML_CUDA_CC_IS_NVIDIA(cc) && ggml_cuda_highest_compiled_arch(cc) == GGML_CUDA_CC_TURING) || (GGML_CUDA_CC_IS_AMD(cc) && DKQ > 256)) { - ggml_cuda_flash_attn_ext_mma_f16_case(ctx, dst); + ggml_cuda_flash_attn_ext_mma_f16_case(ctx, dst); return; } - ggml_cuda_flash_attn_ext_mma_f16_case(ctx, dst); + ggml_cuda_flash_attn_ext_mma_f16_case(ctx, dst); } -template +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; const ggml_tensor * KQV = dst; @@ -75,22 +75,22 @@ static void ggml_cuda_flash_attn_ext_mma_f16_switch_ncols2(ggml_backend_cuda_con // On Volta the GQA optimizations aren't as impactful vs. minimizing wasted compute: if (cc == GGML_CUDA_CC_VOLTA) { if (use_gqa_opt && gqa_ratio % 8 == 0) { - ggml_cuda_flash_attn_ext_mma_f16_switch_ncols1(ctx, dst); + ggml_cuda_flash_attn_ext_mma_f16_switch_ncols1(ctx, dst); return; } if (use_gqa_opt && gqa_ratio % 4 == 0) { - ggml_cuda_flash_attn_ext_mma_f16_switch_ncols1(ctx, dst); + ggml_cuda_flash_attn_ext_mma_f16_switch_ncols1(ctx, dst); return; } if constexpr (DKQ <= 256) { if (use_gqa_opt && gqa_ratio % 2 == 0) { - ggml_cuda_flash_attn_ext_mma_f16_switch_ncols1(ctx, dst); + ggml_cuda_flash_attn_ext_mma_f16_switch_ncols1(ctx, dst); return; } - ggml_cuda_flash_attn_ext_mma_f16_switch_ncols1(ctx, dst); + ggml_cuda_flash_attn_ext_mma_f16_switch_ncols1(ctx, dst); return; } else { GGML_ABORT("fatal error"); @@ -98,22 +98,22 @@ static void ggml_cuda_flash_attn_ext_mma_f16_switch_ncols2(ggml_backend_cuda_con } if (use_gqa_opt && gqa_ratio > 4) { - ggml_cuda_flash_attn_ext_mma_f16_switch_ncols1(ctx, dst); + ggml_cuda_flash_attn_ext_mma_f16_switch_ncols1(ctx, dst); return; } if (use_gqa_opt && gqa_ratio > 2) { - ggml_cuda_flash_attn_ext_mma_f16_switch_ncols1(ctx, dst); + ggml_cuda_flash_attn_ext_mma_f16_switch_ncols1(ctx, dst); return; } if (use_gqa_opt && gqa_ratio > 1) { - ggml_cuda_flash_attn_ext_mma_f16_switch_ncols1(ctx, dst); + ggml_cuda_flash_attn_ext_mma_f16_switch_ncols1(ctx, dst); return; } if constexpr (DKQ <= 256) { - ggml_cuda_flash_attn_ext_mma_f16_switch_ncols1(ctx, dst); + ggml_cuda_flash_attn_ext_mma_f16_switch_ncols1(ctx, dst); } else { GGML_ABORT("fatal error"); } @@ -146,6 +146,14 @@ static void ggml_cuda_flash_attn_ext_mma_f16(ggml_backend_cuda_context & ctx, gg break; case 128: GGML_ASSERT(V->ne[0] == 128); + if (ggml_cuda_fattn_mma_kv_native_supported(dst)) { + if (K->type == GGML_TYPE_Q4_0) { + ggml_cuda_flash_attn_ext_mma_f16_switch_ncols2<128, 128, GGML_TYPE_Q4_0>(ctx, dst); + } else { + ggml_cuda_flash_attn_ext_mma_f16_switch_ncols2<128, 128, GGML_TYPE_Q8_0>(ctx, dst); + } + break; + } ggml_cuda_flash_attn_ext_mma_f16_switch_ncols2<128, 128>(ctx, dst); break; case 192: { @@ -166,6 +174,14 @@ static void ggml_cuda_flash_attn_ext_mma_f16(ggml_backend_cuda_context & ctx, gg } break; case 256: GGML_ASSERT(V->ne[0] == 256); + if (ggml_cuda_fattn_mma_kv_native_supported(dst)) { + if (K->type == GGML_TYPE_Q4_0) { + ggml_cuda_flash_attn_ext_mma_f16_switch_ncols2<256, 256, GGML_TYPE_Q4_0>(ctx, dst); + } else { + ggml_cuda_flash_attn_ext_mma_f16_switch_ncols2<256, 256, GGML_TYPE_Q8_0>(ctx, dst); + } + break; + } ggml_cuda_flash_attn_ext_mma_f16_switch_ncols2<256, 256>(ctx, dst); break; case 320: @@ -474,12 +490,22 @@ static best_fattn_kernel ggml_cuda_get_best_fattn_kernel(const int device, const // guarantee is 1 to 4 queries, see ggml_cuda_batch_invariant) if (ggml_cuda_batch_invariant() && Q->ne[1] <= 8 && Q->ne[3] == 1) { // with an F16 cache the tensor-core kernel is much faster at depth (39 vs 54 ms per token at 64K on - // an RTX 3060); ggml_cuda_flash_attn_ext_mma_f16_switch_ncols1 keeps it batch-invariant up to 4 queries - if (K->type == GGML_TYPE_F16 && V->type == GGML_TYPE_F16 && gqa_opt_applies && Q->ne[1] <= 4) { + // an RTX 3060); ggml_cuda_flash_attn_ext_mma_f16_switch_ncols1 keeps it batch-invariant up to 8 queries + if (K->type == GGML_TYPE_F16 && V->type == GGML_TYPE_F16 && gqa_opt_applies) { + return BEST_FATTN_KERNEL_MMA_F16; + } + // a q4_0/q8_0 cache is read in place by the same kernel + if (ggml_cuda_fattn_mma_kv_native_supported(dst) && mask && max_bias == 0.0f && gqa_ratio > 4) { return BEST_FATTN_KERNEL_MMA_F16; } return BEST_FATTN_KERNEL_VEC; } + // quantized-KV decode with wide GQA: the vector kernel fetches every K/V row once per Q head, the MMA + // kernel packs the GQA heads of a KV head into one tile and reads q4_0/q8_0 in place + if (ggml_cuda_fattn_mma_kv_native_supported(dst) && Q->ne[1] <= 8 && gqa_ratio > 4 && mask && max_bias == 0.0f && + K->ne[1] >= 256) { + return BEST_FATTN_KERNEL_MMA_F16; + } if (!ggml_is_quantized(K->type) && !ggml_is_quantized(V->type)) { if (cc >= GGML_CUDA_CC_ADA_LOVELACE && Q->ne[1] == 1 && Q->ne[3] == 1 && !(gqa_ratio > 4 && K->ne[1] >= 8192)) { return BEST_FATTN_KERNEL_VEC; @@ -568,8 +594,15 @@ size_t ggml_cuda_flash_attn_ext_get_alloc_size(int device, const ggml_tensor * d bool need_f16_V = false; switch (kernel) { - case BEST_FATTN_KERNEL_TILE: case BEST_FATTN_KERNEL_MMA_F16: + if (ggml_cuda_fattn_mma_kv_native_supported(dst)) { + // In-place quantized K/V kernel: nothing to reserve beyond dst. + break; + } + need_f16_K = true; + need_f16_V = true; + break; + case BEST_FATTN_KERNEL_TILE: need_f16_K = true; need_f16_V = true; break; diff --git a/ggml/src/ggml-cuda/template-instances/fattn-mma-f16-instance-q4_0-ncols1_1-ncols2_8.cu b/ggml/src/ggml-cuda/template-instances/fattn-mma-f16-instance-q4_0-ncols1_1-ncols2_8.cu new file mode 100644 index 00000000..355554a1 --- /dev/null +++ b/ggml/src/ggml-cuda/template-instances/fattn-mma-f16-instance-q4_0-ncols1_1-ncols2_8.cu @@ -0,0 +1,6 @@ +// This file has been autogenerated by generate_cu_files.py, do not edit manually. + +#include "../fattn-mma-f16.cuh" + +DECL_FATTN_MMA_F16_CASE_KV(128, 128, 1, 8, GGML_TYPE_Q4_0); +DECL_FATTN_MMA_F16_CASE_KV(256, 256, 1, 8, GGML_TYPE_Q4_0); diff --git a/ggml/src/ggml-cuda/template-instances/fattn-mma-f16-instance-q4_0-ncols1_16-ncols2_1.cu b/ggml/src/ggml-cuda/template-instances/fattn-mma-f16-instance-q4_0-ncols1_16-ncols2_1.cu new file mode 100644 index 00000000..a79fa078 --- /dev/null +++ b/ggml/src/ggml-cuda/template-instances/fattn-mma-f16-instance-q4_0-ncols1_16-ncols2_1.cu @@ -0,0 +1,6 @@ +// This file has been autogenerated by generate_cu_files.py, do not edit manually. + +#include "../fattn-mma-f16.cuh" + +DECL_FATTN_MMA_F16_CASE_KV(128, 128, 16, 1, GGML_TYPE_Q4_0); +DECL_FATTN_MMA_F16_CASE_KV(256, 256, 16, 1, GGML_TYPE_Q4_0); diff --git a/ggml/src/ggml-cuda/template-instances/fattn-mma-f16-instance-q4_0-ncols1_16-ncols2_2.cu b/ggml/src/ggml-cuda/template-instances/fattn-mma-f16-instance-q4_0-ncols1_16-ncols2_2.cu new file mode 100644 index 00000000..709e5de2 --- /dev/null +++ b/ggml/src/ggml-cuda/template-instances/fattn-mma-f16-instance-q4_0-ncols1_16-ncols2_2.cu @@ -0,0 +1,6 @@ +// This file has been autogenerated by generate_cu_files.py, do not edit manually. + +#include "../fattn-mma-f16.cuh" + +DECL_FATTN_MMA_F16_CASE_KV(128, 128, 16, 2, GGML_TYPE_Q4_0); +DECL_FATTN_MMA_F16_CASE_KV(256, 256, 16, 2, GGML_TYPE_Q4_0); diff --git a/ggml/src/ggml-cuda/template-instances/fattn-mma-f16-instance-q4_0-ncols1_16-ncols2_4.cu b/ggml/src/ggml-cuda/template-instances/fattn-mma-f16-instance-q4_0-ncols1_16-ncols2_4.cu new file mode 100644 index 00000000..6f475ec0 --- /dev/null +++ b/ggml/src/ggml-cuda/template-instances/fattn-mma-f16-instance-q4_0-ncols1_16-ncols2_4.cu @@ -0,0 +1,6 @@ +// This file has been autogenerated by generate_cu_files.py, do not edit manually. + +#include "../fattn-mma-f16.cuh" + +DECL_FATTN_MMA_F16_CASE_KV(128, 128, 16, 4, GGML_TYPE_Q4_0); +DECL_FATTN_MMA_F16_CASE_KV(256, 256, 16, 4, GGML_TYPE_Q4_0); diff --git a/ggml/src/ggml-cuda/template-instances/fattn-mma-f16-instance-q4_0-ncols1_2-ncols2_4.cu b/ggml/src/ggml-cuda/template-instances/fattn-mma-f16-instance-q4_0-ncols1_2-ncols2_4.cu new file mode 100644 index 00000000..4144c8e3 --- /dev/null +++ b/ggml/src/ggml-cuda/template-instances/fattn-mma-f16-instance-q4_0-ncols1_2-ncols2_4.cu @@ -0,0 +1,6 @@ +// This file has been autogenerated by generate_cu_files.py, do not edit manually. + +#include "../fattn-mma-f16.cuh" + +DECL_FATTN_MMA_F16_CASE_KV(128, 128, 2, 4, GGML_TYPE_Q4_0); +DECL_FATTN_MMA_F16_CASE_KV(256, 256, 2, 4, GGML_TYPE_Q4_0); diff --git a/ggml/src/ggml-cuda/template-instances/fattn-mma-f16-instance-q4_0-ncols1_2-ncols2_8.cu b/ggml/src/ggml-cuda/template-instances/fattn-mma-f16-instance-q4_0-ncols1_2-ncols2_8.cu new file mode 100644 index 00000000..859de760 --- /dev/null +++ b/ggml/src/ggml-cuda/template-instances/fattn-mma-f16-instance-q4_0-ncols1_2-ncols2_8.cu @@ -0,0 +1,6 @@ +// This file has been autogenerated by generate_cu_files.py, do not edit manually. + +#include "../fattn-mma-f16.cuh" + +DECL_FATTN_MMA_F16_CASE_KV(128, 128, 2, 8, GGML_TYPE_Q4_0); +DECL_FATTN_MMA_F16_CASE_KV(256, 256, 2, 8, GGML_TYPE_Q4_0); diff --git a/ggml/src/ggml-cuda/template-instances/fattn-mma-f16-instance-q4_0-ncols1_32-ncols2_1.cu b/ggml/src/ggml-cuda/template-instances/fattn-mma-f16-instance-q4_0-ncols1_32-ncols2_1.cu new file mode 100644 index 00000000..5b116735 --- /dev/null +++ b/ggml/src/ggml-cuda/template-instances/fattn-mma-f16-instance-q4_0-ncols1_32-ncols2_1.cu @@ -0,0 +1,6 @@ +// This file has been autogenerated by generate_cu_files.py, do not edit manually. + +#include "../fattn-mma-f16.cuh" + +DECL_FATTN_MMA_F16_CASE_KV(128, 128, 32, 1, GGML_TYPE_Q4_0); +DECL_FATTN_MMA_F16_CASE_KV(256, 256, 32, 1, GGML_TYPE_Q4_0); diff --git a/ggml/src/ggml-cuda/template-instances/fattn-mma-f16-instance-q4_0-ncols1_32-ncols2_2.cu b/ggml/src/ggml-cuda/template-instances/fattn-mma-f16-instance-q4_0-ncols1_32-ncols2_2.cu new file mode 100644 index 00000000..861a34c1 --- /dev/null +++ b/ggml/src/ggml-cuda/template-instances/fattn-mma-f16-instance-q4_0-ncols1_32-ncols2_2.cu @@ -0,0 +1,6 @@ +// This file has been autogenerated by generate_cu_files.py, do not edit manually. + +#include "../fattn-mma-f16.cuh" + +DECL_FATTN_MMA_F16_CASE_KV(128, 128, 32, 2, GGML_TYPE_Q4_0); +DECL_FATTN_MMA_F16_CASE_KV(256, 256, 32, 2, GGML_TYPE_Q4_0); diff --git a/ggml/src/ggml-cuda/template-instances/fattn-mma-f16-instance-q4_0-ncols1_4-ncols2_2.cu b/ggml/src/ggml-cuda/template-instances/fattn-mma-f16-instance-q4_0-ncols1_4-ncols2_2.cu new file mode 100644 index 00000000..99c3234e --- /dev/null +++ b/ggml/src/ggml-cuda/template-instances/fattn-mma-f16-instance-q4_0-ncols1_4-ncols2_2.cu @@ -0,0 +1,6 @@ +// This file has been autogenerated by generate_cu_files.py, do not edit manually. + +#include "../fattn-mma-f16.cuh" + +DECL_FATTN_MMA_F16_CASE_KV(128, 128, 4, 2, GGML_TYPE_Q4_0); +DECL_FATTN_MMA_F16_CASE_KV(256, 256, 4, 2, GGML_TYPE_Q4_0); diff --git a/ggml/src/ggml-cuda/template-instances/fattn-mma-f16-instance-q4_0-ncols1_4-ncols2_4.cu b/ggml/src/ggml-cuda/template-instances/fattn-mma-f16-instance-q4_0-ncols1_4-ncols2_4.cu new file mode 100644 index 00000000..d909b720 --- /dev/null +++ b/ggml/src/ggml-cuda/template-instances/fattn-mma-f16-instance-q4_0-ncols1_4-ncols2_4.cu @@ -0,0 +1,6 @@ +// This file has been autogenerated by generate_cu_files.py, do not edit manually. + +#include "../fattn-mma-f16.cuh" + +DECL_FATTN_MMA_F16_CASE_KV(128, 128, 4, 4, GGML_TYPE_Q4_0); +DECL_FATTN_MMA_F16_CASE_KV(256, 256, 4, 4, GGML_TYPE_Q4_0); diff --git a/ggml/src/ggml-cuda/template-instances/fattn-mma-f16-instance-q4_0-ncols1_4-ncols2_8.cu b/ggml/src/ggml-cuda/template-instances/fattn-mma-f16-instance-q4_0-ncols1_4-ncols2_8.cu new file mode 100644 index 00000000..cdee7c6f --- /dev/null +++ b/ggml/src/ggml-cuda/template-instances/fattn-mma-f16-instance-q4_0-ncols1_4-ncols2_8.cu @@ -0,0 +1,6 @@ +// This file has been autogenerated by generate_cu_files.py, do not edit manually. + +#include "../fattn-mma-f16.cuh" + +DECL_FATTN_MMA_F16_CASE_KV(128, 128, 4, 8, GGML_TYPE_Q4_0); +DECL_FATTN_MMA_F16_CASE_KV(256, 256, 4, 8, GGML_TYPE_Q4_0); diff --git a/ggml/src/ggml-cuda/template-instances/fattn-mma-f16-instance-q4_0-ncols1_64-ncols2_1.cu b/ggml/src/ggml-cuda/template-instances/fattn-mma-f16-instance-q4_0-ncols1_64-ncols2_1.cu new file mode 100644 index 00000000..ab677bed --- /dev/null +++ b/ggml/src/ggml-cuda/template-instances/fattn-mma-f16-instance-q4_0-ncols1_64-ncols2_1.cu @@ -0,0 +1,6 @@ +// This file has been autogenerated by generate_cu_files.py, do not edit manually. + +#include "../fattn-mma-f16.cuh" + +DECL_FATTN_MMA_F16_CASE_KV(128, 128, 64, 1, GGML_TYPE_Q4_0); +DECL_FATTN_MMA_F16_CASE_KV(256, 256, 64, 1, GGML_TYPE_Q4_0); diff --git a/ggml/src/ggml-cuda/template-instances/fattn-mma-f16-instance-q4_0-ncols1_8-ncols2_1.cu b/ggml/src/ggml-cuda/template-instances/fattn-mma-f16-instance-q4_0-ncols1_8-ncols2_1.cu new file mode 100644 index 00000000..d9de18d3 --- /dev/null +++ b/ggml/src/ggml-cuda/template-instances/fattn-mma-f16-instance-q4_0-ncols1_8-ncols2_1.cu @@ -0,0 +1,6 @@ +// This file has been autogenerated by generate_cu_files.py, do not edit manually. + +#include "../fattn-mma-f16.cuh" + +DECL_FATTN_MMA_F16_CASE_KV(128, 128, 8, 1, GGML_TYPE_Q4_0); +DECL_FATTN_MMA_F16_CASE_KV(256, 256, 8, 1, GGML_TYPE_Q4_0); diff --git a/ggml/src/ggml-cuda/template-instances/fattn-mma-f16-instance-q4_0-ncols1_8-ncols2_2.cu b/ggml/src/ggml-cuda/template-instances/fattn-mma-f16-instance-q4_0-ncols1_8-ncols2_2.cu new file mode 100644 index 00000000..68f48ae6 --- /dev/null +++ b/ggml/src/ggml-cuda/template-instances/fattn-mma-f16-instance-q4_0-ncols1_8-ncols2_2.cu @@ -0,0 +1,6 @@ +// This file has been autogenerated by generate_cu_files.py, do not edit manually. + +#include "../fattn-mma-f16.cuh" + +DECL_FATTN_MMA_F16_CASE_KV(128, 128, 8, 2, GGML_TYPE_Q4_0); +DECL_FATTN_MMA_F16_CASE_KV(256, 256, 8, 2, GGML_TYPE_Q4_0); diff --git a/ggml/src/ggml-cuda/template-instances/fattn-mma-f16-instance-q4_0-ncols1_8-ncols2_4.cu b/ggml/src/ggml-cuda/template-instances/fattn-mma-f16-instance-q4_0-ncols1_8-ncols2_4.cu new file mode 100644 index 00000000..c7655a7b --- /dev/null +++ b/ggml/src/ggml-cuda/template-instances/fattn-mma-f16-instance-q4_0-ncols1_8-ncols2_4.cu @@ -0,0 +1,6 @@ +// This file has been autogenerated by generate_cu_files.py, do not edit manually. + +#include "../fattn-mma-f16.cuh" + +DECL_FATTN_MMA_F16_CASE_KV(128, 128, 8, 4, GGML_TYPE_Q4_0); +DECL_FATTN_MMA_F16_CASE_KV(256, 256, 8, 4, GGML_TYPE_Q4_0); diff --git a/ggml/src/ggml-cuda/template-instances/fattn-mma-f16-instance-q4_0-ncols1_8-ncols2_8.cu b/ggml/src/ggml-cuda/template-instances/fattn-mma-f16-instance-q4_0-ncols1_8-ncols2_8.cu new file mode 100644 index 00000000..e9cb6703 --- /dev/null +++ b/ggml/src/ggml-cuda/template-instances/fattn-mma-f16-instance-q4_0-ncols1_8-ncols2_8.cu @@ -0,0 +1,6 @@ +// This file has been autogenerated by generate_cu_files.py, do not edit manually. + +#include "../fattn-mma-f16.cuh" + +DECL_FATTN_MMA_F16_CASE_KV(128, 128, 8, 8, GGML_TYPE_Q4_0); +DECL_FATTN_MMA_F16_CASE_KV(256, 256, 8, 8, GGML_TYPE_Q4_0); diff --git a/ggml/src/ggml-cuda/template-instances/fattn-mma-f16-instance-q8_0-ncols1_1-ncols2_8.cu b/ggml/src/ggml-cuda/template-instances/fattn-mma-f16-instance-q8_0-ncols1_1-ncols2_8.cu new file mode 100644 index 00000000..ef09f7b7 --- /dev/null +++ b/ggml/src/ggml-cuda/template-instances/fattn-mma-f16-instance-q8_0-ncols1_1-ncols2_8.cu @@ -0,0 +1,6 @@ +// This file has been autogenerated by generate_cu_files.py, do not edit manually. + +#include "../fattn-mma-f16.cuh" + +DECL_FATTN_MMA_F16_CASE_KV(128, 128, 1, 8, GGML_TYPE_Q8_0); +DECL_FATTN_MMA_F16_CASE_KV(256, 256, 1, 8, GGML_TYPE_Q8_0); diff --git a/ggml/src/ggml-cuda/template-instances/fattn-mma-f16-instance-q8_0-ncols1_16-ncols2_1.cu b/ggml/src/ggml-cuda/template-instances/fattn-mma-f16-instance-q8_0-ncols1_16-ncols2_1.cu new file mode 100644 index 00000000..c7a048ad --- /dev/null +++ b/ggml/src/ggml-cuda/template-instances/fattn-mma-f16-instance-q8_0-ncols1_16-ncols2_1.cu @@ -0,0 +1,6 @@ +// This file has been autogenerated by generate_cu_files.py, do not edit manually. + +#include "../fattn-mma-f16.cuh" + +DECL_FATTN_MMA_F16_CASE_KV(128, 128, 16, 1, GGML_TYPE_Q8_0); +DECL_FATTN_MMA_F16_CASE_KV(256, 256, 16, 1, GGML_TYPE_Q8_0); diff --git a/ggml/src/ggml-cuda/template-instances/fattn-mma-f16-instance-q8_0-ncols1_16-ncols2_2.cu b/ggml/src/ggml-cuda/template-instances/fattn-mma-f16-instance-q8_0-ncols1_16-ncols2_2.cu new file mode 100644 index 00000000..6dd74f06 --- /dev/null +++ b/ggml/src/ggml-cuda/template-instances/fattn-mma-f16-instance-q8_0-ncols1_16-ncols2_2.cu @@ -0,0 +1,6 @@ +// This file has been autogenerated by generate_cu_files.py, do not edit manually. + +#include "../fattn-mma-f16.cuh" + +DECL_FATTN_MMA_F16_CASE_KV(128, 128, 16, 2, GGML_TYPE_Q8_0); +DECL_FATTN_MMA_F16_CASE_KV(256, 256, 16, 2, GGML_TYPE_Q8_0); diff --git a/ggml/src/ggml-cuda/template-instances/fattn-mma-f16-instance-q8_0-ncols1_16-ncols2_4.cu b/ggml/src/ggml-cuda/template-instances/fattn-mma-f16-instance-q8_0-ncols1_16-ncols2_4.cu new file mode 100644 index 00000000..8341b5f9 --- /dev/null +++ b/ggml/src/ggml-cuda/template-instances/fattn-mma-f16-instance-q8_0-ncols1_16-ncols2_4.cu @@ -0,0 +1,6 @@ +// This file has been autogenerated by generate_cu_files.py, do not edit manually. + +#include "../fattn-mma-f16.cuh" + +DECL_FATTN_MMA_F16_CASE_KV(128, 128, 16, 4, GGML_TYPE_Q8_0); +DECL_FATTN_MMA_F16_CASE_KV(256, 256, 16, 4, GGML_TYPE_Q8_0); diff --git a/ggml/src/ggml-cuda/template-instances/fattn-mma-f16-instance-q8_0-ncols1_2-ncols2_4.cu b/ggml/src/ggml-cuda/template-instances/fattn-mma-f16-instance-q8_0-ncols1_2-ncols2_4.cu new file mode 100644 index 00000000..991759e3 --- /dev/null +++ b/ggml/src/ggml-cuda/template-instances/fattn-mma-f16-instance-q8_0-ncols1_2-ncols2_4.cu @@ -0,0 +1,6 @@ +// This file has been autogenerated by generate_cu_files.py, do not edit manually. + +#include "../fattn-mma-f16.cuh" + +DECL_FATTN_MMA_F16_CASE_KV(128, 128, 2, 4, GGML_TYPE_Q8_0); +DECL_FATTN_MMA_F16_CASE_KV(256, 256, 2, 4, GGML_TYPE_Q8_0); diff --git a/ggml/src/ggml-cuda/template-instances/fattn-mma-f16-instance-q8_0-ncols1_2-ncols2_8.cu b/ggml/src/ggml-cuda/template-instances/fattn-mma-f16-instance-q8_0-ncols1_2-ncols2_8.cu new file mode 100644 index 00000000..0cdd4b1e --- /dev/null +++ b/ggml/src/ggml-cuda/template-instances/fattn-mma-f16-instance-q8_0-ncols1_2-ncols2_8.cu @@ -0,0 +1,6 @@ +// This file has been autogenerated by generate_cu_files.py, do not edit manually. + +#include "../fattn-mma-f16.cuh" + +DECL_FATTN_MMA_F16_CASE_KV(128, 128, 2, 8, GGML_TYPE_Q8_0); +DECL_FATTN_MMA_F16_CASE_KV(256, 256, 2, 8, GGML_TYPE_Q8_0); diff --git a/ggml/src/ggml-cuda/template-instances/fattn-mma-f16-instance-q8_0-ncols1_32-ncols2_1.cu b/ggml/src/ggml-cuda/template-instances/fattn-mma-f16-instance-q8_0-ncols1_32-ncols2_1.cu new file mode 100644 index 00000000..b0d5f040 --- /dev/null +++ b/ggml/src/ggml-cuda/template-instances/fattn-mma-f16-instance-q8_0-ncols1_32-ncols2_1.cu @@ -0,0 +1,6 @@ +// This file has been autogenerated by generate_cu_files.py, do not edit manually. + +#include "../fattn-mma-f16.cuh" + +DECL_FATTN_MMA_F16_CASE_KV(128, 128, 32, 1, GGML_TYPE_Q8_0); +DECL_FATTN_MMA_F16_CASE_KV(256, 256, 32, 1, GGML_TYPE_Q8_0); diff --git a/ggml/src/ggml-cuda/template-instances/fattn-mma-f16-instance-q8_0-ncols1_32-ncols2_2.cu b/ggml/src/ggml-cuda/template-instances/fattn-mma-f16-instance-q8_0-ncols1_32-ncols2_2.cu new file mode 100644 index 00000000..e4086fba --- /dev/null +++ b/ggml/src/ggml-cuda/template-instances/fattn-mma-f16-instance-q8_0-ncols1_32-ncols2_2.cu @@ -0,0 +1,6 @@ +// This file has been autogenerated by generate_cu_files.py, do not edit manually. + +#include "../fattn-mma-f16.cuh" + +DECL_FATTN_MMA_F16_CASE_KV(128, 128, 32, 2, GGML_TYPE_Q8_0); +DECL_FATTN_MMA_F16_CASE_KV(256, 256, 32, 2, GGML_TYPE_Q8_0); diff --git a/ggml/src/ggml-cuda/template-instances/fattn-mma-f16-instance-q8_0-ncols1_4-ncols2_2.cu b/ggml/src/ggml-cuda/template-instances/fattn-mma-f16-instance-q8_0-ncols1_4-ncols2_2.cu new file mode 100644 index 00000000..0388387f --- /dev/null +++ b/ggml/src/ggml-cuda/template-instances/fattn-mma-f16-instance-q8_0-ncols1_4-ncols2_2.cu @@ -0,0 +1,6 @@ +// This file has been autogenerated by generate_cu_files.py, do not edit manually. + +#include "../fattn-mma-f16.cuh" + +DECL_FATTN_MMA_F16_CASE_KV(128, 128, 4, 2, GGML_TYPE_Q8_0); +DECL_FATTN_MMA_F16_CASE_KV(256, 256, 4, 2, GGML_TYPE_Q8_0); diff --git a/ggml/src/ggml-cuda/template-instances/fattn-mma-f16-instance-q8_0-ncols1_4-ncols2_4.cu b/ggml/src/ggml-cuda/template-instances/fattn-mma-f16-instance-q8_0-ncols1_4-ncols2_4.cu new file mode 100644 index 00000000..29886f86 --- /dev/null +++ b/ggml/src/ggml-cuda/template-instances/fattn-mma-f16-instance-q8_0-ncols1_4-ncols2_4.cu @@ -0,0 +1,6 @@ +// This file has been autogenerated by generate_cu_files.py, do not edit manually. + +#include "../fattn-mma-f16.cuh" + +DECL_FATTN_MMA_F16_CASE_KV(128, 128, 4, 4, GGML_TYPE_Q8_0); +DECL_FATTN_MMA_F16_CASE_KV(256, 256, 4, 4, GGML_TYPE_Q8_0); diff --git a/ggml/src/ggml-cuda/template-instances/fattn-mma-f16-instance-q8_0-ncols1_4-ncols2_8.cu b/ggml/src/ggml-cuda/template-instances/fattn-mma-f16-instance-q8_0-ncols1_4-ncols2_8.cu new file mode 100644 index 00000000..b0d85792 --- /dev/null +++ b/ggml/src/ggml-cuda/template-instances/fattn-mma-f16-instance-q8_0-ncols1_4-ncols2_8.cu @@ -0,0 +1,6 @@ +// This file has been autogenerated by generate_cu_files.py, do not edit manually. + +#include "../fattn-mma-f16.cuh" + +DECL_FATTN_MMA_F16_CASE_KV(128, 128, 4, 8, GGML_TYPE_Q8_0); +DECL_FATTN_MMA_F16_CASE_KV(256, 256, 4, 8, GGML_TYPE_Q8_0); diff --git a/ggml/src/ggml-cuda/template-instances/fattn-mma-f16-instance-q8_0-ncols1_64-ncols2_1.cu b/ggml/src/ggml-cuda/template-instances/fattn-mma-f16-instance-q8_0-ncols1_64-ncols2_1.cu new file mode 100644 index 00000000..c5c46959 --- /dev/null +++ b/ggml/src/ggml-cuda/template-instances/fattn-mma-f16-instance-q8_0-ncols1_64-ncols2_1.cu @@ -0,0 +1,6 @@ +// This file has been autogenerated by generate_cu_files.py, do not edit manually. + +#include "../fattn-mma-f16.cuh" + +DECL_FATTN_MMA_F16_CASE_KV(128, 128, 64, 1, GGML_TYPE_Q8_0); +DECL_FATTN_MMA_F16_CASE_KV(256, 256, 64, 1, GGML_TYPE_Q8_0); diff --git a/ggml/src/ggml-cuda/template-instances/fattn-mma-f16-instance-q8_0-ncols1_8-ncols2_1.cu b/ggml/src/ggml-cuda/template-instances/fattn-mma-f16-instance-q8_0-ncols1_8-ncols2_1.cu new file mode 100644 index 00000000..ba806751 --- /dev/null +++ b/ggml/src/ggml-cuda/template-instances/fattn-mma-f16-instance-q8_0-ncols1_8-ncols2_1.cu @@ -0,0 +1,6 @@ +// This file has been autogenerated by generate_cu_files.py, do not edit manually. + +#include "../fattn-mma-f16.cuh" + +DECL_FATTN_MMA_F16_CASE_KV(128, 128, 8, 1, GGML_TYPE_Q8_0); +DECL_FATTN_MMA_F16_CASE_KV(256, 256, 8, 1, GGML_TYPE_Q8_0); diff --git a/ggml/src/ggml-cuda/template-instances/fattn-mma-f16-instance-q8_0-ncols1_8-ncols2_2.cu b/ggml/src/ggml-cuda/template-instances/fattn-mma-f16-instance-q8_0-ncols1_8-ncols2_2.cu new file mode 100644 index 00000000..b2fdf963 --- /dev/null +++ b/ggml/src/ggml-cuda/template-instances/fattn-mma-f16-instance-q8_0-ncols1_8-ncols2_2.cu @@ -0,0 +1,6 @@ +// This file has been autogenerated by generate_cu_files.py, do not edit manually. + +#include "../fattn-mma-f16.cuh" + +DECL_FATTN_MMA_F16_CASE_KV(128, 128, 8, 2, GGML_TYPE_Q8_0); +DECL_FATTN_MMA_F16_CASE_KV(256, 256, 8, 2, GGML_TYPE_Q8_0); diff --git a/ggml/src/ggml-cuda/template-instances/fattn-mma-f16-instance-q8_0-ncols1_8-ncols2_4.cu b/ggml/src/ggml-cuda/template-instances/fattn-mma-f16-instance-q8_0-ncols1_8-ncols2_4.cu new file mode 100644 index 00000000..cfe288e5 --- /dev/null +++ b/ggml/src/ggml-cuda/template-instances/fattn-mma-f16-instance-q8_0-ncols1_8-ncols2_4.cu @@ -0,0 +1,6 @@ +// This file has been autogenerated by generate_cu_files.py, do not edit manually. + +#include "../fattn-mma-f16.cuh" + +DECL_FATTN_MMA_F16_CASE_KV(128, 128, 8, 4, GGML_TYPE_Q8_0); +DECL_FATTN_MMA_F16_CASE_KV(256, 256, 8, 4, GGML_TYPE_Q8_0); diff --git a/ggml/src/ggml-cuda/template-instances/fattn-mma-f16-instance-q8_0-ncols1_8-ncols2_8.cu b/ggml/src/ggml-cuda/template-instances/fattn-mma-f16-instance-q8_0-ncols1_8-ncols2_8.cu new file mode 100644 index 00000000..1c35d44c --- /dev/null +++ b/ggml/src/ggml-cuda/template-instances/fattn-mma-f16-instance-q8_0-ncols1_8-ncols2_8.cu @@ -0,0 +1,6 @@ +// This file has been autogenerated by generate_cu_files.py, do not edit manually. + +#include "../fattn-mma-f16.cuh" + +DECL_FATTN_MMA_F16_CASE_KV(128, 128, 8, 8, GGML_TYPE_Q8_0); +DECL_FATTN_MMA_F16_CASE_KV(256, 256, 8, 8, GGML_TYPE_Q8_0); diff --git a/ggml/src/ggml-cuda/template-instances/generate_cu_files.py b/ggml/src/ggml-cuda/template-instances/generate_cu_files.py index c54928a5..073e4154 100755 --- a/ggml/src/ggml-cuda/template-instances/generate_cu_files.py +++ b/ggml/src/ggml-cuda/template-instances/generate_cu_files.py @@ -34,6 +34,11 @@ SOURCE_FATTN_MMA_START = """// This file has been autogenerated by generate_cu_f SOURCE_FATTN_MMA_CASE = "DECL_FATTN_MMA_F16_CASE({head_size_kq}, {head_size_v}, {ncols1}, {ncols2});\n" +SOURCE_FATTN_MMA_CASE_KV = "DECL_FATTN_MMA_F16_CASE_KV({head_size_kq}, {head_size_v}, {ncols1}, {ncols2}, {type_kv});\n" + +TYPES_KV_MMA_NATIVE = ["GGML_TYPE_Q4_0", "GGML_TYPE_Q8_0"] +HEAD_SIZES_MMA_NATIVE = [128, 256] + TYPES_MMQ = [ "GGML_TYPE_Q1_0", "GGML_TYPE_Q2_0", @@ -113,6 +118,17 @@ for ncols in [8, 16, 32, 64]: head_size_v = HEAD_SIZES_V_OVERRIDE.get(head_size_kq, head_size_kq) f.write(SOURCE_FATTN_MMA_CASE.format(ncols1=ncols1, ncols2=ncols2, head_size_kq=head_size_kq, head_size_v=head_size_v)) +# In-place quantized K/V variants of the MMA kernel (no F16 copy of the cache): the head sizes and GQA tilings +# a q4_0 / q8_0 cache is served with in practice. Kept to one file per (type, ncols1, ncols2) for parallel builds. +for type_kv in TYPES_KV_MMA_NATIVE: + for ncols in [8, 16, 32, 64]: + for ncols2 in [1, 2, 4, 8]: + ncols1 = ncols // ncols2 + with open(f"fattn-mma-f16-instance-{get_short_name(type_kv)}-ncols1_{ncols1}-ncols2_{ncols2}.cu", "w") as f: + f.write(SOURCE_FATTN_MMA_START) + for head_size in HEAD_SIZES_MMA_NATIVE: + f.write(SOURCE_FATTN_MMA_CASE_KV.format(ncols1=ncols1, ncols2=ncols2, head_size_kq=head_size, head_size_v=head_size, type_kv=type_kv)) + for type in TYPES_MMQ: with open(f"mmq-instance-{get_short_name(type)}.cu", "w") as f: source = SOURCE_MMQ_CUDA_ONLY if type == "GGML_TYPE_PTQ1_0" else SOURCE_MMQ