CUDA: size routed MoE MMQ N-tiles from typical expert width on RDNA3 (#28552)

Recreated from #24546

---------

Co-authored-by: Carl Philipp Klemm <carl@uvos.xyz>

* CUDA: pick MMQ tile size against ncols_opt set on the host side

Assisted-by: Claude Fable 5.1
Claude-Session: https://claude.ai/code/session_011SYPfRhKoUpU3gMsGxq6go

---------

Co-authored-by: ravel7524 <58877666+ravel7524@users.noreply.github.com>
Co-authored-by: Carl Philipp Klemm <carl@uvos.xyz>
This commit is contained in:
Piotr Wilkin (ilintar)
2026-09-09 13:25:54 +02:00
committed by GitHub
parent 5a4d0fecae
commit d4abd573f6
2 changed files with 11 additions and 3 deletions
+9 -2
View File
@@ -171,7 +171,7 @@ void ggml_cuda_mul_mat_q(
ne00, ne01, ne1, s01, ne11, s1,
ne02, ne12, s02, s12, s2,
ne03, ne13, s03, s13, s3,
ne1};
ne1, ne1};
ggml_cuda_mul_mat_q_switch_type(ctx, args, stream);
return;
}
@@ -244,6 +244,13 @@ void ggml_cuda_mul_mat_q(
ne11 * ne10_padded * sizeof(block_q8_1) / (QK8_1 * sizeof(int));
const int64_t s13 = ne12*s12;
// Each expert only sees ne12*n_expert_used/ne02 tokens on average.
// On RDNA3 and RDNA4 it is faster to pick the tile size against this value instead of ne12.
int64_t ncols_opt = ne12;
if (GGML_CUDA_CC_IS_RDNA3_0(cc) || GGML_CUDA_CC_IS_RDNA4(cc)) {
ncols_opt = (ne12*n_expert_used + ne02 - 1) / ne02;
}
// Note that ne02 is used instead of ne12 because the number of y channels determines the z dimension of the CUDA grid.
const mmq_args args = {
src0_d, src0->type, (const int *) src1_q8_1.get(), ids_dst.get(), expert_bounds.get(), dst_d,
@@ -251,7 +258,7 @@ void ggml_cuda_mul_mat_q(
ne00, ne01, ne_get_rows, s01, ne_get_rows, s1,
ne02, ne02, s02, s12, s2,
ne03, ne13, s03, s13, s3,
ne12};
ne12, ncols_opt};
ggml_cuda_mul_mat_q_switch_type(ctx, args, stream);
}
+2 -1
View File
@@ -1376,6 +1376,7 @@ struct mmq_args {
int64_t nchannels_x; int64_t nchannels_y; int64_t stride_channel_x; int64_t stride_channel_y; int64_t stride_channel_dst;
int64_t nsamples_x; int64_t nsamples_y; int64_t stride_sample_x; int64_t stride_sample_y; int64_t stride_sample_dst;
int64_t ncols_max;
int64_t ncols_opt; // value to optimize the tile size against, launch grid still uses ncols_max
};
static size_t mmq_get_nbytes_shared(const ggml_cuda_mmq_config & config, const int cc) {
@@ -1486,7 +1487,7 @@ void mul_mat_q_switch_J(ggml_backend_cuda_context & ctx, const mmq_args & args,
continue;
}
const int ntiles_x = (args.ncols_max + config.J - 1) / config.J;
const int ntiles_x = (args.ncols_opt + config.J - 1) / config.J;
if (ntiles_x < ntiles_J_best) {
J_best = J;