Skip to content

Commit 0e1efab

Browse files
liangliangchangcursoragent
liangliangchang
and
cursoragent
committed
HIP: batch RDNA3.5 MMQ WMMA issue and pipeline the Q4_K weight tile
On the 64-row tile at J=128 the MMQ inner loop has one MMA per j-tile, so it stalls on that tile's scales before it can issue. mma_half() splits the two WMMA instructions mma() already emits, letting a caller issue every tile's first half, load all the scales, then issue the second halves. It goes into the two existing generic vec_dots rather than a new kernel, so Q4_0 and Q8_0 get it as well. Q4_K additionally gets a split VRAM-load / SRAM-store pair for the weight tile, so the kernel can read the next K-iteration's tile into registers while the current iteration's MMAs run instead of stalling in front of the load. Q8_0 and Q5_K move to the 64-row tile at J=128, which the batched schedule needs. That tile on its own regresses those two, which is why the earlier tile-shape change left them out; combined with batching it pays. test-backend-ops perf on gfx1151, q4_K: +5.4% at m=4096 n=512 k=14336, +4.4% at m=12288 n=128 k=4096, +15.1% at m=1024 n=128 k=4096. Batching alone is worth +2.5% to +4.0% on q4_0 and q8_0. No measured shape regressed. Correctness: 214/214 MUL_MAT, 225/225 MUL_MAT_ID. Assisted-by: Claude Opus 5 Co-authored-by: Cursor <cursoragent@cursor.com>
1 parent 3ac87ad commit 0e1efab

5 files changed

Lines changed: 312 additions & 4 deletions

File tree

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

Lines changed: 29 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1336,6 +1336,35 @@ namespace ggml_cuda_mma {
13361336
#endif // AMD_MFMA_AVAILABLE
13371337
}
13381338

1339+
// Half of the WMMA mma() above: issues only the half'th of its two instructions.
1340+
// The MMA result is only complete once both halves have run. A caller that holds
1341+
// several B tiles can issue every tile's half 0, do unrelated work, then issue the
1342+
// half 1s, which keeps the WMMA pipe busy across that work instead of stalling on
1343+
// it once per tile. WMMA only: the MFMA paths do not decompose the same way.
1344+
template <int half, data_layout dl_d, data_layout dl_ab>
1345+
static __device__ __forceinline__ void mma_half(
1346+
tile<16, 16, int, dl_d> & D, const tile<16, 8, int, dl_ab> & A, const tile<16, 8, int, dl_ab> & B) {
1347+
static_assert(half == 0 || half == 1, "mma_half expects half 0 or 1");
1348+
#if defined(AMD_WMMA_AVAILABLE)
1349+
using int32x8_t = __attribute__((__vector_size__(8 * sizeof(int)))) int;
1350+
int32x8_t * acc = (int32x8_t *) D.x;
1351+
#if defined(RDNA4)
1352+
using int32x2_t = __attribute__((__vector_size__(2 * sizeof(int)))) int;
1353+
const int32x2_t * a_vec = (const int32x2_t *) A.x;
1354+
const int32x2_t * b_vec = (const int32x2_t *) B.x;
1355+
acc[0] = __builtin_amdgcn_wmma_i32_16x16x16_iu8_w32_gfx12(true, a_vec[half], true, b_vec[half], acc[0], true);
1356+
#elif defined(RDNA3)
1357+
using int32x4_t = __attribute__((__vector_size__(4 * sizeof(int)))) int;
1358+
const int32x4_t * a_vec = (const int32x4_t *) A.x;
1359+
const int32x4_t * b_vec = (const int32x4_t *) B.x;
1360+
acc[0] = __builtin_amdgcn_wmma_i32_16x16x16_iu8_w32(true, a_vec[half], true, b_vec[half], acc[0], true);
1361+
#endif // RDNA4
1362+
#else
1363+
GGML_UNUSED_VARS(D, A, B);
1364+
NO_DEVICE_CODE;
1365+
#endif // defined(AMD_WMMA_AVAILABLE)
1366+
}
1367+
13391368
static __device__ __forceinline__ void mma(
13401369
tile<32, 32, int> & D, const tile<32, 4, int> & A, const tile<32, 4, int> & B) {
13411370
#if defined(AMD_MFMA_AVAILABLE)

‎ggml/src/ggml-cuda/mmq-config-rdna3-5.cuh‎

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -79,15 +79,15 @@ static constexpr __host__ __device__ ggml_cuda_mmq_config ggml_cuda_mmq_get_conf
7979
CASE(GGML_TYPE_Q8_0, 128, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
8080
CASE(GGML_TYPE_Q8_0, 128, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
8181
CASE(GGML_TYPE_Q8_0, 128, 2, 64, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
82-
CASE(GGML_TYPE_Q8_0, 256, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
82+
CASE(GGML_TYPE_Q8_0, 128, 2, 64, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
8383
CASE(GGML_TYPE_Q8_0, 128, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
8484
CASE(GGML_TYPE_Q8_0, 128, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
8585
CASE(GGML_TYPE_Q8_0, 128, 2, 64, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
8686
CASE(GGML_TYPE_Q8_0, 128, 2, 64, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
8787
CASE(GGML_TYPE_Q8_0, 128, 2, 64, 80, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
8888
CASE(GGML_TYPE_Q8_0, 256, 2, 128, 96, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
8989
CASE(GGML_TYPE_Q8_0, 128, 2, 64, 112, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
90-
CASE(GGML_TYPE_Q8_0, 256, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
90+
CASE(GGML_TYPE_Q8_0, 128, 2, 64, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
9191

9292
// ---------------------------------------------------------------------------------------------
9393

@@ -129,15 +129,15 @@ static constexpr __host__ __device__ ggml_cuda_mmq_config ggml_cuda_mmq_get_conf
129129
CASE(GGML_TYPE_Q5_K, 128, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true);
130130
CASE(GGML_TYPE_Q5_K, 128, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true);
131131
CASE(GGML_TYPE_Q5_K, 128, 2, 64, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true);
132-
CASE(GGML_TYPE_Q5_K, 256, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true);
132+
CASE(GGML_TYPE_Q5_K, 128, 2, 64, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true);
133133
CASE(GGML_TYPE_Q5_K, 128, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
134134
CASE(GGML_TYPE_Q5_K, 128, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
135135
CASE(GGML_TYPE_Q5_K, 128, 2, 64, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
136136
CASE(GGML_TYPE_Q5_K, 128, 2, 64, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
137137
CASE(GGML_TYPE_Q5_K, 128, 2, 64, 80, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
138138
CASE(GGML_TYPE_Q5_K, 256, 2, 128, 96, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
139139
CASE(GGML_TYPE_Q5_K, 128, 2, 64, 112, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
140-
CASE(GGML_TYPE_Q5_K, 256, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
140+
CASE(GGML_TYPE_Q5_K, 128, 2, 64, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
141141

142142
CASE(GGML_TYPE_Q6_K, 128, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q6_K, MMQ_ITER_K, false, true);
143143
CASE(GGML_TYPE_Q6_K, 128, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q6_K, MMQ_ITER_K, false, true);

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

Lines changed: 84 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -819,6 +819,90 @@ template <ggml_type type, int J, bool fallback> static __device__ __forceinline_
819819
#endif // defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
820820
}
821821

822+
#if defined(RDNA3_5)
823+
// ggml_cuda_mmq_load_tiles_q4_K above reads x from VRAM and writes the SRAM tile in one
824+
// go, so the caller stalls on the VRAM latency before it can issue any MMA. These two
825+
// split that in half: prefetch() does only the reads, into registers, and store() does
826+
// only the SRAM writes. A caller can then issue the reads for the next K-iteration,
827+
// run the current iteration's MMAs while they are in flight, and commit them afterwards.
828+
// The pair must be used together and the tile must not be read between them.
829+
template <ggml_type type, int J, bool fallback>
830+
static __device__ __forceinline__ void ggml_cuda_mmq_prefetch_tiles_q4_K(
831+
const char * __restrict__ x, const int kbx0, const int i_max, const int stride,
832+
int (&qs_cache)[ggml_cuda_mmq_get_I(type, J, fallback)/
833+
(ggml_cuda_mmq_get_nthreads(type, J, fallback)/ggml_cuda_get_physical_warp_size())],
834+
int (&scales_cache)[3], half2 & dm_cache) {
835+
constexpr int warp_size = ggml_cuda_get_physical_warp_size();
836+
constexpr int nthreads = ggml_cuda_mmq_get_nthreads(type, J, fallback);
837+
constexpr int nwarps = nthreads / warp_size;
838+
constexpr int I = ggml_cuda_mmq_get_I(type, J, fallback);
839+
static_assert(warp_size == 32 && nthreads == 128 && I == 64, "unexpected Q4_K MMQ configuration");
840+
841+
#pragma unroll
842+
for (int i0 = 0; i0 < I; i0 += nwarps) {
843+
int i = i0 + threadIdx.y;
844+
if constexpr (fallback) {
845+
i = min(i, i_max);
846+
}
847+
848+
const block_q4_K * bxi = (const block_q4_K *) x + kbx0 + i*stride;
849+
qs_cache[i0/nwarps] = ((const int *) bxi->qs)[threadIdx.x];
850+
}
851+
852+
int i = (threadIdx.y*warp_size + threadIdx.x)/2;
853+
if constexpr (fallback) {
854+
i = min(i, i_max);
855+
}
856+
857+
const block_q4_K * bxi = (const block_q4_K *) x + kbx0 + i*stride;
858+
#pragma unroll
859+
for (int l = 0; l < 3; ++l) {
860+
scales_cache[l] = ((const int *) bxi->scales)[l];
861+
}
862+
dm_cache = bxi->dm;
863+
}
864+
865+
template <ggml_type type, int J, bool fallback>
866+
static __device__ __forceinline__ void ggml_cuda_mmq_store_tiles_q4_K(
867+
int * __restrict__ x_tile,
868+
const int (&qs_cache)[ggml_cuda_mmq_get_I(type, J, fallback)/
869+
(ggml_cuda_mmq_get_nthreads(type, J, fallback)/ggml_cuda_get_physical_warp_size())],
870+
const int (&scales_cache)[3], const half2 dm_cache) {
871+
constexpr int warp_size = ggml_cuda_get_physical_warp_size();
872+
constexpr int nthreads = ggml_cuda_mmq_get_nthreads(type, J, fallback);
873+
constexpr int nwarps = nthreads / warp_size;
874+
constexpr int I = ggml_cuda_mmq_get_I(type, J, fallback);
875+
constexpr int sram_stride = ggml_cuda_mmq_get_sram_stride(type, J, fallback);
876+
static_assert(warp_size == 32 && nthreads == 128 && I == 64, "unexpected Q4_K MMQ configuration");
877+
878+
int * x_qs = x_tile;
879+
#pragma unroll
880+
for (int i0 = 0; i0 < I; i0 += nwarps) {
881+
const int i = i0 + threadIdx.y;
882+
const int qs = qs_cache[i0/nwarps];
883+
int * row_qs = x_qs + i*sram_stride;
884+
const int kqs = 16*(threadIdx.x/8) + threadIdx.x%8;
885+
row_qs[kqs] = qs & 0x0F0F0F0F;
886+
row_qs[kqs+8] = (qs >> 4) & 0x0F0F0F0F;
887+
}
888+
889+
half2 * x_dm = (half2 *) (x_qs + 2*MMQ_TILE_NE_K);
890+
const int linear_tid = threadIdx.y*warp_size + threadIdx.x;
891+
const int i = linear_tid/2;
892+
const int ksc = linear_tid%2;
893+
const int sc32 = unpack_scales_q45_K(scales_cache, ksc);
894+
const int m32 = unpack_scales_q45_K(scales_cache, ksc + 2);
895+
const uint8_t * sc8 = (const uint8_t *) &sc32;
896+
const uint8_t * m8 = (const uint8_t *) &m32;
897+
const half2 dm = dm_cache * make_half2(1.0f, -1.0f);
898+
899+
#pragma unroll
900+
for (int l = 0; l < int(sizeof(int)); ++l) {
901+
x_dm[i*sram_stride + sizeof(int)*ksc + l] = dm*make_half2(sc8[l], m8[l]);
902+
}
903+
}
904+
#endif // defined(RDNA3_5)
905+
822906
template <ggml_type type, int J, bool fallback> static __device__ __forceinline__ void ggml_cuda_mmq_load_tiles_q5_K(
823907
const char * __restrict__ x, int * __restrict__ x_tile, const int kbx0, const int i_max, const int stride) {
824908
constexpr int warp_size = ggml_cuda_get_physical_warp_size();

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

Lines changed: 131 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -148,6 +148,7 @@ static __device__ __forceinline__ void ggml_cuda_mmq_vec_dot_q8_0_q8_1_mma(
148148
typedef tile<16, 8, int, input_layout> tile_B;
149149
typedef tile<16, 16, int, DATA_LAYOUT_J_MAJOR> tile_C;
150150

151+
constexpr int I = ggml_cuda_mmq_get_I(type, J, fallback);
151152
constexpr int sram_stride = ggml_cuda_mmq_get_sram_stride(type, J, fallback);
152153
constexpr int rows_per_warp = ggml_cuda_mmq_get_rows_per_warp(type, J, fallback);
153154
constexpr int ntx = rows_per_warp/tile_C::I; // Number of x minitiles per warp.
@@ -162,6 +163,68 @@ static __device__ __forceinline__ void ggml_cuda_mmq_vec_dot_q8_0_q8_1_mma(
162163

163164
const int i0 = (threadIdx.y / ntx) * rows_per_warp;
164165

166+
#if defined(RDNA3_5)
167+
// A single x minitile per warp leaves the loop below with one MMA per j-tile, so it
168+
// stalls on that tile's scale load before it can issue. With the widest tile there
169+
// are enough j-tiles to hide those loads instead: issue every tile's first mma half,
170+
// load all the scales, then issue the second halves.
171+
if constexpr (I == 64 && J == 128 && ntx == 1) {
172+
constexpr int ntiles = J/tile_C::J;
173+
174+
for (int k01 = 0; k01 < MMQ_TILE_NE_K; k01 += QI8_0) {
175+
const int k0 = k00 + k01;
176+
177+
tile_A A;
178+
load_ldmatrix(A, x_qs + i0*sram_stride + k0, sram_stride);
179+
180+
tile_B B[ntiles];
181+
tile_C C[ntiles];
182+
#pragma unroll
183+
for (int jb = 0; jb < ntiles; ++jb) {
184+
load_ldmatrix(B[jb], y_qs + jb*tile_C::J*MMQ_TILE_Y_K + k01, MMQ_TILE_Y_K);
185+
mma_half<0>(C[jb], A, B[jb]);
186+
}
187+
188+
__builtin_amdgcn_sched_barrier(0);
189+
190+
float dA[tile_C::ne];
191+
float dB[ntiles];
192+
#pragma unroll
193+
for (int l = 0; l < tile_C::ne; ++l) {
194+
const int i = i0 + tile_C::get_i(l);
195+
dA[l] = x_df[i*sram_stride + k0/QI8_0];
196+
}
197+
#pragma unroll
198+
for (int jb = 0; jb < ntiles; ++jb) {
199+
const int j = jb*tile_C::J + tile_C::get_j(0);
200+
if constexpr (ds_layout == MMQ_Q8_1_DS_LAYOUT_D4) {
201+
dB[jb] = y_df[j*MMQ_TILE_Y_K + k01/QI8_1];
202+
} else {
203+
dB[jb] = __low2float(y_ds[j*MMQ_TILE_Y_K + k01/QI8_1]);
204+
}
205+
}
206+
207+
__builtin_amdgcn_sched_barrier(0);
208+
209+
#pragma unroll
210+
for (int jb = 0; jb < ntiles; ++jb) {
211+
mma_half<1>(C[jb], A, B[jb]);
212+
}
213+
214+
__builtin_amdgcn_sched_barrier(0);
215+
216+
#pragma unroll
217+
for (int jb = 0; jb < ntiles; ++jb) {
218+
#pragma unroll
219+
for (int l = 0; l < tile_C::ne; ++l) {
220+
sum[jb*tile_C::ne + l] += C[jb].x[l]*dA[l]*dB[jb];
221+
}
222+
}
223+
}
224+
return;
225+
}
226+
#endif // defined(RDNA3_5)
227+
165228
for (int k01 = 0; k01 < MMQ_TILE_NE_K; k01 += QI8_0) {
166229
const int k0 = k00 + k01;
167230

@@ -318,6 +381,7 @@ template <ggml_type type, int J, bool fallback> static __device__ __forceinline_
318381
typedef tile<16, 8, int, input_layout> tile_B;
319382
typedef tile<16, 16, int, DATA_LAYOUT_J_MAJOR> tile_C;
320383

384+
constexpr int I = ggml_cuda_mmq_get_I(type, J, fallback);
321385
constexpr int sram_stride = ggml_cuda_mmq_get_sram_stride(type, J, fallback);
322386
constexpr int rows_per_warp = ggml_cuda_mmq_get_rows_per_warp(type, J, fallback);
323387
constexpr int ntx = rows_per_warp/tile_C::I; // Number of x minitiles per warp.
@@ -331,6 +395,73 @@ template <ggml_type type, int J, bool fallback> static __device__ __forceinline_
331395

332396
const int i0 = (threadIdx.y / ntx) * rows_per_warp;
333397

398+
#if defined(RDNA3_5)
399+
// See the matching branch in ggml_cuda_mmq_vec_dot_q8_0_q8_1_mma. The scales here are
400+
// half2 rather than float, and the second of their two terms does not depend on the
401+
// MMA result, so that part of the sum is folded in while the first halves are still
402+
// in flight.
403+
if constexpr (I == 64 && J == 128 && ntx == 1) {
404+
constexpr int ntiles = J/tile_C::J;
405+
406+
for (int k01 = 0; k01 < MMQ_TILE_NE_K; k01 += QI8_1) {
407+
const int k0 = k00 + k01;
408+
409+
tile_A A;
410+
load_ldmatrix(A, x_qs + i0*sram_stride + k0, sram_stride);
411+
412+
tile_B B[ntiles];
413+
tile_C C[ntiles];
414+
#pragma unroll
415+
for (int jb = 0; jb < ntiles; ++jb) {
416+
load_ldmatrix(B[jb], y_qs + jb*tile_C::J*MMQ_TILE_Y_K + k01, MMQ_TILE_Y_K);
417+
mma_half<0>(C[jb], A, B[jb]);
418+
}
419+
420+
__builtin_amdgcn_sched_barrier(0);
421+
422+
float2 dmA[tile_C::ne];
423+
float2 dsB[ntiles];
424+
#pragma unroll
425+
for (int l = 0; l < tile_C::ne; ++l) {
426+
const int i = i0 + tile_C::get_i(l);
427+
dmA[l] = __half22float2(x_dm[i*sram_stride + k0/QI8_1]);
428+
}
429+
#pragma unroll
430+
for (int jb = 0; jb < ntiles; ++jb) {
431+
const int j = jb*tile_C::J + tile_C::get_j(0);
432+
dsB[jb] = __half22float2(y_dm[j*MMQ_TILE_Y_K + k01/QI8_1]);
433+
}
434+
435+
// The bias term needs no MMA result, so retire it before the second halves.
436+
#pragma unroll
437+
for (int l = 0; l < tile_C::ne; ++l) {
438+
#pragma unroll
439+
for (int jb = 0; jb < ntiles; ++jb) {
440+
sum[jb*tile_C::ne + l] += dmA[l].y*dsB[jb].y;
441+
}
442+
}
443+
444+
__builtin_amdgcn_sched_barrier(0);
445+
446+
#pragma unroll
447+
for (int jb = 0; jb < ntiles; ++jb) {
448+
mma_half<1>(C[jb], A, B[jb]);
449+
}
450+
451+
__builtin_amdgcn_sched_barrier(0);
452+
453+
#pragma unroll
454+
for (int jb = 0; jb < ntiles; ++jb) {
455+
#pragma unroll
456+
for (int l = 0; l < tile_C::ne; ++l) {
457+
sum[jb*tile_C::ne + l] += dmA[l].x*dsB[jb].x*C[jb].x[l];
458+
}
459+
}
460+
}
461+
return;
462+
}
463+
#endif // defined(RDNA3_5)
464+
334465
for (int k01 = 0; k01 < MMQ_TILE_NE_K; k01 += QI8_1) {
335466
const int k0 = k00 + k01;
336467

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

Lines changed: 64 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -899,6 +899,70 @@ static __device__ __forceinline__ void mul_mat_q_process_tile(
899899

900900
constexpr int sz = sizeof(block_q8_1_mmq) / sizeof(int);
901901

902+
#if defined(RDNA3_5)
903+
// Software-pipelined variant: the x tile for the next K-iteration is read into
904+
// registers before this iteration's second vec_dot and committed to SRAM after it,
905+
// so the VRAM latency overlaps the MMAs instead of stalling in front of them.
906+
// Only Q4_K on the 64-row tile has the split load/store pair for this so far.
907+
if constexpr (type == GGML_TYPE_Q4_K && I == 64 &&
908+
ggml_cuda_mmq_get_nthreads(type, J, fallback) == 128) {
909+
constexpr int qs_cache_size = I/nwarps;
910+
911+
load_tiles(x, tile_x, offset_x + kb0_start, tile_x_max_i, stride_row_x);
912+
913+
for (int kb0 = kb0_start; kb0 < kb0_stop; kb0 += blocks_per_iter) {
914+
{
915+
const int * by0 = y + ncols_y * (kb0 * qk / ne_block) * sz;
916+
#pragma unroll
917+
for (int l0 = 0; l0 < J * MMQ_TILE_Y_K; l0 += nwarps * warp_size) {
918+
int l = l0 + threadIdx.y*warp_size + threadIdx.x;
919+
920+
tile_y[l] = by0[l];
921+
}
922+
}
923+
924+
__syncthreads();
925+
926+
vec_dot(tile_x, tile_y, sum, 0);
927+
928+
__syncthreads();
929+
930+
{
931+
const int * by0 = y + ncols_y * ((kb0 * qk / ne_block) * sz + sz);
932+
#pragma unroll
933+
for (int l0 = 0; l0 < J * MMQ_TILE_Y_K; l0 += nwarps * warp_size) {
934+
int l = l0 + threadIdx.y*warp_size + threadIdx.x;
935+
936+
tile_y[l] = by0[l];
937+
}
938+
}
939+
940+
__syncthreads();
941+
942+
int qs_cache[qs_cache_size];
943+
int scales_cache[3];
944+
half2 dm_cache;
945+
const int kb0_next = kb0 + blocks_per_iter;
946+
const bool have_next = kb0_next < kb0_stop;
947+
if (have_next) {
948+
ggml_cuda_mmq_prefetch_tiles_q4_K<type, J, fallback>(
949+
x, offset_x + kb0_next, tile_x_max_i, stride_row_x,
950+
qs_cache, scales_cache, dm_cache);
951+
}
952+
953+
vec_dot(tile_x, tile_y, sum, MMQ_TILE_NE_K);
954+
955+
__syncthreads();
956+
957+
if (have_next) {
958+
ggml_cuda_mmq_store_tiles_q4_K<type, J, fallback>(
959+
tile_x, qs_cache, scales_cache, dm_cache);
960+
}
961+
962+
__syncthreads();
963+
}
964+
} else
965+
#endif // defined(RDNA3_5)
902966
for (int kb0 = kb0_start; kb0 < kb0_stop; kb0 += blocks_per_iter) {
903967
load_tiles(x, tile_x, offset_x + kb0, tile_x_max_i, stride_row_x);
904968
{

0 commit comments

Comments
 (0)