feat: MoE expert cache — keep hot routed experts resident in VRAM

common/arg.cpp

@@ -2621,6 +2621,23 @@ common_params_context common_params_parser_init(common_params & params, llama_ex
             }
         }
     ).set_env("LLAMA_ARG_N_CPU_MOE"));
+    add_opt(common_arg(
+        {"--moe-cache-profile"}, "FNAME",
+        "routing profile CSV (from llama-moe-trace) used to pick which experts to cache in GPU memory",
+        [](common_params & params, const std::string & value) {
+            params.moe_cache_profile = value;
+        }
+    ).set_env("LLAMA_ARG_MOE_CACHE_PROFILE"));
+    add_opt(common_arg(
+        {"--moe-cache-slots"}, "N",
+        "number of routed experts per layer to keep resident in GPU memory (default: 0 = disabled)",
+        [](common_params & params, int value) {
+            if (value < 0) {
+                throw std::invalid_argument("invalid value");
+            }
+            params.moe_cache_slots = value;
+        }
+    ).set_env("LLAMA_ARG_MOE_CACHE_SLOTS"));
     GGML_ASSERT(params.n_gpu_layers < 0); // string_format would need to be extended for a default >= 0
     add_opt(common_arg(
         {"-ngl", "--gpu-layers", "--n-gpu-layers"}, "N",

common/common.cpp

@@ -1564,6 +1564,11 @@ struct llama_model_params common_model_params_to_llama(common_params & params) {
     mparams.use_extra_bufts = !params.no_extra_bufts;
     mparams.no_host         = params.no_host;
 
+    if (!params.moe_cache_profile.empty()) {
+        mparams.moe_cache_profile = params.moe_cache_profile.c_str();
+    }
+    mparams.moe_cache_slots = params.moe_cache_slots;
+
     if (params.kv_overrides.empty()) {
         mparams.kv_overrides = NULL;
     } else {

common/common.h

@@ -583,6 +583,9 @@ struct common_params {
     bool no_extra_bufts    = false; // disable extra buffer types (used for weight repacking)
     bool no_host           = false; // bypass host buffer allowing extra buffers to be used
 
+    std::string moe_cache_profile = ""; // MoE expert cache routing profile CSV (empty = disabled)
+    int32_t     moe_cache_slots   = 0;  // MoE expert cache slots per layer (0 = disabled)
+
     bool single_turn       = false; // single turn chat conversation
 
     ggml_type cache_type_k = GGML_TYPE_F16; // KV cache data type for the K

ggml/src/ggml-backend.cpp

@@ -1756,7 +1756,10 @@ static enum ggml_status ggml_backend_sched_compute_splits(ggml_backend_sched_t s
                         for (int64_t i1 = 0; i1 < ids_tensor->ne[1]; i1++) {
                             for (int64_t i0 = 0; i0 < ids_tensor->ne[0]; i0++) {
                                 int32_t id = ids[i1 * ids_tensor->nb[1]/sizeof(int32_t) + i0 * ids_tensor->nb[0]/sizeof(int32_t)];
-                                GGML_ASSERT(id >= 0 && id < n_expert);
+                                if (id < 0) {
+                                    continue; // expert not owned by this pack (hot/cold split)
+                                }
+                                GGML_ASSERT(id < n_expert);
                                 ggml_bitset_set(used_ids.data(), id);
                             }
                         }

ggml/src/ggml-cpu/ggml-cpu.c

@@ -1677,7 +1677,14 @@ 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);
+                // id == -1 means "expert not owned by this pack" (hot/cold expert
+                // split): contribute a zero row so the pack outputs merge additively
+                if (i02 < 0) {
+                    memset((char *) dst->data + id*dst->nb[1] + iid1*dst->nb[2], 0, dst->ne[0]*sizeof(float));
+                    continue;
+                }
+
+                assert(i02 < n_as);
 
                 MMID_MATRIX_ROW(i02, matrix_row_counts[i02]) = (struct mmid_row_mapping) {id, iid1};
                 matrix_row_counts[i02] += 1;

ggml/src/ggml-cuda/ggml-cuda.cu

@@ -1882,12 +1882,15 @@ static void ggml_cuda_mul_mat_id(ggml_backend_cuda_context & ctx, ggml_tensor *
             }
         }
 
-        if (ggml_cuda_should_use_mmq(src0->type, cc, ne12, /*n_experts=*/ne02)) {
+        // ids containing -1 (hot/cold expert-pack split, op_params[0] != 0) are
+        // supported by the mmvq and general paths only; mmq/mmf are skipped
+        const bool ids_may_skip = dst->op_params[0] != 0;
+        if (!ids_may_skip && ggml_cuda_should_use_mmq(src0->type, cc, ne12, /*n_experts=*/ne02)) {
             ggml_cuda_mul_mat_q(ctx, src0, src1, ids, dst);
             return;
         }
 
-        if (ggml_cuda_should_use_mmf(src0->type, cc, WARP_SIZE, src0->ne, src0->nb, src1->ne[2], /*mul_mat_id=*/true)) {
+        if (!ids_may_skip && ggml_cuda_should_use_mmf(src0->type, cc, WARP_SIZE, src0->ne, src0->nb, src1->ne[2], /*mul_mat_id=*/true)) {
             ggml_cuda_mul_mat_f(ctx, src0, src1, ids, dst);
             return;
         }
@@ -1911,14 +1914,14 @@ static void ggml_cuda_mul_mat_id(ggml_backend_cuda_context & ctx, ggml_tensor *
 
     std::vector<int32_t> ids_to_sorted_host;
     ids_to_sorted_host.reserve(2*ne_get_rows);
-    std::vector<int32_t> ids_from_sorted_host(ne_get_rows);
+    std::vector<int32_t> ids_from_sorted_host(ne_get_rows, -1); // -1 = slot's expert not in this pack
 
     ggml_cuda_pool_alloc<int32_t> ids_buf_dev(ctx.pool(), 2*ne_get_rows);
 
     std::vector<int32_t> tokens_per_expert(ne02);
 
     ggml_cuda_pool_alloc<char> src1_sorted(ctx.pool(), ne12*n_expert_used*ne10*ts_src1_sorted);
-    ggml_cuda_pool_alloc<char>  dst_sorted(ctx.pool(), ne2 *n_expert_used* ne0*ts_dst_sorted);
+    ggml_cuda_pool_alloc<char>  dst_sorted(ctx.pool(), (ne2*n_expert_used + 1)*ne0*ts_dst_sorted); // +1 zero row for skipped slots
 
     std::vector<char> ids_host(ggml_nbytes(ids));
     CUDA_CHECK(cudaMemcpyAsync(ids_host.data(), ids->data, ggml_nbytes(ids), cudaMemcpyDeviceToHost, stream));
@@ -1928,7 +1931,7 @@ static void ggml_cuda_mul_mat_id(ggml_backend_cuda_context & ctx, ggml_tensor *
         for (int64_t i12 = 0; i12 < ne12; ++i12) { // tokens
             for (int64_t iex = 0; iex < n_expert_used; ++iex) {
                 const int32_t expert_to_use = *(const int32_t *)(ids_host.data() + i12*ids->nb[1] + iex*ids->nb[0]);
-                assert(expert_to_use >= 0 && expert_to_use < ne02);
+                assert(expert_to_use >= -1 && expert_to_use < ne02); // -1 = skip (hot/cold expert split)
                 if (expert_to_use == i02) {
                     ids_from_sorted_host[i12*n_expert_used + iex] = ids_to_sorted_host.size();
                     ids_to_sorted_host.push_back(i12*ne11 + iex % ne11);
@@ -1938,7 +1941,20 @@ static void ggml_cuda_mul_mat_id(ggml_backend_cuda_context & ctx, ggml_tensor *
             }
         }
     }
-    GGML_ASSERT(ids_to_sorted_host.size() == size_t(ne_get_rows));
+    const int64_t ne_rows_used = ids_to_sorted_host.size();
+    GGML_ASSERT(ne_rows_used <= ne_get_rows);
+
+    if (ne_rows_used < ne_get_rows) {
+        // slots whose expert id was -1: scatter from a zeroed row so the
+        // pack outputs merge additively
+        CUDA_CHECK(cudaMemsetAsync(dst_sorted.ptr + ne_rows_used*ne0*ts_dst_sorted, 0, ne0*ts_dst_sorted, stream));
+        for (auto & v : ids_from_sorted_host) {
+            if (v < 0) {
+                v = ne_rows_used;
+            }
+        }
+        ids_to_sorted_host.resize(ne_get_rows, 0); // pad; rows past ne_rows_used are never gathered
+    }
 
     ids_to_sorted_host.insert(ids_to_sorted_host.end(), ids_from_sorted_host.begin(), ids_from_sorted_host.end());
 

ggml/src/ggml-cuda/mmf.cu

@@ -88,6 +88,11 @@ void ggml_cuda_mul_mat_f(ggml_backend_cuda_context & ctx, const ggml_tensor * sr
             static_cast<int>(n_experts), static_cast<int>(n_tokens), static_cast<int>(n_expert_used), static_cast<int>(ne11), si1, sis1, /*write_inverse =*/ false, ctx.stream());
         CUDA_CHECK(cudaGetLastError());
 
+        // slots with expert id -1 (hot/cold expert split) are never scattered to; zero their dst rows
+        ggml_cuda_launch_mm_ids_zero_skipped_rows(ids_d, dst_d,
+            dst->ne[0], static_cast<int>(n_tokens), static_cast<int>(n_expert_used), si1, s1, s2, ctx.stream());
+        CUDA_CHECK(cudaGetLastError());
+
         ids_info.ids_src_compact   = ids_src_compact_dev.get();
         ids_info.ids_dst_compact   = ids_dst_compact_dev.get();
         ids_info.expert_bounds_dev = expert_bounds_dev.get();

ggml/src/ggml-cuda/mmid.cu

@@ -44,7 +44,7 @@ static __global__ void mm_ids_helper(
             int iex_used = -1; // The index at which the expert is used, if any.
             for (int iex = threadIdx.x; iex < n_expert_used; iex += warp_size) {
                 const int expert_used = ids[it*si1 + iex];
-                nex_prev += expert_used < expert;
+                nex_prev += expert_used >= 0 && expert_used < expert; // id -1 (skipped slot) occupies no compact position
                 if (expert_used == expert) {
                     iex_used = iex;
                 }
@@ -69,7 +69,7 @@ static __global__ void mm_ids_helper(
             const int expert_used = (neu_padded == n_expert_used || iex < n_expert_used) && it < n_tokens ?
                 ids[it*si1 + iex] : INT_MAX;
             const int iex_used = expert_used == expert ? iex : -1;
-            nex_prev += expert_used < expert;
+            nex_prev += expert_used >= 0 && expert_used < expert; // id -1 (skipped slot) occupies no compact position
 
             // Whether the threads at this token position have used the expert:
             const int it_compact_add_self = warp_reduce_any<neu_padded>(iex_used != -1);
@@ -140,6 +140,35 @@ static void launch_mm_ids_helper(
         (ids, ids_src1, ids_dst, expert_bounds, n_tokens, n_expert_used_var, nchannels_y, si1, sis1, write_inverse);
 }
 
+// Zero the dst rows of (token, slot) pairs whose expert id is -1 ("expert not owned by
+// this pack", hot/cold expert split). The matrix multiplication kernels skip these slots
+// entirely, so without this the corresponding dst rows would contain garbage. With zeros
+// the outputs of multiple expert packs can be merged additively.
+static __global__ void mm_ids_zero_skipped_rows(
+        const int32_t * __restrict__ ids, float * __restrict__ dst, const int64_t ne0,
+        const int n_tokens, const int si1, const int64_t s_slot, const int64_t s_token) {
+    const int iex = blockIdx.y;
+    for (int it = blockIdx.z; it < n_tokens; it += gridDim.z) {
+        if (ids[it*si1 + iex] >= 0) {
+            continue;
+        }
+        float * dst_row = dst + it*s_token + iex*s_slot;
+        for (int64_t i = blockIdx.x*int64_t(blockDim.x) + threadIdx.x; i < ne0; i += int64_t(gridDim.x)*blockDim.x) {
+            dst_row[i] = 0.0f;
+        }
+    }
+}
+
+void ggml_cuda_launch_mm_ids_zero_skipped_rows(
+        const int32_t * ids, float * dst, const int64_t ne0, const int n_tokens, const int n_expert_used,
+        const int si1, const int64_t s_slot, const int64_t s_token, cudaStream_t stream) {
+    constexpr int block_size = 256;
+    const int blocks_x = (ne0 + block_size - 1) / block_size;
+    const dim3 num_blocks(blocks_x, n_expert_used, n_tokens < 65535 ? n_tokens : 65535);
+    const dim3 block_dims(block_size, 1, 1);
+    mm_ids_zero_skipped_rows<<<num_blocks, block_dims, 0, stream>>>(ids, dst, ne0, n_tokens, si1, s_slot, s_token);
+}
+
 void ggml_cuda_launch_mm_ids_helper(
         const int32_t * __restrict__ ids, int32_t * __restrict__ ids_src1, int32_t * __restrict__ ids_dst, int32_t * __restrict__ expert_bounds,
         const int n_experts, const int n_tokens, const int n_expert_used, const int nchannels_y, const int si1, const int sis1, const bool write_inverse, cudaStream_t stream) {

ggml/src/ggml-cuda/mmid.cuh

@@ -3,3 +3,7 @@
 void ggml_cuda_launch_mm_ids_helper(
         const int32_t * ids, int32_t * ids_src1, int32_t * ids_dst, int32_t * expert_bounds,
         int n_experts, int n_tokens, int n_expert_used, int nchannels_y, int si1, int sis1, bool write_inverse, cudaStream_t stream);
+
+void ggml_cuda_launch_mm_ids_zero_skipped_rows(
+        const int32_t * ids, float * dst, int64_t ne0, int n_tokens, int n_expert_used,
+        int si1, int64_t s_slot, int64_t s_token, cudaStream_t stream);

ggml/src/ggml-cuda/mmq.cu

@@ -194,12 +194,23 @@ void ggml_cuda_mul_mat_q(
 
     {
         GGML_ASSERT(ids->nb[0] == ggml_element_size(ids));
+        // sentinel-fill: compact slots belonging to skipped ids (-1, hot/cold expert
+        // split) are never written by mm_ids_helper; quantize kernels skip on i < 0
+        CUDA_CHECK(cudaMemsetAsync(ids_src1.get(), 0xFF, ne_get_rows*sizeof(int32_t), stream));
+        // ids_dst tail likewise: unwritten compact slots must hold a safe row index,
+        // not pool garbage — tile-padded reads in the mm kernel touch them
+        CUDA_CHECK(cudaMemsetAsync(ids_dst.get(), 0, ne_get_rows*sizeof(int32_t), stream));
         const int si1  = ids->nb[1] / ggml_element_size(ids);
         const int sis1 = nb12 / nb11;
 
         ggml_cuda_launch_mm_ids_helper((const int32_t *) ids->data, ids_src1.get(), ids_dst.get(), expert_bounds.get(),
             ne02, ne12, n_expert_used, ne11, si1, sis1, /*write_inverse =*/ dedup_bcast, stream);
         CUDA_CHECK(cudaGetLastError());
+
+        // slots with expert id -1 (hot/cold expert split) are never scattered to; zero their dst rows
+        ggml_cuda_launch_mm_ids_zero_skipped_rows((const int32_t *) ids->data, (float *) dst->data,
+            dst->ne[0], ne12, n_expert_used, si1, dst->nb[1]/sizeof(float), dst->nb[2]/sizeof(float), stream);
+        CUDA_CHECK(cudaGetLastError());
     }
 
     const size_t nbytes_src1_q8_1 = ne12*n_expert_used*ne10_padded * y_block_size/y_values_per_block +

ggml/src/ggml-cuda/mmvf.cu

@@ -2,6 +2,7 @@
 #include "common.cuh"
 #include "unary.cuh"
 #include "mmvf.cuh"
+#include "mmid.cuh"
 #include "convert.cuh"
 
 template <typename T, typename type_acc, int ncols_dst, int block_size, bool has_fusion = false, bool is_multi_token_id = false>
@@ -39,6 +40,10 @@ static __global__ void mul_mat_vec_f(
         sample_dst = ids ? 0                                                   : blockIdx.z;
     }
 
+    if (ids && channel_x < 0) {
+        return; // expert not owned by this pack; dst row pre-zeroed host-side
+    }
+
     const int sample_x    = fastdiv((uint32_t) sample_dst, sample_ratio);
     const int sample_y    = sample_dst;
 
@@ -651,6 +656,14 @@ void ggml_cuda_mul_mat_vec_f(ggml_backend_cuda_context & ctx, const ggml_tensor
 
     const float   * src1_d =       (const float   *) src1->data;
     const int32_t *  ids_d = ids ? (const int32_t *)  ids->data : nullptr;
+
+    if (ids) {
+        // slots with expert id -1 (hot/cold expert split) are skipped by the kernels; zero their dst rows
+        ggml_cuda_launch_mm_ids_zero_skipped_rows(ids_d, (float *) dst->data,
+            dst->ne[0], ids->ne[1], ids->ne[0], ids->nb[1]/sizeof(int32_t),
+            dst->nb[1]/sizeof(float), dst->nb[2]/sizeof(float), ctx.stream());
+        CUDA_CHECK(cudaGetLastError());
+    }
     float         *  dst_d =       (float         *)  dst->data;
 
     ggml_cuda_mm_fusion_args_device fusion_local{};

ggml/src/ggml-cuda/mmvq.cu

@@ -1,4 +1,5 @@
 #include "mmvq.cuh"
+#include "mmid.cuh"
 #include "quantize.cuh"
 #include "unary.cuh"
 #include "vecdotq.cuh"
@@ -513,6 +514,9 @@ static __global__ void mul_mat_vec_q(
     uint32_t sample_dst;
 
     ggml_cuda_pdl_sync();
+    if (ncols_dst == 1 && ids && ids[channel_dst] < 0) {
+        return; // expert not owned by this pack; dst row pre-zeroed host-side
+    }
     channel_x  = ncols_dst == 1 && ids ? ids[channel_dst]                     : fastdiv(channel_dst, channel_ratio);
     channel_y  = ncols_dst == 1 && ids ? fastmodulo(channel_dst, nchannels_y) : channel_dst;
     sample_dst = blockIdx.z;
@@ -737,7 +741,11 @@ static __global__ void mul_mat_vec_q_moe(
     }
 
     ggml_cuda_pdl_sync();
-    const uint32_t channel_x = ids[channel_dst + token_idx * ids_stride];
+    const int32_t id_used = ids[channel_dst + token_idx * ids_stride];
+    if (id_used < 0) {
+        return; // expert not owned by this pack; dst row pre-zeroed host-side
+    }
+    const uint32_t channel_x = id_used;
     const uint32_t channel_y = fastmodulo(channel_dst, nchannels_y);
 
     const block_q8_1 * y = ((const block_q8_1 *) vy) + channel_y*stride_channel_y + token_idx*stride_col_y;
@@ -1174,6 +1182,14 @@ void ggml_cuda_mul_mat_vec_q(
 
     const float   * src1_d =       (const float   *) src1->data;
     const int32_t *  ids_d = ids ? (const int32_t *)  ids->data : nullptr;
+
+    if (ids) {
+        // slots with expert id -1 (hot/cold expert split) are skipped by the kernels; zero their dst rows
+        ggml_cuda_launch_mm_ids_zero_skipped_rows(ids_d, (float *) dst->data,
+            dst->ne[0], ids->ne[1], ids->ne[0], ids->nb[1]/sizeof(int32_t),
+            dst->nb[1]/sizeof(float), dst->nb[2]/sizeof(float), ctx.stream());
+        CUDA_CHECK(cudaGetLastError());
+    }
     float         *  dst_d =       (float         *)  dst->data;
 
     ggml_cuda_mm_fusion_args_device fusion_local{};

ggml/src/ggml-cuda/quantize.cu

@@ -140,6 +140,9 @@ static __global__ void quantize_mmq_nvfp4(
         const int64_t i2  = blockIdx.y % ne2;
         const int64_t i3  = blockIdx.y / ne2;
         const int64_t i01 = ids ? ids[blockIdx.x] : blockIdx.x;
+        if (i01 < 0) {
+            return; // compact slot unused (hot/cold expert split)
+        }
         base_idx = i3 * s03 + i2 * s02 + i01 * s01;
     }
     const float * __restrict__  x_row = x + base_idx;
@@ -184,6 +187,9 @@ static __global__ void quantize_mmq_nvfp4(
 #pragma unroll
                 for (int slot = 0; slot < n_expert_used; ++slot) {
                     const int64_t i = ids[(int64_t) blockIdx.x * n_expert_used + slot];
+                    if (i < 0) {
+                        continue; // slot skipped (hot/cold expert split)
+                    }
                     scale[i] = warp_amax[0];
                 }
             } else {
@@ -310,6 +316,9 @@ static __global__ void quantize_mmq_nvfp4(
 #pragma unroll
             for (int slot = 0; slot < n_expert_used; ++slot) {
                 const int64_t i = ids[(int64_t) blockIdx.x * n_expert_used + slot];
+                if (i < 0) {
+                    continue; // slot skipped (hot/cold expert split)
+                }
                 block_fp4_mmq * yb = y + (k_block * ne1 + i);
                 uint32_t * yqs = reinterpret_cast<uint32_t *>(yb->qs);
                 yqs[2 * sub + 0] = q0;
@@ -376,6 +385,9 @@ static __global__ void quantize_mmq_mxfp4(const float * __restrict__ x,
         const int64_t i2  = blockIdx.z % ne2;
         const int64_t i3  = blockIdx.z / ne2;
         const int64_t i01 = ids ? ids[blockIdx.x] : blockIdx.x;
+        if (i01 < 0) {
+            return; // compact slot unused (hot/cold expert split)
+        }
         base_pos = i3 * s03 + i2 * s02 + i01 * s01;
     }
 
@@ -428,6 +440,9 @@ static __global__ void quantize_mmq_mxfp4(const float * __restrict__ x,
 #pragma unroll
         for (int slot = 0; slot < n_expert_used; ++slot) {
             const int64_t i = ids[(int64_t) blockIdx.x * n_expert_used + slot];
+            if (i < 0) {
+                continue; // slot skipped (hot/cold expert split)
+            }
             block_fp4_mmq * yb = y + (k_block * ne1 + i);
             char2 * yqs2 = (char2 *) yb->qs;
             if (lane_in_group == 0) {
@@ -479,6 +494,9 @@ static __global__ void quantize_mmq_q8_1(
         const int64_t i2  = blockIdx.z % ne2;
         const int64_t i3  = blockIdx.z / ne2;
         const int64_t i01 = ids ? ids[blockIdx.x] : blockIdx.x;
+        if (i01 < 0) {
+            return; // compact slot unused (hot/cold expert split)
+        }
         base_idx = i3*s03 + i2*s02 + i01*s01;
     }
 
@@ -527,6 +545,9 @@ static __global__ void quantize_mmq_q8_1(
         int64_t ib;
         if constexpr (scatter) {
             const int64_t i = ids[(int64_t) blockIdx.x * n_expert_used + slot];
+            if (i < 0) {
+                continue; // slot skipped (hot/cold expert split)
+            }
             ib = k_block*ne1 + i;
         } else {
             const int64_t ib0 = blockIdx.z*((int64_t)gridDim.x*gridDim.y*blockDim.x/QK8_1); // first block of channel