Commit 701fd25b3 for llama.cpp
commit 701fd25b3a55a90297dec0d925204b29084054c8
Author: Xuan-Son Nguyen <son@huggingface.co>
Date: Sat Oct 24 12:15:36 2026 +0200
ggml : allow MUL_MAT_ID -1 index to skip the computation (#26631)
* ggml : allow MUL_MAT_ID -1 index to skip the computation
* nits
* update vulkan code
* ggml-webgpu: ensure all threads reach workgroup barriers in mul_mat_id_vec
(cherry picked from commit d8a4b9efb6e2b9abf3a866e847756b7e6bbc6b69)
* fix cuda
* openvino
* exclude CUDA from this PR
* add TODO tag
---------
Co-authored-by: yomaytk <yoshimura.masashi.frbs@gmail.com>
diff --git a/ggml/include/ggml.h b/ggml/include/ggml.h
index e6ac1af1b..d62a63104 100644
--- a/ggml/include/ggml.h
+++ b/ggml/include/ggml.h
@@ -926,6 +926,7 @@ extern "C" {
enum ggml_type type);
// dst[i0, i1, i2] = a[i0, i1, i2] + b[i0, ids[i1, i2]]
+ // an id of -1 adds nothing: dst[i0, i1, i2] = a[i0, i1, i2]
GGML_API struct ggml_tensor * ggml_add_id(
struct ggml_context * ctx,
struct ggml_tensor * a,
@@ -1508,6 +1509,7 @@ extern "C" {
enum ggml_op_hint hint);
// indirect matrix multiplication
+ // an id of -1 skips the slot and zeroes the matching row of the result
//
// [TAG_MUL_MAT_ID_DUPLICATE_ID]
// note: this op has undefined behaviour if a given `a` is referenced more than once for a given `b`
diff --git a/ggml/src/ggml-cpu/ggml-cpu.c b/ggml/src/ggml-cpu/ggml-cpu.c
index 4c33369ee..53d39d249 100644
--- a/ggml/src/ggml-cpu/ggml-cpu.c
+++ b/ggml/src/ggml-cpu/ggml-cpu.c
@@ -1652,7 +1652,12 @@ static void ggml_compute_forward_mul_mat_id(
for (int id = 0; id < n_ids; ++id) {
const int32_t i02 = *(const int32_t *) ((const char *) ids->data + iid1*ids->nb[1] + id*ids->nb[0]);
- assert(i02 >= 0 && i02 < n_as);
+ assert(i02 == -1 || (i02 >= 0 && i02 < n_as));
+
+ if (i02 == -1) {
+ memset((char *) dst->data + id*nb1 + iid1*nb2, 0, ne0*sizeof(float));
+ continue;
+ }
MMID_MATRIX_ROW(i02, matrix_row_counts[i02]) = (struct mmid_row_mapping) {id, iid1};
matrix_row_counts[i02] += 1;
diff --git a/ggml/src/ggml-cpu/ops.cpp b/ggml/src/ggml-cpu/ops.cpp
index 02d5efffa..5735f96f8 100644
--- a/ggml/src/ggml-cpu/ops.cpp
+++ b/ggml/src/ggml-cpu/ops.cpp
@@ -743,12 +743,17 @@ static void ggml_compute_forward_add_id_f32(
// src1 indices
const int i11 = *(int32_t *) ((char *) src2->data + i1*nb20 + i2*nb21);
- GGML_ASSERT(i11 >= 0 && i11 < ne11);
+ GGML_ASSERT(i11 == -1 || (i11 >= 0 && i11 < ne11));
- ggml_vec_add_f32(ne0,
- (float *) ((char *) dst->data + i3*nb3 + i2*nb2 + i1*nb1 ),
- (float *) ((char *) src0->data + i3*nb03 + i2*nb02 + i1*nb01),
- (float *) ((char *) src1->data + i11*nb11));
+ float * dst_row = (float *) ((char *) dst->data + i3*nb3 + i2*nb2 + i1*nb1 );
+ const float * src0_row = (float *) ((char *) src0->data + i3*nb03 + i2*nb02 + i1*nb01);
+
+ if (i11 == -1) {
+ memcpy(dst_row, src0_row, ne0*sizeof(float));
+ continue;
+ }
+
+ ggml_vec_add_f32(ne0, dst_row, src0_row, (float *) ((char *) src1->data + i11*nb11));
}
}
diff --git a/ggml/src/ggml-cpu/repack.cpp b/ggml/src/ggml-cpu/repack.cpp
index 2632ae3a0..9282ca299 100644
--- a/ggml/src/ggml-cpu/repack.cpp
+++ b/ggml/src/ggml-cpu/repack.cpp
@@ -4855,7 +4855,12 @@ template <typename BLOC_TYPE, int64_t INTER_SIZE, int64_t NB_COLS, ggml_type PAR
const int32_t i02 =
*(const int32_t *) ((const char *) ids->data + iid1 * ids->nb[1] + id * ids->nb[0]);
- GGML_ASSERT(i02 >= 0 && i02 < n_as);
+ GGML_ASSERT(i02 == -1 || (i02 >= 0 && i02 < n_as));
+
+ if (i02 == -1) {
+ memset((char *) dst->data + id*nb1 + iid1*nb2, 0, ne0*sizeof(float));
+ continue;
+ }
MMID_MATRIX_ROW(i02, matrix_row_counts[i02]) = { id, iid1 };
matrix_row_counts[i02] += 1;
diff --git a/ggml/src/ggml-cpu/spacemit/ime.cpp b/ggml/src/ggml-cpu/spacemit/ime.cpp
index 81e86b9dd..7b8012755 100644
--- a/ggml/src/ggml-cpu/spacemit/ime.cpp
+++ b/ggml/src/ggml-cpu/spacemit/ime.cpp
@@ -682,7 +682,12 @@ template <typename BLOC_TYPE, int64_t INTER_SIZE, int64_t NB_COLS> class tensor_
const int32_t i02 =
*(const int32_t *) ((const char *) ids->data + iid1 * ids->nb[1] + id * ids->nb[0]);
- GGML_ASSERT(i02 >= 0 && i02 < n_as);
+ GGML_ASSERT(i02 == -1 || (i02 >= 0 && i02 < n_as));
+
+ if (i02 == -1) {
+ memset((char *) dst->data + id * nb1 + iid1 * nb2, 0, ne0 * sizeof(float));
+ continue;
+ }
MMID_MATRIX_ROW(i02, matrix_row_counts[i02]) = { id, iid1 };
matrix_row_counts[i02] += 1;
diff --git a/ggml/src/ggml-hexagon/htp/matmul-ops.c b/ggml/src/ggml-hexagon/htp/matmul-ops.c
index e290835fc..eb40fa642 100644
--- a/ggml/src/ggml-hexagon/htp/matmul-ops.c
+++ b/ggml/src/ggml-hexagon/htp/matmul-ops.c
@@ -1413,10 +1413,10 @@ static void hvx_mv_id_tiled(unsigned int nth, unsigned int ith, void * data) {
for (uint32_t ie1 = 0; ie1 < n_aids; ++ie1) { // for each expert
const int32_t eid = *(const int32_t *) ((const uint8_t *) src2->data + ie1 * src2->nb[0]);
- if (eid < 0) {
+ if (eid == -1) {
continue;
}
- assert(eid < (int32_t) n_ids);
+ assert(eid >= 0 && eid < (int32_t) n_ids);
const dma_addr_t src0_row = src0->data + eid * nb02;
const uint8_t * restrict act_col = (const uint8_t *) act_data;
@@ -4166,10 +4166,10 @@ static inline void scan_expert_ids_n(
const int32_t * row_ptr = (const int32_t *) (ids_data + iid1 * ids_nb1);
for (uint32_t id = 0; id < n_ids; ++id) {
const int32_t i02 = row_ptr[id];
- if (i02 < 0) {
+ if (i02 == -1) {
continue;
}
- assert(i02 < n_as);
+ assert(i02 >= 0 && i02 < n_as);
if (matrix_rows) {
matrix_rows[i02 * mapping_stride + counts[i02]] = (struct mmid_row_mapping) { id, iid1 };
@@ -4179,6 +4179,28 @@ static inline void scan_expert_ids_n(
}
}
+// the matmuls only write the rows that an expert owns, so the rows of the skipped slots are zeroed here
+static inline void zero_skipped_dst_rows(
+ const struct htp_tensor * ids,
+ uint32_t n_ids,
+ const struct htp_tensor * dst,
+ uint32_t ne0,
+ size_t nb1,
+ size_t nb2
+) {
+ const uint8_t * ids_data = (const uint8_t *) ids->data;
+
+ for (uint32_t iid1 = 0; iid1 < ids->ne[1]; ++iid1) {
+ for (uint32_t id = 0; id < n_ids; ++id) {
+ const int32_t i02 = *(const int32_t *) (ids_data + iid1 * ids->nb[1] + id * ids->nb[0]);
+ if (i02 != -1) {
+ continue;
+ }
+ memset((uint8_t *) (dst->data + id * nb1 + iid1 * nb2), 0, ne0 * sizeof(float));
+ }
+ }
+}
+
static inline void scan_expert_ids(
const struct htp_tensor * ids,
uint32_t n_ids,
@@ -4204,10 +4226,10 @@ static inline void scan_expert_ids(
const int32_t * row_ptr = (const int32_t *) (ids_data + iid1 * ids_nb1);
for (uint32_t id = 0; id < n_ids; ++id) {
const int32_t i02 = *(const int32_t *) ((const uint8_t *) row_ptr + id * ids_nb0);
- if (i02 < 0) {
+ if (i02 == -1) {
continue;
}
- assert(i02 < n_as);
+ assert(i02 >= 0 && i02 < n_as);
if (matrix_rows) {
matrix_rows[i02 * mapping_stride + counts[i02]] = (struct mmid_row_mapping) { id, iid1 };
@@ -4287,6 +4309,8 @@ int op_matmul_id(struct htp_ops_context * octx) {
scan_expert_ids(ids, n_ids, n_as, matrix_row_counts, matrix_rows, mapping_stride);
}
+ zero_skipped_dst_rows(ids, n_ids, dst, ne0, nb1, nb2);
+
mmctx->matrix_row_counts = matrix_row_counts;
mmctx->matrix_rows = matrix_rows;
mmctx->mapping_stride = mapping_stride;
diff --git a/ggml/src/ggml-metal/ggml-metal-device.cpp b/ggml/src/ggml-metal/ggml-metal-device.cpp
index 65611314a..1a66f112b 100644
--- a/ggml/src/ggml-metal/ggml-metal-device.cpp
+++ b/ggml/src/ggml-metal/ggml-metal-device.cpp
@@ -1283,6 +1283,15 @@ ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_mul_mm_id_map0(g
return res;
}
+ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_mul_mm_id_zero(ggml_metal_library_t lib) {
+ ggml_metal_pipeline_with_params res = ggml_metal_library_get_pipeline(lib, "kernel_mul_mm_id_zero");
+ if (!res.pipeline) {
+ res = ggml_metal_library_compile_pipeline(lib, "kernel_mul_mm_id_zero", "kernel_mul_mm_id_zero", nullptr);
+ }
+
+ return res;
+}
+
ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_mul_mm_id(ggml_metal_library_t lib, const ggml_tensor * op) {
char base[256];
char name[256];
diff --git a/ggml/src/ggml-metal/ggml-metal-device.h b/ggml/src/ggml-metal/ggml-metal-device.h
index 5ccaec6a3..72a8f92f2 100644
--- a/ggml/src/ggml-metal/ggml-metal-device.h
+++ b/ggml/src/ggml-metal/ggml-metal-device.h
@@ -140,6 +140,7 @@ struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_mul_mv_mm
struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_mul_mm (ggml_metal_library_t lib, const struct ggml_tensor * op);
struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_mul_mv (ggml_metal_library_t lib, const struct ggml_tensor * op);
struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_mul_mm_id_map0 (ggml_metal_library_t lib, int ne02, int ne20);
+struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_mul_mm_id_zero (ggml_metal_library_t lib);
struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_mul_mm_id_amax(ggml_metal_library_t lib);
struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_mul_mm_id_amax_part(ggml_metal_library_t lib);
struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_mul_mm_id (ggml_metal_library_t lib, const struct ggml_tensor * op);
diff --git a/ggml/src/ggml-metal/ggml-metal-impl.h b/ggml/src/ggml-metal/ggml-metal-impl.h
index 46b56c9d9..ea6a677f9 100644
--- a/ggml/src/ggml-metal/ggml-metal-impl.h
+++ b/ggml/src/ggml-metal/ggml-metal-impl.h
@@ -572,6 +572,13 @@ typedef struct {
uint64_t nb21;
} ggml_metal_kargs_mul_mm_id_map0;
+typedef struct {
+ int32_t ne0;
+ uint64_t nb1;
+ uint64_t nb2;
+ uint64_t nb21;
+} ggml_metal_kargs_mul_mm_id_zero;
+
typedef struct {
int32_t ne00;
int32_t ne01;
diff --git a/ggml/src/ggml-metal/ggml-metal-ops.cpp b/ggml/src/ggml-metal/ggml-metal-ops.cpp
index fc856b5b2..589f7fafa 100644
--- a/ggml/src/ggml-metal/ggml-metal-ops.cpp
+++ b/ggml/src/ggml-metal/ggml-metal-ops.cpp
@@ -2935,6 +2935,25 @@ int ggml_metal_op_mul_mat_id(ggml_metal_op_t ctx, int idx) {
ggml_metal_encoder_dispatch_threadgroups(enc, 1, 1, 1, ne02, 1, 1);
}
+ // zero the rows that no expert owns - disjoint from what kernel_mul_mm_id writes
+ {
+ ggml_metal_kargs_mul_mm_id_zero args = {
+ /*.ne0 =*/ ne0,
+ /*.nb1 =*/ nb1,
+ /*.nb2 =*/ nb2,
+ /*.nb21 =*/ nb21,
+ };
+
+ auto pipeline = ggml_metal_library_get_pipeline_mul_mm_id_zero(lib);
+
+ ggml_metal_encoder_set_pipeline(enc, pipeline);
+ ggml_metal_encoder_set_bytes (enc, &args, sizeof(args), 0);
+ ggml_metal_encoder_set_buffer (enc, bid_src2, 1);
+ ggml_metal_encoder_set_buffer (enc, bid_dst, 2);
+
+ ggml_metal_encoder_dispatch_threadgroups(enc, ne20, ne21, 1, 32, 1, 1);
+ }
+
ggml_metal_op_concurrency_reset(ctx);
if (use_amax) {
diff --git a/ggml/src/ggml-metal/kernels/binbcast.metal b/ggml/src/ggml-metal/kernels/binbcast.metal
index 7c7ab9b5e..b367b3ea1 100644
--- a/ggml/src/ggml-metal/kernels/binbcast.metal
+++ b/ggml/src/ggml-metal/kernels/binbcast.metal
@@ -185,6 +185,14 @@ kernel void kernel_add_id(
device float * dst_row = (device float *)((device char *)dst + i1*nb1 + i2*nb2);
device const float * src0_row = (device const float *)((device char *)src0 + i1*args.nb01 + i2*args.nb02);
+
+ if (i11 == -1) {
+ for (int i0 = tpitg.x; i0 < args.ne0; i0 += ntg.x) {
+ dst_row[i0] = src0_row[i0];
+ }
+ return;
+ }
+
device const float * src1_row = (device const float *)((device char *)src1 + i11*args.nb11);
for (int i0 = tpitg.x; i0 < args.ne0; i0 += ntg.x) {
diff --git a/ggml/src/ggml-metal/kernels/mul_mm.metal b/ggml/src/ggml-metal/kernels/mul_mm.metal
index c313cb8e2..4acc24e5a 100644
--- a/ggml/src/ggml-metal/kernels/mul_mm.metal
+++ b/ggml/src/ggml-metal/kernels/mul_mm.metal
@@ -391,8 +391,10 @@ kernel void kernel_mul_mm_id_map0(
if (i21 + tpitg < args.ne21) {
device const int32_t * src2_i32 = (device const int32_t *) (src2 + (i21 + tpitg)*args.nb21);
+ // 0xFFFF marks a skipped slot, it never matches an expert since ne02 fits in a threadgroup
FOR_UNROLL (short i20 = 0; i20 < ne20; i20++) {
- sids[tpitg*ne20 + i20] = src2_i32[i20];
+ const int32_t id = src2_i32[i20];
+ sids[tpitg*ne20 + i20] = id == -1 ? 0xFFFF : (uint16_t) id;
}
}
@@ -506,6 +508,30 @@ kernel void kernel_mul_mm_id_map0(
}
}
+// kernel_mul_mm_id only writes the rows that an expert owns, so the skipped rows are zeroed here
+kernel void kernel_mul_mm_id_zero(
+ constant ggml_metal_kargs_mul_mm_id_zero & args,
+ device const char * src2,
+ device char * dst,
+ uint3 tgpig[[threadgroup_position_in_grid]],
+ ushort3 tpitg[[thread_position_in_threadgroup]],
+ ushort3 ntg[[threads_per_threadgroup]]) {
+ const int i20 = tgpig.x; // slot
+ const int i21 = tgpig.y; // token
+
+ const int32_t id = ((device const int32_t *) (src2 + i21*args.nb21))[i20];
+
+ if (id != -1) {
+ return;
+ }
+
+ device float * dst_row = (device float *) (dst + i20*args.nb1 + i21*args.nb2);
+
+ for (int i0 = tpitg.x; i0 < args.ne0; i0 += ntg.x) {
+ dst_row[i0] = 0.0f;
+ }
+}
+
kernel void kernel_mul_mm_id_amax_part_f32(
constant ggml_metal_kargs_mul_mm_id_amax & args,
device const char * src1,
diff --git a/ggml/src/ggml-metal/kernels/mul_mv.metal b/ggml/src/ggml-metal/kernels/mul_mv.metal
index 0d1069f2c..7ad07a29c 100644
--- a/ggml/src/ggml-metal/kernels/mul_mv.metal
+++ b/ggml/src/ggml-metal/kernels/mul_mv.metal
@@ -3302,7 +3302,8 @@ kernel void kernel_mul_mv_id(
uint3 tgpig[[threadgroup_position_in_grid]],
ushort tiitg[[thread_index_in_threadgroup]],
ushort tiisg[[thread_index_in_simdgroup]],
- ushort sgitg[[simdgroup_index_in_threadgroup]]) {
+ ushort sgitg[[simdgroup_index_in_threadgroup]],
+ ushort3 ntg [[threads_per_threadgroup]]) {
const int iid1 = tgpig.z/args.nei0;
const int idx = tgpig.z%args.nei0;
@@ -3316,11 +3317,20 @@ kernel void kernel_mul_mv_id(
const int64_t i1 = idx;
const int64_t i2 = i12;
+ device char * dst_cur = dst + (i1*args.ne0 + i2*args.ne1*args.ne0)*sizeof(float);
+
+ if (i02 == -1) {
+ device float * dst_f32 = (device float *) dst_cur;
+ const ushort nth = ntg.x*ntg.y*ntg.z;
+ for (int i0 = tiitg; i0 < args.ne0; i0 += nth) {
+ dst_f32[i0] = 0.0f;
+ }
+ return;
+ }
+
device const char * src0_cur = src0s + i02*args.nb02;
device const char * src1_cur = src1 + i11*args.nb11 + i12*args.nb12;
- device char * dst_cur = dst + (i1*args.ne0 + i2*args.ne1*args.ne0)*sizeof(float);
-
ggml_metal_kargs_mul_mv args0 = {
/*.ne00 =*/ args.ne00,
/*.ne01 =*/ args.ne01,
diff --git a/ggml/src/ggml-opencl/ggml-opencl.cpp b/ggml/src/ggml-opencl/ggml-opencl.cpp
index 1fafe8e33..fc9877e2a 100644
--- a/ggml/src/ggml-opencl/ggml-opencl.cpp
+++ b/ggml/src/ggml-opencl/ggml-opencl.cpp
@@ -1047,6 +1047,7 @@ struct ggml_backend_opencl_context {
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_zero_dst;
cl_kernel kernel_moe_combine_f32 = nullptr; // fused router-weight mul + cross-expert sum
cl_kernel kernel_moe_combine_bias_f32 = nullptr; // same, with the down-projection bias add folded in
cl_kernel kernel_mul_mv_id_q4_0_f32_8x_flat;
@@ -5547,6 +5548,7 @@ static void load_cl_kernels(ggml_backend_opencl_context *backend_ctx) {
build_program_from_source(backend_ctx, kernel_src.c_str(), CL_moe_compile_opts);
CL_CHECK((backend_ctx->kernel_moe_histogram = clCreateKernel(prog, "kernel_moe_histogram", &err), err));
+ CL_CHECK((backend_ctx->kernel_moe_zero_dst = clCreateKernel(prog, "kernel_moe_zero_dst", &err), err));
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));
@@ -26961,6 +26963,39 @@ static void ggml_cl_mul_mat_id(ggml_backend_t backend, const ggml_tensor * src0,
GGML_ASSERT(ne00 == ne10);
+ // the matmul kernels only write the rows an expert owns, so zero the skipped rows here
+ {
+ cl_int err;
+
+ cl_buffer_region region;
+ region.origin = offset2;
+ region.size = nb21 * ne21;
+ cl_mem buf_ids = clCreateSubBuffer(extra2->data_device, 0, CL_BUFFER_CREATE_TYPE_REGION, ®ion, &err);
+ CL_CHECK(err);
+
+ region.origin = offsetd;
+ region.size = (size_t) ne0 * ne1 * ne2 * sizeof(float);
+ cl_mem buf_dst = clCreateSubBuffer(extrad->data_device, 0, CL_BUFFER_CREATE_TYPE_REGION, ®ion, &err);
+ CL_CHECK(err);
+
+ const cl_uint ids_stride = nb21 / sizeof(int);
+
+ cl_kernel kernel_zero = backend_ctx->kernel_moe_zero_dst;
+ CL_CHECK(clSetKernelArg(kernel_zero, 0, sizeof(cl_mem), &buf_ids));
+ CL_CHECK(clSetKernelArg(kernel_zero, 1, sizeof(cl_mem), &buf_dst));
+ CL_CHECK(clSetKernelArg(kernel_zero, 2, sizeof(cl_uint), &ne21));
+ CL_CHECK(clSetKernelArg(kernel_zero, 3, sizeof(cl_uint), &ne20));
+ CL_CHECK(clSetKernelArg(kernel_zero, 4, sizeof(cl_uint), &ids_stride));
+ CL_CHECK(clSetKernelArg(kernel_zero, 5, sizeof(cl_uint), &ne0));
+
+ size_t zero_local[3] = {1, 1, 64};
+ size_t zero_global[3] = {(size_t) ne21, (size_t) ne20, 64};
+ backend_ctx->enqueue_ndrange_kernel(kernel_zero, 3, zero_global, zero_local, dst);
+
+ CL_CHECK(clReleaseMemObject(buf_ids));
+ CL_CHECK(clReleaseMemObject(buf_dst));
+ }
+
int sgs = 32; // subgroup size
int nsg = 1; // number of subgroups
int nrows = 1; // number of row in src1
diff --git a/ggml/src/ggml-opencl/kernels/add_id.cl b/ggml/src/ggml-opencl/kernels/add_id.cl
index e9c6d55e6..e9aa5a150 100644
--- a/ggml/src/ggml-opencl/kernels/add_id.cl
+++ b/ggml/src/ggml-opencl/kernels/add_id.cl
@@ -34,6 +34,14 @@ kernel void kernel_add_id(
global float * dst_row = (global float *)((global char *)dst + i1*nb1 + i2*nb2);
global float * src0_row = (global float *)((global char *)src0 + i1*nb01 + i2*nb02);
+
+ if (i11 == -1) { // skipped slot, add nothing
+ for (int i0 = get_local_id(0); i0 < ne0; i0 += get_local_size(0)) {
+ dst_row[i0] = src0_row[i0];
+ }
+ return;
+ }
+
global float * src1_row = (global float *)((global char *)src1 + i11*nb11);
for (int i0 = get_local_id(0); i0 < ne0; i0 += get_local_size(0)) {
diff --git a/ggml/src/ggml-opencl/kernels/gemv_moe_mxfp4_f32_ns.cl b/ggml/src/ggml-opencl/kernels/gemv_moe_mxfp4_f32_ns.cl
index ee8b94f44..3f21c3523 100644
--- a/ggml/src/ggml-opencl/kernels/gemv_moe_mxfp4_f32_ns.cl
+++ b/ggml/src/ggml-opencl/kernels/gemv_moe_mxfp4_f32_ns.cl
@@ -89,6 +89,9 @@ __kernel void kernel_gemv_moe_mxfp4_f32_ns(
uint i11 = i20 % ne11;
uint expert_id = src2[i20];
+ if ((int)expert_id == -1) { // skipped slot, the dst row is zeroed by kernel_moe_zero_dst
+ return;
+ }
uint expert_offset = expert_id * ne00 * ne01 / 32;
__private float sum = 0.0f; // each thread calculate partial sum of one output
@@ -188,6 +191,9 @@ __kernel void kernel_gemv_moe_mxfp4_f32_ns_wimg(
uint i11 = i20 % ne11;
uint expert_id = src2[i20];
+ if ((int)expert_id == -1) { // skipped slot, the dst row is zeroed by kernel_moe_zero_dst
+ return;
+ }
uint expert_offset = expert_id * ne00 * ne01 / 32;
__private float sum = 0.0f;
diff --git a/ggml/src/ggml-opencl/kernels/gemv_moe_q4_0_f32_ns.cl b/ggml/src/ggml-opencl/kernels/gemv_moe_q4_0_f32_ns.cl
index 2d28db63e..f90970b21 100644
--- a/ggml/src/ggml-opencl/kernels/gemv_moe_q4_0_f32_ns.cl
+++ b/ggml/src/ggml-opencl/kernels/gemv_moe_q4_0_f32_ns.cl
@@ -44,6 +44,9 @@ __kernel void kernel_gemv_moe_q4_0_f32_ns(
uint i11 = i20 % ne11;
uint expert_id = src2[i20];
+ if ((int)expert_id == -1) { // skipped slot, the dst row is zeroed by kernel_moe_zero_dst
+ return;
+ }
uint expert_offset = expert_id * ne00 * ne01 / 32;
__private float sum = 0.0f; // each thread calculate partial sum of one output
diff --git a/ggml/src/ggml-opencl/kernels/gemv_moe_q4_1_f32_ns.cl b/ggml/src/ggml-opencl/kernels/gemv_moe_q4_1_f32_ns.cl
index b98bdc0f1..44bc15ff9 100644
--- a/ggml/src/ggml-opencl/kernels/gemv_moe_q4_1_f32_ns.cl
+++ b/ggml/src/ggml-opencl/kernels/gemv_moe_q4_1_f32_ns.cl
@@ -45,6 +45,9 @@ __kernel void kernel_gemv_moe_q4_1_f32_ns(
uint i11 = i20 % ne11;
uint expert_id = src2[i20];
+ if ((int)expert_id == -1) { // skipped slot, the dst row is zeroed by kernel_moe_zero_dst
+ return;
+ }
uint expert_offset = expert_id * ne00 * ne01 / 32;
__private float sum = 0.0f; // each thread calculate partial sum of one output
diff --git a/ggml/src/ggml-opencl/kernels/gemv_moe_q4_k_f32_ns.cl b/ggml/src/ggml-opencl/kernels/gemv_moe_q4_k_f32_ns.cl
index d3a3c7db8..290a31593 100644
--- a/ggml/src/ggml-opencl/kernels/gemv_moe_q4_k_f32_ns.cl
+++ b/ggml/src/ggml-opencl/kernels/gemv_moe_q4_k_f32_ns.cl
@@ -61,6 +61,9 @@ __kernel void kernel_gemv_moe_q4_k_f32_ns(
uint i11 = i20 % ne11;
uint expert_id = src2[i20];
+ if ((int)expert_id == -1) { // skipped slot, the dst row is zeroed by kernel_moe_zero_dst
+ return;
+ }
int num_superblocks = ne00 / QK_K;
int num_subblocks = ne00 / 32;
@@ -180,6 +183,9 @@ __kernel void kernel_gemv_moe_q4_k_f32_ns_wimg(
uint i11 = i20 % ne11;
uint expert_id = src2[i20];
+ if ((int)expert_id == -1) { // skipped slot, the dst row is zeroed by kernel_moe_zero_dst
+ return;
+ }
int num_superblocks = ne00 / QK_K;
int num_subblocks = ne00 / 32;
diff --git a/ggml/src/ggml-opencl/kernels/gemv_moe_q5_0_f32_ns.cl b/ggml/src/ggml-opencl/kernels/gemv_moe_q5_0_f32_ns.cl
index b43613638..83b186b44 100644
--- a/ggml/src/ggml-opencl/kernels/gemv_moe_q5_0_f32_ns.cl
+++ b/ggml/src/ggml-opencl/kernels/gemv_moe_q5_0_f32_ns.cl
@@ -45,6 +45,9 @@ __kernel void kernel_gemv_moe_q5_0_f32_ns(
uint i11 = i20 % ne11;
uint expert_id = src2[i20];
+ if ((int)expert_id == -1) { // skipped slot, the dst row is zeroed by kernel_moe_zero_dst
+ return;
+ }
uint expert_offset = expert_id * ne00 * ne01 / 32;
__private float sum = 0.0f; // each thread calculate partial sum of one output
diff --git a/ggml/src/ggml-opencl/kernels/gemv_moe_q5_1_f32_ns.cl b/ggml/src/ggml-opencl/kernels/gemv_moe_q5_1_f32_ns.cl
index 7a666006e..2eba34e68 100644
--- a/ggml/src/ggml-opencl/kernels/gemv_moe_q5_1_f32_ns.cl
+++ b/ggml/src/ggml-opencl/kernels/gemv_moe_q5_1_f32_ns.cl
@@ -46,6 +46,9 @@ __kernel void kernel_gemv_moe_q5_1_f32_ns(
uint i11 = i20 % ne11;
uint expert_id = src2[i20];
+ if ((int)expert_id == -1) { // skipped slot, the dst row is zeroed by kernel_moe_zero_dst
+ return;
+ }
uint expert_offset = expert_id * ne00 * ne01 / 32;
__private float sum = 0.0f; // each thread calculate partial sum of one output
diff --git a/ggml/src/ggml-opencl/kernels/gemv_moe_q5_k_f32_ns.cl b/ggml/src/ggml-opencl/kernels/gemv_moe_q5_k_f32_ns.cl
index 7d868d7ab..a959f72c3 100644
--- a/ggml/src/ggml-opencl/kernels/gemv_moe_q5_k_f32_ns.cl
+++ b/ggml/src/ggml-opencl/kernels/gemv_moe_q5_k_f32_ns.cl
@@ -62,6 +62,9 @@ __kernel void kernel_gemv_moe_q5_k_f32_ns(
uint i11 = i20 % ne11;
uint expert_id = src2[i20];
+ if ((int)expert_id == -1) { // skipped slot, the dst row is zeroed by kernel_moe_zero_dst
+ return;
+ }
int num_superblocks = ne00 / QK_K;
int num_subblocks = ne00 / 32;
diff --git a/ggml/src/ggml-opencl/kernels/gemv_moe_q6_k_f32_ns.cl b/ggml/src/ggml-opencl/kernels/gemv_moe_q6_k_f32_ns.cl
index c166bad5b..dcd34fbc2 100644
--- a/ggml/src/ggml-opencl/kernels/gemv_moe_q6_k_f32_ns.cl
+++ b/ggml/src/ggml-opencl/kernels/gemv_moe_q6_k_f32_ns.cl
@@ -45,6 +45,9 @@ __kernel void kernel_gemv_moe_q6_k_f32_ns(
uint i11 = i20 % ne11;
uint expert_id = src2[i20];
+ if ((int)expert_id == -1) { // skipped slot, the dst row is zeroed by kernel_moe_zero_dst
+ return;
+ }
int num_superblocks = ne00 / QK_K;
int num_subblocks = ne00 / 32; // 8 sub-blocks of 32 per super-block
diff --git a/ggml/src/ggml-opencl/kernels/moe_sort_by_expert.cl b/ggml/src/ggml-opencl/kernels/moe_sort_by_expert.cl
index d0f9e8c87..5ea31afb8 100644
--- a/ggml/src/ggml-opencl/kernels/moe_sort_by_expert.cl
+++ b/ggml/src/ggml-opencl/kernels/moe_sort_by_expert.cl
@@ -18,6 +18,9 @@ __kernel void kernel_moe_histogram(
}
int expert_id = input[n * n_experts + k];
+ if (expert_id == -1) { // skipped slot, owned by no expert
+ return;
+ }
atomic_inc(&hist[expert_id]);
}
@@ -63,6 +66,9 @@ __kernel void kernel_moe_scatter(
}
int val = input[n * n_experts + k];
+ if (val == -1) { // skipped slot, owned by no expert
+ return;
+ }
int local_slot = atomic_inc(&slot_counter[val]);
@@ -74,6 +80,32 @@ __kernel void kernel_moe_scatter(
emap[tile_idx] = val;
}
+// the gemm only writes the rows an expert owns, so the skipped rows (negative ids) are zeroed here
+__kernel void kernel_moe_zero_dst(
+ __global const int * input,
+ __global float * dst,
+ uint N,
+ uint topK,
+ uint n_experts,
+ uint ne0
+) {
+ uint n = get_global_id(0);
+ uint k = get_global_id(1);
+
+ if (n >= N || k >= topK) {
+ return;
+ }
+
+ if (input[n * n_experts + k] != -1) {
+ return;
+ }
+
+ __global float * dst_row = dst + ((ulong)n * topK + k) * ne0;
+ for (uint i = get_local_id(2); i < ne0; i += get_local_size(2)) {
+ dst_row[i] = 0.0f;
+ }
+}
+
// Deterministic replacement for kernel_moe_scatter.
//
// kernel_moe_scatter takes each token's slot from atomic_inc(slot_counter[expert]),
diff --git a/ggml/src/ggml-opencl/kernels/mul_mv_id_mxfp4_f32.cl b/ggml/src/ggml-opencl/kernels/mul_mv_id_mxfp4_f32.cl
index d50bd1fc4..4c0755e14 100644
--- a/ggml/src/ggml-opencl/kernels/mul_mv_id_mxfp4_f32.cl
+++ b/ggml/src/ggml-opencl/kernels/mul_mv_id_mxfp4_f32.cl
@@ -172,6 +172,9 @@ kernel void kernel_mul_mv_id_mxfp4_f32(
const int idx = get_group_id(2)%ne20;
int i02 = ((global int *) (src2 + iid1*nb21))[idx];
+ if (i02 == -1) { // skipped slot, the dst row is zeroed by kernel_moe_zero_dst
+ return;
+ }
int i11 = idx % ne11;
int i12 = iid1;
diff --git a/ggml/src/ggml-opencl/kernels/mul_mv_id_mxfp4_f32_flat.cl b/ggml/src/ggml-opencl/kernels/mul_mv_id_mxfp4_f32_flat.cl
index f65e86ed6..962ddd76d 100644
--- a/ggml/src/ggml-opencl/kernels/mul_mv_id_mxfp4_f32_flat.cl
+++ b/ggml/src/ggml-opencl/kernels/mul_mv_id_mxfp4_f32_flat.cl
@@ -101,6 +101,9 @@ kernel void kernel_mul_mv_id_mxfp4_f32_flat(
const int idx = get_group_id(2) % ne20;
uint i02 = ((global uint *) (src2 + offset2 + iid1 * nb21))[idx];
+ if ((int)i02 == -1) { // skipped slot, the dst row is zeroed by kernel_moe_zero_dst
+ return;
+ }
int i11 = idx % ne11;
diff --git a/ggml/src/ggml-opencl/kernels/mul_mv_id_q4_0_f32_8x_flat.cl b/ggml/src/ggml-opencl/kernels/mul_mv_id_q4_0_f32_8x_flat.cl
index 7ccf41efb..13314689f 100644
--- a/ggml/src/ggml-opencl/kernels/mul_mv_id_q4_0_f32_8x_flat.cl
+++ b/ggml/src/ggml-opencl/kernels/mul_mv_id_q4_0_f32_8x_flat.cl
@@ -267,6 +267,9 @@ kernel void kernel_mul_mv_id_q4_0_f32_8x_flat(
const int idx = get_group_id(2)%ne20;
const int i02 = ((global int *)(src2 + iid1*nb21))[idx];
+ if (i02 == -1) { // skipped slot, the dst row is zeroed by kernel_moe_zero_dst
+ return;
+ }
const int i11 = idx%ne11;
const int i12 = iid1;
diff --git a/ggml/src/ggml-opencl/kernels/mul_mv_id_q8_0_f32.cl b/ggml/src/ggml-opencl/kernels/mul_mv_id_q8_0_f32.cl
index f37e83ee8..2c40a43f1 100644
--- a/ggml/src/ggml-opencl/kernels/mul_mv_id_q8_0_f32.cl
+++ b/ggml/src/ggml-opencl/kernels/mul_mv_id_q8_0_f32.cl
@@ -73,6 +73,9 @@ kernel void kernel_mul_mv_id_q8_0_f32(
int idx = get_group_id(2)%ne20;
int i02 = ((global int *) (src2 + iid1*nb21))[idx];
+ if (i02 == -1) { // skipped slot, the dst row is zeroed by kernel_moe_zero_dst
+ return;
+ }
int i11_ = idx % ne11;
int i12_ = iid1;
diff --git a/ggml/src/ggml-opencl/kernels/mul_mv_id_q8_0_f32_flat.cl b/ggml/src/ggml-opencl/kernels/mul_mv_id_q8_0_f32_flat.cl
index fd3a0710f..3b1b39b62 100644
--- a/ggml/src/ggml-opencl/kernels/mul_mv_id_q8_0_f32_flat.cl
+++ b/ggml/src/ggml-opencl/kernels/mul_mv_id_q8_0_f32_flat.cl
@@ -72,6 +72,9 @@ kernel void kernel_mul_mv_id_q8_0_f32_flat(
int idx = (int)get_group_id(2)%ne20;
int i02 = ((global int *) (src2 + iid1*nb21))[idx];
+ if (i02 == -1) { // skipped slot, the dst row is zeroed by kernel_moe_zero_dst
+ return;
+ }
int i11_ = idx % ne11;
int i12_ = iid1;
diff --git a/ggml/src/ggml-openvino/openvino/op/add_id.cpp b/ggml/src/ggml-openvino/openvino/op/add_id.cpp
index 79bdbe877..04b1d3ff2 100644
--- a/ggml/src/ggml-openvino/openvino/op/add_id.cpp
+++ b/ggml/src/ggml-openvino/openvino/op/add_id.cpp
@@ -9,8 +9,12 @@
#include <openvino/op/constant.hpp>
#include <openvino/op/convert.hpp>
#include <openvino/op/gather.hpp>
+#include <openvino/op/greater_eq.hpp>
+#include <openvino/op/maximum.hpp>
+#include <openvino/op/multiply.hpp>
#include <openvino/op/reshape.hpp>
#include <openvino/op/shape_of.hpp>
+#include <openvino/op/unsqueeze.hpp>
namespace ov {
namespace frontend {
@@ -52,8 +56,16 @@ OutputVector translate_add_id(const NodeContext & context) {
ids = std::make_shared<ov::op::v0::Convert>(ids, ov::element::i32);
}
+ // an id of -1 marks a skipped slot: gather row 0 for it and zero the bias with a mask
+ auto zero = ov::op::v0::Constant::create(ids.get_element_type(), ov::Shape{}, {0});
+ ov::Output<ov::Node> keep = std::make_shared<ov::op::v1::GreaterEqual>(ids, zero);
+ keep = std::make_shared<ov::op::v0::Convert>(keep, bias.get_element_type());
+ keep = std::make_shared<ov::op::v0::Unsqueeze>(keep, ov::op::v0::Constant::create(ov::element::i64, ov::Shape{1}, {-1}));
+ ids = std::make_shared<ov::op::v1::Maximum>(ids, zero);
+
auto gather_axis = ov::op::v0::Constant::create(ov::element::i32, ov::Shape{}, {0});
ov::Output<ov::Node> selected_bias = std::make_shared<ov::op::v8::Gather>(bias, ids, gather_axis);
+ selected_bias = std::make_shared<ov::op::v1::Multiply>(selected_bias, keep);
selected_bias = std::make_shared<ov::op::v1::Reshape>(
selected_bias, std::make_shared<ov::op::v3::ShapeOf>(input, ov::element::i64), false);
diff --git a/ggml/src/ggml-sycl/add-id.cpp b/ggml/src/ggml-sycl/add-id.cpp
index e0adc4fe4..a7901859a 100644
--- a/ggml/src/ggml-sycl/add-id.cpp
+++ b/ggml/src/ggml-sycl/add-id.cpp
@@ -26,6 +26,15 @@ static void add_id_kernel(
float* dst_row = (float*)((char*)dst + i1 * nb1 + i2 * nb2);
const float* src0_row =
(const float*)((const char*)src0 + i1 * nb01 + i2 * nb02);
+
+ if (i11 == -1) { // skipped slot, add nothing
+ for (int64_t i0 = item_ct1.get_local_id(2); i0 < ne0;
+ i0 += item_ct1.get_local_range(2)) {
+ dst_row[i0] = src0_row[i0];
+ }
+ return;
+ }
+
const float* src1_row = (const float*)((const char*)src1 + i11 * nb11);
for (int64_t i0 = item_ct1.get_local_id(2); i0 < ne0;
diff --git a/ggml/src/ggml-sycl/ggml-sycl.cpp b/ggml/src/ggml-sycl/ggml-sycl.cpp
index ed0b27280..f3782c2cb 100644
--- a/ggml/src/ggml-sycl/ggml-sycl.cpp
+++ b/ggml/src/ggml-sycl/ggml-sycl.cpp
@@ -5700,6 +5700,23 @@ __dpct_inline__ static void k_copy_dst_from_contiguous(
}
}
+__dpct_inline__ static void k_zero_dst_rows(
+ char *__restrict__ dst_original,
+ const mmid_row_mapping *__restrict__ row_mapping, int64_t ne0, size_t nb1,
+ size_t nb2, const sycl::nd_item<3> &item_ct1) {
+ int32_t i = item_ct1.get_group(2);
+
+ const int32_t i1 = row_mapping[i].i1;
+ const int32_t i2 = row_mapping[i].i2;
+
+ float * dst_row_original = (float *)(dst_original + i1*nb1 + i2*nb2);
+
+ for (int j = item_ct1.get_local_id(2); j < ne0;
+ j += item_ct1.get_local_range(2)) {
+ dst_row_original[j] = 0.0f;
+ }
+}
+
// returns true when ggml_sycl_mul_mat_id takes the fallback path that requires stream synchronization
static bool ggml_sycl_mul_mat_id_needs_sync(const ggml_tensor * dst) {
const ggml_tensor * src0 = dst->src[0];
@@ -5785,14 +5802,18 @@ static void mmid_counting_sort_rows(
int64_t n_ids, int64_t n_as, int64_t n_routed_rows,
std::vector<int64_t> & expert_counts,
std::vector<int64_t> & expert_row_offsets,
- std::vector<mmid_row_mapping> & routed_row_src) {
+ std::vector<mmid_row_mapping> & routed_row_src,
+ std::vector<mmid_row_mapping> & skipped_row_dst) {
// frequencies: how many routed rows each expert "owns"
expert_counts.assign(n_as, 0);
for (int64_t iid1 = 0; iid1 < ids->ne[1]; iid1++) {
for (int64_t id = 0; id < n_ids; id++) {
const int32_t row_id_i = *(const int32_t *) (ids_host + iid1*ids->nb[1] + id*ids->nb[0]);
- GGML_ASSERT(row_id_i >= 0 && row_id_i < n_as);
+ GGML_ASSERT(row_id_i == -1 || (row_id_i >= 0 && row_id_i < n_as));
+ if (row_id_i == -1) {
+ continue;
+ }
expert_counts[row_id_i]++;
}
}
@@ -5805,10 +5826,17 @@ static void mmid_counting_sort_rows(
std::vector<int64_t> expert_row_next = expert_row_offsets;
routed_row_src.resize(n_routed_rows);
+ skipped_row_dst.clear();
for (int64_t iid1 = 0; iid1 < ids->ne[1]; iid1++) {
for (int64_t id = 0; id < n_ids; id++) {
const int32_t row_id_i = *(const int32_t *) (ids_host + iid1*ids->nb[1] + id*ids->nb[0]);
- GGML_ASSERT(row_id_i >= 0 && row_id_i < n_as);
+ GGML_ASSERT(row_id_i == -1 || (row_id_i >= 0 && row_id_i < n_as));
+
+ // a skipped slot has no expert, its dst row is zeroed instead
+ if (row_id_i == -1) {
+ skipped_row_dst.push_back({(int32_t) id, (int32_t) iid1});
+ continue;
+ }
// find and validate the next free row for a given expert (row_id_i)
const int64_t routed_row = expert_row_next[row_id_i]++;
@@ -5817,6 +5845,7 @@ static void mmid_counting_sort_rows(
routed_row_src[routed_row] = {(int32_t) id, (int32_t) iid1};
}
}
+ routed_row_src.resize(expert_row_offsets[n_as]);
}
static void ggml_sycl_mul_mat_id(ggml_backend_sycl_context & ctx,
@@ -5884,7 +5913,7 @@ static void ggml_sycl_mul_mat_id(ggml_backend_sycl_context & ctx,
for (int64_t iid1 = 0; iid1 < ids->ne[1]; iid1++) {
for (int64_t id = 0; id < n_ids; id++) {
const int32_t i02 = *(const int32_t *) (ids_host.data() + iid1*ids->nb[1] + id*ids->nb[0]);
- GGML_ASSERT(i02 >= 0 && i02 < n_as);
+ GGML_ASSERT(i02 == -1 || (i02 >= 0 && i02 < n_as));
const int64_t i11 = id % ne11;
const int64_t i12 = iid1;
@@ -5892,6 +5921,12 @@ static void ggml_sycl_mul_mat_id(ggml_backend_sycl_context & ctx,
const int64_t i1 = id;
const int64_t i2 = i12;
+ if (i02 == -1) { // skipped slot
+ SYCL_CHECK(CHECK_TRY_ERROR(
+ stream->memset(dst_original + i1*nb1 + i2*nb2, 0, ne0*sizeof(float))));
+ continue;
+ }
+
src0_row.data = src0_original + i02*nb02;
src1_row.data = src1_original + i11*nb11 + i12*nb12;
dst_row.data = dst_original + i1*nb1 + i2*nb2;
@@ -5913,20 +5948,43 @@ static void ggml_sycl_mul_mat_id(ggml_backend_sycl_context & ctx,
std::vector<int64_t> expert_row_offsets;
// the sources (slot/token pairs) of contiguous rows to guide k_copy_src1_to_contiguous
std::vector<mmid_row_mapping> & routed_row_src = ctx.mmid_row_mapping_host;
+ // the slot/token pairs of the skipped slots, whose dst rows are zeroed
+ std::vector<mmid_row_mapping> skipped_row_dst;
mmid_counting_sort_rows(ids, ids_host.data(), n_ids, n_as, n_routed_rows,
- expert_row_counts, expert_row_offsets, routed_row_src);
+ expert_row_counts, expert_row_offsets, routed_row_src, skipped_row_dst);
+
+ const int64_t n_valid_rows = (int64_t) routed_row_src.size();
ggml_sycl_pool_alloc<mmid_row_mapping> dev_row_mapping(ctx.pool(), n_routed_rows);
SYCL_CHECK(CHECK_TRY_ERROR(
- stream->memcpy(dev_row_mapping.get(), routed_row_src.data(), n_routed_rows*sizeof(mmid_row_mapping))));
+ stream->memcpy(dev_row_mapping.get(), routed_row_src.data(), n_valid_rows*sizeof(mmid_row_mapping))));
const unsigned int max_work_group_size = ggml_sycl_info().max_work_group_sizes[ctx.device];
assert(max_work_group_size % (WARP_SIZE * WARP_SIZE) == 0);
- {
+ if (!skipped_row_dst.empty()) {
+ ggml_sycl_pool_alloc<mmid_row_mapping> dev_skipped(ctx.pool(), skipped_row_dst.size());
+ SYCL_CHECK(CHECK_TRY_ERROR(
+ stream->memcpy(dev_skipped.get(), skipped_row_dst.data(), skipped_row_dst.size()*sizeof(mmid_row_mapping))));
+
+ sycl::range<3> block_dims(1, 1, std::min((unsigned int)ne0, max_work_group_size));
+ sycl::range<3> grid_dims(1, 1, skipped_row_dst.size());
+ stream->submit([&](sycl::handler &cgh) {
+ mmid_row_mapping *__restrict dev_skipped_get = dev_skipped.get();
+
+ cgh.parallel_for(
+ sycl::nd_range<3>(grid_dims * block_dims, block_dims),
+ [=](sycl::nd_item<3> item_ct1) {
+ k_zero_dst_rows(dst_original, dev_skipped_get, ne0, nb1, nb2, item_ct1);
+ });
+ });
+ SYCL_CHECK(CHECK_TRY_ERROR(stream->wait()));
+ }
+
+ if (n_valid_rows > 0) {
sycl::range<3> block_dims(1, 1, std::min((unsigned int)ne10, max_work_group_size));
- sycl::range<3> grid_dims(1, 1, n_routed_rows);
+ sycl::range<3> grid_dims(1, 1, n_valid_rows);
stream->submit([&](sycl::handler &cgh) {
char *__restrict src1_contiguous_get =
src1_contiguous.get();
@@ -5992,9 +6050,9 @@ static void ggml_sycl_mul_mat_id(ggml_backend_sycl_context & ctx,
ggml_sycl_mul_mat(ctx, &src0_row, &src1_row, &dst_row);
}
- {
+ if (n_valid_rows > 0) {
sycl::range<3> block_dims(1, 1, std::min((unsigned int)ne0, max_work_group_size));
- sycl::range<3> grid_dims(1, 1, n_routed_rows);
+ sycl::range<3> grid_dims(1, 1, n_valid_rows);
stream->submit([&](sycl::handler &cgh) {
const char *__restrict dst_contiguous_get =
dst_contiguous.get();
diff --git a/ggml/src/ggml-sycl/mmvq.cpp b/ggml/src/ggml-sycl/mmvq.cpp
index 13c6677a8..97ce165c6 100644
--- a/ggml/src/ggml-sycl/mmvq.cpp
+++ b/ggml/src/ggml-sycl/mmvq.cpp
@@ -3149,7 +3149,6 @@ static void mul_mat_vec_q_moe(
const int expert_idx = item_ct1.get_group(1);
const int i02 = ids_dev[expert_idx];
- const char * vx = (const char *) vx_base + (size_t) i02 * expert_weight_stride;
const char * vy = (const char *) vy_base + (size_t) expert_idx * src1_row_stride;
float * dst = (float *) ((char *) dst_base + (size_t) expert_idx * dst_row_stride);
@@ -3159,6 +3158,15 @@ static void mul_mat_vec_q_moe(
return;
}
+ if (i02 == -1) { // skipped slot
+ if (item_ct1.get_local_id(2) == 0) {
+ dst[row] = 0.0f;
+ }
+ return;
+ }
+
+ const char * vx = (const char *) vx_base + (size_t) i02 * expert_weight_stride;
+
const int blocks_per_row = ncols / qk;
constexpr int blocks_per_warp = (vdr * WARP_SIZE + qi - 1) / qi;
@@ -3355,7 +3363,6 @@ static void mul_mat_vec_q_moe_reorder(
const int expert_idx = item_ct1.get_group(1);
const int i02 = ids_dev[expert_idx];
- const char * vx = (const char *) vx_base + (size_t) i02 * expert_weight_stride;
const char * vy = (const char *) vy_base + (size_t) expert_idx * src1_row_stride;
float * dst = (float *) ((char *) dst_base + (size_t) expert_idx * dst_row_stride);
@@ -3366,6 +3373,15 @@ static void mul_mat_vec_q_moe_reorder(
const auto sg = item_ct1.get_sub_group();
+ if (i02 == -1) { // skipped slot
+ if (sg.leader()) {
+ dst[row] = 0.0f;
+ }
+ return;
+ }
+
+ const char * vx = (const char *) vx_base + (size_t) i02 * expert_weight_stride;
+
constexpr int vdr_mmvq = reorder_vec_dot_q_sycl::gtype == GGML_TYPE_Q8_0 ?
VDR_Q8_0_Q8_1_MMVQ : block_traits::vdr_mmvq;
const int blocks_per_row = ncols / block_traits::qk;
diff --git a/ggml/src/ggml-vulkan/ggml-vulkan-push-constants.h b/ggml/src/ggml-vulkan/ggml-vulkan-push-constants.h
index fdce443a1..219c477c6 100644
--- a/ggml/src/ggml-vulkan/ggml-vulkan-push-constants.h
+++ b/ggml/src/ggml-vulkan/ggml-vulkan-push-constants.h
@@ -222,6 +222,9 @@ struct vk_op_count_experts_push_constants {
uint32_t ne00L;
uint32_t row_tile_size;
uint32_t row_ids_offset;
+ uint32_t ne0_dst;
+ uint32_t stride_dst;
+ uint32_t batch_stride_dst;
};
struct vk_op_glu_push_constants {
diff --git a/ggml/src/ggml-vulkan/ggml-vulkan.cpp b/ggml/src/ggml-vulkan/ggml-vulkan.cpp
index 0ea4e063b..4087f9729 100644
--- a/ggml/src/ggml-vulkan/ggml-vulkan.cpp
+++ b/ggml/src/ggml-vulkan/ggml-vulkan.cpp
@@ -3585,9 +3585,9 @@ void ggml_vk_load_shaders(vk_device& device, vk_pipeline requested) {
ggml_vk_create_pipeline(device, device->pipeline_count_equal_i32, "count_equal_i32", count_equal_i32_len, count_equal_i32_data, "main", 3, sizeof(vk_op_push_constants), {512, 1, 1}, { device->subgroup_size }, 1);
if (device->subgroup_arithmetic && device->subgroup_vote && device->subgroup_require_full_support) {
- ggml_vk_create_pipeline(device, device->pipeline_count_experts, "count_experts", count_experts_subgroup_len, count_experts_subgroup_data, "main", 2, sizeof(vk_op_count_experts_push_constants), {1, 1, 1}, {}, 1, true, true);
+ ggml_vk_create_pipeline(device, device->pipeline_count_experts, "count_experts", count_experts_subgroup_len, count_experts_subgroup_data, "main", 3, sizeof(vk_op_count_experts_push_constants), {1, 1, 1}, {}, 1, true, true);
} else {
- ggml_vk_create_pipeline(device, device->pipeline_count_experts, "count_experts", count_experts_len, count_experts_data, "main", 2, sizeof(vk_op_count_experts_push_constants), {1, 1, 1}, {}, 1, true);
+ ggml_vk_create_pipeline(device, device->pipeline_count_experts, "count_experts", count_experts_len, count_experts_data, "main", 3, sizeof(vk_op_count_experts_push_constants), {1, 1, 1}, {}, 1, true);
}
// comb holds a token's 4x4 matrix in one 16-lane slice of a subgroup, so it
@@ -7643,16 +7643,20 @@ static void ggml_vk_mul_mat_id_q_f16(ggml_backend_vk_context * ctx, vk_context&
ggml_vk_sync_buffers(ctx, subctx);
}
{
+ // this also zeroes the dst rows of the skipped slots (negative ids)
vk_op_count_experts_push_constants pc = { (uint32_t)nei0,
(uint32_t)nei1,
(uint32_t)(nbi0 / ggml_type_size(ids->type)),
(uint32_t)(nbi1 / ggml_type_size(ids->type)),
(uint32_t)(get_misalign_bytes(ctx, ids) / ggml_type_size(ids->type)),
(uint32_t)n_as,
- 0, 0, bn, (uint32_t)row_ids_offset };
+ 0, 0, bn, (uint32_t)row_ids_offset,
+ (uint32_t)ne01,
+ (uint32_t)(dst->nb[1] / sizeof(float)),
+ (uint32_t)(dst->nb[2] / sizeof(float)) };
init_pushconst_fastdiv(pc);
ggml_vk_dispatch_pipeline(ctx, subctx, count_experts,
- { d_ids, expert_count_buf }, pc,
+ { d_ids, expert_count_buf, vk_subbuffer{ d_D, d_buf_offset, d_sz } }, pc,
{ 1, 1, 1 });
}
diff --git a/ggml/src/ggml-vulkan/vulkan-shaders/add_id.comp b/ggml/src/ggml-vulkan/vulkan-shaders/add_id.comp
index 495249d5f..bf6206704 100644
--- a/ggml/src/ggml-vulkan/vulkan-shaders/add_id.comp
+++ b/ggml/src/ggml-vulkan/vulkan-shaders/add_id.comp
@@ -27,14 +27,22 @@ void main() {
const uint i1 = gl_WorkGroupID.x;
const uint i2 = gl_WorkGroupID.y;
- const uint i11 = data_c[i1 + i2 * p.s21];
+ const int i11 = data_c[i1 + i2 * p.s21];
const uint s1 = p.ne0;
const uint s2 = p.ne0 * p.ne1;
const uint d0 = i1 * s1 + i2 * s2;
const uint a0 = i1 * p.s01 + i2 * p.s02;
- const uint b0 = i11 * p.s11;
+
+ if (i11 == -1) { // skipped slot, add nothing
+ for (uint i0 = gl_LocalInvocationID.x; i0 < p.ne0; i0 += BLOCK_SIZE) {
+ data_d[d0 + i0] = data_a[a0 + i0];
+ }
+ return;
+ }
+
+ const uint b0 = uint(i11) * p.s11;
for (uint i0 = gl_LocalInvocationID.x; i0 < p.ne0; i0 += BLOCK_SIZE) {
data_d[d0 + i0] = data_a[a0 + i0] + data_b[b0 + i0];
diff --git a/ggml/src/ggml-vulkan/vulkan-shaders/count_experts.comp b/ggml/src/ggml-vulkan/vulkan-shaders/count_experts.comp
index e7e295115..a46a06193 100644
--- a/ggml/src/ggml-vulkan/vulkan-shaders/count_experts.comp
+++ b/ggml/src/ggml-vulkan/vulkan-shaders/count_experts.comp
@@ -23,6 +23,9 @@ layout (push_constant) uniform parameter
uint32_t ne00L;
uint32_t row_tile_size;
uint32_t row_ids_offset;
+ uint32_t ne0_dst;
+ uint32_t stride_dst;
+ uint32_t batch_stride_dst;
} p;
#define BLOCK_SIZE 256
@@ -32,6 +35,7 @@ layout(local_size_x = BLOCK_SIZE, local_size_y = 1, local_size_z = 1) in;
layout (binding = 0) readonly buffer A {uint data_a[];};
layout (binding = 1) writeonly buffer D {uint data_d[];};
+layout (binding = 2) writeonly buffer DST {float data_dst[];};
shared uint vals[EXPERTS_PER_CHUNK];
shared uint offsets[EXPERTS_PER_CHUNK];
@@ -47,6 +51,26 @@ void write_tiles(uint expert, uint count, uint tile_begin, uint n_tiles) {
}
}
+// the matmul writes only the rows that an expert owns, so zero the rows of the skipped slots
+void zero_skipped_rows() {
+ const uint num_elements = p.ne00 * p.ne01;
+ const uint tid = gl_LocalInvocationID.x;
+
+ for (uint idx = 0; idx < num_elements; idx++) {
+ const uint i01 = fastdiv(idx, p.ne00mp, p.ne00L);
+ const uint i00 = idx - i01 * p.ne00;
+
+ if (int(data_a[p.a_offset + i01 * p.nb01 + i00 * p.nb00]) != -1) {
+ continue;
+ }
+
+ const uint d_off = i00 * p.stride_dst + i01 * p.batch_stride_dst;
+ for (uint i = tid; i < p.ne0_dst; i += BLOCK_SIZE) {
+ data_dst[d_off + i] = 0.0f;
+ }
+ }
+}
+
void main() {
const uint num_elements = p.ne00 * p.ne01;
const uint tid = gl_LocalInvocationID.x;
@@ -123,4 +147,6 @@ void main() {
if (tid == 0) {
data_d[0] = total_tiles;
}
+
+ zero_skipped_rows();
}
diff --git a/ggml/src/ggml-vulkan/vulkan-shaders/mul_mat_vec_base.glsl b/ggml/src/ggml-vulkan/vulkan-shaders/mul_mat_vec_base.glsl
index 4aeda68c7..12d9cbf6b 100644
--- a/ggml/src/ggml-vulkan/vulkan-shaders/mul_mat_vec_base.glsl
+++ b/ggml/src/ggml-vulkan/vulkan-shaders/mul_mat_vec_base.glsl
@@ -42,6 +42,7 @@ layout (push_constant) uniform parameter
#ifdef MUL_MAT_ID
uint expert_id;
+bool expert_skip; // an id of -1 skips the slot, the result row is zeroed
#endif
void get_offsets(out uint a_offset, out uint b_offset, out uint d_offset) {
@@ -63,7 +64,9 @@ void get_offsets(out uint a_offset, out uint b_offset, out uint d_offset) {
batch_idx_a = i03 * p.ne02 + i02;
}
#else
- expert_id = data_ids[expert_i0 + p.expert_i1 * p.nbi1];
+ const int expert = data_ids[expert_i0 + p.expert_i1 * p.nbi1];
+ expert_skip = expert == -1;
+ expert_id = expert_skip ? 0 : uint(expert); // clamp, so that the offsets below stay in range
#endif
a_offset =
@@ -113,6 +116,9 @@ void reduce_result(inout FLOAT_TYPE temp[NUM_COLS][NUM_ROWS], const in uint32_t
const uint expert_i0 = gl_GlobalInvocationID.y;
temp[j][n] *= FLOAT_TYPE(data_fuse1[expert_i0]);
}
+ if (expert_skip) {
+ temp[j][n] = FLOAT_TYPE(0);
+ }
#else
if ((p.fusion_flags & MAT_VEC_FUSION_FLAGS_BIAS0) != 0) {
temp[j][n] += FLOAT_TYPE(data_fuse0[j*p.batch_stride_d + d_offset + first_row + n]);
@@ -168,6 +174,9 @@ void reduce_result(FLOAT_TYPE temp[NUM_COLS][NUM_ROWS], const in uint32_t d_offs
const uint expert_i0 = gl_GlobalInvocationID.y;
temp[j][n] *= FLOAT_TYPE(data_fuse1[expert_i0]);
}
+ if (expert_skip) {
+ temp[j][n] = FLOAT_TYPE(0);
+ }
#else
if ((p.fusion_flags & MAT_VEC_FUSION_FLAGS_BIAS0) != 0) {
temp[j][n] += FLOAT_TYPE(data_fuse0[j*p.batch_stride_d + d_offset + first_row + n]);
@@ -213,6 +222,9 @@ void reduce_result(FLOAT_TYPE temp[NUM_COLS][NUM_ROWS], const in uint32_t d_offs
const uint expert_i0 = gl_GlobalInvocationID.y;
tmpsh[j][n][0] *= FLOAT_TYPE(data_fuse1[expert_i0]);
}
+ if (expert_skip) {
+ tmpsh[j][n][0] = FLOAT_TYPE(0);
+ }
#else
if ((p.fusion_flags & MAT_VEC_FUSION_FLAGS_BIAS0) != 0) {
tmpsh[j][n][0] += FLOAT_TYPE(data_fuse0[j*p.batch_stride_d + d_offset + first_row + n]);
diff --git a/ggml/src/ggml-webgpu/ggml-webgpu.cpp b/ggml/src/ggml-webgpu/ggml-webgpu.cpp
index 71d9ae2b9..3643e9d91 100644
--- a/ggml/src/ggml-webgpu/ggml-webgpu.cpp
+++ b/ggml/src/ggml-webgpu/ggml-webgpu.cpp
@@ -1863,10 +1863,12 @@ static webgpu_encoded_op ggml_webgpu_mul_mat_id(webgpu_context & ctx,
// params for mul_mat_id_gather.wgsl
std::vector<uint32_t> gather_params = {
(uint32_t) (ggml_webgpu_tensor_misalignment(ctx, src2) / ggml_type_size(src2->type)),
+ (uint32_t) (ggml_webgpu_tensor_misalignment(ctx, dst) / ggml_type_size(dst->type)),
param_n_expert,
param_n_expert_used,
param_n_tokens,
(uint32_t) (src2->nb[1] / ggml_type_size(src2->type)),
+ (uint32_t) dst->ne[0],
};
const size_t dst_offset = ggml_webgpu_tensor_offset(dst);
@@ -1888,6 +1890,8 @@ static webgpu_encoded_op ggml_webgpu_mul_mat_id(webgpu_context & ctx,
ggml_webgpu_tensor_binding_size(ctx, src2)),
ggml_webgpu_make_bind_group_entry(1, ggml_webgpu_tensor_buf(dst), ids_cache_align_offset, ids_cache_bind_size),
ggml_webgpu_make_bind_group_entry(2, ggml_webgpu_tensor_buf(dst), ids_tiles_align_offset, ids_tiles_bind_size),
+ ggml_webgpu_make_bind_group_entry(3, ggml_webgpu_tensor_buf(dst), ggml_webgpu_tensor_align_offset(ctx, dst),
+ ggml_webgpu_tensor_binding_size(ctx, dst)),
};
dispatches.push_back({
diff --git a/ggml/src/ggml-webgpu/wgsl-shaders/add_id.wgsl b/ggml/src/ggml-webgpu/wgsl-shaders/add_id.wgsl
index 2573926cb..1e377bff5 100644
--- a/ggml/src/ggml-webgpu/wgsl-shaders/add_id.wgsl
+++ b/ggml/src/ggml-webgpu/wgsl-shaders/add_id.wgsl
@@ -46,12 +46,24 @@ fn main(@builtin(workgroup_id) wg_id: vec3<u32>,
let i2 = wg_linear / params.ne1;
let i1 = wg_linear % params.ne1;
- let i11 = u32(ids[params.offset_ids + i1 * params.nb20 + i2 * params.nb21]);
+ let i11 = i32(ids[params.offset_ids + i1 * params.nb20 + i2 * params.nb21]);
let src0_row = params.offset_src0 + i1 * params.nb01 + i2 * params.nb02;
- let src1_row = params.offset_src1 + i11 * params.nb11;
let dst_row = params.offset_dst + i1 * params.ne0 + i2 * (params.ne0 * params.ne1);
+ if (i11 == -1) { // skipped slot, add nothing
+#ifdef INPLACE
+ // src0 already holds the result
+#else
+ for (var i = thread_id;i < params.ne0; i += WG_SIZE) {
+ dst[dst_row + i] = src0[src0_row + i];
+ }
+#endif
+ return;
+ }
+
+ let src1_row = params.offset_src1 + u32(i11) * params.nb11;
+
for (var i = thread_id;i < params.ne0; i += WG_SIZE) {
#ifdef INPLACE
src0[src0_row + i] = src0[src0_row + i] + src1[src1_row + i];
diff --git a/ggml/src/ggml-webgpu/wgsl-shaders/mul_mat_id_gather.wgsl b/ggml/src/ggml-webgpu/wgsl-shaders/mul_mat_id_gather.wgsl
index 1439a12d0..11fcc6ba8 100644
--- a/ggml/src/ggml-webgpu/wgsl-shaders/mul_mat_id_gather.wgsl
+++ b/ggml/src/ggml-webgpu/wgsl-shaders/mul_mat_id_gather.wgsl
@@ -2,19 +2,24 @@ enable f16;
struct MulMatIdGatherParams {
offset_ids: u32,
+ offset_dst: u32,
n_expert: u32,
n_expert_used: u32,
n_tokens: u32,
stride_ids_1: u32,
+
+ m: u32,
};
@group(0) @binding(0) var<storage, read_write> ids: array<i32>; // [n_expert_used, n_tokens]
@group(0) @binding(1) var<storage, read_write> ids_slot_pos: array<u32>; // [n_expert_used * n_tokens]
@group(0) @binding(2) var<storage, read_write> ids_tiles: array<u32>; // [n_tiles, (expert, first slot, n slots) * n_tiles]
-@group(0) @binding(3) var<uniform> params: MulMatIdGatherParams;
+@group(0) @binding(3) var<storage, read_write> dst: array<f32>; // [m, n_expert_used, n_tokens]
+
+@group(0) @binding(4) var<uniform> params: MulMatIdGatherParams;
// per expert: the slot count, replaced by the first slot of the expert in ids_slot_pos by the prefix sum, then advanced as its write cursor
var<workgroup> ids_count: array<atomic<u32>, N_EXPERT>;
@@ -36,8 +41,10 @@ fn main(@builtin(local_invocation_id) local_id: vec3<u32>) {
for (var i = thread_id; i < ids_slots_num; i += WG_SIZE) {
let row = i / params.n_expert_used;
let col = i % params.n_expert_used;
- let expert = u32(ids[params.offset_ids + row * params.stride_ids_1 + col]);
- atomicAdd(&ids_count[expert], 1u);
+ let expert = ids[params.offset_ids + row * params.stride_ids_1 + col];
+ if (expert != -1) { // skipped slot, owned by no expert
+ atomicAdd(&ids_count[u32(expert)], 1u);
+ }
}
workgroupBarrier();
@@ -71,8 +78,23 @@ fn main(@builtin(local_invocation_id) local_id: vec3<u32>) {
for (var i = thread_id; i < ids_slots_num; i += WG_SIZE) {
let row = i / params.n_expert_used;
let col = i % params.n_expert_used;
- let expert = u32(ids[params.offset_ids + row * params.stride_ids_1 + col]);
- let pos = atomicAdd(&ids_count[expert], 1u);
- ids_slot_pos[pos] = i;
+ let expert = ids[params.offset_ids + row * params.stride_ids_1 + col];
+ if (expert != -1) {
+ let pos = atomicAdd(&ids_count[u32(expert)], 1u);
+ ids_slot_pos[pos] = i;
+ }
+ }
+
+ // the matmul writes only the rows that an expert owns, so zero the rows of the skipped slots
+ for (var i: u32 = 0; i < ids_slots_num; i += 1) {
+ let row = i / params.n_expert_used;
+ let col = i % params.n_expert_used;
+ if (ids[params.offset_ids + row * params.stride_ids_1 + col] != -1) {
+ continue;
+ }
+ let dst_row = params.offset_dst + i * params.m;
+ for (var r = thread_id; r < params.m; r += WG_SIZE) {
+ dst[dst_row + r] = 0.0f;
+ }
}
}
diff --git a/ggml/src/ggml-webgpu/wgsl-shaders/mul_mat_id_vec.wgsl b/ggml/src/ggml-webgpu/wgsl-shaders/mul_mat_id_vec.wgsl
index cb95783bc..e87cc6b2d 100644
--- a/ggml/src/ggml-webgpu/wgsl-shaders/mul_mat_id_vec.wgsl
+++ b/ggml/src/ggml-webgpu/wgsl-shaders/mul_mat_id_vec.wgsl
@@ -70,11 +70,18 @@ fn main(
let own_expert = ids[params.offset_ids + slot];
let row_base = (wg_linear % output_groups) * OUTPUTS_PER_WG;
- let src0_batch_offset = params.offset_src0 + own_expert * params.stride_02;
+ // a skipped slot (id -1) keeps the zero accumulator, so the reduction below writes zeros
+ // no early return: the barriers below must stay in uniform control flow
+ let skipped = i32(own_expert) == -1;
+
+ let src0_batch_offset = params.offset_src0 + select(own_expert, 0u, skipped) * params.stride_02;
let src1_idx_base = params.offset_src1 + (slot % params.b_ne1) * params.stride_11;
let dst_idx_base = params.offset_dst + slot * params.m + row_base;
- let acc = accumulate_vec_dot(thread_id, row_base, src0_batch_offset, src1_idx_base);
+ var acc: array<array<f32, OUTPUTS_PER_WG>, NUM_COLS>;
+ if (!skipped) {
+ acc = accumulate_vec_dot(thread_id, row_base, src0_batch_offset, src1_idx_base);
+ }
#ifdef USE_SUBGROUP_REDUCTION
for (var row = 0u; row < OUTPUTS_PER_WG; row++) {
diff --git a/ggml/src/ggml-zendnn/ggml-zendnn.cpp b/ggml/src/ggml-zendnn/ggml-zendnn.cpp
index ec7ce2331..357cfcd99 100644
--- a/ggml/src/ggml-zendnn/ggml-zendnn.cpp
+++ b/ggml/src/ggml-zendnn/ggml-zendnn.cpp
@@ -375,7 +375,12 @@ static void ggml_zendnn_compute_forward_mul_mat_id(
for (int id = 0; id < n_ids; ++id) {
const int32_t i02 = *(const int32_t *)((const char *)ids->data + iid1*ids->nb[1] + id*ids->nb[0]);
- GGML_ASSERT(i02 >= 0 && i02 < n_as);
+ GGML_ASSERT(i02 == -1 || (i02 >= 0 && i02 < n_as));
+
+ if (i02 == -1) {
+ memset((char *) dst->data + id*nb1 + iid1*nb2, 0, ne0*sizeof(float));
+ continue;
+ }
if (matrix_row_counts[i02] == 0) {
n_active_experts++;
diff --git a/ggml/src/ggml.c b/ggml/src/ggml.c
index f33e3f2e4..ed8554b14 100644
--- a/ggml/src/ggml.c
+++ b/ggml/src/ggml.c
@@ -3438,6 +3438,8 @@ void ggml_mul_mat_set_hint(
in b, n_expert_used can be broadcasted to match the n_expert_used of ids
c ~= as[:,:,i] @ b[:,i%r,t], i = ids[e,t] for all e,t in ids
+
+ an id of -1 skips the slot: no matrix is selected, c[:,e,t] is set to zero
*/
struct ggml_tensor * ggml_mul_mat_id(
struct ggml_context * ctx,
diff --git a/tests/test-backend-ops.cpp b/tests/test-backend-ops.cpp
index 5716a2b01..636c61be2 100644
--- a/tests/test-backend-ops.cpp
+++ b/tests/test-backend-ops.cpp
@@ -1194,6 +1194,29 @@ static bool op_filter_entry_matches(std::string_view entry, std::string_view op_
return std::regex_search(op_name.data(), op_name.data() + op_name.size(), re);
}
+// TODO: [TAG_MUL_MAT_ID_SKIP] see https://github.com/ggml-org/llama.cpp/pull/26631
+static bool backend_lacks_skipped_ids(ggml_backend_t backend, ggml_cgraph * gf) {
+ const char * name = ggml_backend_reg_name(ggml_backend_dev_backend_reg(ggml_backend_get_device(backend)));
+ if (strcmp(name, "CUDA") != 0 && strcmp(name, "ROCm") != 0 && strcmp(name, "MUSA") != 0) {
+ return false;
+ }
+ for (int i = 0; i < ggml_graph_n_nodes(gf); i++) {
+ const ggml_tensor * node = ggml_graph_node(gf, i);
+ if (node->op != GGML_OP_MUL_MAT_ID && node->op != GGML_OP_ADD_ID) {
+ continue;
+ }
+ const ggml_tensor * ids = node->src[2];
+ std::vector<int32_t> row(ids->ne[0]);
+ for (int64_t i1 = 0; i1 < ids->ne[1]; i1++) {
+ ggml_backend_tensor_get(ids, row.data(), i1*ids->nb[1], ids->ne[0]*sizeof(int32_t));
+ if (std::find(row.begin(), row.end(), -1) != row.end()) {
+ return true;
+ }
+ }
+ }
+ return false;
+}
+
struct test_case {
virtual ~test_case() {}
@@ -1475,6 +1498,13 @@ struct test_case {
initialize_tensors(ctx_weights.get());
}
+ // TODO: [TAG_MUL_MAT_ID_SKIP] @ngxson : remove this once CUDA implements mul_mat_id with -1 = skipped
+ if (backend_lacks_skipped_ids(backend1, gf)) {
+ test_result result(ggml_backend_name(backend1), current_op_name, vars(), "test", false, false, "not supported");
+ print_test_result_locked(output_printer, result);
+ return test_status_t::NOT_SUPPORTED;
+ }
+
// compare
struct callback_userdata {
bool ok;
@@ -3564,6 +3594,16 @@ struct test_bin_bcast : public test_case {
}
};
+// set some ids to -1, so that MUL_MAT_ID and ADD_ID skip those slots
+// every 4th row skips all of its slots, the other rows skip only some of them
+static void mask_skipped_ids(std::vector<int32_t> & ids, int64_t row) {
+ for (size_t i = 0; i < ids.size(); i++) {
+ if (row % 4 == 1 || (row + (int64_t) i) % 3 == 0) {
+ ids[i] = -1;
+ }
+ }
+}
+
// GGML_OP_ADD_ID
struct test_add_id : public test_case {
const ggml_type type_a;
@@ -3572,9 +3612,10 @@ struct test_add_id : public test_case {
const int64_t n_experts;
const int64_t n_experts_used;
const int64_t n_token;
+ const bool skip; // set some ids to a negative value
std::string vars() override {
- return VARS_TO_STR6(type_a, type_b, n_embd, n_experts, n_experts_used, n_token);
+ return VARS_TO_STR7(type_a, type_b, n_embd, n_experts, n_experts_used, n_token, skip);
}
size_t op_size(ggml_tensor * t) override {
@@ -3586,9 +3627,10 @@ struct test_add_id : public test_case {
int64_t n_embd = 128,
int64_t n_experts = 16,
int64_t n_experts_used = 8,
- int64_t n_token = 10)
+ int64_t n_token = 10,
+ bool skip = false)
: type_a(type_a), type_b(type_b), n_embd(n_embd),
- n_experts(n_experts), n_experts_used(n_experts_used), n_token(n_token) {}
+ n_experts(n_experts), n_experts_used(n_experts_used), n_token(n_token), skip(skip) {}
ggml_tensor * build_graph(ggml_context * ctx) override {
ggml_tensor * a = ggml_new_tensor_3d(ctx, type_a, n_embd, n_experts_used, n_token);
@@ -3617,6 +3659,9 @@ struct test_add_id : public test_case {
data[i] = i % n_experts;
}
std::shuffle(data.begin(), data.end(), rng);
+ if (skip) {
+ mask_skipped_ids(data, r);
+ }
ggml_backend_tensor_set(t, data.data(), r * t->nb[1], t->ne[0] * sizeof(int32_t));
}
} else {
@@ -5550,7 +5595,8 @@ struct test_mul_mat_w4a4 : public test_mul_mat {
}
};
-static void init_mul_mat_id_ids(ggml_context * ctx, int n_mats) {
+static void init_mul_mat_id_ids(ggml_context * ctx, int n_mats, bool skip = false) {
+
std::random_device rd;
std::default_random_engine rng(rd());
for (ggml_tensor * t = ggml_get_first_tensor(ctx); t != NULL; t = ggml_get_next_tensor(ctx, t)) {
@@ -5563,12 +5609,15 @@ static void init_mul_mat_id_ids(ggml_context * ctx, int n_mats) {
data[i] = i % n_mats;
}
std::shuffle(data.begin(), data.end(), rng);
+ if (skip) {
+ mask_skipped_ids(data, r);
+ }
ggml_backend_tensor_set(t, data.data(), r * t->nb[1], t->ne[0] * sizeof(int32_t));
}
}
}
-static void init_mul_mat_id_tensors(ggml_context * ctx, int n_mats, float amax = 1.0f) {
+static void init_mul_mat_id_tensors(ggml_context * ctx, int n_mats, bool skip = false, float amax = 1.0f) {
for (ggml_tensor * t = ggml_get_first_tensor(ctx); t != NULL; t = ggml_get_next_tensor(ctx, t)) {
if (t->type == GGML_TYPE_I32) {
continue;
@@ -5579,7 +5628,7 @@ static void init_mul_mat_id_tensors(ggml_context * ctx, int n_mats, float amax =
init_tensor_uniform(t);
}
}
- init_mul_mat_id_ids(ctx, n_mats);
+ init_mul_mat_id_ids(ctx, n_mats, skip);
}
// GGML_OP_MUL_MAT_ID
@@ -5595,9 +5644,10 @@ struct test_mul_mat_id : public test_case {
const float amax; // magnitude of src1
const int64_t m_v; // rows of as in memory, the experts of as are strided for m_v > m, no view for m_v == 0
const bool ids_offset;
+ const bool skip; // set some ids to a negative value
std::string vars() override {
- return VARS_TO_STR11(type_a, type_b, n_mats, n_used, b, m, n, k, amax, m_v, ids_offset);
+ return VARS_TO_STR12(type_a, type_b, n_mats, n_used, b, m, n, k, amax, m_v, ids_offset, skip);
}
double max_nmse_err() override {
@@ -5622,9 +5672,9 @@ struct test_mul_mat_id : public test_case {
test_mul_mat_id(ggml_type type_a = GGML_TYPE_F32, ggml_type type_b = GGML_TYPE_F32,
int n_mats = 8, int n_used = 2, bool b = false,
int64_t m = 32, int64_t n = 32, int64_t k = 32,
- float amax = 1.0f, int64_t m_v = 0, bool ids_offset = false)
+ float amax = 1.0f, int64_t m_v = 0, bool ids_offset = false, bool skip = false)
: type_a(type_a), type_b(type_b), n_mats(n_mats), n_used(n_used), b(b),
- m(m), n(n), k(k), amax(amax), m_v(m_v), ids_offset(ids_offset) {
+ m(m), n(n), k(k), amax(amax), m_v(m_v), ids_offset(ids_offset), skip(skip) {
GGML_ASSERT(n_used <= n_mats);
GGML_ASSERT(m_v == 0 || m_v > m);
}
@@ -5659,11 +5709,11 @@ struct test_mul_mat_id : public test_case {
}
void initialize_tensors(ggml_context * ctx) override {
- init_mul_mat_id_tensors(ctx, n_mats, amax);
+ init_mul_mat_id_tensors(ctx, n_mats, skip, amax);
}
void reinit_perf_iter(ggml_context * ctx) override {
- init_mul_mat_id_ids(ctx, n_mats);
+ init_mul_mat_id_ids(ctx, n_mats, skip);
}
};
@@ -5738,9 +5788,10 @@ struct test_mul_mat_id_fusion : public test_case {
const int64_t k;
const uint32_t o; // number of outputs
const bool mul;
+ const bool skip; // set some ids to a negative value
std::string vars() override {
- return VARS_TO_STR10(type_a, type_b, n_mats, n_used, b, m, n, k, o, mul);
+ return VARS_TO_STR11(type_a, type_b, n_mats, n_used, b, m, n, k, o, mul, skip);
}
double max_nmse_err() override {
@@ -5754,9 +5805,9 @@ struct test_mul_mat_id_fusion : public test_case {
test_mul_mat_id_fusion(ggml_type type_a = GGML_TYPE_F32, ggml_type type_b = GGML_TYPE_F32,
int n_mats = 8, int n_used = 2, bool b = false,
- int64_t m = 32, int64_t n = 32, int64_t k = 32, uint32_t o = 1, bool mul = false)
+ int64_t m = 32, int64_t n = 32, int64_t k = 32, uint32_t o = 1, bool mul = false, bool skip = false)
: type_a(type_a), type_b(type_b), n_mats(n_mats), n_used(n_used), b(b),
- m(m), n(n), k(k), o(o), mul(mul) {
+ m(m), n(n), k(k), o(o), mul(mul), skip(skip) {
GGML_ASSERT(n_used <= n_mats);
}
@@ -5796,7 +5847,7 @@ struct test_mul_mat_id_fusion : public test_case {
}
void initialize_tensors(ggml_context * ctx) override {
- init_mul_mat_id_tensors(ctx, n_mats);
+ init_mul_mat_id_tensors(ctx, n_mats, skip);
}
bool run_whole_graph() override { return true; }
@@ -7754,19 +7805,20 @@ struct test_mul_mat_vec_fusion : public test_case {
const bool with_gate;
const bool with_lane_scale;
std::array<int64_t, 2> batch_dims;
+ const bool skip; // set some ids to a negative value (only for use_id)
test_mul_mat_vec_fusion(ggml_type type, ggml_glu_op op, int64_t m, int64_t n, int64_t k,
bool use_id = false, int n_mats = 1, int n_used = 1, bool b = false, bool with_bias = false, bool with_gate = true,
- bool with_lane_scale = false, std::array<int64_t, 2> batch_dims = {4, 2})
+ bool with_lane_scale = false, std::array<int64_t, 2> batch_dims = {4, 2}, bool skip = false)
: type(type), glu_op(op), m(m), n(n), k(k), use_id(use_id), n_mats(n_mats), n_used(n_used), b(b), with_bias(with_bias),
- with_gate(with_gate), with_lane_scale(with_lane_scale), batch_dims(batch_dims) {
+ with_gate(with_gate), with_lane_scale(with_lane_scale), batch_dims(batch_dims), skip(skip) {
if (use_id) {
GGML_ASSERT(n_used <= n_mats);
}
}
std::string vars() override {
- return VARS_TO_STR13(type, glu_op, m, n, k, use_id, n_mats, n_used, b, with_bias, with_gate, with_lane_scale, batch_dims);
+ return VARS_TO_STR14(type, glu_op, m, n, k, use_id, n_mats, n_used, b, with_bias, with_gate, with_lane_scale, batch_dims, skip);
}
std::string op_desc(ggml_tensor * t) override {
@@ -7917,7 +7969,7 @@ struct test_mul_mat_vec_fusion : public test_case {
init_tensor_uniform(t);
}
} else {
- init_mul_mat_id_tensors(ctx, n_mats);
+ init_mul_mat_id_tensors(ctx, n_mats, skip);
}
}
@@ -11453,6 +11505,21 @@ static std::vector<std::unique_ptr<test_case>> make_test_cases_eval() {
}
}
+ // skipped slots (negative ids)
+ for (ggml_type type_a : {GGML_TYPE_F32, GGML_TYPE_F16, GGML_TYPE_BF16, GGML_TYPE_Q4_0, GGML_TYPE_Q4_K, GGML_TYPE_Q8_0, GGML_TYPE_MXFP4}) {
+ for (int n_used : {1, 2, 4}) {
+ for (bool b : {false, true}) {
+ for (int n : {1, 4, 33, 129}) {
+ test_cases.emplace_back(new test_mul_mat_id(type_a, GGML_TYPE_F32, 8, n_used, b, 512, n, 256, 1.0f, 0, false, true));
+ }
+ }
+ }
+ }
+ // n_used == n_mats, so every expert can be skipped
+ test_cases.emplace_back(new test_mul_mat_id(GGML_TYPE_F16, GGML_TYPE_F32, 4, 4, false, 512, 64, 256, 1.0f, 0, false, true));
+ test_cases.emplace_back(new test_mul_mat_id(GGML_TYPE_F16, GGML_TYPE_F32, 1, 1, false, 8, 16, 1, 1.0f, 0, false, true));
+ test_cases.emplace_back(new test_mul_mat_id_fusion(GGML_TYPE_F16, GGML_TYPE_F32, 16, 16, false, 32, 32, 32, 3, false, true));
+
for (int bs : {1, 4, 512}) {
for (ggml_type type_a : {GGML_TYPE_F32, GGML_TYPE_F16, GGML_TYPE_Q4_0, GGML_TYPE_Q4_K}) {
for (ggml_type type_b : {GGML_TYPE_F32}) {
@@ -11500,6 +11567,7 @@ static std::vector<std::unique_ptr<test_case>> make_test_cases_eval() {
for (int n_embd : {32, 129}) {
for (int n_token : {1, 32, 129}) {
test_cases.emplace_back(new test_add_id(type_a, type_b, n_embd, n_mats, n_used, n_token));
+ test_cases.emplace_back(new test_add_id(type_a, type_b, n_embd, n_mats, n_used, n_token, true));
}
}
}
@@ -12238,6 +12306,11 @@ static std::vector<std::unique_ptr<test_case>> make_test_cases_eval() {
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}));
}
+ // the lane scale does a GET_ROWS on the ids, which does not accept -1
+ if (use_id && !with_lane_scale) {
+ test_cases.emplace_back(new test_mul_mat_vec_fusion(type, glu_op, 8, 32, 256,
+ use_id, 16, 8, b, with_bias, with_gate, with_lane_scale, {1, 1}, true));
+ }
}
}
}