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

Open
liangliangchang wants to merge 7 commits into
base: gfx11
Choose a base branch
from
+282 −1
Open

q4k mmq optimizations #80

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
7 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
bec4a25
HIP: preserve expert-aware Q4_K MMQ tiling
liangliangchang Aug 3, 2026
f1cd056
HIP: remove Q4_K MMQ test override
liangliangchang Aug 3, 2026
f602551
HIP: extend RDNA3.5 Q4_K prefetch to J128
liangliangchang Aug 12, 2026
96fc67b
HIP: batch RDNA3.5 Q4_K J128 WMMA phases
liangliangchang Aug 12, 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
129 changes: 129 additions & 0 deletions ggml/src/ggml-cuda/mmq-vec-dot.cuh
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -945,6 +945,135 @@ template <ggml_type type, int J, bool fallback> 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 <ggml_type type, int J, bool fallback>
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<type, J, fallback>(x, y, sum, k00);
}

template <ggml_type type, int J, bool fallback> 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();
Expand Down
76 changes: 75 additions & 1 deletion 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) {

@mgehre-amd mgehre-amd Aug 17, 2026

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Why is this fixed to exactly 128?

return 16;
}
return J >= 64 && J % 32 == 0 ? 32 : 16;
#else
return 16;
Expand Down Expand Up @@ -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<type, J, fallback>,
ggml_cuda_mmq_vec_dot_q8_1_q8_1_mma<type, J, fallback>,
ggml_cuda_mmq_vec_dot_q4_K_q8_1_mma_rdna35<type, J, fallback>,

@mgehre-amd mgehre-amd Aug 17, 2026

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

It looks like this function is called for all AMD GPUs, not just rdna35?

ggml_cuda_mmq_write_back_mma<type, J, fallback>);
case GGML_TYPE_Q5_K:
return ggml_cuda_mmq_util_funcs(
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)

@mgehre-amd mgehre-amd Aug 17, 2026

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Doesn't RDNA3_5 imply that WMMA is available and MFMA is not 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<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,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) {
Expand Down
Loading