wylab/llama.cpp
0
0
Fork 1
mirror of https://github.com/ggml-org/llama.cpp.git synced 2026-08-31 17:17:44 +02:00
Code Issues Packages Projects Releases Wiki Activity

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

Browse Source 
* 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
Download Patch File Download Diff File Expand all files Collapse all files
ggml/src/ggml-cuda/ggml-cuda.cu
+6 -5
View File
@@ -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;
}
@@ -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) {
@@ -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;
}
ggml/src/ggml-cuda/mmvq.cu
+104 -11
View File
@@ -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,
@@ -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;
@@ -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);
@@ -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);
}
}
}
}
@@ -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];
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;
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);
}
}
@@ -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,
@@ -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>
@@ -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);
@@ -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
ggml/src/ggml-cuda/topk-moe.cu
+14 -10
View File
@@ -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;
@@ -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);
@@ -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();
ggml/src/ggml-cuda/topk-moe.cuh
+3
View File
@@ -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{};
tests/test-backend-ops.cpp
+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));
}
}
}