opencl: choose the MoE expert matmul by batch size for speculative decoding/MTP (#27637)

* opencl: gate the prebuilt q4_0 MoE GEMM on routing count

* opencl: stop writing zeros into the padded MoE activation slots

* opencl: rephrase claude's comments

---------

Co-authored-by: Li He <lih@qti.qualcomm.com>
This commit is contained in:
Hongqiang Wang
2026-09-15 11:21:05 -07:00
committed by GitHub
parent d1d3c3396a
commit 9f31776c37
2 changed files with 35 additions and 22 deletions
+28 -17
View File
@@ -24566,10 +24566,33 @@ static void ggml_cl_mul_mat_id(ggml_backend_t backend, const ggml_tensor * src0,
CL_CHECK(clReleaseMemObject(buf_src2));
} else { // for gemm
kernel = backend_ctx->kernel_gemm_moe_q4_0_f32_ns;
if (backend_ctx->kernel_gemm_moe_q4_0_f32_ns_bin) {
kernel = backend_ctx->kernel_gemm_moe_q4_0_f32_ns_bin;
}
// dp4a (int8) prefill GEMM variant
static const char * q4_0_moe_dp4a_env = getenv("GGML_OPENCL_Q4_0_MOE_DP4A");
// It turns out that the prebuilt kernel only outperforms the dp4a variant (on X2-90)
// at very large routing counts, so we gate its use accordingly using moe_bin_min,
// which can be overridden via the GGML_OPENCL_MOE_BIN_MIN_ROUTINGS environment variable.
// The routing count is ne20 * ne21 (n_expert_used * n_tokens).
static const char * moe_bin_min_env = getenv("GGML_OPENCL_MOE_BIN_MIN_ROUTINGS");
const int moe_bin_min = moe_bin_min_env ? atoi(moe_bin_min_env) : 4096;
// whether bin kernels are available
const bool bin_available = backend_ctx->kernel_gemm_moe_q4_0_f32_ns_bin != nullptr;
const bool dp4a_bin_available = backend_ctx->kernel_gemm_moe_q4_0_q8_1_dp4a_bin != nullptr;
bool use_moe_dp4a = q4_0_moe_dp4a_env
? (atoi(q4_0_moe_dp4a_env) != 0)
: (backend_ctx->adreno_gen == ADRENO_GPU_GEN::X2E
&& (dp4a_bin_available || !bin_available
|| (int)(ne20 * ne21) < moe_bin_min));
// dot prod has to be available
use_moe_dp4a = backend_ctx->has_integer_dot && use_moe_dp4a;
const bool use_bin_kernel = bin_available && !use_moe_dp4a;
kernel = use_bin_kernel
? backend_ctx->kernel_gemm_moe_q4_0_f32_ns_bin
: backend_ctx->kernel_gemm_moe_q4_0_f32_ns;
// Reorder router if called from test-backend-ops or when new router is generated.
// Otherwise reuse the reordered result from previous mul_mat_id call.
@@ -24582,18 +24605,6 @@ static void ggml_cl_mul_mat_id(ggml_backend_t backend, const ggml_tensor * src0,
cl_mem buf_src1_reordered = nullptr, image_src1_reordered = nullptr;
cl_mem buf_src2, buf_src2_emap;
// dp4a (int8) prefill GEMM variant
static const char * q4_0_moe_dp4a_env = getenv("GGML_OPENCL_Q4_0_MOE_DP4A");
bool use_moe_dp4a = q4_0_moe_dp4a_env
? (atoi(q4_0_moe_dp4a_env) != 0)
: (backend_ctx->adreno_gen == ADRENO_GPU_GEN::X2E);
// dot prod has to be available
use_moe_dp4a = backend_ctx->has_integer_dot && use_moe_dp4a;
// bin kernel takes precedence
if (backend_ctx->kernel_gemm_moe_q4_0_q8_1_dp4a_bin == nullptr) {
use_moe_dp4a = use_moe_dp4a && backend_ctx->kernel_gemm_moe_q4_0_f32_ns_bin == nullptr;
}
cl_buffer_region region;
region.origin = 0;
region.size = sizeof(int) * max_post_router_tile * n_tile_size;
@@ -24632,7 +24643,7 @@ static void ggml_cl_mul_mat_id(ggml_backend_t backend, const ggml_tensor * src0,
cl_image_desc image_desc_buf_src1;
image_format_buf_src1 = {CL_RGBA, CL_FLOAT};
image_desc_buf_src1 = {CL_MEM_OBJECT_IMAGE1D_BUFFER, static_cast<size_t>(ne00 * max_post_router_tile * n_tile_size / 4), 0,0,0,0,0,0,0, {buf_src1_reordered}};
if (backend_ctx->kernel_gemm_moe_q4_0_f32_ns_bin) {
if (use_bin_kernel) {
// bin kernel uses slightly different image format
image_format_buf_src1 = {CL_R, CL_FLOAT};
image_desc_buf_src1.image_width = static_cast<size_t>(ne00 * max_post_router_tile * n_tile_size);
@@ -20,11 +20,13 @@ kernel void kernel_moe_reorder_b(
uint router_idx = router[post_router_idx];
float4 out = (float4)(0);
if (router_idx != 0xFFFFFFFF) {
ushort activation_idx = router_idx / map_ratio;
out = src[activation_idx * K / 4 + k_4];
// Padded slots need not be written at all. The MoE GEMMs accumulate per output
// column and scatter only the real columns, so whatever sits in a padded slot
// never reaches dst
if (router_idx == 0xFFFFFFFF) {
return;
}
dst[post_router_idx * K / 4 + k_4] = out;
ushort activation_idx = router_idx / map_ratio;
dst[post_router_idx * K / 4 + k_4] = src[activation_idx * K / 4 + k_4];
}