Skip to content

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

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

Merged
am17an merged 3 commits into from
Aug 31, 2026
+135 −32
Merged

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

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
7ff3d43
CUDA: extend MOE fusion to specdec, earlier MOE glu fusion and topk-r…
ynankani Aug 23, 2026
0fc74e6
Address review comments
ynankani Aug 25, 2026
2373d10
Add SWIGLU_CLAMP case to multi-token moe fusion
ynankani Aug 31, 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
11 changes: 6 additions & 5 deletions ggml/src/ggml-cuda/ggml-cuda.cu
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1807,7 +1807,7 @@ static bool ggml_cuda_should_fuse_mul_mat_vec_q(const ggml_tensor * tensor) {
return false;
}

if (tensor->op == GGML_OP_MUL_MAT_ID && dst->ne[2] != 1) {
if (tensor->op == GGML_OP_MUL_MAT_ID && dst->ne[2] > get_mmvq_mmid_max_batch(src0->type, cc)) {
return false;
}

Expand Down Expand Up @@ -2983,9 +2983,10 @@ static bool ggml_cuda_check_fusion_memory_ranges(const ggml_cgraph * cgraph,
};

bool is_ok = true;
// exception for topk-moe, as each row is read entirely before writing
if (ggml_nrows(cgraph->nodes[node_idx]) == 1 && is_topk_moe) {
return true;
// one block reads all logits before it writes, so logits may alias the out nodes
const ggml_tensor * logits_may_alias = nullptr;
if (is_topk_moe && ggml_nrows(cgraph->nodes[node_idx]) <= TOPK_MOE_ROWS_PER_BLOCK) {
logits_may_alias = cgraph->nodes[node_idx]->src[0];
}

for (int i = 0; i < out_count; ++i) {
Expand All @@ -2999,7 +3000,7 @@ static bool ggml_cuda_check_fusion_memory_ranges(const ggml_cgraph * cgraph,
for (int src_idx = 0; src_idx < GGML_MAX_SRC; ++src_idx) {
const ggml_tensor * src = cgraph->nodes[j]->src[src_idx];

if (!src || src->op == GGML_OP_NONE) {
if (!src || src->op == GGML_OP_NONE || src == logits_may_alias) {
continue;
}

Expand Down
115 changes: 104 additions & 11 deletions ggml/src/ggml-cuda/mmvq.cu
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -773,10 +773,10 @@ static __global__ void mul_mat_vec_q(
// Grid: (ceil(nrows_x / c_rows_per_block), nchannels_dst)
// Block: (warp_size, ncols_dst) - each warp handles one token independently.
// No shared memory reduction needed since each warp works alone.
template <ggml_type type, int c_rows_per_block>
template <ggml_type type, int c_rows_per_block, bool has_fusion = false>
__launch_bounds__(get_mmvq_mmid_max_batch_for_device<type>()*ggml_cuda_get_physical_warp_size(), 1)
static __global__ void mul_mat_vec_q_moe(
const void * vx_ptr, const void * vy_ptr, const int32_t * ids_ptr,
const void * vx_ptr, const void * vy_ptr, const int32_t * ids_ptr, const ggml_cuda_mm_fusion_args_device fusion,
float * dst_ptr,
const uint32_t ncols_x, const uint3 nchannels_y, const uint32_t nrows_x,
const uint32_t stride_row_x, const uint32_t stride_col_y, const uint32_t stride_col_dst,
Expand All @@ -794,6 +794,29 @@ static __global__ void mul_mat_vec_q_moe(

constexpr vec_dot_q_cuda_t vec_dot_q_cuda = get_vec_dot_q_cuda(type);

// fuse gate, bias, scales, and glu_op into the up projection
bool use_gate = false;
const void * vgate = nullptr;
const float * x_bias = nullptr;
const float * gate_bias = nullptr;
const float * x_scale = nullptr;
const float * gate_scale = nullptr;
ggml_glu_op active_glu = GGML_GLU_OP_SWIGLU;
float glu_limit = 0.0f;

if constexpr (has_fusion) {
use_gate = fusion.gate != nullptr;
vgate = fusion.gate;
x_bias = (const float *) fusion.x_bias;
gate_bias = (const float *) fusion.gate_bias;
active_glu = fusion.glu_op;
glu_limit = fusion.glu_limit;
if constexpr (type == GGML_TYPE_NVFP4) {
x_scale = (const float *) fusion.x_scale;
gate_scale = (const float *) fusion.gate_scale;
}
}

const uint32_t token_idx = threadIdx.y;
const int row0 = c_rows_per_block*blockIdx.x;
const int blocks_per_row_x = ncols_x / qk;
Expand All @@ -814,6 +837,7 @@ static __global__ void mul_mat_vec_q_moe(

// partial sum for each thread
float tmp[c_rows_per_block] = {0.0f};
float tmp_gate[c_rows_per_block] = {0.0f};

for (int kbx = threadIdx.x / (qi/vdr); kbx < blocks_per_row_x; kbx += blocks_per_iter) {
const int kby = kbx * (qk/QK8_1);
Expand All @@ -822,6 +846,11 @@ static __global__ void mul_mat_vec_q_moe(
#pragma unroll
for (int i = 0; i < c_rows_per_block; ++i) {
tmp[i] += vec_dot_q_cuda(vx, &y[kby], kbx_offset + i*stride_row_x + kbx, kqs);
if constexpr (has_fusion) {
if (use_gate) {
tmp_gate[i] += vec_dot_q_cuda(vgate, &y[kby], kbx_offset + i*stride_row_x + kbx, kqs);
}
}
}
}

Expand All @@ -831,11 +860,63 @@ static __global__ void mul_mat_vec_q_moe(
#pragma unroll
for (int i = 0; i < c_rows_per_block; ++i) {
tmp[i] = warp_reduce_sum<warp_size>(tmp[i]);
if constexpr (has_fusion) {
if (use_gate) {
tmp_gate[i] = warp_reduce_sum<warp_size>(tmp_gate[i]);
}
}
}

// Write results
if (threadIdx.x < c_rows_per_block && (c_rows_per_block == 1 || uint32_t(row0 + threadIdx.x) < nrows_x)) {
dst[channel_dst*stride_channel_dst + token_idx*stride_col_dst + row0 + threadIdx.x] = tmp[threadIdx.x];
float result = tmp[threadIdx.x];
if constexpr (has_fusion) {
const uint32_t bias_idx = channel_x*stride_channel_dst + row0 + threadIdx.x;

if constexpr (type == GGML_TYPE_NVFP4) {
if (x_scale) {
result *= x_scale[channel_x];
}
}
if (x_bias) {
result += x_bias[bias_idx];
}
if (use_gate) {
float gate_value = tmp_gate[threadIdx.x];
Comment thread
ORippler marked this conversation as resolved.
if constexpr (type == GGML_TYPE_NVFP4) {
if (gate_scale) {
gate_value *= gate_scale[channel_x];
}
}
if (gate_bias) {
gate_value += gate_bias[bias_idx];
}
switch (active_glu) {
case GGML_GLU_OP_SWIGLU:
result *= ggml_cuda_op_silu_single(gate_value);
break;
case GGML_GLU_OP_GEGLU:
result *= ggml_cuda_op_gelu_single(gate_value);
break;
case GGML_GLU_OP_SWIGLU_OAI:
result = ggml_cuda_op_swiglu_oai_single(gate_value, result);
break;
case GGML_GLU_OP_SWIGLU_CLAMP:
result = ggml_cuda_op_swiglu_clamp_single(gate_value, result, glu_limit);
break;
default:
result = result * gate_value;
Comment on lines +894 to +908

@am17an am17an Aug 24, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

let's also add sqrtsoftplus used for deepseek4

@ynankani ynankani Aug 24, 2026 •
edited
Loading

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I might be missing something here, I don't see sqrtsoftplus in the glu op enum

llama.cpp/ggml/include/ggml.h

Line 623 in c060ca9

enum ggml_glu_op {
. Also sqrt_softplus is already fused in topk-moe right?

llama.cpp/ggml/src/ggml-cuda/ggml-cuda.cu

Line 2825 in c060ca9

} else if (unary_op == GGML_UNARY_OP_SOFTPLUS && node_idx + 1 < n_nodes &&

@am17an am17an Aug 24, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Ah okay, then that can be done in a follow-up PR. I was thinking I already added it in #25896

break;
}
}
}
dst[channel_dst*stride_channel_dst + token_idx*stride_col_dst + row0 + threadIdx.x] = result;
}

if constexpr (!has_fusion) {
GGML_UNUSED_VARS(use_gate, tmp_gate, vgate, x_bias, gate_bias, active_glu, glu_limit, x_scale, gate_scale);
} else if constexpr (type != GGML_TYPE_NVFP4) {
GGML_UNUSED_VARS(x_scale, gate_scale);
}
}

Expand Down Expand Up @@ -885,7 +966,7 @@ static void mul_mat_vec_q_switch_fusion(

template <ggml_type type>
static void mul_mat_vec_q_moe_launch(
const void * vx, const void * vy, const int32_t * ids, float * dst,
const void * vx, const void * vy, const int32_t * ids, const ggml_cuda_mm_fusion_args_device fusion, float * dst,
const uint32_t ncols_x, const uint3 nchannels_y, const uint32_t nrows_x,
const uint32_t stride_row_x, const uint32_t stride_col_y, const uint32_t stride_col_dst,
const uint32_t stride_channel_x, const uint32_t stride_channel_y, const uint32_t stride_channel_dst,
Expand All @@ -898,11 +979,22 @@ static void mul_mat_vec_q_moe_launch(
const dim3 block_dims(warp_size, ncols_dst);
const ggml_cuda_kernel_launch_params launch_params = ggml_cuda_kernel_launch_params(block_nums, block_dims, 0, stream);

ggml_cuda_kernel_launch(mul_mat_vec_q_moe<type, rows_per_block>, launch_params,
vx, vy, ids, dst, ncols_x, nchannels_y, nrows_x,
stride_row_x, stride_col_y, stride_col_dst,
stride_channel_x, stride_channel_y, stride_channel_dst,
ncols_dst, ids_stride);
const bool has_fusion = fusion.gate != nullptr || fusion.x_bias != nullptr || fusion.gate_bias != nullptr ||
fusion.x_scale != nullptr || fusion.gate_scale != nullptr;

if (has_fusion) {
ggml_cuda_kernel_launch(mul_mat_vec_q_moe<type, rows_per_block, true>, launch_params,
vx, vy, ids, fusion, dst, ncols_x, nchannels_y, nrows_x,
stride_row_x, stride_col_y, stride_col_dst,
stride_channel_x, stride_channel_y, stride_channel_dst,
ncols_dst, ids_stride);
} else {
ggml_cuda_kernel_launch(mul_mat_vec_q_moe<type, rows_per_block, false>, launch_params,
vx, vy, ids, fusion, dst, ncols_x, nchannels_y, nrows_x,
stride_row_x, stride_col_y, stride_col_dst,
stride_channel_x, stride_channel_y, stride_channel_dst,
ncols_dst, ids_stride);
}
}

template <ggml_type type>
Expand Down Expand Up @@ -998,7 +1090,7 @@ static void mul_mat_vec_q_switch_ncols_dst(
if (has_ids && ncols_dst > 1) {
// Multi-token MUL_MAT_ID path - dedicated MoE kernel
mul_mat_vec_q_moe_launch<type>(
vx, vy, ids, dst, ncols_x, nchannels_y_fd, nrows_x,
vx, vy, ids, fusion, dst, ncols_x, nchannels_y_fd, nrows_x,
stride_row_x, stride_col_y, stride_col_dst,
stride_channel_x, stride_channel_y, stride_channel_dst,
ncols_dst, ids_stride, warp_size, nchannels_dst, stream);
Expand Down Expand Up @@ -1280,7 +1372,8 @@ void ggml_cuda_mul_mat_vec_q(
ggml_cuda_mm_fusion_args_device fusion_local{};

if (fusion) {
GGML_ASSERT( !ids || dst->ne[2] == 1);
const int cc = ggml_cuda_info().devices[ggml_cuda_get_device()].cc;
GGML_ASSERT( !ids || dst->ne[2] <= get_mmvq_mmid_max_batch(src0->type, cc));
GGML_ASSERT( ids || dst->ne[1] == 1);
// Scale fusion is only allowed for NVFP4 currently as the cost of checking this at run-time in the prologue is
// non-negligible for some models such as gpt-oss-20b
Expand Down
24 changes: 14 additions & 10 deletions ggml/src/ggml-cuda/topk-moe.cu
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -88,15 +88,16 @@ __device__ void sqrt_softplus_warp_inplace(float (&vals)[experts_per_thread], co
It is intended as fusion of softmax->top-k->get_rows pipeline for MoE models
*/
template <int n_experts, bool has_bias>
__launch_bounds__(4 * WARP_SIZE, 1) __global__ void topk_moe_cuda(const float * logits,
float * weights,
int32_t * ids,
float * bias,
const int n_rows,
const int n_expert_used,
const float clamp_val,
const float scale_val,
const topk_moe_config config) {
__launch_bounds__(TOPK_MOE_ROWS_PER_BLOCK * WARP_SIZE, 1)
__global__ void topk_moe_cuda(const float * logits,
float * weights,
int32_t * ids,
float * bias,
const int n_rows,
const int n_expert_used,
const float clamp_val,
const float scale_val,
const topk_moe_config config) {
const int row = blockIdx.x * blockDim.y + threadIdx.y;
if (row >= n_rows) {
return;
Expand All @@ -123,6 +124,9 @@ __launch_bounds__(4 * WARP_SIZE, 1) __global__ void topk_moe_cuda(const float *
wt[i / WARP_SIZE] = (n_experts % WARP_SIZE == 0 || expert < n_experts) ? logits[expert] : -INFINITY;
}

// Weights and IDs can alias logits, so wait until every row in the block reads its logits.
__syncthreads();

if (!config.delayed_softmax) {
if (config.use_sigmoid) {
sigmoid_warp_inplace<experts_per_thread, false>(wt, n_experts, threadIdx.x);
Expand Down Expand Up @@ -282,7 +286,7 @@ static void launch_topk_moe_cuda(ggml_backend_cuda_context & ctx,
const topk_moe_config config) {
GGML_ASSERT(!(config.with_norm && config.delayed_softmax) &&
"delayed softmax is not supported with weight normalization");
const int rows_per_block = 4;
const int rows_per_block = TOPK_MOE_ROWS_PER_BLOCK;
dim3 grid_dims((n_rows + rows_per_block - 1) / rows_per_block, 1, 1);
dim3 block_dims(WARP_SIZE, rows_per_block, 1);
cudaStream_t stream = ctx.stream();
Expand Down
3 changes: 3 additions & 0 deletions ggml/src/ggml-cuda/topk-moe.cuh
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -3,6 +3,9 @@

#include <initializer_list>

// Rows that one CUDA block handles.
#define TOPK_MOE_ROWS_PER_BLOCK 8

struct ggml_cuda_topk_moe_args {
bool sigmoid{};
bool sqrt_softplus{};
Expand Down
14 changes: 8 additions & 6 deletions tests/test-backend-ops.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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}));
}
}
}
Expand Down Expand Up @@ -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));
Comment thread
ynankani marked this conversation as resolved.
test_cases.emplace_back(new test_topk_moe({32, 9, 1, 1}, 8, with_norm, bias_probs, gate, scale_w));
}
}
}
Expand Down
Loading