From 5de3e40802497cb61ada8a2673d41aa161805c1c Mon Sep 17 00:00:00 2001 From: Piotr Wilkin Date: Thu, 2 Jul 2026 17:35:49 +0200 Subject: [PATCH 1/3] cuda: enable topk-moe fusion for 288 experts MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The topk-moe fusion only accepted power-of-2 expert counts (or the special-cased 576), so models with 288 experts (e.g. Step-3.7-Flash) fell back to the unfused per-layer routing chain: softmax/sigmoid, argsort, get_rows, sum_rows, div, clamp, scale. At batch size 1 that is ~330 extra tiny graph nodes per token. 288 is a multiple of the warp size, so the existing kernel already handles it; this adds the missing template instantiation and accepts 288 in the eligibility check. Measured on gfx1151 with Step-3.7-Flash IQ4_XS (llama-bench, -b 4096 -ub 4096 -fa 1 -dio 1 -ctk q8_0 -ctv q8_0; machine idle, before/after paired so pp4096 stays matched as a load control): test | before | after ----------------+----------------+---------------- pp4096 | 460.99 ± 0.45 | 462.47 ± 0.34 (unchanged) tg128 | 19.10 ± 0.04 | 19.56 ± 0.03 (+2.4%) tg128 @ d30000 | 12.68 ± 0.04 | 12.69 ± 0.03 (unchanged) Prompt processing is unaffected (the fusion only touches decode routing). The decode gain is ~+2.4% at shallow context and fades with depth: by 30k tokens each step is attention-bound over the KV cache, so removing the fixed routing overhead is no longer visible. Assisted-By: Claude Fable 5 --- ggml/src/ggml-cuda/topk-moe.cu | 8 +++++++- tests/test-backend-ops.cpp | 1 + 2 files changed, 8 insertions(+), 1 deletion(-) diff --git a/ggml/src/ggml-cuda/topk-moe.cu b/ggml/src/ggml-cuda/topk-moe.cu index c4253bfa43b1..8f9137bea02d 100644 --- a/ggml/src/ggml-cuda/topk-moe.cu +++ b/ggml/src/ggml-cuda/topk-moe.cu @@ -312,6 +312,10 @@ static void launch_topk_moe_cuda(ggml_backend_cuda_context & ctx, ggml_cuda_kernel_launch(topk_moe_cuda<256, has_bias>, launch_params, logits, weights, ids, bias, n_rows, n_expert_used, clamp_val, scale_val, config); break; + case 288: + ggml_cuda_kernel_launch(topk_moe_cuda<288, has_bias>, launch_params, + logits, weights, ids, bias, n_rows, n_expert_used, clamp_val, scale_val, config); + break; case 512: ggml_cuda_kernel_launch(topk_moe_cuda<512, has_bias>, launch_params, logits, weights, ids, bias, n_rows, n_expert_used, clamp_val, scale_val, config); @@ -377,8 +381,10 @@ bool ggml_cuda_should_use_topk_moe(const ggml_tensor * gating_op, const ggml_tensor * weights, const ggml_tensor * logits, const ggml_tensor * ids) { + // must match an instantiation of launch_topk_moe_cuda: a power of 2 up to 512, + // or one of the non-power-of-2 expert counts of supported models const int n_expert = ids->nb[1] / ids->nb[0]; - if (((n_expert & (n_expert - 1)) != 0 || n_expert > 512) && n_expert != 576) { + if (((n_expert & (n_expert - 1)) != 0 || n_expert > 512) && n_expert != 288 && n_expert != 576) { return false; } diff --git a/tests/test-backend-ops.cpp b/tests/test-backend-ops.cpp index 34811bd50491..e14e125ecb53 100644 --- a/tests/test-backend-ops.cpp +++ b/tests/test-backend-ops.cpp @@ -9219,6 +9219,7 @@ static std::vector> make_test_cases_eval() { test_cases.emplace_back(new test_topk_moe({128, 1, 1, 1}, 128, with_norm, bias_probs, gate, scale_w)); test_cases.emplace_back(new test_topk_moe({129, 1, 1, 1}, 128, with_norm, bias_probs, gate, scale_w)); test_cases.emplace_back(new test_topk_moe({160, 4, 1, 1}, 160, with_norm, bias_probs, gate, scale_w)); + test_cases.emplace_back(new test_topk_moe({288, 22, 1, 1}, 8, with_norm, bias_probs, gate, scale_w)); } } } From c46bc26fe96dfd92c73c1764f8dcdf838005e152 Mon Sep 17 00:00:00 2001 From: "Piotr Wilkin (ilintar)" Date: Fri, 3 Jul 2026 15:14:46 +0200 Subject: [PATCH 2/3] Update tests/test-backend-ops.cpp Co-authored-by: Oliver Simons --- tests/test-backend-ops.cpp | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/test-backend-ops.cpp b/tests/test-backend-ops.cpp index e14e125ecb53..5d1f7d8ad576 100644 --- a/tests/test-backend-ops.cpp +++ b/tests/test-backend-ops.cpp @@ -9219,7 +9219,7 @@ static std::vector> make_test_cases_eval() { test_cases.emplace_back(new test_topk_moe({128, 1, 1, 1}, 128, with_norm, bias_probs, gate, scale_w)); test_cases.emplace_back(new test_topk_moe({129, 1, 1, 1}, 128, with_norm, bias_probs, gate, scale_w)); test_cases.emplace_back(new test_topk_moe({160, 4, 1, 1}, 160, with_norm, bias_probs, gate, scale_w)); - test_cases.emplace_back(new test_topk_moe({288, 22, 1, 1}, 8, with_norm, bias_probs, gate, scale_w)); + test_cases.emplace_back(new test_topk_moe({288, 22, 1, 1}, 8, with_norm, bias_probs, gate, scale_w)); // Used by StepFun 3.7 } } } From c685c0337587a98888156020df497481edffd318 Mon Sep 17 00:00:00 2001 From: "Piotr Wilkin (ilintar)" Date: Fri, 3 Jul 2026 15:16:16 +0200 Subject: [PATCH 3/3] Add comment for case 288 in topk-moe.cu --- ggml/src/ggml-cuda/topk-moe.cu | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/ggml/src/ggml-cuda/topk-moe.cu b/ggml/src/ggml-cuda/topk-moe.cu index 8f9137bea02d..c80394e31ff3 100644 --- a/ggml/src/ggml-cuda/topk-moe.cu +++ b/ggml/src/ggml-cuda/topk-moe.cu @@ -312,7 +312,7 @@ static void launch_topk_moe_cuda(ggml_backend_cuda_context & ctx, ggml_cuda_kernel_launch(topk_moe_cuda<256, has_bias>, launch_params, logits, weights, ids, bias, n_rows, n_expert_used, clamp_val, scale_val, config); break; - case 288: + case 288: // StepFun 3.7 ggml_cuda_kernel_launch(topk_moe_cuda<288, has_bias>, launch_params, logits, weights, ids, bias, n_rows, n_expert_used, clamp_val, scale_val, config); break;