diff --git a/ggml/src/ggml-cuda/mmq-load-tiles.cuh b/ggml/src/ggml-cuda/mmq-load-tiles.cuh index acf62501e07..12a84bb3d9f 100644 --- a/ggml/src/ggml-cuda/mmq-load-tiles.cuh +++ b/ggml/src/ggml-cuda/mmq-load-tiles.cuh @@ -263,6 +263,90 @@ template static __device__ __forceinline_ } } +#if defined(GGML_USE_HIP) && defined(RDNA3_5) +template +static __device__ __forceinline__ void ggml_cuda_mmq_prefetch_tiles_q4_0_rdna35( + const char * __restrict__ x, const int kbx0, const int i_max, const int stride, + int (&qs_cache)[ggml_cuda_mmq_get_I(type, J, fallback)/ + (ggml_cuda_mmq_get_nthreads(type, J, fallback)/ggml_cuda_get_physical_warp_size())], + float (&d_cache)[4]) { + constexpr int warp_size = ggml_cuda_get_physical_warp_size(); + constexpr int nthreads = ggml_cuda_mmq_get_nthreads(type, J, fallback); + constexpr int nwarps = nthreads / warp_size; + constexpr int I = ggml_cuda_mmq_get_I(type, J, fallback); + static_assert(warp_size == 32 && nthreads == 128 && I == 64, "unexpected RDNA3.5 Q4_0 MMQ configuration"); + + const int kbx = threadIdx.x / QI4_0; + const int kqsx = threadIdx.x % QI4_0; + +#pragma unroll + for (int i0 = 0; i0 < I; i0 += nwarps) { + int i = i0 + threadIdx.y; + if constexpr (fallback) { + i = min(i, i_max); + } + + const block_q4_0 * bxi = (const block_q4_0 *) x + kbx0 + i*stride + kbx; + qs_cache[i0/nwarps] = get_int_b2(bxi->qs, kqsx); + } + + constexpr int blocks_per_tile_x_row = MMQ_TILE_NE_K / QI4_0; + constexpr int scale_rows_per_warp = warp_size / blocks_per_tile_x_row; + const int kbxd = threadIdx.x % blocks_per_tile_x_row; + int d_idx = 0; +#pragma unroll + for (int i0 = 0; i0 < I; i0 += nwarps * scale_rows_per_warp) { + int i = i0 + threadIdx.y * scale_rows_per_warp + threadIdx.x / blocks_per_tile_x_row; + if constexpr (fallback) { + i = min(i, i_max); + } + + const block_q4_0 * bxi = (const block_q4_0 *) x + kbx0 + i*stride + kbxd; + d_cache[d_idx++] = bxi->d; + } + + asm volatile("" ::: "memory"); +} + +template +static __device__ __forceinline__ void ggml_cuda_mmq_store_tiles_q4_0_rdna35( + int * __restrict__ x_tile, + const int (&qs_cache)[ggml_cuda_mmq_get_I(type, J, fallback)/ + (ggml_cuda_mmq_get_nthreads(type, J, fallback)/ggml_cuda_get_physical_warp_size())], + const float (&d_cache)[4]) { + constexpr int warp_size = ggml_cuda_get_physical_warp_size(); + constexpr int nthreads = ggml_cuda_mmq_get_nthreads(type, J, fallback); + constexpr int nwarps = nthreads / warp_size; + constexpr int I = ggml_cuda_mmq_get_I(type, J, fallback); + constexpr int sram_stride = ggml_cuda_mmq_get_sram_stride(type, J, fallback); + static_assert(warp_size == 32 && nthreads == 128 && I == 64, "unexpected RDNA3.5 Q4_0 MMQ configuration"); + + int * x_qs = x_tile; + float * x_df = (float *) (x_qs + 2*MMQ_TILE_NE_K); + const int txi = threadIdx.x; + const int kbx = txi / QI4_0; + const int kqsx = txi % QI4_0; + +#pragma unroll + for (int i0 = 0; i0 < I; i0 += nwarps) { + const int i = i0 + threadIdx.y; + const int qs0 = qs_cache[i0/nwarps]; + x_qs[i*sram_stride + kbx*(2*QI4_0) + kqsx + 0] = __vsubss4((qs0 >> 0) & 0x0F0F0F0F, 0x08080808); + x_qs[i*sram_stride + kbx*(2*QI4_0) + kqsx + QI4_0] = __vsubss4((qs0 >> 4) & 0x0F0F0F0F, 0x08080808); + } + + constexpr int blocks_per_tile_x_row = MMQ_TILE_NE_K / QI4_0; + constexpr int scale_rows_per_warp = warp_size / blocks_per_tile_x_row; + const int kbxd = threadIdx.x % blocks_per_tile_x_row; + int d_idx = 0; +#pragma unroll + for (int i0 = 0; i0 < I; i0 += nwarps * scale_rows_per_warp) { + const int i = i0 + threadIdx.y * scale_rows_per_warp + threadIdx.x / blocks_per_tile_x_row; + x_df[i*sram_stride + kbxd] = d_cache[d_idx++]; + } +} +#endif // defined(GGML_USE_HIP) && defined(RDNA3_5) + template static __device__ __forceinline__ void ggml_cuda_mmq_load_tiles_q4_1( const char * __restrict__ x, int * __restrict__ x_tile, const int kbx0, const int i_max, const int stride) { constexpr int warp_size = ggml_cuda_get_physical_warp_size(); @@ -548,6 +632,89 @@ template static __device__ __forceinline_ } } +#if defined(GGML_USE_HIP) && defined(RDNA3_5) +template +static __device__ __forceinline__ void ggml_cuda_mmq_prefetch_tiles_q8_0_rdna35( + const char * __restrict__ x, const int kbx0, const int i_max, const int stride, + int (&qs_cache)[2 * ggml_cuda_mmq_get_I(type, J, fallback)/ + (ggml_cuda_mmq_get_nthreads(type, J, fallback)/ggml_cuda_get_physical_warp_size())], + float (&d_cache)[4]) { + constexpr int warp_size = ggml_cuda_get_physical_warp_size(); + constexpr int nthreads = ggml_cuda_mmq_get_nthreads(type, J, fallback); + constexpr int nwarps = nthreads / warp_size; + constexpr int I = ggml_cuda_mmq_get_I(type, J, fallback); + static_assert(warp_size == 32 && nthreads == 128 && I == 64, "unexpected RDNA3.5 Q8_0 MMQ configuration"); + + const int txi = threadIdx.x; + const int kbx = txi / QI8_0; + const int kqsx = txi % QI8_0; + +#pragma unroll + for (int i0 = 0; i0 < I; i0 += nwarps) { + int i = i0 + threadIdx.y; + if constexpr (fallback) { + i = min(i, i_max); + } + + const block_q8_0 * bxi = (const block_q8_0 *) x + kbx0 + i*stride + kbx; + qs_cache[2*(i0/nwarps) + 0] = get_int_b2(bxi[0].qs, kqsx); + qs_cache[2*(i0/nwarps) + 1] = get_int_b2(bxi[MMQ_TILE_NE_K/QI8_0].qs, kqsx); + } + + constexpr int blocks_per_tile_x_row = 2*MMQ_TILE_NE_K / QI8_0; + constexpr int scale_rows_per_warp = warp_size / blocks_per_tile_x_row; + const int kbxd = threadIdx.x % blocks_per_tile_x_row; + int d_idx = 0; +#pragma unroll + for (int i0 = 0; i0 < I; i0 += nwarps * scale_rows_per_warp) { + int i = i0 + threadIdx.y * scale_rows_per_warp + threadIdx.x / blocks_per_tile_x_row; + if constexpr (fallback) { + i = min(i, i_max); + } + + const block_q8_0 * bxi = (const block_q8_0 *) x + kbx0 + i*stride + kbxd; + d_cache[d_idx++] = bxi->d; + } + + asm volatile("" ::: "memory"); +} + +template +static __device__ __forceinline__ void ggml_cuda_mmq_store_tiles_q8_0_rdna35( + int * __restrict__ x_tile, + const int (&qs_cache)[2 * ggml_cuda_mmq_get_I(type, J, fallback)/ + (ggml_cuda_mmq_get_nthreads(type, J, fallback)/ggml_cuda_get_physical_warp_size())], + const float (&d_cache)[4]) { + constexpr int warp_size = ggml_cuda_get_physical_warp_size(); + constexpr int nthreads = ggml_cuda_mmq_get_nthreads(type, J, fallback); + constexpr int nwarps = nthreads / warp_size; + constexpr int I = ggml_cuda_mmq_get_I(type, J, fallback); + constexpr int sram_stride = ggml_cuda_mmq_get_sram_stride(type, J, fallback); + static_assert(warp_size == 32 && nthreads == 128 && I == 64, "unexpected RDNA3.5 Q8_0 MMQ configuration"); + + int * x_qs = x_tile; + float * x_df = (float *) (x_tile + 2*MMQ_TILE_NE_K); + const int txi = threadIdx.x; + +#pragma unroll + for (int i0 = 0; i0 < I; i0 += nwarps) { + const int i = i0 + threadIdx.y; + x_qs[i*sram_stride + 0 + txi] = qs_cache[2*(i0/nwarps) + 0]; + x_qs[i*sram_stride + MMQ_TILE_NE_K + txi] = qs_cache[2*(i0/nwarps) + 1]; + } + + constexpr int blocks_per_tile_x_row = 2*MMQ_TILE_NE_K / QI8_0; + constexpr int scale_rows_per_warp = warp_size / blocks_per_tile_x_row; + const int kbxd = threadIdx.x % blocks_per_tile_x_row; + int d_idx = 0; +#pragma unroll + for (int i0 = 0; i0 < I; i0 += nwarps * scale_rows_per_warp) { + const int i = i0 + threadIdx.y * scale_rows_per_warp + threadIdx.x / blocks_per_tile_x_row; + x_df[i*sram_stride + kbxd] = d_cache[d_idx++]; + } +} +#endif // defined(GGML_USE_HIP) && defined(RDNA3_5) + // --------------------------------------------------------------------------------------------- template static __device__ __forceinline__ void ggml_cuda_mmq_load_tiles_q2_K( @@ -845,6 +1012,166 @@ static __device__ __forceinline__ void ggml_cuda_mmq_store_tiles_q4_K_rdna35( const uint8_t * m8 = (const uint8_t *) &m32; const half2 dm = dm_cache * make_half2(1.0f, -1.0f); +#pragma unroll + for (int l = 0; l < int(sizeof(int)); ++l) { + x_dm[i*sram_stride + sizeof(int)*ksc + l] = dm*make_half2(sc8[l], m8[l]); + } +} + +template +static __device__ __forceinline__ void ggml_cuda_mmq_load_tiles_q5_K_rdna35( + const char * __restrict__ x, int * __restrict__ x_tile, const int kbx0, const int i_max, const int stride) { + constexpr int warp_size = ggml_cuda_get_physical_warp_size(); + constexpr int nthreads = ggml_cuda_mmq_get_nthreads(type, J, fallback); + constexpr int I = ggml_cuda_mmq_get_I(type, J, fallback); + constexpr int sram_stride = ggml_cuda_mmq_get_sram_stride(type, J, fallback); + static_assert(warp_size == 32 && nthreads == 128 && I == 64, "unexpected RDNA3.5 Q5_K MMQ configuration"); + + int * x_qs = (int *) x_tile; + half2 * x_dm = (half2 *) (x_qs + 2*MMQ_TILE_NE_K); + + const int linear_tid = threadIdx.y*warp_size + threadIdx.x; + int i = linear_tid/2; + if (fallback) { + i = min(i, i_max); + } + + const block_q5_K * bxi = (const block_q5_K *) x + kbx0 + i*stride; + const int * scales = (const int *) bxi->scales; + const int ksc = linear_tid % 2; + const int sc32 = unpack_scales_q45_K(scales, ksc); + const int m32 = unpack_scales_q45_K(scales, ksc + 2); + const uint8_t * sc8 = (const uint8_t *) &sc32; + const uint8_t * m8 = (const uint8_t *) &m32; + const half2 dm = bxi->dm * make_half2(1.0f, -1.0f); + +#pragma unroll + for (int l = 0; l < int(sizeof(int)); ++l) { + x_dm[i*sram_stride + sizeof(int)*ksc + l] = dm*make_half2(sc8[l], m8[l]); + } + + const int txi = threadIdx.x; + const int kqs = 16*(txi/8) + txi%8; + const int qh_shift0 = 2*(txi/8); + +#pragma unroll + for (int i0 = 0; i0 < I; i0 += nthreads/warp_size) { + int row = i0 + threadIdx.y; + if (fallback) { + row = min(row, i_max); + } + + const block_q5_K * bxq = (const block_q5_K *) x + kbx0 + row*stride; + const int qs = ((const int *) bxq->qs)[txi]; + const int qh = ((const int *) bxq->qh)[txi % (QI5_K/4)]; + int * row_qs = x_qs + row*sram_stride; + row_qs[kqs] = (qs & 0x0F0F0F0F) | (((qh >> (qh_shift0 + 0)) << 4) & 0x10101010); + row_qs[kqs+8] = ((qs >> 4) & 0x0F0F0F0F) | (((qh >> (qh_shift0 + 1)) << 4) & 0x10101010); + } +} + +template +static __device__ __forceinline__ void ggml_cuda_mmq_prefetch_tiles_q5_K_rdna35( + const char * __restrict__ x, const int kbx0, const int i_max, const int stride, + int (&qs_cache)[ggml_cuda_mmq_get_I(type, J, fallback)/ + (ggml_cuda_mmq_get_nthreads(type, J, fallback)/ggml_cuda_get_physical_warp_size())], + int (&qh_cache)[ggml_cuda_mmq_get_I(type, J, fallback)*(QI5_K/4)/ + ggml_cuda_mmq_get_nthreads(type, J, fallback)], + int (&scales_cache)[3], half2 & dm_cache) { + constexpr int warp_size = ggml_cuda_get_physical_warp_size(); + constexpr int nthreads = ggml_cuda_mmq_get_nthreads(type, J, fallback); + constexpr int nwarps = nthreads / warp_size; + constexpr int I = ggml_cuda_mmq_get_I(type, J, fallback); + constexpr int qh_words_per_row = QI5_K/4; + constexpr int qh_cache_size = I*qh_words_per_row/nthreads; + constexpr int rows_per_warp = I/nwarps; + static_assert(warp_size == 32 && nthreads == 128 && I == 64, "unexpected RDNA3.5 Q5_K MMQ configuration"); + static_assert(qh_cache_size*warp_size == rows_per_warp*qh_words_per_row, + "Q5_K high bits must be distributed evenly across the warp"); + +#pragma unroll + for (int i0 = 0; i0 < I; i0 += nwarps) { + int i = i0 + threadIdx.y; + if constexpr (fallback) { + i = min(i, i_max); + } + + const block_q5_K * bxi = (const block_q5_K *) x + kbx0 + i*stride; + qs_cache[i0/nwarps] = ((const int *) bxi->qs)[threadIdx.x]; + } + +#pragma unroll + for (int l = 0; l < qh_cache_size; ++l) { + const int qh_linear = l*warp_size + threadIdx.x; + int i = (qh_linear/qh_words_per_row)*nwarps + threadIdx.y; + if constexpr (fallback) { + i = min(i, i_max); + } + + const block_q5_K * bxi = (const block_q5_K *) x + kbx0 + i*stride; + qh_cache[l] = ((const int *) bxi->qh)[qh_linear % qh_words_per_row]; + } + + int i = (threadIdx.y*warp_size + threadIdx.x)/2; + if constexpr (fallback) { + i = min(i, i_max); + } + + const block_q5_K * bxi = (const block_q5_K *) x + kbx0 + i*stride; +#pragma unroll + for (int l = 0; l < 3; ++l) { + scales_cache[l] = ((const int *) bxi->scales)[l]; + } + dm_cache = bxi->dm; + + asm volatile("" ::: "memory"); +} + +template +static __device__ __forceinline__ void ggml_cuda_mmq_store_tiles_q5_K_rdna35( + int * __restrict__ x_tile, + const int (&qs_cache)[ggml_cuda_mmq_get_I(type, J, fallback)/ + (ggml_cuda_mmq_get_nthreads(type, J, fallback)/ggml_cuda_get_physical_warp_size())], + const int (&qh_cache)[ggml_cuda_mmq_get_I(type, J, fallback)*(QI5_K/4)/ + ggml_cuda_mmq_get_nthreads(type, J, fallback)], + const int (&scales_cache)[3], const half2 dm_cache) { + constexpr int warp_size = ggml_cuda_get_physical_warp_size(); + constexpr int nthreads = ggml_cuda_mmq_get_nthreads(type, J, fallback); + constexpr int nwarps = nthreads / warp_size; + constexpr int I = ggml_cuda_mmq_get_I(type, J, fallback); + constexpr int sram_stride = ggml_cuda_mmq_get_sram_stride(type, J, fallback); + constexpr int qh_words_per_row = QI5_K/4; + constexpr int qh_rows_per_slot = warp_size/qh_words_per_row; + static_assert(warp_size == 32 && nthreads == 128 && I == 64, "unexpected RDNA3.5 Q5_K MMQ configuration"); + + int * x_qs = x_tile; + const int txi = threadIdx.x; + const int kqs = 16*(txi/8) + txi%8; + const int qh_shift0 = 2*(txi/8); + +#pragma unroll + for (int i0 = 0; i0 < I; i0 += nwarps) { + const int row_in_warp = i0/nwarps; + const int qh_slot = row_in_warp/qh_rows_per_slot; + const int qh_src_lane = (row_in_warp % qh_rows_per_slot)*qh_words_per_row + txi%qh_words_per_row; + const int qs = qs_cache[row_in_warp]; + const int qh = __shfl_sync(0xFFFFFFFF, qh_cache[qh_slot], qh_src_lane, warp_size); + const int i = i0 + threadIdx.y; + int * row_qs = x_qs + i*sram_stride; + row_qs[kqs] = (qs & 0x0F0F0F0F) | (((qh >> (qh_shift0 + 0)) << 4) & 0x10101010); + row_qs[kqs+8] = ((qs >> 4) & 0x0F0F0F0F) | (((qh >> (qh_shift0 + 1)) << 4) & 0x10101010); + } + + half2 * x_dm = (half2 *) (x_qs + 2*MMQ_TILE_NE_K); + const int linear_tid = threadIdx.y*warp_size + threadIdx.x; + const int i = linear_tid/2; + const int ksc = linear_tid%2; + const int sc32 = unpack_scales_q45_K(scales_cache, ksc); + const int m32 = unpack_scales_q45_K(scales_cache, ksc + 2); + const uint8_t * sc8 = (const uint8_t *) &sc32; + const uint8_t * m8 = (const uint8_t *) &m32; + const half2 dm = dm_cache * make_half2(1.0f, -1.0f); + #pragma unroll for (int l = 0; l < int(sizeof(int)); ++l) { x_dm[i*sram_stride + sizeof(int)*ksc + l] = dm*make_half2(sc8[l], m8[l]); @@ -970,6 +1297,11 @@ template static __device__ __forceinline_ template static __device__ __forceinline__ void ggml_cuda_mmq_load_tiles_q5_K( const char * __restrict__ x, int * __restrict__ x_tile, const int kbx0, const int i_max, const int stride) { +#if defined(RDNA3_5) + ggml_cuda_mmq_load_tiles_q5_K_rdna35(x, x_tile, kbx0, i_max, stride); + return; +#endif + constexpr int warp_size = ggml_cuda_get_physical_warp_size(); constexpr int nwarps = ggml_cuda_mmq_get_nthreads(type, J, fallback) / warp_size; constexpr int I = ggml_cuda_mmq_get_I(type, J, fallback); diff --git a/ggml/src/ggml-cuda/mmq-vec-dot.cuh b/ggml/src/ggml-cuda/mmq-vec-dot.cuh index 9f55335ef8e..90aa3c15880 100644 --- a/ggml/src/ggml-cuda/mmq-vec-dot.cuh +++ b/ggml/src/ggml-cuda/mmq-vec-dot.cuh @@ -975,7 +975,7 @@ static __device__ __forceinline__ void ggml_cuda_mmq_mma_q4_K_rdna35_high( template static __device__ __forceinline__ void ggml_cuda_mmq_vec_dot_q4_K_q8_1_mma_rdna35( const int * __restrict__ x, const int * __restrict__ y, float * __restrict__ sum, const int k00) { - if constexpr (type == GGML_TYPE_Q4_K && J == 128) { + if constexpr ((type == GGML_TYPE_Q4_K || type == GGML_TYPE_Q5_K) && J == 128) { constexpr data_layout input_layout = get_input_data_layout(); typedef tile<16, 8, int, input_layout> tile_A; typedef tile<16, 8, int, input_layout> tile_B; @@ -986,7 +986,7 @@ static __device__ __forceinline__ void ggml_cuda_mmq_vec_dot_q4_K_q8_1_mma_rdna3 constexpr int rows_per_warp = ggml_cuda_mmq_get_rows_per_warp(type, J, fallback); constexpr int ntx = rows_per_warp/tile_C::I; constexpr int ntiles = J/tile_C::J; - static_assert(I == 64 && ntx == 1, "unexpected RDNA3.5 Q4_K J128 configuration"); + static_assert(I == 64 && ntx == 1, "unexpected RDNA3.5 Q4_K/Q5_K ntx=1 configuration"); const int * x_qs = (const int *) x; const half2 * x_dm = (const half2 *) x_qs + 2*MMQ_TILE_NE_K; @@ -1070,6 +1070,154 @@ static __device__ __forceinline__ void ggml_cuda_mmq_vec_dot_q4_K_q8_1_mma_rdna3 } ggml_cuda_mmq_vec_dot_q8_1_q8_1_mma(x, y, sum, k00); } + +template +static __device__ __forceinline__ void ggml_cuda_mmq_vec_dot_q4_0_q8_1_mma_rdna35( + const int * __restrict__ x, const int * __restrict__ y, float * __restrict__ sum, const int k00) { + if constexpr (type == GGML_TYPE_Q4_0 && J == 128) { + constexpr data_layout input_layout = get_input_data_layout(); + typedef tile<16, 8, int, input_layout> tile_A; + typedef tile<16, 8, int, input_layout> tile_B; + typedef tile<16, 16, int, DATA_LAYOUT_J_MAJOR> tile_C; + + constexpr int I = ggml_cuda_mmq_get_I(type, J, fallback); + constexpr int sram_stride = ggml_cuda_mmq_get_sram_stride(type, J, fallback); + constexpr int rows_per_warp = ggml_cuda_mmq_get_rows_per_warp(type, J, fallback); + constexpr int ntx = rows_per_warp/tile_C::I; + constexpr int ntiles = J/tile_C::J; + static_assert(I == 64 && ntx == 1, "unexpected RDNA3.5 Q4_0 J128 configuration"); + + const int * x_qs = (const int *) x; + const float * x_df = (const float *) x_qs + 2*MMQ_TILE_NE_K; + const int * y_qs = (const int *) y + 4; + const half2 * y_ds = (const half2 *) y; + + const int i0 = threadIdx.y*rows_per_warp; + + for (int k01 = 0; k01 < MMQ_TILE_NE_K; k01 += QI8_0) { + const int k0 = k00 + k01; + + tile_A A; + load_ldmatrix(A, x_qs + i0*sram_stride + k0, sram_stride); + + tile_B B[ntiles]; + tile_C C[ntiles]; +#pragma unroll + for (int jb = 0; jb < ntiles; ++jb) { + load_ldmatrix(B[jb], y_qs + jb*tile_C::J*MMQ_TILE_Y_K + k01, MMQ_TILE_Y_K); + ggml_cuda_mmq_mma_q4_K_rdna35_low(C[jb], A, B[jb]); + } + + __builtin_amdgcn_sched_barrier(0); + + float dA[tile_C::ne]; + float dB[ntiles]; +#pragma unroll + for (int l = 0; l < tile_C::ne; ++l) { + const int i = i0 + tile_C::get_i(l); + dA[l] = x_df[i*sram_stride + k0/QI8_0]; + } +#pragma unroll + for (int jb = 0; jb < ntiles; ++jb) { + const int j = jb*tile_C::J + tile_C::get_j(0); + dB[jb] = __low2float(y_ds[j*MMQ_TILE_Y_K + k01/QI8_1]); + } + + __builtin_amdgcn_sched_barrier(0); + +#pragma unroll + for (int jb = 0; jb < ntiles; ++jb) { + ggml_cuda_mmq_mma_q4_K_rdna35_high(C[jb], A, B[jb]); + } + + __builtin_amdgcn_sched_barrier(0); + +#pragma unroll + for (int jb = 0; jb < ntiles; ++jb) { +#pragma unroll + for (int l = 0; l < tile_C::ne; ++l) { + sum[jb*tile_C::ne + l] += C[jb].x[l]*dA[l]*dB[jb]; + } + } + } + return; + } + ggml_cuda_mmq_vec_dot_q8_0_q8_1_mma(x, y, sum, k00); +} + +template +static __device__ __forceinline__ void ggml_cuda_mmq_vec_dot_q8_0_q8_1_mma_rdna35( + const int * __restrict__ x, const int * __restrict__ y, float * __restrict__ sum, const int k00) { + if constexpr (type == GGML_TYPE_Q8_0 && J == 128) { + constexpr data_layout input_layout = get_input_data_layout(); + typedef tile<16, 8, int, input_layout> tile_A; + typedef tile<16, 8, int, input_layout> tile_B; + typedef tile<16, 16, int, DATA_LAYOUT_J_MAJOR> tile_C; + + constexpr int I = ggml_cuda_mmq_get_I(type, J, fallback); + constexpr int sram_stride = ggml_cuda_mmq_get_sram_stride(type, J, fallback); + constexpr int rows_per_warp = ggml_cuda_mmq_get_rows_per_warp(type, J, fallback); + constexpr int ntx = rows_per_warp/tile_C::I; + constexpr int ntiles = J/tile_C::J; + static_assert(I == 64 && ntx == 1, "unexpected RDNA3.5 Q8_0 J128 configuration"); + + const int * x_qs = (const int *) x; + const float * x_df = (const float *) x_qs + 2*MMQ_TILE_NE_K; + const int * y_qs = (const int *) y + 4; + const float * y_df = (const float *) y; + + const int i0 = threadIdx.y*rows_per_warp; + + for (int k01 = 0; k01 < MMQ_TILE_NE_K; k01 += QI8_0) { + const int k0 = k00 + k01; + + tile_A A; + load_ldmatrix(A, x_qs + i0*sram_stride + k0, sram_stride); + + tile_B B[ntiles]; + tile_C C[ntiles]; +#pragma unroll + for (int jb = 0; jb < ntiles; ++jb) { + load_ldmatrix(B[jb], y_qs + jb*tile_C::J*MMQ_TILE_Y_K + k01, MMQ_TILE_Y_K); + ggml_cuda_mmq_mma_q4_K_rdna35_low(C[jb], A, B[jb]); + } + + __builtin_amdgcn_sched_barrier(0); + + float dA[tile_C::ne]; + float dB[ntiles]; +#pragma unroll + for (int l = 0; l < tile_C::ne; ++l) { + const int i = i0 + tile_C::get_i(l); + dA[l] = x_df[i*sram_stride + k0/QI8_0]; + } +#pragma unroll + for (int jb = 0; jb < ntiles; ++jb) { + const int j = jb*tile_C::J + tile_C::get_j(0); + dB[jb] = y_df[j*MMQ_TILE_Y_K + k01/QI8_1]; + } + + __builtin_amdgcn_sched_barrier(0); + +#pragma unroll + for (int jb = 0; jb < ntiles; ++jb) { + ggml_cuda_mmq_mma_q4_K_rdna35_high(C[jb], A, B[jb]); + } + + __builtin_amdgcn_sched_barrier(0); + +#pragma unroll + for (int jb = 0; jb < ntiles; ++jb) { +#pragma unroll + for (int l = 0; l < tile_C::ne; ++l) { + sum[jb*tile_C::ne + l] += C[jb].x[l]*dA[l]*dB[jb]; + } + } + } + return; + } + ggml_cuda_mmq_vec_dot_q8_0_q8_1_mma(x, y, sum, k00); +} #endif template static __device__ __forceinline__ void ggml_cuda_mmq_vec_dot_q5_K_q8_1_dp4a( diff --git a/ggml/src/ggml-cuda/mmq.cuh b/ggml/src/ggml-cuda/mmq.cuh index 42c0ef5542f..107009f07bc 100644 --- a/ggml/src/ggml-cuda/mmq.cuh +++ b/ggml/src/ggml-cuda/mmq.cuh @@ -181,7 +181,14 @@ struct ggml_cuda_mmq_config { constexpr __device__ int rows_per_warp() const { #if defined(AMD_MFMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE) #if defined(RDNA3_5) - if (type == GGML_TYPE_Q4_K && J == 128) { + if ((type == GGML_TYPE_Q4_K || type == GGML_TYPE_Q4_0 || type == GGML_TYPE_Q8_0 || + type == GGML_TYPE_Q5_K) && J == 128) { + return 16; + } + // Block quants only have a batched WMMA vec_dot at J=128; at the other two widths + // that would take ntx=2 (J=64 and J=96, the only J>=64 multiples of 32) they fall + // back to the generic schedule, which is several times slower there. Keep ntx=1. + if ((type == GGML_TYPE_Q4_0 || type == GGML_TYPE_Q8_0) && (J == 64 || J == 96)) { return 16; } return J >= 64 && J % 32 == 0 ? 32 : 16; @@ -814,7 +821,11 @@ static constexpr __device__ ggml_cuda_mmq_util_funcs ggml_cuda_mmq_get_util_func return ggml_cuda_mmq_util_funcs( -1, ggml_cuda_mmq_load_tiles_q4_0, +#if defined(RDNA3_5) + ggml_cuda_mmq_vec_dot_q4_0_q8_1_mma_rdna35, +#else ggml_cuda_mmq_vec_dot_q8_0_q8_1_mma, +#endif ggml_cuda_mmq_write_back_mma); case GGML_TYPE_Q4_1: return ggml_cuda_mmq_util_funcs( @@ -838,7 +849,11 @@ static constexpr __device__ ggml_cuda_mmq_util_funcs ggml_cuda_mmq_get_util_func return ggml_cuda_mmq_util_funcs( -1, ggml_cuda_mmq_load_tiles_q8_0, +#if defined(RDNA3_5) + ggml_cuda_mmq_vec_dot_q8_0_q8_1_mma_rdna35, +#else ggml_cuda_mmq_vec_dot_q8_0_q8_1_mma, +#endif ggml_cuda_mmq_write_back_mma); // --------------------------------------------------------------------------------------------- case GGML_TYPE_Q2_K: @@ -867,7 +882,11 @@ static constexpr __device__ ggml_cuda_mmq_util_funcs ggml_cuda_mmq_get_util_func return ggml_cuda_mmq_util_funcs( -1, ggml_cuda_mmq_load_tiles_q5_K, +#if defined(RDNA3_5) + ggml_cuda_mmq_vec_dot_q4_K_q8_1_mma_rdna35, +#else ggml_cuda_mmq_vec_dot_q8_1_q8_1_mma, +#endif ggml_cuda_mmq_write_back_mma); case GGML_TYPE_Q6_K: return ggml_cuda_mmq_util_funcs( @@ -1096,6 +1115,150 @@ static __device__ __forceinline__ void mul_mat_q_process_tile( } __syncthreads(); } + } else if constexpr (type == GGML_TYPE_Q5_K && (J == 64 || J == 128)) { + constexpr int qs_cache_size = I/nwarps; + constexpr int qh_cache_size = I*(QI5_K/4)/(nwarps*warp_size); + + __syncthreads(); + load_tiles(x, tile_x, offset_x + kb0_start, tile_x_max_i, stride_row_x); + __syncthreads(); + + for (int kb0 = kb0_start; kb0 < kb0_stop; kb0 += blocks_per_iter) { + const int yk = kb0*qk/ne_block; + const int * by0 = y + ncols_y*yk*sz; + const int * by1 = y + ncols_y*(yk + 1)*sz; + +#pragma unroll + for (int l0 = 0; l0 < J*MMQ_TILE_Y_K; l0 += nwarps*warp_size) { + const int l = l0 + threadIdx.y*warp_size + threadIdx.x; + tile_y[l] = by0[l]; + } + __syncthreads(); + vec_dot(tile_x, tile_y, sum, 0); + + __syncthreads(); +#pragma unroll + for (int l0 = 0; l0 < J*MMQ_TILE_Y_K; l0 += nwarps*warp_size) { + const int l = l0 + threadIdx.y*warp_size + threadIdx.x; + tile_y[l] = by1[l]; + } + __syncthreads(); + + int qs_cache[qs_cache_size]; + int qh_cache[qh_cache_size]; + int scales_cache[3]; + half2 dm_cache; + const int kb0_next = kb0 + blocks_per_iter; + const bool have_next = kb0_next < kb0_stop; + if (have_next) { + ggml_cuda_mmq_prefetch_tiles_q5_K_rdna35( + x, offset_x + kb0_next, tile_x_max_i, stride_row_x, + qs_cache, qh_cache, scales_cache, dm_cache); + } + + vec_dot(tile_x, tile_y, sum, MMQ_TILE_NE_K); + __syncthreads(); + + if (have_next) { + ggml_cuda_mmq_store_tiles_q5_K_rdna35( + tile_x, qs_cache, qh_cache, scales_cache, dm_cache); + } + __syncthreads(); + } + } else if constexpr (type == GGML_TYPE_Q4_0 && (J == 64 || J == 128)) { + constexpr int qs_cache_size = I/nwarps; + constexpr int d_cache_size = 4; + + __syncthreads(); + load_tiles(x, tile_x, offset_x + kb0_start, tile_x_max_i, stride_row_x); + __syncthreads(); + + for (int kb0 = kb0_start; kb0 < kb0_stop; kb0 += blocks_per_iter) { + const int yk = kb0*qk/ne_block; + const int * by0 = y + ncols_y*yk*sz; + const int * by1 = y + ncols_y*(yk + 1)*sz; + +#pragma unroll + for (int l0 = 0; l0 < J*MMQ_TILE_Y_K; l0 += nwarps*warp_size) { + const int l = l0 + threadIdx.y*warp_size + threadIdx.x; + tile_y[l] = by0[l]; + } + __syncthreads(); + vec_dot(tile_x, tile_y, sum, 0); + + __syncthreads(); +#pragma unroll + for (int l0 = 0; l0 < J*MMQ_TILE_Y_K; l0 += nwarps*warp_size) { + const int l = l0 + threadIdx.y*warp_size + threadIdx.x; + tile_y[l] = by1[l]; + } + __syncthreads(); + + int qs_cache[qs_cache_size]; + float d_cache[d_cache_size]; + const int kb0_next = kb0 + blocks_per_iter; + const bool have_next = kb0_next < kb0_stop; + if (have_next) { + ggml_cuda_mmq_prefetch_tiles_q4_0_rdna35( + x, offset_x + kb0_next, tile_x_max_i, stride_row_x, qs_cache, d_cache); + } + + vec_dot(tile_x, tile_y, sum, MMQ_TILE_NE_K); + __syncthreads(); + + if (have_next) { + ggml_cuda_mmq_store_tiles_q4_0_rdna35( + tile_x, qs_cache, d_cache); + } + __syncthreads(); + } + } else if constexpr (type == GGML_TYPE_Q8_0 && J == 128) { + constexpr int qs_cache_size = 2*(I/nwarps); + constexpr int d_cache_size = 4; + + __syncthreads(); + load_tiles(x, tile_x, offset_x + kb0_start, tile_x_max_i, stride_row_x); + __syncthreads(); + + for (int kb0 = kb0_start; kb0 < kb0_stop; kb0 += blocks_per_iter) { + const int yk = kb0*qk/ne_block; + const int * by0 = y + ncols_y*yk*sz; + const int * by1 = y + ncols_y*(yk + 1)*sz; + +#pragma unroll + for (int l0 = 0; l0 < J*MMQ_TILE_Y_K; l0 += nwarps*warp_size) { + const int l = l0 + threadIdx.y*warp_size + threadIdx.x; + tile_y[l] = by0[l]; + } + __syncthreads(); + vec_dot(tile_x, tile_y, sum, 0); + + __syncthreads(); +#pragma unroll + for (int l0 = 0; l0 < J*MMQ_TILE_Y_K; l0 += nwarps*warp_size) { + const int l = l0 + threadIdx.y*warp_size + threadIdx.x; + tile_y[l] = by1[l]; + } + __syncthreads(); + + int qs_cache[qs_cache_size]; + float d_cache[d_cache_size]; + const int kb0_next = kb0 + blocks_per_iter; + const bool have_next = kb0_next < kb0_stop; + if (have_next) { + ggml_cuda_mmq_prefetch_tiles_q8_0_rdna35( + x, offset_x + kb0_next, tile_x_max_i, stride_row_x, qs_cache, d_cache); + } + + vec_dot(tile_x, tile_y, sum, MMQ_TILE_NE_K); + __syncthreads(); + + if (have_next) { + ggml_cuda_mmq_store_tiles_q8_0_rdna35( + tile_x, qs_cache, d_cache); + } + __syncthreads(); + } } else { #endif #if defined(RDNA3_5) @@ -1645,12 +1808,11 @@ struct mmq_args { #if defined(GGML_USE_HIP) // RDNA3.5 dual-WG (mmq_x=64, nbytes <= smpbo/2) helps K-quants with large per-tile LDS -// (Q6_K WMMA tuning). Block quants (Q5_0, Q8_0, Q4_0) and Q4_K prefill are faster at mmq_x=128 ntx=1. +// (Q6_K WMMA tuning). Block quants (Q5_0, Q8_0, Q4_0) and Q4_K/Q5_K prefill are faster at mmq_x=128 ntx=1. static bool mmq_rdna35_dual_wg_eligible(const ggml_type type) { switch (type) { case GGML_TYPE_Q2_K: case GGML_TYPE_Q3_K: - case GGML_TYPE_Q5_K: case GGML_TYPE_Q6_K: return true; default: @@ -1660,10 +1822,14 @@ static bool mmq_rdna35_dual_wg_eligible(const ggml_type type) { // J values that measure faster on RDNA3.5 than the width the occupancy rule picks. // -// J=64 is 6-7x slower than the neighbouring widths for q8_0; the other types stay within -// 1.05x of their best there and keep it. For MoE the compacted per-expert grid already supplies -// the parallelism that column tiling supplies for a dense GEMM, so anything above a 32-wide tile -// only adds padding - q8_0 excepted, whose narrow tiles are slow. +// Dense shapes now keep the width the occupancy rule picks: q8_0 used to be widened from +// J=64 to J=96 because J=64 was several times slower, but that was the generic ntx=2 +// schedule, which rows_per_warp() no longer selects for block quants. Without the widening +// q8_0 avoids padding 64 valid columns into a 96-wide tile and is 27% faster there. +// +// For MoE the compacted per-expert grid already supplies the parallelism that column tiling +// supplies for a dense GEMM, so anything above a 32-wide tile only adds padding - q8_0 +// excepted, whose narrow tiles are slow. static int mmq_rdna35_tuned_J(const ggml_type type, const bool moe, const int J_occupancy) { if (moe) { if (type == GGML_TYPE_Q8_0) { @@ -1671,9 +1837,6 @@ static int mmq_rdna35_tuned_J(const ggml_type type, const bool moe, const int J_ } return J_occupancy < 32 ? J_occupancy : 32; } - if (J_occupancy == 64 && type == GGML_TYPE_Q8_0) { - return 96; - } return J_occupancy; } #endif // GGML_USE_HIP @@ -1873,10 +2036,7 @@ void mul_mat_q_switch_J(ggml_backend_cuda_context & ctx, const mmq_args & args, } #if defined(GGML_USE_HIP) - const bool dual_wg_type = - type == GGML_TYPE_Q2_K || type == GGML_TYPE_Q3_K || - type == GGML_TYPE_Q5_K || type == GGML_TYPE_Q6_K; - if (GGML_CUDA_CC_IS_RDNA3_5(cc) && J_best > 0 && dual_wg_type) { + if (GGML_CUDA_CC_IS_RDNA3_5(cc) && J_best > 0 && mmq_rdna35_dual_wg_eligible(type)) { const size_t lds_dual_wg = smpbo/2; if (mmq_get_nbytes_shared(ggml_cuda_mmq_get_config(type, J_best, fallback, cc), cc) > lds_dual_wg) { int J_dual = 0; @@ -1918,20 +2078,20 @@ void mul_mat_q_switch_J(ggml_backend_cuda_context & ctx, const mmq_args & args, } } - if constexpr (type == GGML_TYPE_Q4_K) { - constexpr int q4_k_J_default = 128; - constexpr int q4_k_J_small = 64; - constexpr int q4_k_m_small_max = 1024; - constexpr int q4_k_ncols_pipeline = 128; - const bool use_q4_k_pipeline = - args.expert_bounds == nullptr && args.ncols_max == q4_k_ncols_pipeline; - if (use_q4_k_pipeline) { - const bool use_small = args.nrows_x <= q4_k_m_small_max; - const int q4_k_J = use_small ? q4_k_J_small : q4_k_J_default; - const ggml_cuda_mmq_config config = ggml_cuda_mmq_get_config(type, q4_k_J, fallback, cc); + if constexpr (type == GGML_TYPE_Q4_K || type == GGML_TYPE_Q5_K) { + constexpr int qk_J_default = 128; + constexpr int qk_J_small = 64; + constexpr int qk_m_small_max = 1024; + constexpr int qk_ncols_pipeline = 128; + const bool use_qk_pipeline = + args.expert_bounds == nullptr && args.ncols_max == qk_ncols_pipeline; + if (use_qk_pipeline) { + const bool use_small = args.nrows_x <= qk_m_small_max; + const int qk_J = use_small ? qk_J_small : qk_J_default; + const ggml_cuda_mmq_config config = ggml_cuda_mmq_get_config(type, qk_J, fallback, cc); if (GGML_CUDA_CC_IS_RDNA3_5(cc) && config.type != GGML_TYPE_COUNT && mmq_get_nbytes_shared(config, cc) <= smpbo) { - J_best = q4_k_J; + J_best = qk_J; } } } diff --git a/tests/test-backend-ops.cpp b/tests/test-backend-ops.cpp index a62d781f677..58a917b1df3 100644 --- a/tests/test-backend-ops.cpp +++ b/tests/test-backend-ops.cpp @@ -8209,6 +8209,12 @@ static void add_rdna35_mmq_cases(std::vector> & test_ test_cases.emplace_back(new test_mul_mat(GGML_TYPE_Q4_K, GGML_TYPE_F32, 4096, 16, 4096, {1, 1}, {1, 1})); test_cases.emplace_back(new test_mul_mat(GGML_TYPE_Q4_K, GGML_TYPE_F32, 4096, 1024, 12288, {1, 1}, {1, 1})); test_cases.emplace_back(new test_mul_mat(GGML_TYPE_Q6_K, GGML_TYPE_F32, 4096, 128, 12288, {1, 1}, {1, 1})); + test_cases.emplace_back(new test_mul_mat(GGML_TYPE_Q6_K, GGML_TYPE_F32, 2560, 128, 9216, {1, 1}, {1, 1})); + test_cases.emplace_back(new test_mul_mat(GGML_TYPE_Q4_K, GGML_TYPE_F32, 8192, 128, 2560, {1, 1}, {1, 1})); + test_cases.emplace_back(new test_mul_mat(GGML_TYPE_Q5_K, GGML_TYPE_F32, 8192, 128, 2560, {1, 1}, {1, 1})); + test_cases.emplace_back(new test_mul_mat(GGML_TYPE_Q5_K, GGML_TYPE_F32, 8192, 512, 2560, {1, 1}, {1, 1})); + test_cases.emplace_back(new test_mul_mat(GGML_TYPE_Q5_K, GGML_TYPE_F32, 2560, 128, 4096, {1, 1}, {1, 1})); + test_cases.emplace_back(new test_mul_mat(GGML_TYPE_Q5_K, GGML_TYPE_F32, 2560, 512, 4096, {1, 1}, {1, 1})); test_cases.emplace_back(new test_mul_mat(GGML_TYPE_Q8_0, GGML_TYPE_F32, 32, 128, 4096, {1, 1}, {1, 1})); test_cases.emplace_back(new test_mul_mat(GGML_TYPE_Q4_0, GGML_TYPE_F32, 32, 128, 4096, {1, 1}, {1, 1})); for (int64_t gate_n : {8, 16, 64}) {