Repository: ggml-org/llama.cpp
Pull request: ggml-zendnn : add MUL_MAT_ID op support for MoE models

This PR adds support for the `MUL_MAT_ID` op in the ZenDNN backend. The `MUL_MAT_ID` op is used in Mixture-of-Experts (MoE) models to perform matrix multiplication with expert selection based on input IDs.

- Acceleration of the `MUL_MAT_ID` op for MoE models using ZenDNN.
- Fallback to the CPU backend when the total number of experts exceeds 32, as ZenDNN may not be efficient in these scenarios (for now).
- Updated the ZenDNN lib to the latest bits (`ZenDNN-2026-WW13`).

### Test Configuration
- **Hardware**: AMD EPYC 9004 Series (Zen 4)
- **Threads**: 96
- **Tool**: llama-bench and llama-batched-bench
- **Precision**: BF16
- **k/v cache**: bf16
- **llama.cpp version**: 8690
- **ZenDNN version**: 1.0.0 (ZenDNN-2026-WW13)
- **Environment**: `ZENDNNL_MATMUL_ALGO=1` (Blocked AOCL BLIS)

### Benchmark Results

#### Optimal `ubatch` Selection - CPU Backend (ggml-cpu)

The heatmaps below show prompt-processing speed (t/s) across all valid `(prompt length, ubatch)` combinations for the CPU backend.
The optimal ubatch for the CPU backend varies by model - `512` for gpt-oss-20B, `768` for Mixtral-8x7B and Phi-3.5-MoE. For ZenDNN, `4096` is consistently the best ubatch across all three models.

*gpt-oss-20B · Best CPU ubatch: **512***
<img width="1522" height="1270" alt="heatmap_CPU_gpt-oss-20B-BF16 gguf_96t" src="https://github.com/user-attachments/assets/e9cbd3df-25eb-4175-9a2e-b8c92073e091" />

*Mixtral-8x7B · Best CPU ubatch: **768***
<img width="1522" height="1270" alt="heatmap_CPU_Mixtral-8x7B-Instruct-v0 1-BF16 gguf_96t" src="https://github.com/user-attachments/assets/b4c1b3a6-85d5-44a4-aa75-bb18fe2a4ce2" />

*Phi-3.5-MoE · Best CPU ubatch: **768***
<img width="1522" height="1270" alt="heatmap_CPU_Phi-3 5-MoE-16x4 1B-instruct-BF16 gguf_96t" src="https://github.com/user-attachments/assets/6be47aea-a45f-44c2-8772-89f2636cbfa7" />

### `llama-bench` - Prompt Processing Throughput (tokens/sec)

> Config: 96 threads · BF16 · bf16 k/v cache
> CPU numbers use the best ubatch per model (512 / 768 / 768). ZenDNN numbers use ubatch=4096.

#### Mixtral-8x7B-Instruct-v0.1 - 86.99 GiB · 46.70B params

| Test   | CPU (t/s)      | ZenDNN (t/s)   | Speedup |
|--------|---------------:|---------------:|:-------:|
| pp128  | 171.30 ± 0.36  | 162.33 ± 0.19  | 0.95×   |
| pp256  | 181.34 ± 0.10  | 235.79 ± 0.88  | **1.30×** |
| pp384  | 193.57 ± 0.06  | 276.02 ± 0.84  | **1.43×** |
| pp512  | 184.60 ± 0.06  | 308.42 ± 0.22  | **1.67×** |
| pp768  | 185.62 ± 0.12  | 360.08 ± 0.43  | **1.94×** |
| pp1024 | 181.51 ± 0.08  | 374.11 ± 0.48  | **2.06×** |
| pp1536 | 180.46 ± 0.13  | 392.31 ± 0.53  | **2.17×** |
| pp2048 | 177.43 ± 0.07  | 411.17 ± 0.28  | **2.32×** |
| pp3072 | 171.63 ± 0.02  | 413.27 ± 0.07  | **2.41×** |
| pp4096 | 166.41 ± 0.03  | 393.64 ± 0.48  | **2.37×** |

#### Phi-3.5-MoE-instruct - 78.00 GiB · 41.87B params

| Test   | CPU (t/s)      | ZenDNN (t/s)    | Speedup |
|--------|---------------:|----------------:|:-------:|
| pp128  | 282.71 ± 0.82  | 210.67 ± 11.31  | 0.74×   |
| pp256  | 347.71 ± 0.40  | 298.33 ± 3.35   | 0.86×   |
| pp384  | 359.16 ± 0.63  | 396.00 ± 4.39   | **1.10×** |
| pp512  | 371.69 ± 0.26  | 436.85 ± 2.89   | **1.18×** |
| pp768  | 375.51 ± 0.68  | 512.55 ± 4.00   | **1.37×** |
| pp1024 | 360.19 ± 0.39  | 536.58 ± 7.34   | **1.49×** |
| pp1536 | 356.10 ± 0.48  | 561.33 ± 3.22   | **1.58×** |
| pp2048 | 344.96 ± 0.62  | 583.05 ± 1.22   | **1.69×** |
| pp3072 | 323.80 ± 1.30  | 600.95 ± 3.63   | **1.86×** |
| pp4096 | 307.02 ± 0.08  | 556.82 ± 9.53   | **1.81×** |

#### unsloth/gpt-oss-20b - 38.97 GiB · 20.91B params

| Test   | CPU (t/s)      | ZenDNN (t/s)    | Speedup |
|--------|---------------:|----------------:|:-------:|
| pp128  | 471.56 ± 2.90  | 284.14 ± 6.81   | 0.60×   |
| pp256  | 589.39 ± 1.80  | 426.90 ± 7.60   | 0.72×   |
| pp384  | 599.90 ± 1.32  | 512.07 ± 5.06   | 0.85×   |
| pp512  | 618.43 ± 0.38  | 563.97 ± 1.11   | 0.91×   |
| pp768  | 587.94 ± 0.81  | 633.89 ± 0.87   | **1.08×** |
| pp1024 | 577.87 ± 0.54  | 609.45 ± 1.67   | **1.05×** |
| pp1536 | 553.02 ± 0.35  | 692.40 ± 0.63   | **1.25×** |
| pp2048 | 533.92 ± 0.38  | 764.32 ± 1.61   | **1.43×** |
| pp3072 | 498.19 ± 0.50  | 672.33 ± 37.24  | **1.35×** |
| pp4096 | 467.38 ± 5.60  | 579.32 ± 31.03  | **1.24×** |


### `llama-batched-bench` - PP + TG Performance

> **Configuration:** 96 threads · BF16 compute · BF16 KV cache
> **Prompt length:** 512 tokens
> **Generation length:** 128 tokens
> **Context length per request:** 512 + 128 = 640 tokens
>
> **Total context size (batch-wise):**
>
> * Batch 16 -> 640 × 16 = 10,240 tokens
> * Batch 32 -> 640 × 32 = 20,480 tokens

#### Mixtral-8x7B-Instruct-v0.1

| Batch | CPU uBatch | ZenDNN uBatch | CPU PP (t/s) | ZenDNN PP (t/s) | CPU TG (t/s) | ZenDNN TG (t/s) | PP Speedup |
| ----: | ---------: | ------------: | -----------: | --------------: | -----------: | --------------: | :--------: |
|    16 |       4096 |          4096 |       194.48 |          468.08 |        56.48 |           42.87 | **2.41x**  |
|    32 |       4096 |          8192 |       199.08 |          515.11 |        92.38 |           69.00 | **2.59x**  |


#### Phi-3.5-MoE-instruct

| Batch | CPU uBatch | ZenDNN uBatch | CPU PP (t/s) | ZenDNN PP (t/s) | CPU TG (t/s) | ZenDNN TG (t/s) | PP Speedup |
| ----: | ---------: | ------------: | -----------: | --------------: | -----------: | --------------: | :--------: |
|    16 |       1024 |          8192 |       412.05 |          689.85 |       108.46 |           65.02 |    **1.67x**   |
|    32 |       1024 |          8192 |       415.51 |          825.44 |       163.35 |           98.43 |    **1.99x**   |


#### unsloth/gpt-oss-20b

| Batch | CPU uBatch | ZenDNN uBatch | CPU PP (t/s) | ZenDNN PP (t/s) | CPU TG (t/s) | ZenDNN TG (t/s) | PP Speedup |
| ----: | ---------: | ------------: | -----------: | --------------: | -----------: | --------------: | :--------: |
|    16 |        512 |         16384 |       800.39 |         1338.16 |       225.68 |           96.59 |    **1.67x**   |
|    32 |        512 |         16384 |       810.13 |         1629.12 |       325.10 |          152.93 |    **2.01x**   |


> **Note:** TG (token generation) throughput may be lower with ZenDNN in some configurations with small batch sizes. The ZenDNN team is actively working on optimizations for smaller batch sizes, and improvements are expected in upcoming releases.

cc: @amukho @avinashcpandey @taronaeo @danbev

AI usage disclosure: AI assistance was used for scripts that generate graph and text refractor only.

Changed files:
filename: docs/backend/ZenDNN.md
patch: @@ -57,13 +57,14 @@ ZenDNN is optimized for AMD EPYC™ processors and AMD Ryzen™ processors based
 
 ## Supported Operations
 
-The ZenDNN backend currently accelerates **matrix multiplication (MUL_MAT)** operations only. Other operations are handled by the standard CPU backend.
+The ZenDNN backend accelerates **matrix multiplication (MUL_MAT)** and **expert-based matrix multiplication (MUL_MAT_ID)** operations. Other operations are handled by the standard CPU backend.
 
 | Operation    | Status  | Notes                                          |
 |:-------------|:-------:|:----------------------------------------------:|
 | MUL_MAT      | Support | Accelerated via ZenDNN LowOHA MatMul           |
+| MUL_MAT_ID   | Support | Accelerated via ZenDNN LowOHA MatMul (MoE)     |
 
-*Note:* Since only MUL_MAT is accelerated, models will benefit most from ZenDNN when matrix multiplications dominate the computational workload (which is typical for transformer-based LLMs).
+*Note:* Since MUL_MAT and MUL_MAT_ID are accelerated, models will benefit most from ZenDNN when matrix multiplications dominate the computational workload (which is typical for transformer-based LLMs and Mixture-of-Experts models).
 
 ## DataType Supports
 
@@ -181,7 +182,7 @@ For detailed profiling and logging options, refer to the [ZenDNN Logging Documen
 
 ## Known Issues
 
-- **Limited operation support**: Currently only matrix multiplication (MUL_MAT) is accelerated via ZenDNN. Other operations fall back to the standard CPU backend.
+- **Limited operation support**: Currently matrix multiplication (MUL_MAT) and expert-based matrix multiplication (MUL_MAT_ID) are accelerated via ZenDNN. Other operations fall back to the standard CPU backend. Future updates may expand supported operations.
 - **BF16 support**: BF16 operations require AMD Zen 4 or Zen 5 architecture (EPYC 9004/9005 series). On older CPUs, operations will use FP32.
 - **NUMA awareness**: For multi-socket systems, manual NUMA binding may be required for optimal performance.
 
@@ -216,4 +217,4 @@ Please add the **[ZenDNN]** prefix/tag in issues/PRs titles to help the ZenDNN-t
 
 ## TODO
 
-- Expand operation support beyond MUL_MAT (attention operations, activations, etc.)
+- Expand operation support beyond MUL_MAT and MUL_MAT_ID (attention operations, activations, etc.)

filename: docs/ops.md
patch: @@ -68,7 +68,7 @@ Legend:
 |                             MEAN | ❌ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ | ❌ |
 |                              MUL | ❌ | ✅ | ✅ | ✅ | 🟡 | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ |
 |                          MUL_MAT | 🟡 | 🟡 | 🟡 | 🟡 | 🟡 | 🟡 | 🟡 | 🟡 | 🟡 | 🟡 | 🟡 |
-|                       MUL_MAT_ID | ❌ | 🟡 | ✅ | ✅ | 🟡 | 🟡 | 🟡 | ✅ | ❌ | ❌ | ❌ |
+|                       MUL_MAT_ID | ❌ | 🟡 | ✅ | ✅ | 🟡 | 🟡 | 🟡 | ✅ | ❌ | 🟡 | ❌ |
 |                              NEG | ❌ | ✅ | ✅ | 🟡 | ✅ | ❌ | ✅ | 🟡 | ✅ | ❌ | ❌ |
 |                             NORM | ❌ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | 🟡 | ❌ | ❌ | ❌ |
 |                   OPT_STEP_ADAMW | ❌ | ❌ | ✅ | ✅ | ✅ | ❌ | ❌ | ✅ | ❌ | ❌ | ❌ |

filename: ggml/src/ggml-zendnn/CMakeLists.txt
patch: @@ -28,7 +28,7 @@ if (NOT ZENDNN_ROOT OR ZENDNN_ROOT STREQUAL "" OR ZENDNN_ROOT STREQUAL "OFF")
     ExternalProject_Add(
         zendnn
         GIT_REPOSITORY https://github.com/amd/ZenDNN.git
-        GIT_TAG a18adf8c605fb5f5e52cefd7eda08a7b18febbaf    # ZenDNN-2026-WW08
+        GIT_TAG f79f7321a1add65ced6397a6bfab7edba6e3e14e    # ZenDNN-2026-WW13
         PREFIX      ${ZENDNN_PREFIX}
         SOURCE_DIR  ${ZENDNN_SOURCE_DIR}
         BINARY_DIR  ${ZENDNN_BUILD_DIR}

filename: ggml/src/ggml-zendnn/ggml-zendnn.cpp
patch: @@ -190,6 +190,170 @@ static void ggml_zendnn_compute_forward_mul_mat(
     }
 }
 
+struct mmid_row_mapping {
+    int32_t i1;
+    int32_t i2;
+};
+
+static void ggml_zendnn_compute_forward_mul_mat_id(
+    ggml_backend_zendnn_context * ctx,
+    ggml_tensor * dst) {
+
+    const ggml_tensor * src0 = dst->src[0];  // expert weights
+    const ggml_tensor * src1 = dst->src[1];  // inputs
+    const ggml_tensor * ids  = dst->src[2];  // expert ids
+
+    GGML_TENSOR_BINARY_OP_LOCALS
+
+    // exit for no tokens to process
+    if (ne2 == 0 || ne11 == 0) {
+        return;
+    }
+
+    ggml_type         const vec_dot_type = src0->type;
+    ggml_from_float_t const from_float = ggml_get_type_traits(vec_dot_type)->from_float_ref;
+
+    // we don't support permuted src0 or src1
+    GGML_ASSERT(nb00 == ggml_type_size(src0->type));
+    GGML_ASSERT(nb10 == ggml_type_size(src1->type));
+
+    // dst cannot be transposed or permuted
+    GGML_ASSERT(nb0 == sizeof(float));
+    GGML_ASSERT(nb0 <= nb1);
+    GGML_ASSERT(nb1 <= nb2);
+    GGML_ASSERT(nb2 <= nb3);
+
+    GGML_ASSERT(ne03 == 1);
+    GGML_ASSERT(ne13 == 1);
+    GGML_ASSERT(ne3  == 1);
+
+    // row groups
+    const int n_ids = ids->ne[0]; // n_expert_used
+    const int n_as  = ne02;       // n_experts
+
+    std::vector<int64_t> matrix_row_counts(n_as, 0);
+    std::vector<std::vector<mmid_row_mapping>> matrix_rows(n_as);
+
+    int64_t max_rows = 0;
+    // group rows by expert (preprocessing step)
+    for (int64_t iid1 = 0; iid1 < ids->ne[1]; ++iid1) {
+        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);
+
+            matrix_rows[i02].push_back({id, iid1});
+            matrix_row_counts[i02]++;
+            if (matrix_row_counts[i02] > max_rows) {
+                max_rows = matrix_row_counts[i02];
+            }
+        }
+    }
+
+    if (max_rows == 0) {
+        return; // no rows to process
+    }
+
+    const size_t row_size = ggml_row_size(vec_dot_type, ne10);
+
+    // size for converting src1 rows to vec_dot_type if needed
+    const size_t nbw1 = row_size;
+    const size_t nbw2 = nbw1 * ne11;
+    const size_t nbw3 = nbw2 * ne12;
+    const size_t src1_conv_size = (src1->type != vec_dot_type) ? ne13 * nbw3 : 0;
+
+    // size for MoE gather/scatter buffers
+    const size_t wdata_cur_size = max_rows * row_size;
+    const size_t dst_cur_size = max_rows * ggml_row_size(dst->type, ne01);
+
+    // allocate single buffer for all needs
+    const size_t total_size = src1_conv_size + wdata_cur_size + dst_cur_size;
+    if (ctx->work_size < total_size) {
+        ctx->work_data.reset(new char[total_size]);
+        ctx->work_size = total_size;
+    }
+
+    // partition the buffer
+    char * work_data = ctx->work_data.get();
+    char * wdata_cur = work_data + src1_conv_size;
+    char * dst_cur = wdata_cur + wdata_cur_size;
+
+    if (src1->type != vec_dot_type) {
+        GGML_ASSERT(src1->type == GGML_TYPE_F32);
+
+        #pragma omp parallel for collapse(3) num_threads(ctx->n_threads) schedule(static)
+        for (int64_t i13 = 0; i13 < ne13; ++i13) {
+            for (int64_t i12 = 0; i12 < ne12; ++i12) {
+                for (int64_t i11 = 0; i11 < ne11; ++i11) {
+                    const float * src1_f32 = (float *)((char *)src1->data + i11*nb11 + i12*nb12 + i13*nb13);
+                    void * src1_conv = (char *)work_data + i11*nbw1 + i12*nbw2 + i13*nbw3;
+                    from_float(src1_f32, src1_conv, ne10);
+                }
+            }
+        }
+    }
+
+    const void * wdata = src1->type == vec_dot_type ? src1->data : work_data;
+
+    // process each expert with gather -> gemm -> scatter pattern
+    for (int64_t cur_a = 0; cur_a < n_as; ++cur_a) {
+        const int64_t cne1 = matrix_row_counts[cur_a];
+
+        if (cne1 == 0) {
+            continue;
+        }
+
+        const char * src0_cur = (const char *) src0->data + cur_a*nb02;
+
+        // gather input rows for this expert
+        #pragma omp parallel for num_threads(ctx->n_threads) schedule(static)
+        for (int64_t ir1 = 0; ir1 < cne1; ++ir1) {
+            const mmid_row_mapping & row_mapping = matrix_rows[cur_a][ir1];
+            const int64_t id = row_mapping.i1;
+            const int64_t i11 = id % ne11;
+            const int64_t i12 = row_mapping.i2;
+
+            std::memcpy(
+                wdata_cur + ir1 * row_size,
+                (const char *) wdata + (i11 + i12*ne11) * row_size,
+                row_size
+            );
+        }
+
+        // batched gemm for all tokens in this expert
+        if (!ggml_zendnn_sgemm(ctx,
+                              ne01,       // m
+                              cne1,       // n
+                              ne10,       // k
+                              src0_cur,
+                              ne00,       // lda
+                              wdata_cur,
+                              ne10,       // ldb
+                              dst_cur,
+                              ne01,       // ldc
+                              src0->type,
+                              vec_dot_type,
+                              dst->type)) {
+            GGML_ABORT("%s: ZenDNN sgemm failed\n", __func__);
+        }
+
+        // scatter output rows to destination
+        #pragma omp parallel for num_threads(ctx->n_threads) schedule(static)
+        for (int64_t ir1 = 0; ir1 < cne1; ++ir1) {
+            const mmid_row_mapping & row_mapping = matrix_rows[cur_a][ir1];
+            const int64_t id = row_mapping.i1;
+            const int64_t i1 = id;
+            const int64_t i2 = row_mapping.i2;
+
+            std::memcpy(
+                (char *) dst->data + i1*nb1 + i2*nb2,
+                dst_cur + ir1 * ggml_row_size(dst->type, ne01),
+                ggml_row_size(dst->type, ne01)
+            );
+        }
+    }
+}
+
 // backend interface
 
 static const char * ggml_backend_zendnn_get_name(ggml_backend_t backend) {
@@ -218,6 +382,9 @@ static ggml_status ggml_backend_zendnn_graph_compute(ggml_backend_t backend, ggm
             case GGML_OP_MUL_MAT:
                 ggml_zendnn_compute_forward_mul_mat(ctx, node);
                 break;
+            case GGML_OP_MUL_MAT_ID:
+                ggml_zendnn_compute_forward_mul_mat_id(ctx, node);
+                break;
             case GGML_OP_NONE:
             case GGML_OP_RESHAPE:
             case GGML_OP_VIEW:
@@ -361,6 +528,7 @@ static bool ggml_backend_zendnn_device_supports_op(ggml_backend_dev_t dev, const
             return true;
 
         case GGML_OP_MUL_MAT:
+        case GGML_OP_MUL_MAT_ID:
         {
             const ggml_tensor * weights = op->src[0];
             const ggml_tensor * inputs = op->src[1];
@@ -374,6 +542,17 @@ static bool ggml_backend_zendnn_device_supports_op(ggml_backend_dev_t dev, const
                 ne0 < min_batch || ne1 < min_batch || ne10 < min_batch) {
                     return false;
             }
+            // MUL_MAT_ID performs best with a moderate number of experts due to its
+            // gather + batched matmul + scatter approach. Future versions will leverage
+            // ZenDNN's grouped_gemm for better scalability with larger expert counts:
+            // https://github.com/amd/ZenDNN/blob/main/docs/operator/lowoha_group_gemm_operator.md
+            if (op->op == GGML_OP_MUL_MAT_ID) {
+                const int64_t n_experts = weights->ne[2];
+                const int64_t max_experts = 32;
+                if (n_experts > max_experts) {
+                    return false;
+                }
+            }
             switch (weights->type) {
                 case GGML_TYPE_F32:
                 case GGML_TYPE_BF16: