Skip to content

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

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
ravel7524 wants to merge 5 commits into
base: master
Choose a base branch
from
Open

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

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
5 commits
Select commit Hold shift + click to select a range
d590d22
adjust ncols_picker for routed MoE in mul_mat_q_case function
ravel7524 Jun 12, 2026
3087cea
Adding CDNA, RDNA2 and RDNA4
ravel7524 Jun 19, 2026
a57b8ea
fix: update mmq_use_routed_moe_ncols_picker to include NVIDIA + Volta…
ravel7524 Aug 5, 2026
f3a7a80
feat: enhance mmq configuration for various architectures with moe_nc…
ravel7524 Aug 5, 2026
2adc35b
refactor: replace moe_ncols_min_cc with use_typical_moe_ncols in mmq …
ravel7524 Aug 6, 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
3 changes: 2 additions & 1 deletion ggml/src/ggml-cuda/mmq-config-ampere.cuh
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,4 +1,5 @@
static constexpr __host__ __device__ ggml_cuda_mmq_config ggml_cuda_mmq_get_config_ampere(ggml_type type, int J, bool fallback) {
constexpr bool use_typical_moe_ncols = false;
CASE(GGML_TYPE_Q1_0, 256, 1, 128, 8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, true);
CASE(GGML_TYPE_Q1_0, 256, 1, 128, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, true);
CASE(GGML_TYPE_Q1_0, 256, 1, 128, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, true);
Expand Down Expand Up @@ -379,5 +380,5 @@ static constexpr __host__ __device__ ggml_cuda_mmq_config ggml_cuda_mmq_get_conf
CASE(GGML_TYPE_NVFP4, 256, 1, 128, 112, GGML_CUDA_MMQ_SRAM_LAYOUT_NVFP4, MMQ_ITER_K, true, false);
CASE(GGML_TYPE_NVFP4, 256, 1, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_NVFP4, MMQ_ITER_K, true, false);

return ggml_cuda_mmq_config(GGML_TYPE_COUNT, 256, 1, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, 256, false, true);
return ggml_cuda_mmq_config(GGML_TYPE_COUNT, 256, 1, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, 256, use_typical_moe_ncols, false, true);
}
1 change: 1 addition & 0 deletions ggml/src/ggml-cuda/mmq-config-blackwell.cuh
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,4 +1,5 @@
static constexpr __host__ __device__ ggml_cuda_mmq_config ggml_cuda_mmq_get_config_blackwell(ggml_type type, int J, bool fallback) {
constexpr bool use_typical_moe_ncols = false;
CASE(GGML_TYPE_MXFP4, 256, 1, 128, 8, GGML_CUDA_MMQ_SRAM_LAYOUT_FP4, MMQ_ITER_K_FP4, true, true);
CASE(GGML_TYPE_MXFP4, 256, 1, 128, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_FP4, MMQ_ITER_K_FP4, true, true);
CASE(GGML_TYPE_MXFP4, 256, 1, 128, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_FP4, MMQ_ITER_K_FP4, true, true);
Expand Down
3 changes: 2 additions & 1 deletion ggml/src/ggml-cuda/mmq-config-cdna.cuh
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,4 +1,5 @@
static constexpr __host__ __device__ ggml_cuda_mmq_config ggml_cuda_mmq_get_config_cdna(ggml_type type, int J, bool fallback) {
constexpr bool use_typical_moe_ncols = false;
CASE(GGML_TYPE_Q1_0, 512, 1, 128, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, true);
CASE(GGML_TYPE_Q1_0, 512, 1, 128, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, true);
CASE(GGML_TYPE_Q1_0, 512, 1, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, true);
Expand Down Expand Up @@ -181,5 +182,5 @@ static constexpr __host__ __device__ ggml_cuda_mmq_config ggml_cuda_mmq_get_conf
CASE(GGML_TYPE_NVFP4, 512, 1, 128, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_NVFP4, MMQ_ITER_K, true, false);
CASE(GGML_TYPE_NVFP4, 512, 1, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_NVFP4, MMQ_ITER_K, true, false);

return ggml_cuda_mmq_config(GGML_TYPE_COUNT, 512, 1, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, 256, false, true);
return ggml_cuda_mmq_config(GGML_TYPE_COUNT, 512, 1, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, 256, use_typical_moe_ncols, false, true);
}
3 changes: 2 additions & 1 deletion ggml/src/ggml-cuda/mmq-config-pascal.cuh
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,4 +1,5 @@
static constexpr __host__ __device__ ggml_cuda_mmq_config ggml_cuda_mmq_get_config_pascal(ggml_type type, int J, bool fallback) {
constexpr bool use_typical_moe_ncols = false;
CASE(GGML_TYPE_Q1_0, 256, 2, 64, 8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
CASE(GGML_TYPE_Q1_0, 256, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
CASE(GGML_TYPE_Q1_0, 256, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
Expand Down Expand Up @@ -269,5 +270,5 @@ static constexpr __host__ __device__ ggml_cuda_mmq_config ggml_cuda_mmq_get_conf
CASE(GGML_TYPE_NVFP4, 256, 2, 64, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_NVFP4, MMQ_ITER_K, false, false);
CASE(GGML_TYPE_NVFP4, 256, 2, 64, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_NVFP4, MMQ_ITER_K, false, false);

return ggml_cuda_mmq_config(GGML_TYPE_COUNT, 256, 2, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, 256, false, true);
return ggml_cuda_mmq_config(GGML_TYPE_COUNT, 256, 2, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, 256, use_typical_moe_ncols, false, true);
}
3 changes: 2 additions & 1 deletion ggml/src/ggml-cuda/mmq-config-rdna2.cuh
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,4 +1,5 @@
static constexpr __host__ __device__ ggml_cuda_mmq_config ggml_cuda_mmq_get_config_rdna2(ggml_type type, int J, bool fallback) {
constexpr bool use_typical_moe_ncols = false;
CASE(GGML_TYPE_Q1_0, 256, 2, 128, 8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
CASE(GGML_TYPE_Q1_0, 256, 2, 128, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
CASE(GGML_TYPE_Q1_0, 256, 2, 128, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
Expand Down Expand Up @@ -269,5 +270,5 @@ static constexpr __host__ __device__ ggml_cuda_mmq_config ggml_cuda_mmq_get_conf
CASE(GGML_TYPE_NVFP4, 256, 2, 128, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_NVFP4, MMQ_ITER_K, false, false);
CASE(GGML_TYPE_NVFP4, 256, 2, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_NVFP4, MMQ_ITER_K, false, false);

return ggml_cuda_mmq_config(GGML_TYPE_COUNT, 256, 2, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, 256, false, true);
return ggml_cuda_mmq_config(GGML_TYPE_COUNT, 256, 2, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, 256, use_typical_moe_ncols, false, true);
}
3 changes: 2 additions & 1 deletion ggml/src/ggml-cuda/mmq-config-rdna3-5.cuh
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,4 +1,5 @@
static constexpr __host__ __device__ ggml_cuda_mmq_config ggml_cuda_mmq_get_config_rdna3_5(ggml_type type, int J, bool fallback) {
constexpr bool use_typical_moe_ncols = false;
CASE(GGML_TYPE_Q1_0, 128, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
CASE(GGML_TYPE_Q1_0, 128, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
CASE(GGML_TYPE_Q1_0, 256, 2, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
Expand Down Expand Up @@ -286,5 +287,5 @@ static constexpr __host__ __device__ ggml_cuda_mmq_config ggml_cuda_mmq_get_conf
CASE(GGML_TYPE_NVFP4, 256, 2, 128, 112, GGML_CUDA_MMQ_SRAM_LAYOUT_NVFP4, MMQ_ITER_K, false, false);
CASE(GGML_TYPE_NVFP4, 256, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_NVFP4, MMQ_ITER_K, false, false);

return ggml_cuda_mmq_config(GGML_TYPE_COUNT, 256, 2, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, 256, false, true);
return ggml_cuda_mmq_config(GGML_TYPE_COUNT, 256, 2, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, 256, use_typical_moe_ncols, false, true);
}
3 changes: 2 additions & 1 deletion ggml/src/ggml-cuda/mmq-config-rdna3.cuh
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,4 +1,5 @@
static constexpr __host__ __device__ ggml_cuda_mmq_config ggml_cuda_mmq_get_config_rdna3(ggml_type type, int J, bool fallback) {
constexpr bool use_typical_moe_ncols = true;
CASE(GGML_TYPE_Q1_0, 128, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
CASE(GGML_TYPE_Q1_0, 128, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
CASE(GGML_TYPE_Q1_0, 256, 2, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
Expand Down Expand Up @@ -286,5 +287,5 @@ static constexpr __host__ __device__ ggml_cuda_mmq_config ggml_cuda_mmq_get_conf
CASE(GGML_TYPE_NVFP4, 256, 2, 128, 112, GGML_CUDA_MMQ_SRAM_LAYOUT_NVFP4, MMQ_ITER_K, false, false);
CASE(GGML_TYPE_NVFP4, 256, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_NVFP4, MMQ_ITER_K, false, false);

return ggml_cuda_mmq_config(GGML_TYPE_COUNT, 256, 2, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, 256, false, true);
return ggml_cuda_mmq_config(GGML_TYPE_COUNT, 256, 2, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, 256, use_typical_moe_ncols, false, true);
}
3 changes: 2 additions & 1 deletion ggml/src/ggml-cuda/mmq-config-rdna4.cuh
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,4 +1,5 @@
static constexpr __host__ __device__ ggml_cuda_mmq_config ggml_cuda_mmq_get_config_rdna4(ggml_type type, int J, bool fallback) {
constexpr bool use_typical_moe_ncols = false;
CASE(GGML_TYPE_Q1_0, 128, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
CASE(GGML_TYPE_Q1_0, 128, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
CASE(GGML_TYPE_Q1_0, 128, 2, 64, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
Expand Down Expand Up @@ -286,5 +287,5 @@ static constexpr __host__ __device__ ggml_cuda_mmq_config ggml_cuda_mmq_get_conf
CASE(GGML_TYPE_NVFP4, 256, 2, 128, 112, GGML_CUDA_MMQ_SRAM_LAYOUT_NVFP4, MMQ_ITER_K, false, false);
CASE(GGML_TYPE_NVFP4, 256, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_NVFP4, MMQ_ITER_K, false, false);

return ggml_cuda_mmq_config(GGML_TYPE_COUNT, 256, 2, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, 256, false, true);
return ggml_cuda_mmq_config(GGML_TYPE_COUNT, 256, 2, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, 256, use_typical_moe_ncols, false, true);
}
23 changes: 19 additions & 4 deletions ggml/src/ggml-cuda/mmq.cuh
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -170,12 +170,13 @@ struct ggml_cuda_mmq_config {
int J; // SRAM tile width in src1->ne[1]/dst->ne[1] direction.
ggml_cuda_mmq_sram_layout sram_layout; // SRAM tile length in src0->ne[0]/src1->ne[0] direction (physical 32 bit elements).
int K_vram; // VRAM tile length in src0->ne[0]/src1->ne[0] direction (logical elements).
bool use_typical_moe_ncols;
bool stream_k; // Whether or not to use stream-k decomposition.
bool fallback; // Whether a fallback for out-of-bounds check in src0->ne[1] direction is needed.

constexpr __host__ __device__ ggml_cuda_mmq_config(
ggml_type type, int nthreads, int occupancy, int I, int J, ggml_cuda_mmq_sram_layout sram_layout, int K_vram, bool stream_k, bool fallback) :
type(type), nthreads(nthreads), occupancy(occupancy), I(I), J(J), sram_layout(sram_layout), K_vram(K_vram), stream_k(stream_k), fallback(fallback) {}
ggml_type type, int nthreads, int occupancy, int I, int J, ggml_cuda_mmq_sram_layout sram_layout, int K_vram, bool use_typical_moe_ncols, bool stream_k, bool fallback) :
type(type), nthreads(nthreads), occupancy(occupancy), I(I), J(J), sram_layout(sram_layout), K_vram(K_vram), use_typical_moe_ncols(use_typical_moe_ncols), stream_k(stream_k), fallback(fallback) {}

constexpr __device__ int rows_per_warp() const {
#if defined(AMD_MFMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
Expand Down Expand Up @@ -210,7 +211,7 @@ struct ggml_cuda_mmq_config {
static_assert((I_) % 32 == 0, "bad I"); \
static_assert((J_) % 8 == 0, "bad J"); \
static_assert((K_vram_) % 256 == 0, "bad K_vram"); \
return ggml_cuda_mmq_config((type_), (nthreads_), (occupancy_), (I_), (J_), (sram_layout_), (K_vram_), (stream_k_), (fallback_)); \
return ggml_cuda_mmq_config((type_), (nthreads_), (occupancy_), (I_), (J_), (sram_layout_), (K_vram_), use_typical_moe_ncols, (stream_k_), (fallback_)); \
} \

#include "mmq-config-pascal.cuh"
Expand Down Expand Up @@ -1472,6 +1473,20 @@ void mul_mat_q_switch_J(ggml_backend_cuda_context & ctx, const mmq_args & args,
const int cc = ggml_cuda_info().devices[id].cc;
const size_t smpbo = ggml_cuda_info().devices[id].smpbo;

int64_t ncols_picker = args.ncols_max;
if (args.expert_bounds != nullptr && args.nchannels_x > 0) {
const int J_max = ggml_cuda_mmq_get_J_max(type, fallback, cc, 128);
const ggml_cuda_mmq_config config_max = ggml_cuda_mmq_get_config(type, J_max, fallback, cc);
if (config_max.use_typical_moe_ncols) {
// Use the typical expert width only for tile selection.
// The launch grid still uses args.ncols_max.
const int64_t ncols_typical = (args.ncols_dst + args.nchannels_x - 1) / args.nchannels_x;
if (ncols_typical >= 1 && ncols_typical < J_max && ncols_typical < ncols_picker) {
ncols_picker = ncols_typical;
}
}
}

int J_best = 0;
int ntiles_J_best = INT_MAX;

Expand All @@ -1485,7 +1500,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 = (ncols_picker + config.J - 1) / config.J;

if (ntiles_x < ntiles_J_best) {
J_best = J;
Expand Down