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>

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`

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;

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));
     }
 }
 

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;

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;

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;

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];

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);

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;

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) {

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) {

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,

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,

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, &region, &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, &region, &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

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)) {

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;

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

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

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;

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

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

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;

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

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]),

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;

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;
 

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;

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;

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;

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);
 

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;

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();

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;

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 {

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 });
     }
 

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];

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();
 }

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]);

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({

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];

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;
+        }
     }
 }

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++) {

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++;

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,

src/llama-hparams.h

@@ -61,7 +61,7 @@ struct llama_hparams {
     //       for example: n_embd_out() vs n_embd_out_impl
 
     bool vocab_only;
-    bool no_alloc; // TODO: no_alloc is not a model hyperparameter; move it out of llama_hparams
+    bool no_alloc;
     bool rope_finetuned;
     bool use_par_res;
     bool swin_norm;

src/llama-kv-cache-dsv4.cpp

@@ -905,12 +905,13 @@ llama_dsv4_comp_state::llama_dsv4_comp_state(
             uint32_t        n_rs_seq,
         const char    * name,
         const llama_memory_i::layer_filter_cb & filter) :
-    hparams(model.hparams),
     ratio(ratio),
     state_size(state_size),
     n_embd_state(n_embd_state),
     n_stream(unified ? 1 : n_seq_max),
     n_rs_seq(n_rs_seq) {
+    const llama_hparams & hparams = model.hparams;
+
     struct ggml_backend_buft_comparator {
         bool operator()(const ggml_backend_buffer_type_t & lhs, const ggml_backend_buffer_type_t & rhs) const {
             return strcmp(ggml_backend_buft_name(lhs), ggml_backend_buft_name(rhs)) < 0;
@@ -985,15 +986,7 @@ llama_dsv4_comp_state::llama_dsv4_comp_state(
     }
 
     for (auto & [buft, ctx] : ctx_map) {
-        ggml_backend_buffer_t buf;
-        if (hparams.no_alloc) {
-            buf = ggml_backend_buft_alloc_buffer(buft, /*size =*/ 0); // dummy buffer
-            for (ggml_tensor * t = ggml_get_first_tensor(ctx.get()); t != nullptr; t = ggml_get_next_tensor(ctx.get(), t)) {
-                t->buffer = buf; // set dummy buffer for DSV4 state so that the backend scheduler won't try to allocate it
-            }
-        } else {
-            buf = ggml_backend_alloc_ctx_tensors_from_buft(ctx.get(), buft); // real buffer
-        }
+        ggml_backend_buffer_t buf = ggml_backend_alloc_ctx_tensors_from_buft(ctx.get(), buft);
         if (!buf) {
             throw std::runtime_error("failed to allocate buffer for DSV4 compressor state");
         }
@@ -1081,15 +1074,9 @@ uint32_t llama_dsv4_comp_state::get_n_rows() const {
 
 std::map<ggml_backend_buffer_type_t, size_t> llama_dsv4_comp_state::memory_breakdown() const {
     std::map<ggml_backend_buffer_type_t, size_t> ret;
-    for (const auto & [ctx, buf] : ctxs_bufs) {
+    for (const auto & [_, buf] : ctxs_bufs) {
         ggml_backend_buffer_type_t buft = ggml_backend_buffer_get_type(buf.get());
-
-        if (hparams.no_alloc) {
-            GGML_ASSERT(ggml_backend_buffer_get_base(buf.get()) == nullptr);
-            ret[buft] += ggml_backend_alloc_ctx_tensors_from_buft_size(ctx.get(), buft);
-        } else {
-            ret[buft] += ggml_backend_buffer_get_size(buf.get());
-        }
+        ret[buft] += ggml_backend_buffer_get_size(buf.get());
     }
     return ret;
 }

src/llama-kv-cache-dsv4.h

@@ -60,8 +60,6 @@ private:
         std::vector<ggml_tensor *> score_stream;
     };
 
-    const llama_hparams & hparams;
-
     const uint32_t ratio;
     const uint32_t state_size;
     const uint32_t n_embd_state;

src/llama-memory-recurrent.cpp

@@ -116,15 +116,7 @@ llama_memory_recurrent::llama_memory_recurrent(
 
     // allocate tensors and initialize the buffers to avoid NaNs in the padding
     for (auto & [buft, ctx] : ctx_map) {
-        ggml_backend_buffer_t buf;
-        if (hparams.no_alloc) {
-            buf = ggml_backend_buft_alloc_buffer(buft, /*size =*/ 0); // dummy buffer
-            for (ggml_tensor * t = ggml_get_first_tensor(ctx.get()); t != nullptr; t = ggml_get_next_tensor(ctx.get(), t)) {
-                t->buffer = buf; // set dummy buffer for RS cache so that the backend scheduler won't try to allocate it
-            }
-        } else {
-            buf = ggml_backend_alloc_ctx_tensors_from_buft(ctx.get(), buft); // real buffer
-        }
+        ggml_backend_buffer_t buf = ggml_backend_alloc_ctx_tensors_from_buft(ctx.get(), buft);
         if (!buf) {
             throw std::runtime_error("failed to allocate buffer for rs cache");
         }
@@ -443,15 +435,8 @@ void llama_memory_recurrent::set_rs_idx(llama_seq_id seq_id, uint32_t idx) {
 
 std::map<ggml_backend_buffer_type_t, size_t> llama_memory_recurrent::memory_breakdown() const {
     std::map<ggml_backend_buffer_type_t, size_t> ret;
-    for (const auto & [ctx, buf] : ctxs_bufs) {
-        ggml_backend_buffer_type_t buft = ggml_backend_buffer_get_type(buf.get());
-
-        if (hparams.no_alloc) {
-            GGML_ASSERT(ggml_backend_buffer_get_base(buf.get()) == nullptr);
-            ret[buft] += ggml_backend_alloc_ctx_tensors_from_buft_size(ctx.get(), buft);
-        } else {
-            ret[buft] += ggml_backend_buffer_get_size(buf.get());
-        }
+    for (const auto & [_, buf] : ctxs_bufs) {
+        ret[ggml_backend_buffer_get_type(buf.get())] += ggml_backend_buffer_get_size(buf.get());
     }
     return ret;
 }

tests/CMakeLists.txt

@@ -240,6 +240,26 @@ if (NOT WIN32 OR NOT BUILD_SHARED_LIBS)
         FIXTURES_SETUP generate-models
     )
 
+    # Target model for test-speculative-rejection-sampling
+    set(SPEC_TARGET_DIR "${CMAKE_CURRENT_BINARY_DIR}/test-speculative-models/target")
+    file(MAKE_DIRECTORY "${SPEC_TARGET_DIR}")
+    llama_test(
+        test-llama-archs
+        NAME test-generate-speculative-target
+        LABEL main
+        ARGS -a llama -s 999 -o "${SPEC_TARGET_DIR}"
+    )
+    set_tests_properties(
+        test-generate-speculative-target
+        PROPERTIES FIXTURES_SETUP speculative-model
+    )
+
+    # To avoid recreating the model use the fixture exclude any option:
+    # -FA "^speculative-model$"
+    llama_build_and_test(test-speculative-rejection-sampling.cpp ARGS
+        -m "${SPEC_TARGET_DIR}/llama-dense.gguf")
+    set_tests_properties(test-speculative-rejection-sampling PROPERTIES FIXTURES_REQUIRED speculative-model)
+
     # Test context state handling across all architectures, using the generated dummy models
     llama_test(
         test-llama-context

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));
+                                }
                             }
                         }
                     }

tests/test-speculative-rejection-sampling.cpp

@@ -0,0 +1,229 @@
+#include "sampling.h"
+#include "llama.h"
+#include "common.h"
+#include "arg.h"
+#include "log.h"
+
+#include <cmath>
+#include <cstdlib>
+#include <cstring>
+#include <optional>
+#include <random>
+#include <unordered_map>
+
+// Tests that common_sampler_sample_and_accept_n_rejection preserves the target
+// distribution regardless of what the draft distribution Q looks like.
+
+static int test_speculative_rejection_sampling(common_params & params, int n_trials) {
+    params.sampling.backend_sampling = false;
+    auto llama_init = common_init_from_params(params);
+
+    auto * model_tgt = llama_init->model();
+    auto * ctx_tgt   = llama_init->context();
+    if (!model_tgt || !ctx_tgt) {
+        LOG_ERR("failed to initialize target model and context\n");
+        return 1;
+    }
+
+    const llama_tokens prompt = { 1, 2, 3, 1, 2 };
+
+    // Process the prompt with the target model to get its true distribution for
+    // the position right after the prompt.
+    common_batch batch = common_batch_get_one(ctx_tgt, prompt);
+    if (llama_process(ctx_tgt, LLAMA_PROCESS_TYPE_DECODE, batch.get()) != 0) {
+        LOG_ERR("failed to decode the test prompt\n");
+        return 1;
+    }
+
+    common_params_sampling sampling;
+    sampling.seed     = 18;
+    sampling.samplers = { COMMON_SAMPLER_TYPE_TEMPERATURE };
+    sampling.temp     = 1.0f;
+
+    // Initialize the sampler for the target model.
+    common_sampler_ptr smpl(common_sampler_init(model_tgt, sampling));
+    if (!smpl) {
+        LOG_ERR("failed to initialize target sampler\n");
+        return 1;
+    }
+
+    // Sample once so the sampler chain runs and populates cur_p. We don't need
+    // the returned token id itself so the return value is ignored.
+    common_sampler_sample(smpl.get(), ctx_tgt, -1);
+
+    // Get the target candidates distribution in sorted order.
+    const llama_token_data_array * dist_tgt = common_sampler_get_candidates(smpl.get(), true);
+
+    // Clone smpl before the trial loop starts. We need the dist_tgt pointer to stay valid.
+    common_sampler_ptr smpl_verify(common_sampler_clone(smpl.get()));
+    if (!smpl_verify) {
+        LOG_ERR("failed to clone target sampler\n");
+        return 1;
+    }
+
+    const llama_vocab * vocab = llama_model_get_vocab(model_tgt);
+    const int n_vocab = llama_vocab_n_tokens(vocab);
+
+    std::mt19937 rng(18);
+    std::uniform_real_distribution<float> uni(0.0f, 1.0f);
+
+    std::unordered_map<llama_token, int> token_counts;
+
+    std::vector<llama_token_data> q(n_vocab);
+    for (int i = 0; i < n_vocab; ++i) {
+        q[i].id = i;
+        q[i].logit = 0.0f;
+    }
+
+    for (int trial = 0; trial < n_trials; ++trial) {
+        float q_sum = 0.0f;
+
+        for (auto & td : q) {
+            td.p = uni(rng);
+            q_sum += td.p;
+        }
+
+        const float inv_sum = 1.0f / q_sum;
+        for (auto & td : q) {
+            td.p *= inv_sum;
+        }
+
+        float u = uni(rng);
+
+        // Fallback to the last bucket to absorb residual probability from FP precision loss
+        llama_token draft_tok = q.back().id;
+        for (const auto & e : q) {
+            u -= e.p;
+            if (u <= 0.0f) {
+                draft_tok = e.id;
+                break;
+            }
+        }
+
+        // Both indices reuse the prompt's last output logits. The second entry
+        // is the bonus draw slot; it is always stale and its result is ignored.
+        const llama_tokens draft = { draft_tok };
+        const std::vector<std::vector<llama_token_data>> draft_q = { q };
+        const std::vector<int> idxs = { -1, -1 };
+
+        const std::vector<llama_token> result =
+            common_sampler_sample_and_accept_n_rejection(smpl_verify.get(), ctx_tgt, idxs, draft, draft_q, false);
+
+        token_counts[result[0]]++;
+    }
+
+    LOG("\n");
+    LOG("Empirical vs target distribution after %d trials:\n", n_trials);
+
+    bool   ok      = true;
+    size_t covered = 0;
+
+    auto get_count_for = [&](llama_token id) -> int {
+        const auto it = token_counts.find(id);
+        return it != token_counts.end() ? it->second : 0;
+    };
+
+    // Iterate over the target distribution
+    for (size_t i = 0; i < dist_tgt->size; ++i) {
+        const llama_token id  = dist_tgt->data[i].id;
+        const float       p_t = dist_tgt->data[i].p;
+
+        const int   count = get_count_for(id);
+        const float p_e   = (float) count / n_trials;
+        const float diff  = std::fabs(p_t - p_e);
+
+        // We only deal with ids in dist_tgt, and if each trial result was one
+        // of dist_tgt's candidates then covered ends up equal to the number of
+        // trials. If not this inconsistency tells us how many trials produced
+        // a token outside of the target's distribution.
+        covered += count;
+
+        const float standard_error = std::sqrt(p_t * (1.0f - p_t) / n_trials);
+        const float tolerance      = 6.0f * standard_error;
+
+        LOG("token %3d : target = %f, empirical = %f, diff = %f, tol = %f\n", id, p_t, p_e, diff, tolerance);
+
+        if (diff > tolerance) {
+            LOG_ERR("token %d: empirical probability %f too far from target %f (diff %f > tol %f)\n",
+                    id, p_e, p_t, diff, tolerance);
+            ok = false;
+        }
+    }
+
+    if (covered != (size_t) n_trials) {
+        LOG_ERR("%zu of %d trials produced a token outside dist_tgt's candidates\n",
+                (size_t) n_trials - covered, n_trials);
+        ok = false;
+    }
+
+    if (!ok) {
+        return 1;
+    }
+
+    return 0;
+}
+
+struct parsed_args {
+    int n_trials;
+    std::vector<char *> argv;
+};
+
+static std::optional<parsed_args> parse_n_trials_arg(int argc, char ** argv, int default_n_trials) {
+    parsed_args result;
+    result.n_trials = default_n_trials;
+    result.argv.push_back(argv[0]);
+
+    for (int i = 1; i < argc; ++i) {
+        if (strcmp(argv[i], "--n-trials") == 0) {
+            if (i + 1 >= argc) {
+                LOG_ERR("%s: --n-trials requires a number argument\n", __func__);
+                return std::nullopt;
+            }
+            int value = std::atoi(argv[i + 1]);
+            if (value <= 0) {
+                LOG_ERR("%s: --n-trials requires a positive number argument\n", __func__);
+                return std::nullopt;
+            }
+            result.n_trials = value;
+            i++;
+        } else {
+            result.argv.push_back(argv[i]);
+        }
+    }
+    result.argv.push_back(nullptr);
+
+    return result;
+}
+
+int main(int argc, char ** argv) {
+    std::setlocale(LC_NUMERIC, "C");
+
+    common_params params;
+    params.n_ctx = 256;
+
+    common_init();
+
+    auto parsed = parse_n_trials_arg(argc, argv, 5000);
+    if (!parsed) {
+        return 1;
+    }
+    const int fargc = (int) parsed->argv.size() - 1;
+
+    if (!common_params_parse(fargc, parsed->argv.data(), params, LLAMA_EXAMPLE_COMMON)) {
+        return 1;
+    }
+
+    LOG("Using target model '%s'\n", params.model.path.c_str());
+
+    llama_backend_init();
+
+    const int result = test_speculative_rejection_sampling(params, parsed->n_trials);
+    if (result == 0) {
+        LOG("\nAll tests passed.\n");
+    } else {
+        LOG("\nSome tests failed.\n");
+    }
+
+    llama_backend_free();
+    return result;
+}