Skip to content

Commit 8a6872b

Browse files
liangliangchangcursoragent
liangliangchang
and
cursoragent
committed
HIP: pipeline RDNA3.5 Q5_K tile loads
Prefetch Q5_K low- and high-bit data during WMMA while distributing high bits across lanes to keep J128 spill-free. Co-authored-by: Cursor <cursoragent@cursor.com>
1 parent 7050612 commit 8a6872b

2 files changed

Lines changed: 215 additions & 0 deletions

File tree

‎ggml/src/ggml-cuda/mmq-load-tiles.cuh‎

Lines changed: 165 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1012,6 +1012,166 @@ static __device__ __forceinline__ void ggml_cuda_mmq_store_tiles_q4_K_rdna35(
10121012
const uint8_t * m8 = (const uint8_t *) &m32;
10131013
const half2 dm = dm_cache * make_half2(1.0f, -1.0f);
10141014

1015+
#pragma unroll
1016+
for (int l = 0; l < int(sizeof(int)); ++l) {
1017+
x_dm[i*sram_stride + sizeof(int)*ksc + l] = dm*make_half2(sc8[l], m8[l]);
1018+
}
1019+
}
1020+
1021+
template <ggml_type type, int J, bool fallback>
1022+
static __device__ __forceinline__ void ggml_cuda_mmq_load_tiles_q5_K_rdna35(
1023+
const char * __restrict__ x, int * __restrict__ x_tile, const int kbx0, const int i_max, const int stride) {
1024+
constexpr int warp_size = ggml_cuda_get_physical_warp_size();
1025+
constexpr int nthreads = ggml_cuda_mmq_get_nthreads(type, J, fallback);
1026+
constexpr int I = ggml_cuda_mmq_get_I(type, J, fallback);
1027+
constexpr int sram_stride = ggml_cuda_mmq_get_sram_stride(type, J, fallback);
1028+
static_assert(warp_size == 32 && nthreads == 128 && I == 64, "unexpected RDNA3.5 Q5_K MMQ configuration");
1029+
1030+
int * x_qs = (int *) x_tile;
1031+
half2 * x_dm = (half2 *) (x_qs + 2*MMQ_TILE_NE_K);
1032+
1033+
const int linear_tid = threadIdx.y*warp_size + threadIdx.x;
1034+
int i = linear_tid/2;
1035+
if (fallback) {
1036+
i = min(i, i_max);
1037+
}
1038+
1039+
const block_q5_K * bxi = (const block_q5_K *) x + kbx0 + i*stride;
1040+
const int * scales = (const int *) bxi->scales;
1041+
const int ksc = linear_tid % 2;
1042+
const int sc32 = unpack_scales_q45_K(scales, ksc);
1043+
const int m32 = unpack_scales_q45_K(scales, ksc + 2);
1044+
const uint8_t * sc8 = (const uint8_t *) &sc32;
1045+
const uint8_t * m8 = (const uint8_t *) &m32;
1046+
const half2 dm = bxi->dm * make_half2(1.0f, -1.0f);
1047+
1048+
#pragma unroll
1049+
for (int l = 0; l < int(sizeof(int)); ++l) {
1050+
x_dm[i*sram_stride + sizeof(int)*ksc + l] = dm*make_half2(sc8[l], m8[l]);
1051+
}
1052+
1053+
const int txi = threadIdx.x;
1054+
const int kqs = 16*(txi/8) + txi%8;
1055+
const int qh_shift0 = 2*(txi/8);
1056+
1057+
#pragma unroll
1058+
for (int i0 = 0; i0 < I; i0 += nthreads/warp_size) {
1059+
int row = i0 + threadIdx.y;
1060+
if (fallback) {
1061+
row = min(row, i_max);
1062+
}
1063+
1064+
const block_q5_K * bxq = (const block_q5_K *) x + kbx0 + row*stride;
1065+
const int qs = ((const int *) bxq->qs)[txi];
1066+
const int qh = ((const int *) bxq->qh)[txi % (QI5_K/4)];
1067+
int * row_qs = x_qs + row*sram_stride;
1068+
row_qs[kqs] = (qs & 0x0F0F0F0F) | (((qh >> (qh_shift0 + 0)) << 4) & 0x10101010);
1069+
row_qs[kqs+8] = ((qs >> 4) & 0x0F0F0F0F) | (((qh >> (qh_shift0 + 1)) << 4) & 0x10101010);
1070+
}
1071+
}
1072+
1073+
template <ggml_type type, int J, bool fallback>
1074+
static __device__ __forceinline__ void ggml_cuda_mmq_prefetch_tiles_q5_K_rdna35(
1075+
const char * __restrict__ x, const int kbx0, const int i_max, const int stride,
1076+
int (&qs_cache)[ggml_cuda_mmq_get_I(type, J, fallback)/
1077+
(ggml_cuda_mmq_get_nthreads(type, J, fallback)/ggml_cuda_get_physical_warp_size())],
1078+
int (&qh_cache)[ggml_cuda_mmq_get_I(type, J, fallback)*(QI5_K/4)/
1079+
ggml_cuda_mmq_get_nthreads(type, J, fallback)],
1080+
int (&scales_cache)[3], half2 & dm_cache) {
1081+
constexpr int warp_size = ggml_cuda_get_physical_warp_size();
1082+
constexpr int nthreads = ggml_cuda_mmq_get_nthreads(type, J, fallback);
1083+
constexpr int nwarps = nthreads / warp_size;
1084+
constexpr int I = ggml_cuda_mmq_get_I(type, J, fallback);
1085+
constexpr int qh_words_per_row = QI5_K/4;
1086+
constexpr int qh_cache_size = I*qh_words_per_row/nthreads;
1087+
constexpr int rows_per_warp = I/nwarps;
1088+
static_assert(warp_size == 32 && nthreads == 128 && I == 64, "unexpected RDNA3.5 Q5_K MMQ configuration");
1089+
static_assert(qh_cache_size*warp_size == rows_per_warp*qh_words_per_row,
1090+
"Q5_K high bits must be distributed evenly across the warp");
1091+
1092+
#pragma unroll
1093+
for (int i0 = 0; i0 < I; i0 += nwarps) {
1094+
int i = i0 + threadIdx.y;
1095+
if constexpr (fallback) {
1096+
i = min(i, i_max);
1097+
}
1098+
1099+
const block_q5_K * bxi = (const block_q5_K *) x + kbx0 + i*stride;
1100+
qs_cache[i0/nwarps] = ((const int *) bxi->qs)[threadIdx.x];
1101+
}
1102+
1103+
#pragma unroll
1104+
for (int l = 0; l < qh_cache_size; ++l) {
1105+
const int qh_linear = l*warp_size + threadIdx.x;
1106+
int i = (qh_linear/qh_words_per_row)*nwarps + threadIdx.y;
1107+
if constexpr (fallback) {
1108+
i = min(i, i_max);
1109+
}
1110+
1111+
const block_q5_K * bxi = (const block_q5_K *) x + kbx0 + i*stride;
1112+
qh_cache[l] = ((const int *) bxi->qh)[qh_linear % qh_words_per_row];
1113+
}
1114+
1115+
int i = (threadIdx.y*warp_size + threadIdx.x)/2;
1116+
if constexpr (fallback) {
1117+
i = min(i, i_max);
1118+
}
1119+
1120+
const block_q5_K * bxi = (const block_q5_K *) x + kbx0 + i*stride;
1121+
#pragma unroll
1122+
for (int l = 0; l < 3; ++l) {
1123+
scales_cache[l] = ((const int *) bxi->scales)[l];
1124+
}
1125+
dm_cache = bxi->dm;
1126+
1127+
asm volatile("" ::: "memory");
1128+
}
1129+
1130+
template <ggml_type type, int J, bool fallback>
1131+
static __device__ __forceinline__ void ggml_cuda_mmq_store_tiles_q5_K_rdna35(
1132+
int * __restrict__ x_tile,
1133+
const int (&qs_cache)[ggml_cuda_mmq_get_I(type, J, fallback)/
1134+
(ggml_cuda_mmq_get_nthreads(type, J, fallback)/ggml_cuda_get_physical_warp_size())],
1135+
const int (&qh_cache)[ggml_cuda_mmq_get_I(type, J, fallback)*(QI5_K/4)/
1136+
ggml_cuda_mmq_get_nthreads(type, J, fallback)],
1137+
const int (&scales_cache)[3], const half2 dm_cache) {
1138+
constexpr int warp_size = ggml_cuda_get_physical_warp_size();
1139+
constexpr int nthreads = ggml_cuda_mmq_get_nthreads(type, J, fallback);
1140+
constexpr int nwarps = nthreads / warp_size;
1141+
constexpr int I = ggml_cuda_mmq_get_I(type, J, fallback);
1142+
constexpr int sram_stride = ggml_cuda_mmq_get_sram_stride(type, J, fallback);
1143+
constexpr int qh_words_per_row = QI5_K/4;
1144+
constexpr int qh_rows_per_slot = warp_size/qh_words_per_row;
1145+
static_assert(warp_size == 32 && nthreads == 128 && I == 64, "unexpected RDNA3.5 Q5_K MMQ configuration");
1146+
1147+
int * x_qs = x_tile;
1148+
const int txi = threadIdx.x;
1149+
const int kqs = 16*(txi/8) + txi%8;
1150+
const int qh_shift0 = 2*(txi/8);
1151+
1152+
#pragma unroll
1153+
for (int i0 = 0; i0 < I; i0 += nwarps) {
1154+
const int row_in_warp = i0/nwarps;
1155+
const int qh_slot = row_in_warp/qh_rows_per_slot;
1156+
const int qh_src_lane = (row_in_warp % qh_rows_per_slot)*qh_words_per_row + txi%qh_words_per_row;
1157+
const int qs = qs_cache[row_in_warp];
1158+
const int qh = __shfl_sync(0xFFFFFFFF, qh_cache[qh_slot], qh_src_lane, warp_size);
1159+
const int i = i0 + threadIdx.y;
1160+
int * row_qs = x_qs + i*sram_stride;
1161+
row_qs[kqs] = (qs & 0x0F0F0F0F) | (((qh >> (qh_shift0 + 0)) << 4) & 0x10101010);
1162+
row_qs[kqs+8] = ((qs >> 4) & 0x0F0F0F0F) | (((qh >> (qh_shift0 + 1)) << 4) & 0x10101010);
1163+
}
1164+
1165+
half2 * x_dm = (half2 *) (x_qs + 2*MMQ_TILE_NE_K);
1166+
const int linear_tid = threadIdx.y*warp_size + threadIdx.x;
1167+
const int i = linear_tid/2;
1168+
const int ksc = linear_tid%2;
1169+
const int sc32 = unpack_scales_q45_K(scales_cache, ksc);
1170+
const int m32 = unpack_scales_q45_K(scales_cache, ksc + 2);
1171+
const uint8_t * sc8 = (const uint8_t *) &sc32;
1172+
const uint8_t * m8 = (const uint8_t *) &m32;
1173+
const half2 dm = dm_cache * make_half2(1.0f, -1.0f);
1174+
10151175
#pragma unroll
10161176
for (int l = 0; l < int(sizeof(int)); ++l) {
10171177
x_dm[i*sram_stride + sizeof(int)*ksc + l] = dm*make_half2(sc8[l], m8[l]);
@@ -1137,6 +1297,11 @@ template <ggml_type type, int J, bool fallback> static __device__ __forceinline_
11371297

11381298
template <ggml_type type, int J, bool fallback> static __device__ __forceinline__ void ggml_cuda_mmq_load_tiles_q5_K(
11391299
const char * __restrict__ x, int * __restrict__ x_tile, const int kbx0, const int i_max, const int stride) {
1300+
#if defined(RDNA3_5)
1301+
ggml_cuda_mmq_load_tiles_q5_K_rdna35<type, J, fallback>(x, x_tile, kbx0, i_max, stride);
1302+
return;
1303+
#endif
1304+
11401305
constexpr int warp_size = ggml_cuda_get_physical_warp_size();
11411306
constexpr int nwarps = ggml_cuda_mmq_get_nthreads(type, J, fallback) / warp_size;
11421307
constexpr int I = ggml_cuda_mmq_get_I(type, J, fallback);

‎ggml/src/ggml-cuda/mmq.cuh‎

Lines changed: 50 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1109,6 +1109,56 @@ static __device__ __forceinline__ void mul_mat_q_process_tile(
11091109
}
11101110
__syncthreads();
11111111
}
1112+
} else if constexpr (type == GGML_TYPE_Q5_K && (J == 64 || J == 128)) {
1113+
constexpr int qs_cache_size = I/nwarps;
1114+
constexpr int qh_cache_size = I*(QI5_K/4)/(nwarps*warp_size);
1115+
1116+
__syncthreads();
1117+
load_tiles(x, tile_x, offset_x + kb0_start, tile_x_max_i, stride_row_x);
1118+
__syncthreads();
1119+
1120+
for (int kb0 = kb0_start; kb0 < kb0_stop; kb0 += blocks_per_iter) {
1121+
const int yk = kb0*qk/ne_block;
1122+
const int * by0 = y + ncols_y*yk*sz;
1123+
const int * by1 = y + ncols_y*(yk + 1)*sz;
1124+
1125+
#pragma unroll
1126+
for (int l0 = 0; l0 < J*MMQ_TILE_Y_K; l0 += nwarps*warp_size) {
1127+
const int l = l0 + threadIdx.y*warp_size + threadIdx.x;
1128+
tile_y[l] = by0[l];
1129+
}
1130+
__syncthreads();
1131+
vec_dot(tile_x, tile_y, sum, 0);
1132+
1133+
__syncthreads();
1134+
#pragma unroll
1135+
for (int l0 = 0; l0 < J*MMQ_TILE_Y_K; l0 += nwarps*warp_size) {
1136+
const int l = l0 + threadIdx.y*warp_size + threadIdx.x;
1137+
tile_y[l] = by1[l];
1138+
}
1139+
__syncthreads();
1140+
1141+
int qs_cache[qs_cache_size];
1142+
int qh_cache[qh_cache_size];
1143+
int scales_cache[3];
1144+
half2 dm_cache;
1145+
const int kb0_next = kb0 + blocks_per_iter;
1146+
const bool have_next = kb0_next < kb0_stop;
1147+
if (have_next) {
1148+
ggml_cuda_mmq_prefetch_tiles_q5_K_rdna35<type, J, fallback>(
1149+
x, offset_x + kb0_next, tile_x_max_i, stride_row_x,
1150+
qs_cache, qh_cache, scales_cache, dm_cache);
1151+
}
1152+
1153+
vec_dot(tile_x, tile_y, sum, MMQ_TILE_NE_K);
1154+
__syncthreads();
1155+
1156+
if (have_next) {
1157+
ggml_cuda_mmq_store_tiles_q5_K_rdna35<type, J, fallback>(
1158+
tile_x, qs_cache, qh_cache, scales_cache, dm_cache);
1159+
}
1160+
__syncthreads();
1161+
}
11121162
} else if constexpr (type == GGML_TYPE_Q4_0 && (J == 64 || J == 128)) {
11131163
constexpr int qs_cache_size = I/nwarps;
11141164
constexpr int d_cache_size = 4;

0 commit comments

Comments
 (0)