diff --git a/ggml/src/ggml-cuda/mmq.cuh b/ggml/src/ggml-cuda/mmq.cuh index b58ac9e7b428..1a29334d8c17 100644 --- a/ggml/src/ggml-cuda/mmq.cuh +++ b/ggml/src/ggml-cuda/mmq.cuh @@ -6,6 +6,7 @@ #include #include +#include using namespace ggml_cuda_mma; @@ -275,8 +276,26 @@ static constexpr __host__ __device__ int mmq_get_mma_tile_x_k(ggml_type type) { #define MMQ_TILE_Y_K (MMQ_TILE_NE_K + MMQ_TILE_NE_K / QI8_1) #define MMQ_TILE_Y_FP4_K MMQ_TILE_Y_K +static int mmq_get_tile_y_k_padded_host(const int mmq_x, const int cc, const int warp_size, const int nwarps) { + const int pad = GGML_CUDA_CC_IS_RDNA3_5(cc) ? 2*nwarps*warp_size : nwarps*warp_size; + return GGML_PAD(mmq_x*MMQ_TILE_Y_K, pad); +} + +#if defined(RDNA3_5) +static constexpr __device__ int mmq_get_tile_y_k_padded_device(const int mmq_x, const int nwarps, const int warp_size) { + return GGML_PAD(mmq_x*MMQ_TILE_Y_K, 2*nwarps*warp_size); +} +#else +static constexpr __device__ int mmq_get_tile_y_k_padded_device(const int mmq_x, const int nwarps, const int warp_size) { + return GGML_PAD(mmq_x*MMQ_TILE_Y_K, nwarps*warp_size); +} +#endif // RDNA3_5 + static int mmq_get_granularity_host(const int mmq_x, const int cc) { if (amd_mfma_available(cc) || amd_wmma_available(cc)) { + if (GGML_CUDA_CC_IS_RDNA3_5(cc) && mmq_x >= 64) { + return 32; + } return mmq_x >= 128 ? 32 : 16; } else if (turing_mma_available(cc) && mmq_x >= 48) { return 16; @@ -287,7 +306,11 @@ static int mmq_get_granularity_host(const int mmq_x, const int cc) { #if defined(AMD_MFMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE) static constexpr __device__ int mmq_get_granularity_device(const int mmq_x) { +#if defined(RDNA3_5) + return mmq_x >= 64 ? 32 : 16; +#else return mmq_x >= 128 ? 32 : 16; +#endif // RDNA3_5 } #elif defined(TURING_MMA_AVAILABLE) static constexpr __device__ int mmq_get_granularity_device(const int mmq_x) { @@ -1303,6 +1326,133 @@ static __device__ __forceinline__ void vec_dot_q8_0_q8_1_mma( #endif // defined(AMD_MFMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE) } +template +static __device__ __forceinline__ void vec_dot_q4_0_q8_1_mma( + 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) + 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 granularity = mmq_get_granularity_device(mmq_x); + constexpr int rows_per_warp = granularity; + constexpr int ntx = rows_per_warp/tile_C::I; + + y += (threadIdx.y % ntx) * (tile_C::J*MMQ_TILE_Y_K); + + const int * x_qs = (const int *) x; + const float * x_df = (const float *) x_qs + 2*MMQ_TILE_NE_K; + const int * y_qs = (const int *) y + 4; + const half2 * y_ds = (const half2 *) y; + + const int i0 = (threadIdx.y / ntx) * rows_per_warp; + + for (int k01 = 0; k01 < MMQ_TILE_NE_K; k01 += QI8_0) { + const int k0 = k00 + k01; + + tile_A A[ntx]; + float dA[ntx][tile_C::ne]; + +#pragma unroll + for (int n = 0; n < ntx; ++n) { + load_ldmatrix(A[n], x_qs + (i0 + n*tile_A::I)*MMQ_MMA_TILE_X_K_Q8_0 + k0, MMQ_MMA_TILE_X_K_Q8_0); + +#pragma unroll + for (int l = 0; l < tile_C::ne; ++l) { + const int i = i0 + n*tile_A::I + tile_C::get_i(l); + dA[n][l] = x_df[i*MMQ_MMA_TILE_X_K_Q8_0 + k0/QI8_0]; + } + } + +#pragma unroll + for (int j0 = 0; j0 < mmq_x; j0 += ntx*tile_C::J) { + tile_B B; + load_ldmatrix(B, y_qs + j0*MMQ_TILE_Y_K + k01, MMQ_TILE_Y_K); + + const int j = j0 + tile_C::get_j(0); + const float dB = __low2float(y_ds[j*MMQ_TILE_Y_K + k01/QI8_1]); + +#pragma unroll + for (int n = 0; n < ntx; ++n) { + tile_C C; + mma(C, A[n], B); + +#pragma unroll + for (int l = 0; l < tile_C::ne; ++l) { + sum[(j0/tile_C::J + n)*tile_C::ne + l] += C.x[l]*dA[n][l]*dB; + } + } + } + } +#else + vec_dot_q8_0_q8_1_mma(x, y, sum, k00); +#endif +} + +template +static __device__ __forceinline__ void vec_dot_q8_0_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) + 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 granularity = mmq_get_granularity_device(mmq_x); + constexpr int rows_per_warp = granularity; + constexpr int ntx = rows_per_warp/tile_C::I; + + y += (threadIdx.y % ntx) * (tile_C::J*MMQ_TILE_Y_K); + + const int * x_qs = (const int *) x; + const float * x_df = (const float *) x_qs + 2*MMQ_TILE_NE_K; + const int * y_qs = (const int *) y + 4; + const float * y_df = (const float *) y; + + const int i0 = (threadIdx.y / ntx) * rows_per_warp; + + for (int k01 = 0; k01 < MMQ_TILE_NE_K; k01 += QI8_0) { + const int k0 = k00 + k01; + + tile_A A[ntx]; + float dA[ntx][tile_C::ne]; + +#pragma unroll + for (int n = 0; n < ntx; ++n) { + load_ldmatrix(A[n], x_qs + (i0 + n*tile_A::I)*MMQ_MMA_TILE_X_K_Q8_0 + k0, MMQ_MMA_TILE_X_K_Q8_0); + +#pragma unroll + for (int l = 0; l < tile_C::ne; ++l) { + const int i = i0 + n*tile_A::I + tile_C::get_i(l); + dA[n][l] = x_df[i*MMQ_MMA_TILE_X_K_Q8_0 + k0/QI8_0]; + } + } + +#pragma unroll + for (int j0 = 0; j0 < mmq_x; j0 += ntx*tile_C::J) { + tile_B B; + load_ldmatrix(B, y_qs + j0*MMQ_TILE_Y_K + k01, MMQ_TILE_Y_K); + + const int j = j0 + tile_C::get_j(0); + const float dB = y_df[j*MMQ_TILE_Y_K + k01/QI8_1]; + +#pragma unroll + for (int n = 0; n < ntx; ++n) { + tile_C C; + mma(C, A[n], B); + +#pragma unroll + for (int l = 0; l < tile_C::ne; ++l) { + sum[(j0/tile_C::J + n)*tile_C::ne + l] += C.x[l]*dA[n][l]*dB; + } + } + } + } +#else + vec_dot_q8_0_q8_1_mma(x, y, sum, k00); +#endif +} template static __device__ __forceinline__ void vec_dot_q8_1_q8_1_dp4a( @@ -2100,6 +2250,78 @@ static __device__ __forceinline__ int unpack_scales_q45_K(const int * scales, co ((scales[ksc/2] >> (2 * (ksc % 2))) & 0x30303030); // upper 2 bits } +// LDS int offset for pseudo-q8_1 qs within one weight row (matches Q4_1 / WMMA ldmatrix layout). +static __device__ __forceinline__ int mmq_q4_K_qs_lds_k(const int txi, const int nibble) { + return ((txi >> 3) << 4) + (txi & 7) + (nibble << 3); +} + +#if defined(RDNA3_5) +// gfx115x mmq_y=64: 128 threads map 2:1 to rows for scale/dm prep (no warp divergence). +template +static __device__ __forceinline__ void load_tiles_q4_K_dm_rdna35( + const char * __restrict__ x, half2 * __restrict__ x_dm, const int kbx0, const int i_max, const int stride) { + constexpr int warp_size = ggml_cuda_get_physical_warp_size(); + + const int dm_tid = threadIdx.y*warp_size + threadIdx.x; + const int dm_row = dm_tid/2; + const int dm_ksc = dm_tid%2; + + if (dm_row >= mmq_y) { + return; + } + + int i = dm_row; + if (need_check) { + i = min(i, i_max); + } + + const block_q4_K * bxi = (const block_q4_K *) x + kbx0 + i*stride; + const int * scales = (const int *) bxi->scales; + + const int sc32 = unpack_scales_q45_K(scales, dm_ksc + 0); + const int m32 = unpack_scales_q45_K(scales, dm_ksc + 2); + + const uint8_t * sc8 = (const uint8_t *) &sc32; + const uint8_t * m8 = (const uint8_t *) &m32; + + const half2 dm = bxi->dm * make_half2(1.0f, -1.0f); + +#pragma unroll + for (int l = 0; l < 4; ++l) { + x_dm[i*MMQ_MMA_TILE_X_K_Q8_1 + sizeof(int)*dm_ksc + l] = dm*make_half2(sc8[l], m8[l]); + } +} + +// gfx115x WMMA: one warp per row, coalesced qs global load, nibble expand to ldmatrix-ready LDS. +template +static __device__ __forceinline__ void load_tiles_q4_K_qs_wmma_rdna35( + const char * __restrict__ x, int * __restrict__ x_qs, const int kbx0, const int i_max, const int stride) { + constexpr int nwarps = mmq_get_nwarps_device(); + constexpr int warp_size = ggml_cuda_get_physical_warp_size(); + constexpr int qs_per_row = MMQ_ITER_K / (4 * QR4_K); // 32 ints per block_q4_K::qs + static_assert(qs_per_row == 32, "bad Q4_K qs_per_row"); + constexpr int nrows = warp_size / qs_per_row; + static_assert(nrows == 1, "Q4_K RDNA3.5 WMMA qs path expects one row per warp"); + +#pragma unroll + for (int i0 = 0; i0 < mmq_y; i0 += nrows*nwarps) { + int i = i0 + threadIdx.y; + + if (need_check) { + i = min(i, i_max); + } + + const block_q4_K * bxi = (const block_q4_K *) x + kbx0 + i*stride; + const int qs0 = ((const int *) bxi->qs)[threadIdx.x]; + + int * row_qs = x_qs + i*MMQ_MMA_TILE_X_K_Q8_1; + const int kqs = mmq_q4_K_qs_lds_k(threadIdx.x, 0); + row_qs[kqs] = qs0 & 0x0F0F0F0F; + row_qs[kqs+8] = (qs0 >> 4) & 0x0F0F0F0F; + } +} +#endif // RDNA3_5 + template static __device__ __forceinline__ void load_tiles_q4_K( const char * __restrict__ x, int * __restrict__ x_tile, const int kbx0, const int i_max, const int stride) { constexpr int nwarps = mmq_get_nwarps_device(); @@ -2115,6 +2337,11 @@ template static __device__ __forceinline__ void loa int * x_sc = (int *) (x_dm + txs.dm); #endif // defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE) +#if defined(RDNA3_5) && defined(AMD_WMMA_AVAILABLE) && !defined(AMD_MFMA_AVAILABLE) + // dm first: warms cache lines before qs reads from same block_q4_K. + load_tiles_q4_K_dm_rdna35(x, x_dm, kbx0, i_max, stride); + load_tiles_q4_K_qs_wmma_rdna35(x, x_qs, kbx0, i_max, stride); +#else constexpr int threads_per_row = MMQ_ITER_K / (4 * QR4_K); constexpr int nrows = warp_size / threads_per_row; const int txi = warp_size > threads_per_row ? threadIdx.x % threads_per_row : threadIdx.x; @@ -2128,30 +2355,36 @@ template static __device__ __forceinline__ void loa } const block_q4_K * bxi = (const block_q4_K *) x + kbx0 + i*stride; +#if defined(RDNA3_5) + const int qs0 = ((const int *) bxi->qs)[txi]; +#else const int qs0 = get_int_b4(bxi->qs, txi); +#endif #if defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE) - x_qs[i*MMQ_MMA_TILE_X_K_Q8_1 + 16*(txi/8) + txi % 8 + 0] = (qs0 >> 0) & 0x0F0F0F0F; - x_qs[i*MMQ_MMA_TILE_X_K_Q8_1 + 16*(txi/8) + txi % 8 + 8] = (qs0 >> 4) & 0x0F0F0F0F; + x_qs[i*MMQ_MMA_TILE_X_K_Q8_1 + mmq_q4_K_qs_lds_k(txi, 0)] = (qs0 >> 0) & 0x0F0F0F0F; + x_qs[i*MMQ_MMA_TILE_X_K_Q8_1 + mmq_q4_K_qs_lds_k(txi, 1)] = (qs0 >> 4) & 0x0F0F0F0F; #else x_qs[i*(MMQ_TILE_NE_K + 1) + txi] = qs0; #endif // defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE) } +#endif // RDNA3_5 WMMA qs path #if defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE) +#if defined(RDNA3_5) && defined(AMD_WMMA_AVAILABLE) && !defined(AMD_MFMA_AVAILABLE) + // dm handled above. +#else constexpr int rows_per_warp = warp_size / 2; #pragma unroll for (int i0 = 0; i0 < mmq_y; i0 += nwarps*rows_per_warp) { -#if defined(AMD_MFMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE) - // Need if on AMD instead of % because warp_size == 64 - // This causes double work and throughput loss (MI300X) - // H100 loses about 100 t/s with 'if' condition over '%' +#if defined(AMD_MFMA_AVAILABLE) + // Need if on CDNA (warp_size == 64) instead of %. int i = i0 + threadIdx.y*rows_per_warp + threadIdx.x/2; if (i < mmq_y) { #else int i = (i0 + threadIdx.y*rows_per_warp + threadIdx.x/2) % mmq_y; { -#endif // defined(AMD_MFMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE) +#endif // defined(AMD_MFMA_AVAILABLE) if (need_check) { i = min(i, i_max); } @@ -2175,6 +2408,7 @@ template static __device__ __forceinline__ void loa } } } +#endif // RDNA3_5 RDNA WMMA dm path #else #pragma unroll for (int i0 = 0; i0 < mmq_y; i0 += nwarps*warp_size) { @@ -2244,6 +2478,79 @@ static __device__ __forceinline__ void vec_dot_q4_K_q8_1_dp4a( } } +template +static __device__ __forceinline__ void vec_dot_q4_K_q8_1_mma( + 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) + // Phase WMMA before scale reads: int tiles first, then dm/ds from LDS (HIP codegen). + 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 granularity = mmq_get_granularity_device(mmq_x); + constexpr int rows_per_warp = granularity; + constexpr int ntx = rows_per_warp/tile_C::I; + + y += (threadIdx.y % ntx) * (tile_C::J*MMQ_TILE_Y_K); + + 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 / ntx) * rows_per_warp; + + for (int k01 = 0; k01 < MMQ_TILE_NE_K; k01 += QI8_1) { + const int k0 = k00 + k01; + + tile_A A[ntx]; + +#pragma unroll + for (int n = 0; n < ntx; ++n) { + load_ldmatrix(A[n], x_qs + (i0 + n*tile_A::I)*MMQ_MMA_TILE_X_K_Q8_1 + k0, MMQ_MMA_TILE_X_K_Q8_1); + } + +#pragma unroll + for (int j0 = 0; j0 < mmq_x; j0 += ntx*tile_C::J) { + tile_B B; + load_ldmatrix(B, y_qs + j0*MMQ_TILE_Y_K + k01, MMQ_TILE_Y_K); + +#if defined(GGML_USE_HIP) + __builtin_amdgcn_sched_barrier(0); +#endif + + tile_C C[ntx]; + +#pragma unroll + for (int n = 0; n < ntx; ++n) { +#if defined(GGML_USE_HIP) + __builtin_amdgcn_sched_barrier(0); +#endif + mma(C[n], A[n], B); + } + + const int j = j0 + tile_C::get_j(0); + const float2 dsB = __half22float2(y_dm[j*MMQ_TILE_Y_K + k01/QI8_1]); + +#pragma unroll + for (int n = 0; n < ntx; ++n) { +#pragma unroll + for (int l = 0; l < tile_C::ne; ++l) { + const int i = i0 + n*tile_A::I + tile_C::get_i(l); + const float2 dm = __half22float2(x_dm[i*MMQ_MMA_TILE_X_K_Q8_1 + k0/QI8_1]); + const int si = (j0/tile_C::J + n)*tile_C::ne + l; + sum[si] += dm.x*dsB.x*C[n].x[l]; + sum[si] += dm.y*dsB.y; + } + } + } + } +#else + vec_dot_q8_1_q8_1_mma(x, y, sum, k00); +#endif +} + template static __device__ __forceinline__ void load_tiles_q5_K( const char * __restrict__ x, int * __restrict__ x_tile, const int kbx0, const int i_max, const int stride) { constexpr int nwarps = mmq_get_nwarps_device(); @@ -2525,7 +2832,85 @@ static __device__ __forceinline__ void vec_dot_q6_K_q8_1_dp4a( template static __device__ __forceinline__ void vec_dot_q6_K_q8_1_mma( const int * __restrict__ x, const int * __restrict__ y, float * __restrict__ sum, const int k00) { -#if defined(AMD_MFMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE) +#if defined(RDNA3_5) && defined(AMD_WMMA_AVAILABLE) && !defined(AMD_MFMA_AVAILABLE) + // Hoist sclA = sc[k01/4]*d per k-slice — scales were loaded per (n,l) in j-loop. + constexpr data_layout input_layout = get_input_data_layout(); + typedef tile<16, 4, int, input_layout> tile_A; + typedef tile<16, 4, int, input_layout> tile_B; + typedef tile<16, 16, int, DATA_LAYOUT_J_MAJOR> tile_C; + + constexpr int granularity = mmq_get_granularity_device(mmq_x); + constexpr int rows_per_warp = granularity; + constexpr int ntx = rows_per_warp/tile_C::I; + + y += (threadIdx.y % ntx) * (tile_C::J*MMQ_TILE_Y_K); + + const int * x_qs = (const int *) x; + const float * x_df = (const float *) x_qs + MMQ_TILE_NE_K*2; + const int * x_sc = (const int *) x_df + MMQ_TILE_NE_K/QI6_K; + const int * y_qs = (const int *) y + 4; + const float * y_df = (const float *) y; + + const int i0 = (threadIdx.y / ntx) * rows_per_warp; + + for (int k01 = 0; k01 < MMQ_TILE_NE_K; k01 += 4) { + const int k0 = k00 + k01; + + tile_A A[ntx]; + float sclA[ntx][tile_C::ne]; + +#pragma unroll + for (int n = 0; n < ntx; ++n) { + load_ldmatrix(A[n], x_qs + (i0 + n*tile_A::I)*MMQ_MMA_TILE_X_K_Q6_K + k0, MMQ_MMA_TILE_X_K_Q6_K); + +#pragma unroll + for (int l = 0; l < tile_C::ne; ++l) { + const int i = i0 + n*tile_A::I + tile_C::get_i(l); + const int8_t * sc = (const int8_t *) (x_sc + i*MMQ_MMA_TILE_X_K_Q6_K + k00/16); + sclA[n][l] = (float) sc[k01/4] * x_df[i*MMQ_MMA_TILE_X_K_Q6_K]; + } + } + + constexpr int j_step = ntx*tile_C::J; + + tile_B B0; + load_ldmatrix(B0, y_qs + 0, MMQ_TILE_Y_K); + float dB0 = y_df[tile_C::get_j(0)*MMQ_TILE_Y_K + k01/QI8_1]; + +#pragma unroll + for (int j0 = 0; j0 < mmq_x; j0 += j_step) { + const int j0_next = j0 + j_step; + + tile_B B1; + float dB1 = 0.0f; + if (j0_next < mmq_x) { + load_ldmatrix(B1, y_qs + j0_next*MMQ_TILE_Y_K + k01, MMQ_TILE_Y_K); + const int jn = j0_next + tile_C::get_j(0); + dB1 = y_df[jn*MMQ_TILE_Y_K + k01/QI8_1]; + } + + tile_C C[ntx]; + +#pragma unroll + for (int n = 0; n < ntx; ++n) { + mma(C[n], A[n], B0); + } + +#pragma unroll + for (int n = 0; n < ntx; ++n) { +#pragma unroll + for (int l = 0; l < tile_C::ne; ++l) { + sum[(j0/tile_C::J + n)*tile_C::ne + l] += C[n].x[l]*sclA[n][l]*dB0; + } + } + + if (j0_next < mmq_x) { + B0 = B1; + dB0 = dB1; + } + } + } +#elif defined(AMD_MFMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE) constexpr data_layout input_layout = get_input_data_layout(); typedef tile<16, 4, int, input_layout> tile_A; typedef tile<16, 4, int, input_layout> tile_B; @@ -3287,7 +3672,7 @@ template struct mmq_type_traits { static constexpr int vdr = VDR_Q4_0_Q8_1_MMQ; static constexpr load_tiles_mmq_t load_tiles = load_tiles_q4_0; - static constexpr vec_dot_mmq_t vec_dot_mma = vec_dot_q8_0_q8_1_mma; + static constexpr vec_dot_mmq_t vec_dot_mma = vec_dot_q4_0_q8_1_mma; static constexpr vec_dot_mmq_t vec_dot_dp4a = vec_dot_q4_0_q8_1_dp4a; }; @@ -3319,7 +3704,7 @@ template struct mmq_type_traits { static constexpr int vdr = VDR_Q8_0_Q8_1_MMQ; static constexpr load_tiles_mmq_t load_tiles = load_tiles_q8_0; - static constexpr vec_dot_mmq_t vec_dot_mma = vec_dot_q8_0_q8_1_mma; + static constexpr vec_dot_mmq_t vec_dot_mma = vec_dot_q8_0_q8_1_mma_rdna35; static constexpr vec_dot_mmq_t vec_dot_dp4a = vec_dot_q8_0_q8_1_dp4a; }; @@ -3369,7 +3754,7 @@ template struct mmq_type_traits { static constexpr int vdr = VDR_Q4_K_Q8_1_MMQ; static constexpr load_tiles_mmq_t load_tiles = load_tiles_q4_K; - static constexpr vec_dot_mmq_t vec_dot_mma = vec_dot_q8_1_q8_1_mma; + static constexpr vec_dot_mmq_t vec_dot_mma = vec_dot_q4_K_q8_1_mma; static constexpr vec_dot_mmq_t vec_dot_dp4a = vec_dot_q4_K_q8_1_dp4a; }; @@ -3453,6 +3838,49 @@ struct mmq_type_traits { static constexpr vec_dot_mmq_t vec_dot_dp4a = vec_dot_q8_0_q8_1_dp4a; }; + +#if defined(RDNA3_5) +// Software-pipeline activation tile loads: issue global loads into registers, run WMMA, +// then store to LDS. Overlaps memory latency with vec_dot on gfx115x (no ISA prefetch). +template +static __device__ __forceinline__ void mmq_tile_y_load_global( + int * __restrict__ tile_y, const int * __restrict__ by) { +#pragma unroll + for (int c = 0; c < nchunks; ++c) { + const int l = c*(nwarps*warp_size) + threadIdx.y*warp_size + threadIdx.x; + tile_y[l] = by[l]; + } +} + +template +static __device__ __forceinline__ void mmq_tile_y_load_global_to_regs( + const int * __restrict__ by, int (&cache)[nchunks]) { +#pragma unroll + for (int c = 0; c < nchunks; ++c) { + const int l = c*(nwarps*warp_size) + threadIdx.y*warp_size + threadIdx.x; + cache[c] = by[l]; + } +} + +template +static __device__ __forceinline__ void mmq_tile_y_store_regs( + int * __restrict__ tile_y, const int (&cache)[nchunks]) { +#pragma unroll + for (int c = 0; c < nchunks; ++c) { + const int l = c*(nwarps*warp_size) + threadIdx.y*warp_size + threadIdx.x; + tile_y[l] = cache[c]; + } +} +#endif // RDNA3_5 + +#if defined(GGML_USE_HIP) && defined(RDNA3_5) +// Extra tile-scope barriers give the AMDGPU backend the same DS / WMMA / FMA_MIX +// scheduling regions as the cold/hot CFG split (threadIdx.z) without a dead branch. +#define MMQ_HIP_TILE_BARRIER() __syncthreads() +#else +#define MMQ_HIP_TILE_BARRIER() ((void)0) +#endif + template static __device__ __forceinline__ void mul_mat_q_process_tile( const char * __restrict__ x, const int offset_x, const int * __restrict__ y, @@ -3468,7 +3896,7 @@ static __device__ __forceinline__ void mul_mat_q_process_tile( extern __shared__ int data_mul_mat_q[]; int * tile_y = data_mul_mat_q + mmq_x; - int * tile_x = tile_y + GGML_PAD(mmq_x*MMQ_TILE_Y_K, nwarps*warp_size); + int * tile_x = tile_y + mmq_get_tile_y_k_padded_device(mmq_x, nwarps, warp_size); #if defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE) constexpr vec_dot_mmq_t vec_dot = mmq_type_traits::vec_dot_mma; @@ -3492,39 +3920,101 @@ static __device__ __forceinline__ void mul_mat_q_process_tile( constexpr int sz = sizeof(block_q8_1_mmq) / sizeof(int); - for (int kb0 = kb0_start; kb0 < kb0_stop; kb0 += blocks_per_iter) { - load_tiles(x, tile_x, offset_x + kb0, tile_x_max_i, stride_row_x); - { - const int * by0 = y + ncols_y * (kb0 * qk / ne_block) * sz; -#pragma unroll - for (int l0 = 0; l0 < mmq_x * MMQ_TILE_Y_K; l0 += nwarps * warp_size) { - int l = l0 + threadIdx.y*warp_size + threadIdx.x; - - tile_y[l] = by0[l]; +#if defined(RDNA3_5) + constexpr int tile_y_elems = mmq_x*MMQ_TILE_Y_K; + constexpr int tile_y_load_stride = nwarps*warp_size; + if constexpr (mmq_x <= 64 && tile_y_elems % tile_y_load_stride == 0) { + constexpr int tile_y_nchunks = tile_y_elems/tile_y_load_stride; + + int y0_next_cache[tile_y_nchunks]; + bool have_y0_prefetch = false; + + for (int kb0 = kb0_start; kb0 < kb0_stop; kb0 += blocks_per_iter) { + { + MMQ_HIP_TILE_BARRIER(); + load_tiles(x, tile_x, offset_x + kb0, tile_x_max_i, stride_row_x); + MMQ_HIP_TILE_BARRIER(); } - } - __syncthreads(); + const int yk = kb0 * qk / ne_block; + const int * by0 = y + ncols_y * yk * sz; + const int * by1 = y + ncols_y * (yk + 1) * sz; - vec_dot(tile_x, tile_y, sum, 0); + { + MMQ_HIP_TILE_BARRIER(); - __syncthreads(); + if (have_y0_prefetch) { + mmq_tile_y_store_regs(tile_y, y0_next_cache); + have_y0_prefetch = false; + } else { + mmq_tile_y_load_global(tile_y, by0); + } - { - const int * by0 = y + ncols_y * ((kb0 * qk / ne_block) * sz + sz); -#pragma unroll - for (int l0 = 0; l0 < mmq_x * MMQ_TILE_Y_K; l0 += nwarps * warp_size) { - int l = l0 + threadIdx.y*warp_size + threadIdx.x; + __syncthreads(); + + int y1_cache[tile_y_nchunks]; + mmq_tile_y_load_global_to_regs(by1, y1_cache); + + vec_dot(tile_x, tile_y, sum, 0); + + __syncthreads(); + + mmq_tile_y_store_regs(tile_y, y1_cache); + + __syncthreads(); - tile_y[l] = by0[l]; + const int kb0_next = kb0 + blocks_per_iter; + if (kb0_next < kb0_stop) { + const int * by0_next = y + ncols_y * (kb0_next * qk / ne_block) * sz; + mmq_tile_y_load_global_to_regs(by0_next, y0_next_cache); + have_y0_prefetch = true; + } + + vec_dot(tile_x, tile_y, sum, MMQ_TILE_NE_K); + + MMQ_HIP_TILE_BARRIER(); } } + } else +#endif // RDNA3_5 + { + for (int kb0 = kb0_start; kb0 < kb0_stop; kb0 += blocks_per_iter) { + { + MMQ_HIP_TILE_BARRIER(); + load_tiles(x, tile_x, offset_x + kb0, tile_x_max_i, stride_row_x); + MMQ_HIP_TILE_BARRIER(); + } - __syncthreads(); + const int yk = kb0 * qk / ne_block; + const int * by0 = y + ncols_y * yk * sz; + const int * by1 = y + ncols_y * (yk + 1) * sz; - vec_dot(tile_x, tile_y, sum, MMQ_TILE_NE_K); + { + MMQ_HIP_TILE_BARRIER(); +#pragma unroll + for (int l0 = 0; l0 < mmq_x * MMQ_TILE_Y_K; l0 += nwarps * warp_size) { + int l = l0 + threadIdx.y*warp_size + threadIdx.x; - __syncthreads(); + tile_y[l] = by0[l]; + } + __syncthreads(); + + vec_dot(tile_x, tile_y, sum, 0); + + __syncthreads(); +#pragma unroll + for (int l0 = 0; l0 < mmq_x * MMQ_TILE_Y_K; l0 += nwarps * warp_size) { + int l = l0 + threadIdx.y*warp_size + threadIdx.x; + + tile_y[l] = by1[l]; + } + __syncthreads(); + + vec_dot(tile_x, tile_y, sum, MMQ_TILE_NE_K); + + MMQ_HIP_TILE_BARRIER(); + } + } } if (fixup) { @@ -3535,6 +4025,7 @@ static __device__ __forceinline__ void mul_mat_q_process_tile( } + // The mul_mat_q kernel implements "stream-k" work partitioning as described in https://arxiv.org/abs/2301.03598 template @@ -3939,6 +4430,23 @@ struct mmq_args { bool use_stream_k; int64_t ncols_max; }; +#if defined(GGML_USE_HIP) +// RDNA3.5 dual-WG (mmq_x=64, nbytes <= smpbo/2) helps K-quants with large per-tile LDS +// (Q6_K WMMA tuning). Block quants (Q5_0, Q8_0, Q4_0) are faster at mmq_x=128 ntx=1. +static bool mmq_rdna35_dual_wg_eligible(const ggml_type type) { + switch (type) { + case GGML_TYPE_Q2_K: + case GGML_TYPE_Q3_K: + case GGML_TYPE_Q4_K: + case GGML_TYPE_Q5_K: + case GGML_TYPE_Q6_K: + return true; + default: + return false; + } +} +#endif // GGML_USE_HIP + template static size_t mmq_get_nbytes_shared(const int mmq_x, const int mmq_y, const int cc, const int warp_size, const int nwarps) { const tile_x_sizes txs = mmq_get_dp4a_tile_x_sizes(type, mmq_y); @@ -3946,7 +4454,11 @@ static size_t mmq_get_nbytes_shared(const int mmq_x, const int mmq_y, const int const size_t nbs_ids = mmq_x*sizeof(int); const size_t nbs_x = (turing_mma_available(cc) || amd_mfma_available(cc) || amd_wmma_available(cc)) ? mmq_y*mmq_tile_x_k*sizeof(int) : txs.qs*sizeof(int) + txs.dm*sizeof(half2) + txs.sc*sizeof(int); const size_t nbs_y = mmq_x * (sizeof(block_q8_1_mmq)); - return nbs_ids + nbs_x + GGML_PAD(nbs_y, nwarps*warp_size*sizeof(int)); + const int tile_y_k_padded = mmq_get_tile_y_k_padded_host(mmq_x, cc, warp_size, nwarps); + const size_t nbs_y_padded = std::max(nbs_y, (size_t) tile_y_k_padded*sizeof(int)); + const int pad = GGML_CUDA_CC_IS_RDNA3_5(cc) ? 2*nwarps*warp_size : nwarps*warp_size; + size_t nbytes = nbs_ids + nbs_x + GGML_PAD(nbs_y_padded, pad*sizeof(int)); + return nbytes; } template @@ -3957,6 +4469,7 @@ static void launch_mul_mat_q(ggml_backend_cuda_context & ctx, const mmq_args & a const int warp_size = ggml_cuda_info().devices[id].warp_size; const int nwarps = mmq_get_nwarps_host(cc, warp_size); const int mmq_y = get_mmq_y_host(cc); + const size_t smpbo = ggml_cuda_info().devices[id].smpbo; const dim3 block_dims(warp_size, nwarps, 1); @@ -4091,6 +4604,76 @@ void mul_mat_q_case(ggml_backend_cuda_context & ctx, const mmq_args & args, cuda } } +#if defined(GGML_USE_HIP) + // RDNA3.5 (gfx115x): mmq_x=128 uses ~59% of LDS/CU → 1 WG/CU. For K-quants only, + // pick the smallest mmq_x with nbytes_shared <= smpbo/2 (2 WGs/CU). + if (GGML_CUDA_CC_IS_RDNA3_5(cc) && mmq_x_best > 0 && mmq_rdna35_dual_wg_eligible(type)) { + const size_t lds_dual_wg = smpbo / 2; + const size_t nbytes_best = mmq_get_nbytes_shared(mmq_x_best, mmq_y, cc, warp_size, nwarps); + if (nbytes_best > lds_dual_wg) { + int mmq_x_lds = 0; + size_t nbytes_lds_min = SIZE_MAX; + for (int mmq_x = 8; mmq_x <= mmq_x_max; mmq_x += 8) { + const int granularity = mmq_get_granularity_host(mmq_x, cc); + if (mmq_x % granularity != 0) { + continue; + } + const size_t nbytes = mmq_get_nbytes_shared(mmq_x, mmq_y, cc, warp_size, nwarps); + if (nbytes > smpbo || nbytes > lds_dual_wg) { + continue; + } + const int ntiles_x = (args.ncols_max + mmq_x - 1) / mmq_x; + if (ntiles_x > ntiles_x_best * 2) { + continue; + } + if (nbytes < nbytes_lds_min) { + mmq_x_lds = mmq_x; + nbytes_lds_min = nbytes; + } + } + if (mmq_x_lds > 0) { + mmq_x_best = mmq_x_lds; + } + } + } +#endif // GGML_USE_HIP + +#if defined(GGML_USE_HIP) + // Tiny M (e.g. gate 128×32×4096 Q8_0): prefer largest mmq_x within LDS so N is not + // split; occupancy is irrelevant when only a handful of blocks launch. + if (GGML_CUDA_CC_IS_RDNA3_5(cc) && args.nrows_x <= 32) { + const int nty = (args.nrows_x + mmq_y - 1) / mmq_y; + const int ntzw = static_cast(args.nchannels_y * args.nsamples_y); + + int mmq_x_tiny = 0; + int ntiles_tiny = INT_MAX; + for (int mmq_x = mmq_x_max; mmq_x >= 8; mmq_x -= 8) { + const int granularity = mmq_get_granularity_host(mmq_x, cc); + if (mmq_x % granularity != 0) { + continue; + } + const size_t nbytes = mmq_get_nbytes_shared(mmq_x, mmq_y, cc, warp_size, nwarps); + if (nbytes > smpbo) { + continue; + } + const int ntx = (args.ncols_max + mmq_x - 1) / mmq_x; + const int grid_blocks = nty * ntx * ntzw; + if (grid_blocks <= 4 && ntx <= ntiles_tiny) { + mmq_x_tiny = mmq_x; + ntiles_tiny = ntx; + if (ntiles_tiny == 1) { + break; + } + } + } + if (mmq_x_tiny > 0) { + mmq_x_best = mmq_x_tiny; + } + } + + // Q6_K / K-quants narrow-N: dual-WG path (mmq_x=64, ntx=2) from smpbo/2 selection. +#endif // GGML_USE_HIP + switch (mmq_x_best) { case 8: launch_mul_mat_q(ctx, args, stream);