diff --git a/ggml/src/ggml-cuda/mmq-load-tiles.cuh b/ggml/src/ggml-cuda/mmq-load-tiles.cuh index 688cf72c5506..3653434092b8 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.cuh b/ggml/src/ggml-cuda/mmq.cuh index fdc1b8ca8505..a728d818148e 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; @@ -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) { + 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,26 @@ 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 char * q4_k_J_env = getenv("GGML_HIP_Q4K_MMQ_X"); + if (q4_k_J_env != nullptr || args.ncols_max == q4_k_ncols_pipeline) { + const bool force_small = + q4_k_J_env != nullptr && q4_k_J_env[0] == '6' && q4_k_J_env[1] == '4' && q4_k_J_env[2] == '\0'; + const bool use_small = + q4_k_J_env != nullptr ? force_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) {