diff --git a/ggml/src/ggml-cuda/ggml-cuda.cu b/ggml/src/ggml-cuda/ggml-cuda.cu index a6fc655c41..53eccdd650 100644 --- a/ggml/src/ggml-cuda/ggml-cuda.cu +++ b/ggml/src/ggml-cuda/ggml-cuda.cu @@ -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; } diff --git a/ggml/src/ggml-cuda/mmvq.cu b/ggml/src/ggml-cuda/mmvq.cu index 79f7a3f6fe..2be2f24910 100644 --- a/ggml/src/ggml-cuda/mmvq.cu +++ b/ggml/src/ggml-cuda/mmvq.cu @@ -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 +template __launch_bounds__(get_mmvq_mmid_max_batch_for_device()*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(tmp[i]); + if constexpr (has_fusion) { + if (use_gate) { + tmp_gate[i] = warp_reduce_sum(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 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, 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, 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, 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 @@ -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( - 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 diff --git a/ggml/src/ggml-cuda/topk-moe.cu b/ggml/src/ggml-cuda/topk-moe.cu index c8cec70bb3..dadcd601cb 100644 --- a/ggml/src/ggml-cuda/topk-moe.cu +++ b/ggml/src/ggml-cuda/topk-moe.cu @@ -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 -__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(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(); diff --git a/ggml/src/ggml-cuda/topk-moe.cuh b/ggml/src/ggml-cuda/topk-moe.cuh index 091ef02a41..061b37e297 100644 --- a/ggml/src/ggml-cuda/topk-moe.cuh +++ b/ggml/src/ggml-cuda/topk-moe.cuh @@ -3,6 +3,9 @@ #include +// Rows that one CUDA block handles. +#define TOPK_MOE_ROWS_PER_BLOCK 8 + struct ggml_cuda_topk_moe_args { bool sigmoid{}; bool sqrt_softplus{}; diff --git a/tests/test-backend-ops.cpp b/tests/test-backend-ops.cpp index becf6a5552..a1875c8ed9 100644 --- a/tests/test-backend-ops.cpp +++ b/tests/test-backend-ops.cpp @@ -10205,12 +10205,10 @@ static std::vector> 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> 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)); } } }