CUDA: extend MOE fusion to specdec, earlier MOE glu fusion and topk-router fusion were restricted to 1 token (#27621)

* CUDA: extend MOE fusion to specdec, earlier MOE glu fusion and topk-router fusion were resticted to 1 token

Signed-off-by: ynankani <ynankani@nvidia.com>

* Address review comments

Signed-off-by: ynankani <ynankani@nvidia.com>

* Add SWIGLU_CLAMP case to multi-token moe fusion

Signed-off-by: ynankani <ynankani@nvidia.com>

---------

Signed-off-by: ynankani <ynankani@nvidia.com>
This commit is contained in:
ynankani
2026-08-31 19:22:28 +08:00
committed by GitHub
parent a32af33de2
commit 41ef91f7c8
5 changed files with 135 additions and 32 deletions
+8 -6
View File
@@ -10205,12 +10205,10 @@ static std::vector<std::unique_ptr<test_case>> make_test_cases_eval() {
use_id, 16, 8, b, with_bias, with_gate, with_lane_scale));
test_cases.emplace_back(new test_mul_mat_vec_fusion(type, glu_op, 1, 32, 256,
use_id, 16, 8, b, with_bias, with_gate, with_lane_scale, {1, 1}));
if (!use_id && with_gate && !with_bias && glu_op != GGML_GLU_OP_SWIGLU_CLAMP) {
// small multi-token batches (speculative decoding / MTP verify)
for (int64_t m_batch : { 2, 4, 8 }) {
test_cases.emplace_back(new test_mul_mat_vec_fusion(type, glu_op, m_batch, 32, 256,
use_id, 16, 8, b, with_bias, with_gate, with_lane_scale, {1, 1}));
}
// multi-token batches (spec decoding)
for (int64_t m_batch : { 2, 4, 8 }) {
test_cases.emplace_back(new test_mul_mat_vec_fusion(type, glu_op, m_batch, 32, 256,
use_id, 16, 8, b, with_bias, with_gate, with_lane_scale, {1, 1}));
}
}
}
@@ -10239,6 +10237,10 @@ static std::vector<std::unique_ptr<test_case>> make_test_cases_eval() {
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({256, 22, 1, 1}, 6, with_norm, bias_probs, gate, scale_w)); // Used by DeepSeek-V4
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
// rows at and just past the limit where one block still covers all rows
test_cases.emplace_back(new test_topk_moe({32, 8, 1, 1}, 4, with_norm, bias_probs, gate, scale_w));
test_cases.emplace_back(new test_topk_moe({32, 8, 1, 1}, 8, with_norm, bias_probs, gate, scale_w));
test_cases.emplace_back(new test_topk_moe({32, 9, 1, 1}, 8, with_norm, bias_probs, gate, scale_w));
}
}
}