Skip to content

q4k mmq optimizations #80

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Draft
liangliangchang wants to merge 3 commits into
base: gfx11
Choose a base branch
from
+154 −0
Draft

q4k mmq optimizations #80

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
2815b3e
HIP: optimize RDNA3.5 Q4_K MMQ pipeline
liangliangchang Jul 31, 2026
a06b975
HIP: fix RDNA3.5 Q4_K pipeline integration
liangliangchang Jul 31, 2026
c83bdac
HIP: tune Q4_K wave mapping by tile width
liangliangchang Jul 31, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
78 changes: 78 additions & 0 deletions ggml/src/ggml-cuda/mmq-load-tiles.cuh
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -691,6 +691,84 @@ static __device__ __forceinline__ void ggml_cuda_mmq_load_tiles_q4_K_rdna35(
row_qs[kqs+8] = (qs >> 4) & 0x0F0F0F0F;
}
}

template <ggml_type type, int J, bool fallback>
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 <ggml_type type, int J, bool fallback>
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 <ggml_type type, int J, bool fallback> static __device__ __forceinline__ void ggml_cuda_mmq_load_tiles_q4_K(
Expand Down
76 changes: 76 additions & 0 deletions ggml/src/ggml-cuda/mmq.cuh
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down Expand Up @@ -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<type, J, fallback>(
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<type, J, fallback>(
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;
Expand Down Expand Up @@ -1097,6 +1150,9 @@ static __device__ __forceinline__ void mul_mat_q_process_tile(
mmq_hip_tile_barrier<J>();
}
}
#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);
Expand Down Expand Up @@ -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) {
Expand Down
Loading