diff --git a/benches/prism-helios-ampere-decode-perf.svg b/benches/prism-helios-ampere-decode-perf.svg new file mode 100644 index 000000000000..eee448563af7 --- /dev/null +++ b/benches/prism-helios-ampere-decode-perf.svg @@ -0,0 +1,61 @@ + + +Bonsai 2 27B (Prism PQ2_0) - decode throughput, RTX 3060 (SM86, Ampere) +llama-batched-bench, pp512 tg128, f16 KV, clocks locked, mean of 2 passes (t/s) + +0 + +10 + +20 + +30 + +40 + +50 + +60 + +70 + +80 +tokens / sec + +34.6 + +34.6 + +34.5 +1 stream + +58.4 + +59.9 + +59.1 +2 streams + +62.8 + +68.5 + +62.3 +3 streams + +61.5 + +76.0 + +75.4 +4 streams + +base (pre-PR) + +this PR + +this PR, reuse=0 ++2.5% ++9.2% ++23.6% + \ No newline at end of file diff --git a/ggml/src/ggml-cuda/fattn-mma-f16.cuh b/ggml/src/ggml-cuda/fattn-mma-f16.cuh index 7f4cfd5511ff..7e083b0bbc5e 100644 --- a/ggml/src/ggml-cuda/fattn-mma-f16.cuh +++ b/ggml/src/ggml-cuda/fattn-mma-f16.cuh @@ -527,9 +527,94 @@ static __device__ __forceinline__ void flash_attn_ext_f16_load_mask( } } +static __device__ __forceinline__ half2 ggml_cuda_u32_as_half2(const uint32_t x) { + half2 r; + memcpy(&r, &x, sizeof(r)); + return r; +} + +// Decode q8_0 K/V rows directly into the SRAM tile, same layout as flash_attn_ext_f16_load_tile. +// Each lane decodes 8 values: the byte q+128 is put below the high byte 0x64 which gives fp16 1024+q+128, subtracting 1152 is exact. +template +static __device__ __forceinline__ void flash_attn_ext_q8_0_load_tile( + const char * const __restrict__ KV_raw, half2 * const __restrict__ tile_KV, + const int stride_bytes, const int col_offset, const int i_sup) { + constexpr int warp_size = ggml_cuda_get_physical_warp_size(); + constexpr int h2_per_chunk = 4; + static_assert(D2 % h2_per_chunk == 0, "D2 must be a multiple of 4"); + constexpr int chunks_per_row = D2 / h2_per_chunk; + constexpr int nchunks = nbatch_fa * chunks_per_row; + const int tid = threadIdx.y * warp_size + threadIdx.x; + + const half2 bias = ggml_cuda_u32_as_half2(0x64806480u); + + for (int linear = tid; linear < nchunks; linear += nthreads) { + const int row = linear / chunks_per_row; + const int chunk = linear - row * chunks_per_row; + const int col = chunk * h2_per_chunk; + + static_assert(!swz, "smem swizzle is not available in this tree"); + uint4 * dst = (uint4 *) (tile_KV + row * stride_tile + col); + + if (oob_check && row >= i_sup) { + *dst = {}; + continue; + } + + // eight values starting at an index that is a multiple of 8, so they are inside one block + const int i0 = 2 * (col_offset + col); + const block_q8_0 & block = ((const block_q8_0 *) (KV_raw + int64_t(row) * stride_bytes))[i0 / QK8_0]; + const half2 d2 = __half2half2(block.d); + const uint16_t * qp = (const uint16_t *) (block.qs + i0 % QK8_0); + const uint32_t w0 = __byte_perm((uint32_t) qp[0], (uint32_t) qp[1], 0x5410) ^ 0x80808080u; + const uint32_t w1 = __byte_perm((uint32_t) qp[2], (uint32_t) qp[3], 0x5410) ^ 0x80808080u; + uint4 decoded; + half2 * values = (half2 *) &decoded; + values[0] = d2 * __hsub2(ggml_cuda_u32_as_half2(__byte_perm(w0, 0x64646464u, 0x5140)), bias); + values[1] = d2 * __hsub2(ggml_cuda_u32_as_half2(__byte_perm(w0, 0x64646464u, 0x7362)), bias); + values[2] = d2 * __hsub2(ggml_cuda_u32_as_half2(__byte_perm(w1, 0x64646464u, 0x5140)), bias); + values[3] = d2 * __hsub2(ggml_cuda_u32_as_half2(__byte_perm(w1, 0x64646464u, 0x7362)), bias); + *dst = decoded; + } +} + +// Raw q8_0 rows of the next tile are copied to shared memory with cp.async while the current tile is computed. +template +static constexpr __host__ __device__ bool ggml_cuda_fattn_q8_0_stage() { + return type_K == GGML_TYPE_Q8_0 && ncols2 > 1 && DKQ == 256 && DV == 256; +} + +template +static constexpr __host__ __device__ int ggml_cuda_fattn_q8_0_row_bytes() { + return (D/QK8_0) * (int) sizeof(block_q8_0); +} + +// byte offset of the q8_0 staging buffers in shared memory, tile strides are in half2 +static constexpr __host__ __device__ int ggml_cuda_fattn_mma_q8_0_stage_offset( + const int ncols, const int ncols1, const int nbatch_fa, const int stride_tile_Q, const int stride_tile_KV, const bool Q_in_reg) { + const int n_KV_mask = nbatch_fa * stride_tile_KV + ncols1 * (nbatch_fa/2 + 4); + const int n_Q = ncols * stride_tile_Q; + const int n = Q_in_reg ? (n_KV_mask > n_Q ? n_KV_mask : n_Q) : n_Q + n_KV_mask; + return (n * (int) sizeof(half2) + 15) & ~15; +} + +template +static __device__ __forceinline__ void flash_attn_ext_q8_0_stage_issue(const unsigned int dst32, const char * src, const int stride_bytes) { + static_assert(row_bytes % 16 == 0, "q8_0 staging needs rows aligned to 16 bytes"); + constexpr int chunks_per_row = row_bytes / 16; + constexpr int nchunks = nrows * chunks_per_row; + const int tid = threadIdx.y * ggml_cuda_get_physical_warp_size() + threadIdx.x; + for (int c = tid; c < nchunks; c += nthreads) { + const int row = c / chunks_per_row; + const int off = (c - row * chunks_per_row) * 16; + cp_async_cg_16<0>(dst32 + row * row_bytes + off, src + int64_t(row) * stride_bytes + off); + } +} + 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 = GGML_TYPE_F16> static __device__ __forceinline__ void flash_attn_ext_f16_iter( const float2 * const __restrict__ Q_f2, const half2 * const __restrict__ K_h2, @@ -555,7 +640,9 @@ 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, + char * const __restrict__ raw_K, + char * const __restrict__ raw_V) { #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; @@ -566,7 +653,10 @@ 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); + // q8_0 KV is decoded synchronously into SRAM, cp.async multi-stage loading would copy raw bytes as half2 + constexpr bool is_q8_0 = type_KV == GGML_TYPE_Q8_0; + constexpr bool q8_0_stage = ggml_cuda_fattn_q8_0_stage() && !oob_check; + constexpr int nstages = is_q8_0 ? 0 : ggml_cuda_fattn_mma_get_nstages(DKQ, DV, ncols1, ncols2); constexpr int stride_tile_K = nbatch_K2 + 4; @@ -604,7 +694,27 @@ static __device__ __forceinline__ void flash_attn_ext_f16_iter( for (int k0_start = (DKQ/2-1) - (DKQ/2-1) % nbatch_K2; k0_start >= 0; k0_start -= nbatch_K2) { const int k0_stop = k0_start + nbatch_K2 < DKQ/2 ? k0_start + nbatch_K2 : DKQ/2; - if constexpr (nstages <= 1) { + if constexpr (is_q8_0) { + static_assert(nbatch_K2 == DKQ/2, "q8_0 MMA load needs full K rows"); + static_assert(!V_is_K_view, "q8_0 MMA load does not support K views"); + constexpr int nthreads = nwarps * warp_size; + const char * K_raw = (const char *) K_h2 + int64_t(k_VKQ_0) * stride_K; + int K_pitch = stride_K; + if constexpr (q8_0_stage) { + cp_async_wait_all(); + __syncthreads(); + K_raw = raw_K; + K_pitch = ggml_cuda_fattn_q8_0_row_bytes(); + } + flash_attn_ext_q8_0_load_tile + (K_raw, tile_K, K_pitch, k0_start, k_VKQ_sup); + __syncthreads(); + if constexpr (q8_0_stage && !last_iter) { + // raw_K is consumed, stage the K rows of the next tile + flash_attn_ext_q8_0_stage_issue()> + (ggml_cuda_cvta_generic_to_shared(raw_K), (const char *) K_h2 + int64_t(k_VKQ_0 + nbatch_fa) * stride_K, stride_K); + } + } else 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 @@ -955,7 +1065,24 @@ static __device__ __forceinline__ void flash_attn_ext_f16_iter( static_assert(DV % (2*nbatch_V2) == 0, "bad loop size"); const int i0_stop = i0_start + 2*nbatch_V2; - if constexpr (nstages <= 1) { + if constexpr (is_q8_0) { + static_assert(nbatch_V2 == DV/2, "q8_0 MMA load needs full V rows"); + constexpr int nthreads = nwarps * warp_size; + const char * V_raw = (const char *) V_h2 + int64_t(k_VKQ_0) * stride_V; + int V_pitch = stride_V; + if constexpr (q8_0_stage) { + // staged together with K and already waited for + V_raw = raw_V; + V_pitch = ggml_cuda_fattn_q8_0_row_bytes(); + } + flash_attn_ext_q8_0_load_tile + (V_raw, tile_V, V_pitch, i0_start/2, k_VKQ_sup); + __syncthreads(); + if constexpr (q8_0_stage && !last_iter) { + flash_attn_ext_q8_0_stage_issue()> + (ggml_cuda_cvta_generic_to_shared(raw_V), (const char *) V_h2 + int64_t(k_VKQ_0 + nbatch_fa) * stride_V, stride_V); + } + } else 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; @@ -1019,7 +1146,7 @@ static __device__ __forceinline__ void flash_attn_ext_f16_iter( scale, slope, logit_softcap, ne01, ne02, stride_K, stride_V, stride_mask, tile_Q, tile_K, tile_V, tile_mask, - Q_B, VKQ_C, KQ_max, KQ_rowsum, kb0); + Q_B, VKQ_C, KQ_max, KQ_rowsum, kb0, raw_K, raw_V); NO_DEVICE_CODE; #endif // defined(VOLTA_MMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE) || defined(AMD_MFMA_AVAILABLE) } @@ -1113,7 +1240,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, @@ -1158,7 +1286,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 = type_KV == GGML_TYPE_Q8_0 ? 0 : ggml_cuda_fattn_mma_get_nstages(DKQ, DV, ncols1, ncols2); if (cols_per_warp > ncols) { NO_DEVICE_CODE; @@ -1178,6 +1306,12 @@ static __device__ __forceinline__ void flash_attn_ext_f16_process_tile( half2 * tile_V = nstages > 1 ? tile_K + nbatch_fa * stride_tile_K : tile_K; half * tile_mask = (half *) (nstages > 1 ? tile_V + nbatch_fa * stride_tile_V : tile_V + nbatch_fa * stride_tile_KV_max); + // staging buffers for raw q8_0 rows follow the Q/KV/mask data, the launcher uses the same offset + constexpr bool q8_0_stage = ggml_cuda_fattn_q8_0_stage(); + constexpr int stage_off = ggml_cuda_fattn_mma_q8_0_stage_offset(ncols, ncols1, nbatch_fa, stride_tile_Q, stride_tile_KV_max, Q_in_reg); + char * raw_K = q8_0_stage ? (char *) tile_Q + stage_off : nullptr; + char * raw_V = q8_0_stage ? raw_K + nbatch_fa * ggml_cuda_fattn_q8_0_row_bytes() : nullptr; + T_B_KQ Q_B[(Q_in_reg ? DKQ/(2*T_B_KQ::J) : 1)]; #if defined(TURING_MMA_AVAILABLE) T_C_VKQ VKQ_C[cols_per_warp == 8 ? DV/T_C_VKQ::I : DV/(2*T_C_VKQ::J)]; @@ -1269,6 +1403,14 @@ static __device__ __forceinline__ void flash_attn_ext_f16_process_tile( (K_h2 + int64_t(kb0)*nbatch_fa*stride_K, tile_K, nbatch_K2, stride_K, k_VKQ_sup); } + if constexpr (q8_0_stage) { + constexpr int nthreads = nwarps * warp_size; + flash_attn_ext_q8_0_stage_issue()> + (ggml_cuda_cvta_generic_to_shared(raw_K), (const char *) K_h2 + int64_t(kb0) * nbatch_fa * stride_K, stride_K); + flash_attn_ext_q8_0_stage_issue()> + (ggml_cuda_cvta_generic_to_shared(raw_V), (const char *) V_h2 + int64_t(kb0) * nbatch_fa * stride_V, stride_V); + } + // kb0_start is always < kb0_stop so the last iter can be executed unconditionally. if constexpr (ncols2 == 1) { constexpr bool oob_check = true; @@ -1277,19 +1419,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, 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); + KQ_max, KQ_rowsum, jt, kb0, k_VKQ_sup, raw_K, raw_V); } 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, 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); + KQ_max, KQ_rowsum, jt, kb0, k_VKQ_sup, raw_K, raw_V); } else { constexpr bool oob_check = false; for (; kb0 < kb0_stop-1; ++kb0) { @@ -1297,19 +1439,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, 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); + KQ_max, KQ_rowsum, jt, kb0, k_VKQ_sup, raw_K, raw_V); } 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, 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); + KQ_max, KQ_rowsum, jt, kb0, k_VKQ_sup, raw_K, raw_V); } // With multi-stage loading there is no __syncthreads at the end of the iter, @@ -1700,7 +1842,8 @@ 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 +1925,12 @@ 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); + // q8_0 KV is passed as raw bytes, its strides are byte pitches + constexpr bool is_q8_0 = type_KV == GGML_TYPE_Q8_0; + const int stride_K = is_q8_0 ? nb11 : nb11 / sizeof(half2); 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 : (is_q8_0 ? nb21 : nb21 / sizeof(half2)); const int iter_k = (ne11 + (nbatch_fa - 1)) / nbatch_fa; const int iter_j = (ne01.z + (ncols1 - 1)) / ncols1; @@ -1829,12 +1974,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 +2020,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 @@ -1963,6 +2108,88 @@ void ggml_cuda_flash_attn_ext_mma_f16_case(ggml_backend_cuda_context & ctx, ggml (ctx, dst, fattn_kernel, nwarps, nbytes_shared_total, nbatch_fa, true, true, true, warp_size_host); } +// q8_0 K and V are decoded inside the kernel, no f16 copy of the cache is made. +template +void ggml_cuda_flash_attn_ext_mma_q8_0_case(ggml_backend_cuda_context & ctx, ggml_tensor * dst) { + const ggml_tensor * KQV = dst; + const int id = ggml_cuda_get_device(); + const int cc = ggml_cuda_info().devices[id].cc; + + constexpr int ncols = ncols1 * ncols2; + + 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 cols_per_warp = std::min(ncols, get_cols_per_warp(cc)); + const int warp_size_host = ggml_cuda_info().devices[ctx.device].warp_size; + const int nwarps = nthreads / warp_size_host; + + constexpr bool V_is_K_view = false; + + // single-stage loading, KV tile strides must match flash_attn_ext_f16_iter / _process_tile + const int stride_tile_K = nbatch_K2 + 4; + const int stride_tile_V = nbatch_V2 + 4; + const int stride_tile_KV = std::max(stride_tile_K, stride_tile_V); + + const size_t nbytes_shared_KV = nbatch_fa * stride_tile_KV * sizeof(half2); + const size_t nbytes_shared_Q = ncols * (DKQ/2 + 4) * sizeof(half2); + 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); + + size_t nbytes_shared_total = std::max(nbytes_shared_combine, Q_in_reg ? + std::max(nbytes_shared_Q, nbytes_shared_KV + nbytes_shared_mask) : + nbytes_shared_Q + nbytes_shared_KV + nbytes_shared_mask); + + if constexpr (ggml_cuda_fattn_q8_0_stage()) { + const size_t stage_off = ggml_cuda_fattn_mma_q8_0_stage_offset(ncols, ncols1, nbatch_fa, DKQ/2 + 4, stride_tile_KV, Q_in_reg); + const size_t stage_bytes = nbatch_fa * (ggml_cuda_fattn_q8_0_row_bytes() + ggml_cuda_fattn_q8_0_row_bytes()); + nbytes_shared_total = std::max(nbytes_shared_total, stage_off + stage_bytes); + } + + float logit_softcap; + memcpy(&logit_softcap, (const float *) KQV->op_params + 2, sizeof(float)); + +#if defined(GGML_USE_HIP) + using fattn_kernel_ptr_t = const void*; +#else + using fattn_kernel_ptr_t = fattn_kernel_t; +#endif // defined(GGML_USE_HIP) + fattn_kernel_t fattn_kernel; + if (logit_softcap == 0.0f) { + constexpr bool use_logit_softcap = false; + fattn_kernel = flash_attn_ext_f16; + +#if !defined(GGML_USE_MUSA) + static bool shared_memory_limit_raised[GGML_CUDA_MAX_DEVICES] = {false}; + if (!shared_memory_limit_raised[id]) { + CUDA_CHECK(cudaFuncSetAttribute(reinterpret_cast(fattn_kernel), cudaFuncAttributeMaxDynamicSharedMemorySize, nbytes_shared_total)); + shared_memory_limit_raised[id] = true; + } +#endif // !defined(GGML_USE_MUSA) + } else { + constexpr bool use_logit_softcap = true; + fattn_kernel = flash_attn_ext_f16; + +#if !defined(GGML_USE_MUSA) + static bool shared_memory_limit_raised[GGML_CUDA_MAX_DEVICES] = {false}; + if (!shared_memory_limit_raised[id]) { + CUDA_CHECK(cudaFuncSetAttribute(reinterpret_cast(fattn_kernel), cudaFuncAttributeMaxDynamicSharedMemorySize, nbytes_shared_total)); + shared_memory_limit_raised[id] = true; + } +#endif // !defined(GGML_USE_MUSA) + } + + const bool need_f16_K = false; + const bool need_f16_V = false; + const bool stream_k = true; + launch_fattn + (ctx, dst, fattn_kernel, nwarps, nbytes_shared_total, nbatch_fa, need_f16_K, need_f16_V, stream_k, warp_size_host); +} + #define DECL_FATTN_MMA_F16_CASE(DKQ, DV, ncols1, ncols2) \ template void ggml_cuda_flash_attn_ext_mma_f16_case \ @@ -2031,3 +2258,12 @@ extern DECL_FATTN_MMA_F16_CASE(576, 512, 8, 4); extern DECL_FATTN_MMA_F16_CASE(576, 512, 16, 4); extern DECL_FATTN_MMA_F16_CASE(576, 512, 1, 32); extern DECL_FATTN_MMA_F16_CASE(576, 512, 2, 32); + +#define DECL_FATTN_MMA_Q8_0_CASE(DKQ, DV, ncols1, ncols2) \ + template void ggml_cuda_flash_attn_ext_mma_q8_0_case \ + (ggml_backend_cuda_context & ctx, ggml_tensor * dst) \ + +extern DECL_FATTN_MMA_Q8_0_CASE(256, 256, 1, 8); +extern DECL_FATTN_MMA_Q8_0_CASE(256, 256, 2, 8); +extern DECL_FATTN_MMA_Q8_0_CASE(256, 256, 4, 8); +extern DECL_FATTN_MMA_Q8_0_CASE(256, 256, 8, 8); diff --git a/ggml/src/ggml-cuda/fattn.cu b/ggml/src/ggml-cuda/fattn.cu index ab7a3b297c07..8f377328253e 100644 --- a/ggml/src/ggml-cuda/fattn.cu +++ b/ggml/src/ggml-cuda/fattn.cu @@ -533,6 +533,8 @@ static best_fattn_kernel ggml_cuda_get_best_fattn_kernel(const int device, const return BEST_FATTN_KERNEL_TILE; } +static bool ggml_cuda_flash_attn_ext_mma_q8_0_eligible(int device, const ggml_tensor * dst); + size_t ggml_cuda_flash_attn_ext_get_alloc_size(int device, const ggml_tensor * dst) { GGML_ASSERT(dst->op == GGML_OP_FLASH_ATTN_EXT); @@ -561,15 +563,104 @@ size_t ggml_cuda_flash_attn_ext_get_alloc_size(int device, const ggml_tensor * d break; } + // the fused q8_0 path reads the quantized cache directly, so skip the f16 K/V reservation + if (ggml_cuda_flash_attn_ext_mma_q8_0_eligible(device, dst)) { + need_f16_K = false; + need_f16_V = false; + } + const ggml_cuda_flash_attn_ext_f16_extra_data f16_extra = ggml_cuda_flash_attn_ext_get_f16_extra_data(dst, need_f16_K, need_f16_V); return f16_extra.end - (uintptr_t) dst->data; } +// GGML_CUDA_FA_Q8_0_MMA=0 disables the fused q8_0 KV decode path +static bool ggml_cuda_fattn_q8_0_mma_enabled() { + static const bool value = [] { + const char * env = getenv("GGML_CUDA_FA_Q8_0_MMA"); + return env == nullptr || env[0] != '0'; + }(); + return value; +} + +// smallest ncols1 for single-token queries, 2 by default because the padded (2,8) kernel fills its blocks better +static int ggml_cuda_fattn_q8_0_mma_ncols1_min() { + static const int value = [] { + const char * env = getenv("GGML_CUDA_FA_Q8_0_MMA_NCOLS1_MIN"); + return env != nullptr && env[0] == '1' ? 1 : 2; + }(); + return value; +} + +// eligibility for the fused q8_0 path, shared by the dispatcher and get_alloc_size +// so the f16 K/V reservation is skipped exactly when the fused kernel runs +static bool ggml_cuda_flash_attn_ext_mma_q8_0_eligible(int device, const ggml_tensor * dst) { + const ggml_tensor * KQV = dst; + const ggml_tensor * Q = dst->src[0]; + const ggml_tensor * K = dst->src[1]; + const ggml_tensor * V = dst->src[2]; + const ggml_tensor * mask = dst->src[3]; + + const int cc = ggml_cuda_info().devices[device].cc; + // only tuned and tested on Ampere + const bool is_ampere = GGML_CUDA_CC_IS_NVIDIA(cc) && ggml_cuda_highest_compiled_arch(cc) >= GGML_CUDA_CC_AMPERE && ggml_cuda_highest_compiled_arch(cc) < GGML_CUDA_CC_ADA_LOVELACE; + if (!ggml_cuda_fattn_q8_0_mma_enabled() || !turing_mma_available(cc) || !is_ampere) { + return false; + } + if (K->type != GGML_TYPE_Q8_0 || V->type != GGML_TYPE_Q8_0 || Q->ne[0] != 256 || V->ne[0] != 256 || Q->ne[1] > 5) { + return false; + } + + float max_bias = 0.0f; + memcpy(&max_bias, (const float *) KQV->op_params + 1, sizeof(float)); + if (!mask || max_bias != 0.0f || K->ne[1] % FATTN_KQ_STRIDE != 0 || Q->ne[2] / K->ne[2] <= 4) { + return false; + } + + const bool V_is_K_view = V->view_src && (V->view_src == K || (V->view_src == K->view_src && V->view_offs == K->view_offs)); + if (V_is_K_view) { + return false; + } + + // the raw rows are staged with 16 byte copies + for (const ggml_tensor * t : {K, V}) { + if (t->nb[1] % 16 != 0 || t->nb[2] % 16 != 0 || t->nb[3] % 16 != 0 || ((uintptr_t) t->data) % 16 != 0) { + return false; + } + } + + return true; +} + +// Decode with q8_0 K and V: the kernel reads the quantized cache and packs the GQA heads (ncols2 = 8), +// the generic path would first convert the whole cache to f16. +static bool ggml_cuda_flash_attn_ext_mma_q8_0(ggml_backend_cuda_context & ctx, ggml_tensor * dst) { + if (!ggml_cuda_flash_attn_ext_mma_q8_0_eligible(ggml_cuda_get_device(), dst)) { + return false; + } + + const ggml_tensor * Q = dst->src[0]; + + if (Q->ne[1] == 1 && ggml_cuda_fattn_q8_0_mma_ncols1_min() == 1) { + ggml_cuda_flash_attn_ext_mma_q8_0_case<256, 256, 1, 8>(ctx, dst); + } else if (Q->ne[1] <= 2) { + ggml_cuda_flash_attn_ext_mma_q8_0_case<256, 256, 2, 8>(ctx, dst); + } else if (Q->ne[1] <= 4) { + ggml_cuda_flash_attn_ext_mma_q8_0_case<256, 256, 4, 8>(ctx, dst); + } else { + ggml_cuda_flash_attn_ext_mma_q8_0_case<256, 256, 8, 8>(ctx, dst); + } + return true; +} + void ggml_cuda_flash_attn_ext(ggml_backend_cuda_context & ctx, ggml_tensor * dst) { ggml_cuda_set_device(ctx.device); - switch (ggml_cuda_get_best_fattn_kernel(ggml_cuda_get_device(), dst)) { + const best_fattn_kernel kernel = ggml_cuda_get_best_fattn_kernel(ggml_cuda_get_device(), dst); + if (kernel != BEST_FATTN_KERNEL_NONE && ggml_cuda_flash_attn_ext_mma_q8_0(ctx, dst)) { + return; + } + switch (kernel) { case BEST_FATTN_KERNEL_NONE: GGML_ABORT("fatal error"); case BEST_FATTN_KERNEL_TILE: diff --git a/ggml/src/ggml-cuda/mmvq.cu b/ggml/src/ggml-cuda/mmvq.cu index bf51b61e17b1..e8d4ee454753 100644 --- a/ggml/src/ggml-cuda/mmvq.cu +++ b/ggml/src/ggml-cuda/mmvq.cu @@ -68,8 +68,140 @@ static constexpr __host__ __device__ int get_vdr_mmvq(ggml_type type) { } } +// Decode each weight group once and dot it with all ncols_dst columns. +// Integer and FP ops per column match vec_dot_iq4_xs_q8_1, so results are bit-exact. +template +static __device__ __forceinline__ void vec_dot_iq4_xs_q8_1_multi( + const void * __restrict__ vbq, const block_q8_1 * __restrict__ y, + const int stride_col_y, const int kby, const int kbx, const int iqs, + float (&dots)[ncols_dst]) { + const block_iq4_xs * bq4 = (const block_iq4_xs *) vbq + kbx; + + int sumi[ncols_dst]; +#pragma unroll + for (int col = 0; col < ncols_dst; ++col) { + sumi[col] = 0; + } +#pragma unroll + for (int j = 0; j < 4; ++j) { + const int aux_q4 = get_int_b4(bq4->qs, iqs + j); + const int2 v = get_int_from_table_16_interleaved(aux_q4, kvalues_iq4nl); +#pragma unroll + for (int col = 0; col < ncols_dst; ++col) { + const block_q8_1 * bq8 = y + col*stride_col_y + kby + iqs/4; + sumi[col] = ggml_cuda_dp4a(v.x, get_int_b4(bq8->qs, 2*j + 0), sumi[col]); + sumi[col] = ggml_cuda_dp4a(v.y, get_int_b4(bq8->qs, 2*j + 1), sumi[col]); + } + } + + const int ls = vec_dot_iq4_xs_q8_1_scale(bq4, iqs); +#pragma unroll + for (int col = 0; col < ncols_dst; ++col) { + int s = sumi[col]; + s *= ls - 32; + const float d = __half2float(bq4->d) * __low2float((y + col*stride_col_y + kby + iqs/4)->ds); + dots[col] = d * s; + } +} + +// Same as above for PQ2_0 (group 128, iqs selects one 32-element chunk), the ops per column match vec_dot_pq2_0_q8_1. +template +static __device__ __forceinline__ void vec_dot_pq2_0_q8_1_multi( + const void * __restrict__ vbq, const block_q8_1 * __restrict__ y, + const int stride_col_y, const int kby, const int kbx, const int iqs, + float (&dots)[ncols_dst]) { + const block_pq2_0 * bq2_0 = (const block_pq2_0 *) vbq + kbx; + const int16_t * qs = (const int16_t *) bq2_0->qs + iqs * 4; + + int sumi[ncols_dst]; +#pragma unroll + for (int col = 0; col < ncols_dst; ++col) { + sumi[col] = 0; + } +#pragma unroll + for (int j = 0; j < 4; ++j) { + const int q = qs[j]; + const int qe = __byte_perm(0x020100FF, 0x020100FF, q >> 0); + const int qo = __byte_perm(0x020100FF, 0x020100FF, q >> 2); + const int qx = __byte_perm(qe, qo, 0x5140); + const int qy = __byte_perm(qe, qo, 0x7362); +#pragma unroll + for (int col = 0; col < ncols_dst; ++col) { + const block_q8_1 * bq8 = y + col*stride_col_y + kby + iqs; + sumi[col] = ggml_cuda_dp4a(get_int_b4(bq8->qs, j*2 + 0), qx, sumi[col]); + sumi[col] = ggml_cuda_dp4a(get_int_b4(bq8->qs, j*2 + 1), qy, sumi[col]); + } + } + + const float d2 = bq2_0->d; +#pragma unroll + for (int col = 0; col < ncols_dst; ++col) { + const float d8 = __low2float((y + col*stride_col_y + kby + iqs)->ds); + dots[col] = d2 * d8 * sumi[col]; + } +} + +// Same as above for Q5_0, the ops per column match vec_dot_q5_0_q8_1_impl. +template +static __device__ __forceinline__ void vec_dot_q5_0_q8_1_multi( + const void * __restrict__ vbq, const block_q8_1 * __restrict__ y, + const int stride_col_y, const int kby, const int kbx, const int iqs, + float (&dots)[ncols_dst]) { + const block_q5_0 * bq5_0 = (const block_q5_0 *) vbq + kbx; + + int sumi[ncols_dst]; +#pragma unroll + for (int col = 0; col < ncols_dst; ++col) { + sumi[col] = 0; + } + + const int qh = get_int_b2(bq5_0->qh, 0); +#pragma unroll + for (int i = 0; i < VDR_Q5_0_Q8_1_MMVQ; ++i) { + const int vl = get_int_b2(bq5_0->qs, iqs + i); + const int vh = qh >> (4 * (iqs + i)); + + int vi0 = (vl >> 0) & 0x0F0F0F0F; // lower 4 qs bits, still need qh as 5th bits + vi0 |= (vh << 4) & 0x00000010; // 0 -> 4 + vi0 |= (vh << 11) & 0x00001000; // 1 -> 12 + vi0 |= (vh << 18) & 0x00100000; // 2 -> 20 + vi0 |= (vh << 25) & 0x10000000; // 3 -> 28 + + int vi1 = (vl >> 4) & 0x0F0F0F0F; // upper 4 qs bits, still need qh as 5th bits + vi1 |= (vh >> 12) & 0x00000010; // 16 -> 4 + vi1 |= (vh >> 5) & 0x00001000; // 17 -> 12 + vi1 |= (vh << 2) & 0x00100000; // 18 -> 20 + vi1 |= (vh << 9) & 0x10000000; // 19 -> 28 + +#pragma unroll + for (int col = 0; col < ncols_dst; ++col) { + const block_q8_1 * bq8 = y + col*stride_col_y + kby; + sumi[col] = ggml_cuda_dp4a(vi0, get_int_b4(bq8->qs, iqs + i), sumi[col]); + sumi[col] = ggml_cuda_dp4a(vi1, get_int_b4(bq8->qs, iqs + i + QI5_0), sumi[col]); + } + } + + const float d5 = bq5_0->d; +#pragma unroll + for (int col = 0; col < ncols_dst; ++col) { + const float2 ds8f = __half22float2((y + col*stride_col_y + kby)->ds); + // second part effectively subtracts 16 from each quant value + dots[col] = d5 * (sumi[col] * ds8f.x - (16*VDR_Q5_0_Q8_1_MMVQ/QI5_0) * ds8f.y); + } +} + +// GGML_CUDA_MMVQ_REUSE=0 disables the multi-column vec_dot path +static bool ggml_cuda_mmvq_reuse_enabled() { + static const bool value = [] { + const char * env = getenv("GGML_CUDA_MMVQ_REUSE"); + return env == nullptr || env[0] != '0'; + }(); + return value; +} + enum mmvq_parameter_table_id { MMVQ_PARAMETERS_GENERIC = 0, + MMVQ_PARAMETERS_AMPERE, MMVQ_PARAMETERS_TURING, MMVQ_PARAMETERS_GCN, MMVQ_PARAMETERS_RDNA2, @@ -89,6 +221,8 @@ static constexpr __device__ mmvq_parameter_table_id get_device_table_id() { return MMVQ_PARAMETERS_GCN; #elif defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= GGML_CUDA_CC_TURING && __CUDA_ARCH__ < GGML_CUDA_CC_AMPERE return MMVQ_PARAMETERS_TURING; +#elif defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= GGML_CUDA_CC_AMPERE && __CUDA_ARCH__ < GGML_CUDA_CC_ADA_LOVELACE + return MMVQ_PARAMETERS_AMPERE; #elif defined(__CUDA_ARCH__) && __CUDA_ARCH__ == GGML_CUDA_CC_DGX_SPARK return MMVQ_PARAMETERS_GB10; #else @@ -112,6 +246,9 @@ static __host__ mmvq_parameter_table_id get_device_table_id(int cc) { if (GGML_CUDA_CC_IS_NVIDIA(cc) && ggml_cuda_highest_compiled_arch(cc) >= GGML_CUDA_CC_TURING && ggml_cuda_highest_compiled_arch(cc) < GGML_CUDA_CC_AMPERE) { return MMVQ_PARAMETERS_TURING; } + if (GGML_CUDA_CC_IS_NVIDIA(cc) && ggml_cuda_highest_compiled_arch(cc) >= GGML_CUDA_CC_AMPERE && ggml_cuda_highest_compiled_arch(cc) < GGML_CUDA_CC_ADA_LOVELACE) { + return MMVQ_PARAMETERS_AMPERE; + } if (GGML_CUDA_CC_IS_NVIDIA(cc) && ggml_cuda_highest_compiled_arch(cc) == GGML_CUDA_CC_DGX_SPARK) { return MMVQ_PARAMETERS_GB10; } @@ -404,7 +541,7 @@ static constexpr __device__ int get_mmvq_mmid_max_batch_for_device() { } static constexpr __host__ __device__ int calc_nwarps(ggml_type type, int ncols_dst, mmvq_parameter_table_id table_id, bool small_k = false, bool halve_iters = false) { - if (table_id == MMVQ_PARAMETERS_GENERIC) { + if (table_id == MMVQ_PARAMETERS_GENERIC || table_id == MMVQ_PARAMETERS_AMPERE) { switch (ncols_dst) { case 1: case 2: @@ -534,8 +671,18 @@ static constexpr __host__ __device__ int calc_nwarps(ggml_type type, int ncols_d return 1; } -static constexpr __host__ __device__ int calc_rows_per_block(int ncols_dst, int table_id, bool small_k = false, int nwarps = 1) { - if (table_id == MMVQ_PARAMETERS_GENERIC || table_id == MMVQ_PARAMETERS_GCN || table_id == MMVQ_PARAMETERS_TURING || table_id == MMVQ_PARAMETERS_GB10) { +// PQ2_0 at 2 columns: 4 rows per block measured faster than 8 on Bonsai 2 27B shapes (RTX 3060) +#ifndef GGML_CUDA_MMVQ_AMPERE_PQ2_0_RPB_NCOLS2 +#define GGML_CUDA_MMVQ_AMPERE_PQ2_0_RPB_NCOLS2 4 +#endif + +static constexpr __host__ __device__ int calc_rows_per_block(int ncols_dst, int table_id, bool small_k = false, int nwarps = 1, + ggml_type type = GGML_TYPE_COUNT) { + // Ampere: 8 rows per block for the speculative verify widths + if (table_id == MMVQ_PARAMETERS_AMPERE && ncols_dst >= 2 && ncols_dst <= 4) { + return type == GGML_TYPE_PQ2_0 && ncols_dst == 2 ? GGML_CUDA_MMVQ_AMPERE_PQ2_0_RPB_NCOLS2 : 8; + } + if (table_id == MMVQ_PARAMETERS_GENERIC || table_id == MMVQ_PARAMETERS_AMPERE ||table_id == MMVQ_PARAMETERS_GCN || table_id == MMVQ_PARAMETERS_TURING || table_id == MMVQ_PARAMETERS_GB10) { switch (ncols_dst) { case 1: return small_k ? nwarps : 1; @@ -554,7 +701,7 @@ static constexpr __host__ __device__ int calc_rows_per_block(int ncols_dst, int return 1; } -template +template __launch_bounds__(calc_nwarps(type, ncols_dst, get_device_table_id(), small_k, halve_iters)*ggml_cuda_get_physical_warp_size(), 1) static __global__ void mul_mat_vec_q( const void * vx_ptr, const void * vy_ptr, const int32_t * ids_ptr, const ggml_cuda_mm_fusion_args_device fusion, float * dst_ptr, @@ -573,7 +720,7 @@ static __global__ void mul_mat_vec_q( constexpr int vdr = get_vdr_mmvq(type); constexpr mmvq_parameter_table_id table_id = get_device_table_id(); constexpr int nwarps = calc_nwarps(type, ncols_dst, table_id, small_k, halve_iters); - constexpr int rows_per_cuda_block = calc_rows_per_block(ncols_dst, table_id, small_k, nwarps); + constexpr int rows_per_cuda_block = calc_rows_per_block(ncols_dst, table_id, small_k, nwarps, type); constexpr int warp_size = ggml_cuda_get_physical_warp_size(); constexpr vec_dot_q_cuda_t vec_dot_q_cuda = get_vec_dot_q_cuda(type); @@ -715,6 +862,23 @@ static __global__ void mul_mat_vec_q( // x block quant index when casting the quants to int const int kqs = vdr * (tid % (qi/vdr)); + if constexpr (reuse_weights && (type == GGML_TYPE_IQ4_XS || type == GGML_TYPE_Q5_0 || type == GGML_TYPE_PQ2_0)) { +#pragma unroll + for (int i = 0; i < rows_per_cuda_block; ++i) { + float dots[ncols_dst]; + if constexpr (type == GGML_TYPE_IQ4_XS) { + vec_dot_iq4_xs_q8_1_multi(vx, y, stride_col_y, kby, kbx_offset + i*stride_row_x + kbx, kqs, dots); + } else if constexpr (type == GGML_TYPE_PQ2_0) { + vec_dot_pq2_0_q8_1_multi(vx, y, stride_col_y, kby, kbx_offset + i*stride_row_x + kbx, kqs, dots); + } else { + vec_dot_q5_0_q8_1_multi(vx, y, stride_col_y, kby, kbx_offset + i*stride_row_x + kbx, kqs, dots); + } +#pragma unroll + for (int j = 0; j < ncols_dst; ++j) { + tmp[j][i] += dots[j]; + } + } + } else #if !defined(GGML_USE_HIP) if constexpr (type == GGML_TYPE_PTQ1_0 && ncols_dst > 1 && ncols_dst <= 3) { # pragma unroll @@ -922,7 +1086,7 @@ static std::pair calc_launch_params( const int ncols_dst, const int nrows_x, const int nchannels_dst, const int nsamples_or_ntokens, const int warp_size, const mmvq_parameter_table_id table_id, const bool small_k = false, const bool halve_iters = false) { const int nwarps = calc_nwarps(type, ncols_dst, table_id, small_k, halve_iters); - const int rpb = calc_rows_per_block(ncols_dst, table_id, small_k, nwarps); + const int rpb = calc_rows_per_block(ncols_dst, table_id, small_k, nwarps, type); const int64_t nblocks = (nrows_x + rpb - 1) / rpb; const dim3 block_nums(nblocks, nchannels_dst, nsamples_or_ntokens); const dim3 block_dims(warp_size, nwarps, 1); @@ -962,6 +1126,16 @@ static void mul_mat_vec_q_switch_fusion( GGML_ASSERT(!has_fusion && "fusion only supported for ncols_dst=1"); const ggml_cuda_kernel_launch_params launch_params = ggml_cuda_kernel_launch_params(block_nums, block_dims, nbytes_shared, stream); + if constexpr ((type == GGML_TYPE_IQ4_XS || type == GGML_TYPE_Q5_0 || type == GGML_TYPE_PQ2_0) && c_ncols_dst >= 2 && c_ncols_dst <= 5) { + const int cc = ggml_cuda_info().devices[ggml_cuda_get_device()].cc; + if (get_device_table_id(cc) == MMVQ_PARAMETERS_AMPERE && ggml_cuda_mmvq_reuse_enabled()) { + ggml_cuda_kernel_launch(mul_mat_vec_q, launch_params, + vx, vy, ids, fusion, dst, ncols_x, nchannels_y, stride_row_x, stride_col_y, stride_col_dst, + channel_ratio, stride_channel_x, stride_channel_y, stride_channel_dst, + sample_ratio, stride_sample_x, stride_sample_y, stride_sample_dst, ids_stride); + return; + } + } ggml_cuda_kernel_launch(mul_mat_vec_q, launch_params, vx, vy, ids, fusion, dst, ncols_x, nchannels_y, stride_row_x, stride_col_y, stride_col_dst, channel_ratio, stride_channel_x, stride_channel_y, stride_channel_dst, diff --git a/ggml/src/ggml-cuda/quantize.cu b/ggml/src/ggml-cuda/quantize.cu index fbbc6314ab33..9edaff3a8d27 100644 --- a/ggml/src/ggml-cuda/quantize.cu +++ b/ggml/src/ggml-cuda/quantize.cu @@ -51,6 +51,8 @@ static __device__ __forceinline__ float nvfp4_native_scale_error( #endif // CUDART_VERSION >= 12080 #endif // defined(BLACKWELL_MMA_AVAILABLE) +// swizzle_iq4: store element k at byte 2*(k%16) + k/16, the layout vec_dot_iq4_xs_q8_1 reads +template __launch_bounds__(CUDA_QUANTIZE_BLOCK_SIZE, 1) static __global__ void quantize_q8_1( const float * x_ptr, void * vy_ptr, @@ -92,7 +94,11 @@ static __global__ void quantize_q8_1( const float d = amax / 127.0f; const int8_t q = amax == 0.0f ? 0 : roundf(xi / d); - y[ib].qs[iqs] = q; + if constexpr (swizzle_iq4) { + y[ib].qs[2*(iqs % 16) + iqs/16] = q; + } else { + y[ib].qs[iqs] = q; + } if (iqs > 0) { return; @@ -648,8 +654,11 @@ void quantize_row_q8_1_cuda( const dim3 num_blocks(block_num_x, ne1, ne2*ne3); const dim3 block_size(CUDA_QUANTIZE_BLOCK_SIZE, 1, 1); const ggml_cuda_kernel_launch_params launch_params = ggml_cuda_kernel_launch_params(num_blocks, block_size, 0, stream); - ggml_cuda_kernel_launch(quantize_q8_1, launch_params, x, vy, ne00, s01, s02, s03, ne0, ne1, ne2_fastdiv); - GGML_UNUSED(type_src0); + if (type_src0 == GGML_TYPE_IQ4_XS) { + ggml_cuda_kernel_launch(quantize_q8_1, launch_params, x, vy, ne00, s01, s02, s03, ne0, ne1, ne2_fastdiv); + } else { + ggml_cuda_kernel_launch(quantize_q8_1, launch_params, x, vy, ne00, s01, s02, s03, ne0, ne1, ne2_fastdiv); + } } void quantize_mmq_q8_1_cuda( diff --git a/ggml/src/ggml-cuda/template-instances/fattn-mma-f16-instance-ncols1_1-ncols2_8.cu b/ggml/src/ggml-cuda/template-instances/fattn-mma-f16-instance-ncols1_1-ncols2_8.cu index 6ae77bec8958..8e1c8eaab884 100644 --- a/ggml/src/ggml-cuda/template-instances/fattn-mma-f16-instance-ncols1_1-ncols2_8.cu +++ b/ggml/src/ggml-cuda/template-instances/fattn-mma-f16-instance-ncols1_1-ncols2_8.cu @@ -10,3 +10,4 @@ DECL_FATTN_MMA_F16_CASE(128, 128, 1, 8); DECL_FATTN_MMA_F16_CASE(192, 128, 1, 8); DECL_FATTN_MMA_F16_CASE(256, 256, 1, 8); DECL_FATTN_MMA_F16_CASE(512, 512, 1, 8); +DECL_FATTN_MMA_Q8_0_CASE(256, 256, 1, 8); diff --git a/ggml/src/ggml-cuda/template-instances/fattn-mma-f16-instance-ncols1_2-ncols2_8.cu b/ggml/src/ggml-cuda/template-instances/fattn-mma-f16-instance-ncols1_2-ncols2_8.cu index 9f4bef11a443..c2b935ca4302 100644 --- a/ggml/src/ggml-cuda/template-instances/fattn-mma-f16-instance-ncols1_2-ncols2_8.cu +++ b/ggml/src/ggml-cuda/template-instances/fattn-mma-f16-instance-ncols1_2-ncols2_8.cu @@ -10,3 +10,4 @@ DECL_FATTN_MMA_F16_CASE(128, 128, 2, 8); DECL_FATTN_MMA_F16_CASE(192, 128, 2, 8); DECL_FATTN_MMA_F16_CASE(256, 256, 2, 8); DECL_FATTN_MMA_F16_CASE(512, 512, 2, 8); +DECL_FATTN_MMA_Q8_0_CASE(256, 256, 2, 8); diff --git a/ggml/src/ggml-cuda/template-instances/fattn-mma-f16-instance-ncols1_4-ncols2_8.cu b/ggml/src/ggml-cuda/template-instances/fattn-mma-f16-instance-ncols1_4-ncols2_8.cu index 859bea5c5253..1d40596ba4d9 100644 --- a/ggml/src/ggml-cuda/template-instances/fattn-mma-f16-instance-ncols1_4-ncols2_8.cu +++ b/ggml/src/ggml-cuda/template-instances/fattn-mma-f16-instance-ncols1_4-ncols2_8.cu @@ -10,3 +10,4 @@ DECL_FATTN_MMA_F16_CASE(128, 128, 4, 8); DECL_FATTN_MMA_F16_CASE(192, 128, 4, 8); DECL_FATTN_MMA_F16_CASE(256, 256, 4, 8); DECL_FATTN_MMA_F16_CASE(512, 512, 4, 8); +DECL_FATTN_MMA_Q8_0_CASE(256, 256, 4, 8); diff --git a/ggml/src/ggml-cuda/template-instances/fattn-mma-f16-instance-ncols1_8-ncols2_8.cu b/ggml/src/ggml-cuda/template-instances/fattn-mma-f16-instance-ncols1_8-ncols2_8.cu index c975ce6b9b74..89b8325c4fa1 100644 --- a/ggml/src/ggml-cuda/template-instances/fattn-mma-f16-instance-ncols1_8-ncols2_8.cu +++ b/ggml/src/ggml-cuda/template-instances/fattn-mma-f16-instance-ncols1_8-ncols2_8.cu @@ -10,3 +10,4 @@ DECL_FATTN_MMA_F16_CASE(128, 128, 8, 8); DECL_FATTN_MMA_F16_CASE(192, 128, 8, 8); DECL_FATTN_MMA_F16_CASE(256, 256, 8, 8); DECL_FATTN_MMA_F16_CASE(512, 512, 8, 8); +DECL_FATTN_MMA_Q8_0_CASE(256, 256, 8, 8); 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 c54928a50ae6..901407378446 100755 --- a/ggml/src/ggml-cuda/template-instances/generate_cu_files.py +++ b/ggml/src/ggml-cuda/template-instances/generate_cu_files.py @@ -33,6 +33,7 @@ """ SOURCE_FATTN_MMA_CASE = "DECL_FATTN_MMA_F16_CASE({head_size_kq}, {head_size_v}, {ncols1}, {ncols2});\n" +SOURCE_FATTN_MMA_Q8_0_CASE = "DECL_FATTN_MMA_Q8_0_CASE({head_size_kq}, {head_size_v}, {ncols1}, {ncols2});\n" TYPES_MMQ = [ "GGML_TYPE_Q1_0", @@ -113,6 +114,10 @@ def get_short_name(long_quant_name): 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)) + # q8_0 KV decoded inside the kernel, used by the fused decode path in fattn.cu + if ncols2 == 8: + f.write(SOURCE_FATTN_MMA_Q8_0_CASE.format(ncols1=ncols1, ncols2=ncols2, head_size_kq=256, head_size_v=256)) + 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 diff --git a/ggml/src/ggml-cuda/vecdotq.cuh b/ggml/src/ggml-cuda/vecdotq.cuh index 6c70b6cb5a39..7232292c222b 100644 --- a/ggml/src/ggml-cuda/vecdotq.cuh +++ b/ggml/src/ggml-cuda/vecdotq.cuh @@ -1564,25 +1564,49 @@ static __device__ __forceinline__ float vec_dot_iq4_nl_q8_1( #define VDR_IQ4_XS_Q8_1_MMVQ 4 #define VDR_IQ4_XS_Q8_1_MMQ 4 -static __device__ __forceinline__ float vec_dot_iq4_xs_q8_1( - const void * __restrict__ vbq, const block_q8_1 * __restrict__ bq8_1, const int & kbx, const int & iqs) { - - const block_iq4_xs * bq4 = (const block_iq4_xs *) vbq + kbx; +// Nibble lookup in the interleaved order that quantize_q8_1 uses for IQ4_XS activations. +// .x = {lo(b0), hi(b0), lo(b1), hi(b1)}, .y = {lo(b2), hi(b2), lo(b3), hi(b3)}. +static __device__ __forceinline__ int2 get_int_from_table_16_interleaved(const int & q4, const int8_t * table) { + const uint32_t * values = (const uint32_t *) table; + const uint32_t q = (uint32_t) q4; + const uint32_t selectors = 0x32103210u | ((q & 0x88888888u) >> 1); + const uint32_t lo0 = __byte_perm(values[0], values[1], q); + const uint32_t hi0 = __byte_perm(values[2], values[3], q); + const uint32_t lo1 = __byte_perm(values[0], values[1], q >> 16); + const uint32_t hi1 = __byte_perm(values[2], values[3], q >> 16); + return make_int2(__byte_perm(lo0, hi0, selectors), __byte_perm(lo1, hi1, selectors >> 16)); +} +// bq8 must use the interleaved IQ4_XS activation layout written by quantize_q8_1. +static __device__ __forceinline__ int vec_dot_iq4_xs_q8_1_sumi( + const block_iq4_xs * __restrict__ bq4, const block_q8_1 * __restrict__ bq8, const int & iqs) { int sumi = 0; #pragma unroll for (int j = 0; j < 4; ++j) { const int aux_q4 = get_int_b4(bq4->qs, iqs + j); - const int2 v = get_int_from_table_16(aux_q4, kvalues_iq4nl); + const int2 v = get_int_from_table_16_interleaved(aux_q4, kvalues_iq4nl); - const int u0 = get_int_b4(bq8_1[iqs/4].qs, j + 0); - const int u1 = get_int_b4(bq8_1[iqs/4].qs, j + 4); + const int u0 = get_int_b4(bq8->qs, 2*j + 0); + const int u1 = get_int_b4(bq8->qs, 2*j + 1); sumi = ggml_cuda_dp4a(v.x, u0, sumi); sumi = ggml_cuda_dp4a(v.y, u1, sumi); } + return sumi; +} + +static __device__ __forceinline__ int vec_dot_iq4_xs_q8_1_scale(const block_iq4_xs * __restrict__ bq4, const int & iqs) { + return ((bq4->scales_l[iqs/8] >> (iqs & 0x04)) & 0x0F) | (((bq4->scales_h >> (iqs/2)) & 0x03) << 4); +} + +static __device__ __forceinline__ float vec_dot_iq4_xs_q8_1( + const void * __restrict__ vbq, const block_q8_1 * __restrict__ bq8_1, const int & kbx, const int & iqs) { + + const block_iq4_xs * bq4 = (const block_iq4_xs *) vbq + kbx; + + int sumi = vec_dot_iq4_xs_q8_1_sumi(bq4, &bq8_1[iqs/4], iqs); - const int ls = ((bq4->scales_l[iqs/8] >> (iqs & 0x04)) & 0x0F) | (((bq4->scales_h >> (iqs/2)) & 0x03) << 4); + const int ls = vec_dot_iq4_xs_q8_1_scale(bq4, iqs); sumi *= ls - 32; const float d = __half2float(bq4->d) * __low2float(bq8_1[iqs/4].ds); diff --git a/tests/test-backend-ops.cpp b/tests/test-backend-ops.cpp index 14f654f966ca..4f836f1ee082 100644 --- a/tests/test-backend-ops.cpp +++ b/tests/test-backend-ops.cpp @@ -9285,6 +9285,20 @@ static std::vector> make_test_cases_eval() { test_cases.emplace_back(new test_mul_mat(GGML_TYPE_Q8_0, GGML_TYPE_F32, 8192, 512, 5120, {128, 1}, {1, 1})); #endif + // MMVQ at ncols_dst 2..5 (multi-column vec_dot and Ampere rows per block), opt-in + if (getenv("GGML_QWEN38_REUSE_TEST") != nullptr) { + for (ggml_type type_a : {GGML_TYPE_Q5_0, GGML_TYPE_Q6_K, GGML_TYPE_IQ4_XS, GGML_TYPE_Q4_K, GGML_TYPE_Q5_K, GGML_TYPE_PQ2_0}) { + for (int64_t n : {2, 3, 4, 5}) { + for (int64_t k : {256, 512, 1024, 1280}) { + test_cases.emplace_back(new test_mul_mat(type_a, GGML_TYPE_F32, 512, n, k, {1, 1}, {1, 1})); + } + // row count not a multiple of rows per block + test_cases.emplace_back(new test_mul_mat(type_a, GGML_TYPE_F32, 517, n, 1024, {1, 1}, {1, 1})); + test_cases.emplace_back(new test_mul_mat(type_a, GGML_TYPE_F32, 517, n, 1024, {2, 3}, {1, 1})); + } + } + } + for (ggml_type type_a : all_types) { for (int i = 1; i < 10; ++i) { test_cases.emplace_back(new test_mul_mat(type_a, GGML_TYPE_F32, 16, i, 1*256, { 1, 1}, {1, 1})); @@ -10022,6 +10036,21 @@ static std::vector> make_test_cases_eval() { } } + // q8_0 KV decode with head size 256 and GQA > 4 (fused q8_0 MMA path on Ampere), opt-in + if (getenv("GGML_QWEN38_FA_Q8_0_TEST") != nullptr) { + for (int nr2 : { 6, 8 }) { + for (int kv : { 512, 1024, 4096, 1000 }) { + for (int nb : { 1, 2, 3, 4, 5 }) { + for (bool sinks : { false, true }) { + for (float logit_softcap : { 0.0f, 10.0f }) { + test_cases.emplace_back(new test_flash_attn_ext(256, 256, 4, {nr2, 1}, kv, nb, true, sinks, 0.0f, logit_softcap, GGML_PREC_F32, GGML_TYPE_Q8_0, GGML_TYPE_Q8_0)); + } + } + } + } + } + } + for (int hsk : { 40, 64, 72, 80, 96, 128, 192, 256, 320, 512, 576 }) { for (int hsv : { 40, 64, 72, 80, 96, 128, 192, 256, 512 }) { if (hsk != 192 && hsk != 320 && hsk != 576 && hsk != hsv) continue; @@ -10302,6 +10331,19 @@ static std::vector> make_test_cases_perf() { } } + // Bonsai 2 27B PQ2_0 decode GEMV shapes at 1..5 columns, opt-in + if (getenv("GGML_BONSAI_PERF") != nullptr) { + const std::array, 7> mk = {{ + {17408, 5120}, {12288, 5120}, {10240, 5120}, {6144, 5120}, {1024, 5120}, {5120, 17408}, {5120, 6144}, + }}; + for (const auto & s : mk) { + for (int64_t n : {1, 2, 3, 4, 5}) { + test_cases.emplace_back(new test_mul_mat(GGML_TYPE_PQ2_0, GGML_TYPE_F32, s[0], n, s[1], {1, 1}, {1, 1})); + } + } + return test_cases; + } + // SWIGLU at a 27B-class FFN width, fused [gate|up] vs split operands // note: same bytes either way, so a backend that indexes them differently shows it here for (ggml_type type : {GGML_TYPE_F16, GGML_TYPE_F32}) {