diff --git a/ggml/src/ggml-cuda/mmq-load-tiles.cuh b/ggml/src/ggml-cuda/mmq-load-tiles.cuh index 688cf72c550..3653434092b 100644 --- a/ggml/src/ggml-cuda/mmq-load-tiles.cuh +++ b/ggml/src/ggml-cuda/mmq-load-tiles.cuh @@ -691,6 +691,84 @@ static __device__ __forceinline__ void ggml_cuda_mmq_load_tiles_q4_K_rdna35( row_qs[kqs+8] = (qs >> 4) & 0x0F0F0F0F; } } + +template +static __device__ __forceinline__ void ggml_cuda_mmq_prefetch_tiles_q4_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 (&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); + static_assert(warp_size == 32 && nthreads == 128 && I == 64, "unexpected RDNA3.5 Q4_K MMQ configuration"); + +#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_K * bxi = (const block_q4_K *) x + kbx0 + i*stride; + qs_cache[i0/nwarps] = ((const int *) bxi->qs)[threadIdx.x]; + } + + int i = (threadIdx.y*warp_size + threadIdx.x)/2; + if constexpr (fallback) { + i = min(i, i_max); + } + + const block_q4_K * bxi = (const block_q4_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_q4_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 (&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); + static_assert(warp_size == 32 && nthreads == 128 && I == 64, "unexpected RDNA3.5 Q4_K MMQ configuration"); + + int * x_qs = x_tile; +#pragma unroll + for (int i0 = 0; i0 < I; i0 += nwarps) { + const int i = i0 + threadIdx.y; + const int qs = qs_cache[i0/nwarps]; + int * row_qs = x_qs + i*sram_stride; + const int kqs = 16*(threadIdx.x/8) + threadIdx.x%8; + row_qs[kqs] = qs & 0x0F0F0F0F; + row_qs[kqs+8] = (qs >> 4) & 0x0F0F0F0F; + } + + 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]); + } +} #endif template static __device__ __forceinline__ void ggml_cuda_mmq_load_tiles_q4_K( diff --git a/ggml/src/ggml-cuda/mmq-vec-dot.cuh b/ggml/src/ggml-cuda/mmq-vec-dot.cuh index d573433865f..5cd239acea6 100644 --- a/ggml/src/ggml-cuda/mmq-vec-dot.cuh +++ b/ggml/src/ggml-cuda/mmq-vec-dot.cuh @@ -945,6 +945,135 @@ template static __device__ __forceinline_ } } +#if defined(RDNA3_5) && defined(AMD_WMMA_AVAILABLE) && !defined(AMD_MFMA_AVAILABLE) +static __device__ __forceinline__ void ggml_cuda_mmq_mma_q4_K_rdna35_low( + tile<16, 16, int, DATA_LAYOUT_J_MAJOR> & D, + const tile<16, 8, int, DATA_LAYOUT_I_MAJOR_MIRRORED> & A, + const tile<16, 8, int, DATA_LAYOUT_I_MAJOR_MIRRORED> & B) { + using int32x4_t = __attribute__((__vector_size__(4 * sizeof(int)))) int; + using int32x8_t = __attribute__((__vector_size__(8 * sizeof(int)))) int; + + int32x8_t * acc = (int32x8_t *) D.x; + const int32x4_t * a_vec = (const int32x4_t *) A.x; + const int32x4_t * b_vec = (const int32x4_t *) B.x; + acc[0] = __builtin_amdgcn_wmma_i32_16x16x16_iu8_w32(true, a_vec[0], true, b_vec[0], acc[0], true); +} + +static __device__ __forceinline__ void ggml_cuda_mmq_mma_q4_K_rdna35_high( + tile<16, 16, int, DATA_LAYOUT_J_MAJOR> & D, + const tile<16, 8, int, DATA_LAYOUT_I_MAJOR_MIRRORED> & A, + const tile<16, 8, int, DATA_LAYOUT_I_MAJOR_MIRRORED> & B) { + using int32x4_t = __attribute__((__vector_size__(4 * sizeof(int)))) int; + using int32x8_t = __attribute__((__vector_size__(8 * sizeof(int)))) int; + + int32x8_t * acc = (int32x8_t *) D.x; + const int32x4_t * a_vec = (const int32x4_t *) A.x; + const int32x4_t * b_vec = (const int32x4_t *) B.x; + acc[0] = __builtin_amdgcn_wmma_i32_16x16x16_iu8_w32(true, a_vec[1], true, b_vec[1], acc[0], true); +} +#endif + +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 defined(RDNA3_5) && defined(AMD_WMMA_AVAILABLE) && !defined(AMD_MFMA_AVAILABLE) + if constexpr (type == GGML_TYPE_Q4_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; + 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_K J128 configuration"); + + const int * x_qs = (const int *) x; + const half2 * x_dm = (const half2 *) x_qs + 2*MMQ_TILE_NE_K; + const int * y_qs = (const int *) y + 4; + const half2 * y_dm = (const half2 *) y; + + const int i0 = threadIdx.y*rows_per_warp; + + for (int k01 = 0; k01 < MMQ_TILE_NE_K; k01 += QI8_1) { + 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); + + half2 dmA_half[tile_C::ne]; + half2 dsB_half[ntiles]; +#pragma unroll + for (int l = 0; l < tile_C::ne; ++l) { + const int i = i0 + tile_C::get_i(l); + dmA_half[l] = x_dm[i*sram_stride + k0/QI8_1]; + } +#pragma unroll + for (int jb = 0; jb < ntiles; ++jb) { + const int j = jb*tile_C::J + tile_C::get_j(0); + dsB_half[jb] = y_dm[j*MMQ_TILE_Y_K + k01/QI8_1]; + } + + __builtin_amdgcn_sched_barrier(0); + + float dmA_scale[tile_C::ne]; + float2 dsB[ntiles]; +#pragma unroll + for (int jb = 0; jb < ntiles; ++jb) { + dsB[jb] = __half22float2(dsB_half[jb]); + asm volatile("" : "+v"(dsB[jb].x), "+v"(dsB[jb].y)); + } + +#pragma unroll + for (int l = 0; l < tile_C::ne; l += 2) { + float2 dm0 = __half22float2(dmA_half[l + 0]); + float2 dm1 = __half22float2(dmA_half[l + 1]); + asm volatile("" : "+v"(dm0.x), "+v"(dm0.y), "+v"(dm1.x), "+v"(dm1.y)); + dmA_scale[l + 0] = dm0.x; + dmA_scale[l + 1] = dm1.x; +#pragma unroll + for (int jb = 0; jb < ntiles; ++jb) { + sum[jb*tile_C::ne + l + 0] += dm0.y*dsB[jb].y; + sum[jb*tile_C::ne + l + 1] += dm1.y*dsB[jb].y; + } + } + + __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) { + const int si = jb*tile_C::ne + l; + sum[si] += dmA_scale[l]*dsB[jb].x*C[jb].x[l]; + } + } + } + return; + } +#endif + ggml_cuda_mmq_vec_dot_q8_1_q8_1_mma(x, y, sum, k00); +} + template static __device__ __forceinline__ void ggml_cuda_mmq_vec_dot_q5_K_q8_1_dp4a( const int * __restrict__ x, const int * __restrict__ y, float * __restrict__ sum, const int k00) { constexpr int warp_size = ggml_cuda_get_physical_warp_size(); diff --git a/ggml/src/ggml-cuda/mmq.cuh b/ggml/src/ggml-cuda/mmq.cuh index fdc1b8ca850..5f85f4d5a28 100644 --- a/ggml/src/ggml-cuda/mmq.cuh +++ b/ggml/src/ggml-cuda/mmq.cuh @@ -180,6 +180,9 @@ 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) { + return 16; + } return J >= 64 && J % 32 == 0 ? 32 : 16; #else return 16; @@ -828,7 +831,7 @@ 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_K, - ggml_cuda_mmq_vec_dot_q8_1_q8_1_mma, + ggml_cuda_mmq_vec_dot_q4_K_q8_1_mma_rdna35, ggml_cuda_mmq_write_back_mma); case GGML_TYPE_Q5_K: return ggml_cuda_mmq_util_funcs( @@ -1015,6 +1018,56 @@ static __device__ __forceinline__ void mul_mat_q_process_tile( constexpr int sz = sizeof(block_q8_1_mmq) / sizeof(int); +#if defined(RDNA3_5) && defined(AMD_WMMA_AVAILABLE) && !defined(AMD_MFMA_AVAILABLE) + if constexpr (type == GGML_TYPE_Q4_K && (J == 64 || J == 128)) { + constexpr int qs_cache_size = I/nwarps; + + __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 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_q4_K_rdna35( + x, offset_x + kb0_next, tile_x_max_i, stride_row_x, qs_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_q4_K_rdna35( + tile_x, qs_cache, scales_cache, dm_cache); + } + __syncthreads(); + } + } else { +#endif #if defined(RDNA3_5) constexpr int tile_y_elems = J*MMQ_TILE_Y_K; constexpr int tile_y_load_stride = nwarps*warp_size; @@ -1097,6 +1150,9 @@ static __device__ __forceinline__ void mul_mat_q_process_tile( mmq_hip_tile_barrier(); } } +#if defined(RDNA3_5) && defined(AMD_WMMA_AVAILABLE) && !defined(AMD_MFMA_AVAILABLE) + } +#endif if (fixup) { write_back(sum, ids_dst, tmp_fixup + blockIdx.x*(J*I), y_scale, I, I, J); @@ -1831,6 +1887,24 @@ void mul_mat_q_switch_J(ggml_backend_cuda_context & ctx, const mmq_args & args, J_best = J_tuned; } } + + 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 (GGML_CUDA_CC_IS_RDNA3_5(cc) && + config.type != GGML_TYPE_COUNT && mmq_get_nbytes_shared(config, cc) <= smpbo) { + J_best = q4_k_J; + } + } + } #endif // GGML_USE_HIP switch (J_best) {