Lumpiasty/llama.cpp
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

opencl: make the MoE expert scatter deterministic (#26464)

Browse Source
This commit is contained in:
Hongqiang Wang
2026-08-19 20:40:19 -07:00
committed by GitHub
parent d59d455fd8
commit 9ee9fc04c1
2 changed files with 110 additions and 11 deletions
Download Patch File Download Diff File Expand all files Collapse all files
ggml/src/ggml-opencl/ggml-opencl.cpp
+37 -11
View File
@@ -895,6 +895,7 @@ struct ggml_backend_opencl_context {
cl_kernel kernel_gemm_moe_q4_0_q8_1_dp4a = nullptr; // dp4a (int8) q4_0 MoE prefill GEMM
cl_kernel kernel_moe_reorder_b;
cl_kernel kernel_moe_histogram, kernel_moe_scan, kernel_moe_fill, kernel_moe_scatter;
cl_kernel kernel_moe_scatter_stable = nullptr; // deterministic slot assignment
cl_kernel kernel_moe_combine_f32 = nullptr; // fused router-weight mul + cross-expert sum
cl_kernel kernel_mul_mv_id_q4_0_f32_8x_flat;
cl_kernel kernel_mul_mv_id_q8_0_f32, kernel_mul_mv_id_q8_0_f32_flat;
@@ -4463,6 +4464,7 @@ static void load_cl_kernels(ggml_backend_opencl_context *backend_ctx) {
CL_CHECK((backend_ctx->kernel_moe_scan = clCreateKernel(prog, "kernel_moe_scan", &err), err));
CL_CHECK((backend_ctx->kernel_moe_fill = clCreateKernel(prog, "kernel_moe_fill", &err), err));
CL_CHECK((backend_ctx->kernel_moe_scatter = clCreateKernel(prog, "kernel_moe_scatter", &err), err));
CL_CHECK((backend_ctx->kernel_moe_scatter_stable = clCreateKernel(prog, "kernel_moe_scatter_stable", &err), err));
CL_CHECK(clReleaseProgram(prog));
GGML_LOG_CONT(".");
}
@@ -20863,18 +20865,42 @@ static void moe_router_reoerder(ggml_backend_t backend, const ggml_tensor * src,
size_t fill_local_size[] = {64, 1, 1};
backend_ctx->enqueue_ndrange_kernel(kernel, 3, fill_global_size, fill_local_size, src);
// Scatter
kernel = backend_ctx->kernel_moe_scatter;
CL_CHECK(clSetKernelArg(kernel, 0, sizeof(cl_mem), &original_router_buf));
CL_CHECK(clSetKernelArg(kernel, 1, sizeof(cl_mem), &post_router_buf));
CL_CHECK(clSetKernelArg(kernel, 2, sizeof(cl_mem), &emap_buf));
CL_CHECK(clSetKernelArg(kernel, 3, sizeof(cl_mem), &tile_offset_buf));
CL_CHECK(clSetKernelArg(kernel, 4, sizeof(cl_mem), &slot_counter_buf));
CL_CHECK(clSetKernelArg(kernel, 5, sizeof(int), &ne21));
CL_CHECK(clSetKernelArg(kernel, 6, sizeof(int), &ne20));
CL_CHECK(clSetKernelArg(kernel, 7, sizeof(int), &ne02));
// Scatter. The deterministic variant is the default: kernel_moe_scatter derives
// each token's slot from an atomic counter, so the packing inside an expert - and
// with it the output of the ragged prefill GEMM - changes from run to run. Set
// GGML_OPENCL_MOE_STABLE_SCATTER=0 to restore the atomic version.
static const bool stable_scatter = []{
const char * e = getenv("GGML_OPENCL_MOE_STABLE_SCATTER");
return !e || e[0] == '\0' || e[0] != '0';
}();
backend_ctx->enqueue_ndrange_kernel(kernel, 3, histogram_global_size, histogram_local_size, src);
if (stable_scatter) {
kernel = backend_ctx->kernel_moe_scatter_stable;
CL_CHECK(clSetKernelArg(kernel, 0, sizeof(cl_mem), &original_router_buf));
CL_CHECK(clSetKernelArg(kernel, 1, sizeof(cl_mem), &post_router_buf));
CL_CHECK(clSetKernelArg(kernel, 2, sizeof(cl_mem), &emap_buf));
CL_CHECK(clSetKernelArg(kernel, 3, sizeof(cl_mem), &tile_offset_buf));
CL_CHECK(clSetKernelArg(kernel, 4, sizeof(int), &ne21));
CL_CHECK(clSetKernelArg(kernel, 5, sizeof(int), &ne20));
CL_CHECK(clSetKernelArg(kernel, 6, sizeof(int), &ne02));
// one workgroup (one wave) per expert; each ranks its own tokens
size_t scatter_global_size[] = {64, (size_t)ne02};
size_t scatter_local_size[] = {64, 1};
backend_ctx->enqueue_ndrange_kernel(kernel, 2, scatter_global_size, scatter_local_size, src);
} else {
kernel = backend_ctx->kernel_moe_scatter;
CL_CHECK(clSetKernelArg(kernel, 0, sizeof(cl_mem), &original_router_buf));
CL_CHECK(clSetKernelArg(kernel, 1, sizeof(cl_mem), &post_router_buf));
CL_CHECK(clSetKernelArg(kernel, 2, sizeof(cl_mem), &emap_buf));
CL_CHECK(clSetKernelArg(kernel, 3, sizeof(cl_mem), &tile_offset_buf));
CL_CHECK(clSetKernelArg(kernel, 4, sizeof(cl_mem), &slot_counter_buf));
CL_CHECK(clSetKernelArg(kernel, 5, sizeof(int), &ne21));
CL_CHECK(clSetKernelArg(kernel, 6, sizeof(int), &ne20));
CL_CHECK(clSetKernelArg(kernel, 7, sizeof(int), &ne02));
backend_ctx->enqueue_ndrange_kernel(kernel, 3, histogram_global_size, histogram_local_size, src);
}
// [MOE_TILES] env-gated padding probe: read back total_tiles (= Sum_e
// ceil(k_e/n_tile_size)) and compare to the ideal tile count for the real