[ngxson]

Overview

Sorry for the big PR, but unfortunately this change cannot be gated by a graph build-time condition, so I need to add it to all backends at once. Important to note: code is 100% AI-generated.

This change allow mul_mat_id to read -1 index as a sentiment for "skip computation". This allow implement an "expert mask" similar to pytorch impl of MoE.

It's useful for these 2 models:

• Feature Request: Support for ERNIE-4.5-VL #15512 --> model has "small" and "big" experts, so it will be split into 2 mul_mat_id after this PR, one masked small exp and one masked big exp
• model: support Longcat-Flash (need testing) #19182 --> implemented via a correct by sub-optimal solution, need this PR to work properly

Details about masking logic

The masking impl wll be like this:

Example on x = [2, 0, 1, 3] ; mask out all exp >= 2 --> expected [-1, 0, 1, -1]
---------------------------------------
x = [2, 0, 1, 3]

Step 1: y = sub(x, 1)
        y = [1, -1, 0, 2]

Step 2: c = clamp(y, 0, 1)
        c = [1, 0, 0, 1]

Step 3: two_c = sub(c, neg(c))
        neg(c) = [-1, 0, 0, -1]
        two_c  = [2, 0, 0, 2]

Step 4: z = clamp(x, -1, 1)
        z = [1, 0, 1, 1]

Step 5: result = sub(z, two_c)
        result = [1-2, 0-0, 1-0, 1-2]
               = [-1, 0, 1, -1]

Final output: [-1, 0, 1, -1]

Generalized formula for masking exp >= N:

result = clamp(x, -1, N-1) - N * clamp(x - (N-1), 0, 1)

Additional info

Most backends are covered, but I only tested on Metal.

Some are real skip (computation is skipped), and some is fake mul mat with 0

Important: CUDA is not included in this PR, see comments below

Backend	Path	Skips computation?	How
CPU	ggml-cpu.c mul_mat_id	Real skip	slot dropped from matrix_rows, dst row memset
CPU (repack extra buffer)	repack.cpp	Real skip	same grouping loop
CPU (spacemit IME)	spacemit/ime.cpp	Real skip	same grouping loop
Metal	kernel_mul_mm_id (batch >= 32)	Real skip	map0 marks the slot 0xFFFF so no expert claims it, kernel_mul_mm_id_zero writes 0
Metal	kernel_mul_mv_id (batch < 32)	Real skip	threadgroup returns before reading src0, writes 0
Vulkan	mul_mm / mul_mmq / mul_mm_cm2 (batch)	Real skip	excluded by count_experts + load_row_ids, count_experts also writes 0 to the dst row
Vulkan	mul_mat_vec (bs=1)	Compute + discard	shared get_offsets clamps the expert to 0, reduce_result overwrites with 0
SYCL	batched sort path	Real skip	slot dropped from the counting sort, k_zero_dst_rows kernel writes 0
SYCL	ne12 == 1 per-slot path	Real skip	ggml_sycl_mul_mat not called, memset instead
SYCL	mmvq MoE GEMV (+ reorder)	Real skip	warp returns before reading src0, writes 0
OpenCL	gemm MoE (sorted, batch)	Real skip	slot dropped in kernel_moe_histogram/kernel_moe_scatter, kernel_moe_zero_dst writes 0
OpenCL	gemv_moe_* and mul_mv_id_*	Real skip	workgroup returns before reading src0, kernel_moe_zero_dst writes 0
WebGPU	mul_mat_id (batch)	Real skip	slot dropped in mul_mat_id_gather, same shader writes 0 to the dst row
WebGPU	mul_mat_id_vec (bs=1)	Real skip	slot never enters gathered_count_ids, workgroup 0 writes 0, extra workgroups return
Hexagon	op_matmul_id (HMX + HVX)	Real skip	scan_expert_ids / hvx_mv_id already skipped -1, zero_skipped_dst_rows writes 0
zendnn	mul_mat_id	Real skip	slot dropped from the grouping loop, dst row memset
ggml-et	mul_mat_id_* kernels	No work saved (already correct)	already wrote 0 for out-of-range ids before this change; the kernel is per-output-element, so nothing is wasted either
CANN	ggml_cann_mul_mat_id_fp / _quant	NOT IMPLEMENTED	uses aclnn IndexSelect with the ids verbatim; a negative index is invalid
OpenVINO	translate_mul_mat_id	NOT IMPLEMENTED	uses ov::op::Gather, where a negative index means "from the end" -> silently picks the wrong expert
RPC	passthrough	n/a	forwards to the remote backend

Requirements

• I have read and agree with the contributing guidelines
• AI usage disclosure: yes

[arthw]

It's good job!
It's OK for SYCL part!
All related UT cases are passed.

Thank you!

[ggerganov]

The Metal changes are good.

[ngxson]

I think CUDA / vulkan / webgpu changes are still not correct though, but still keeping this PR if backend maintainers are interested in fixing it (you can push to another branch then I will cherry-pick it)

[jeffbolznv]

Should these last two be strides? But this might work if the output is guaranteed contiguous

[ngxson]

I attempted to address this in 35dbed7 , PTAL

(feel free to let me know if vulkan code is ok here, or it's better to be a follow-up PR, cc @0cc4m too)

[jeffbolznv]

I think the vulkan code is probably fine as-is. I had one small suggestion and one non-issue that I resolved.

[yomaytk]

@ngxson I pushed a fix for the webgpu backend to the following branch, so can you cherry-pick the latest commit? I confirmed that mul_mat_id ops test passes with this commit on my M5 and NVIDIA machine.

https://github.com/yomaytk/llama.cpp/tree/pr-26631-fix-webgpu

The commit ensures that all threads reach the workgroup barriers in mul_mat_id_vec, which fixes the error 'workgroupBarrier' must only be called from uniform control flow. The other changes in this PR for the webgpu backend look good to me.

[max-krasnyansky]

For hexagon we're going to need to replace memset(0) with hvx_splat_f32(0) for better performance, but we can do that in a follow up PR once the rest of the stuff is merged.

[Titaniumtown]

@ggml-org/ggml-cuda can one of you review this PR? Thank you!

[Titaniumtown]

@ngxson arthw reviewed it for sycl, I don't need to review. I just noticed this PR was sitting here for a while so wanted to push it along

[0cc4m]

Vulkan changes are fine

[am17an]

n_tokens here is blockIdx.y which will overflow for n_tokens > 65535.

[ngxson]

can you have a look at 23d8ac0 ?

[ggerganov]

After the #30373 changes, the Metal patch here will likely need to be adjusted a bit.

[ngxson]

@ggerganov @JohannesGaessler @am17an could you please have a look?

I'm waiting for final approval for metal and CUDA

[ggerganov]

The Metal change is good.

[ngxson]

OK so as @JohannesGaessler requested, I excluded CUDA from this PR. For test-backend-ops, I included a small code (a bit hacky) to skip CUDA tests. Remind that we cannot automatically fallback to CPU in this case, because it's impossible to tell if backend support it without looking at tensor data. That code is to be reverted once CUDA support this feature.

I will merge this PR once the CI passes. Won't ask for another round of review, otherwise it will take forever 😓

[ggerganov]

Add a TODO tag so we keep track of this in the https://github.com/ggml-org/llama.cpp-dev/blob/master/README-todo-counts.md