diff --git a/common/arg.cpp b/common/arg.cpp index 08921810b212..c4c1cdbb70e5 100644 --- a/common/arg.cpp +++ b/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", diff --git a/common/common.cpp b/common/common.cpp index a68766cbbbc8..2dff9454bfaf 100644 --- a/common/common.cpp +++ b/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 { diff --git a/common/common.h b/common/common.h index cb1827f5dd24..cc05835fd488 100644 --- a/common/common.h +++ b/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 diff --git a/ggml/src/ggml-backend.cpp b/ggml/src/ggml-backend.cpp index f54c87a2ab0a..5a5d9151e2fe 100644 --- a/ggml/src/ggml-backend.cpp +++ b/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); } } diff --git a/ggml/src/ggml-cpu/ggml-cpu.c b/ggml/src/ggml-cpu/ggml-cpu.c index 23433cd5392a..1cb28bb128cb 100644 --- a/ggml/src/ggml-cpu/ggml-cpu.c +++ b/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; diff --git a/ggml/src/ggml-cuda/ggml-cuda.cu b/ggml/src/ggml-cuda/ggml-cuda.cu index 37aaa5bec72e..0ccf04fdf37e 100644 --- a/ggml/src/ggml-cuda/ggml-cuda.cu +++ b/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 ids_to_sorted_host; ids_to_sorted_host.reserve(2*ne_get_rows); - std::vector ids_from_sorted_host(ne_get_rows); + std::vector ids_from_sorted_host(ne_get_rows, -1); // -1 = slot's expert not in this pack ggml_cuda_pool_alloc ids_buf_dev(ctx.pool(), 2*ne_get_rows); std::vector tokens_per_expert(ne02); ggml_cuda_pool_alloc src1_sorted(ctx.pool(), ne12*n_expert_used*ne10*ts_src1_sorted); - ggml_cuda_pool_alloc dst_sorted(ctx.pool(), ne2 *n_expert_used* ne0*ts_dst_sorted); + ggml_cuda_pool_alloc dst_sorted(ctx.pool(), (ne2*n_expert_used + 1)*ne0*ts_dst_sorted); // +1 zero row for skipped slots std::vector 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()); diff --git a/ggml/src/ggml-cuda/mmf.cu b/ggml/src/ggml-cuda/mmf.cu index 646a5899c803..1c428d0339f7 100644 --- a/ggml/src/ggml-cuda/mmf.cu +++ b/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(n_experts), static_cast(n_tokens), static_cast(n_expert_used), static_cast(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(n_tokens), static_cast(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(); diff --git a/ggml/src/ggml-cuda/mmid.cu b/ggml/src/ggml-cuda/mmid.cu index f80442fbe4e8..06674d2bf103 100644 --- a/ggml/src/ggml-cuda/mmid.cu +++ b/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(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<<>>(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) { diff --git a/ggml/src/ggml-cuda/mmid.cuh b/ggml/src/ggml-cuda/mmid.cuh index 74c2db43385e..b6b669fa8d6a 100644 --- a/ggml/src/ggml-cuda/mmid.cuh +++ b/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); diff --git a/ggml/src/ggml-cuda/mmq.cu b/ggml/src/ggml-cuda/mmq.cu index 485a5c06757e..e35c5bdcb817 100644 --- a/ggml/src/ggml-cuda/mmq.cu +++ b/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 + diff --git a/ggml/src/ggml-cuda/mmvf.cu b/ggml/src/ggml-cuda/mmvf.cu index d7dbc8b99282..826223337537 100644 --- a/ggml/src/ggml-cuda/mmvf.cu +++ b/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 @@ -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{}; diff --git a/ggml/src/ggml-cuda/mmvq.cu b/ggml/src/ggml-cuda/mmvq.cu index 0589e65bdf8c..189e7231ca01 100644 --- a/ggml/src/ggml-cuda/mmvq.cu +++ b/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{}; diff --git a/ggml/src/ggml-cuda/quantize.cu b/ggml/src/ggml-cuda/quantize.cu index 2bd9b6262390..95e50ffe643c 100644 --- a/ggml/src/ggml-cuda/quantize.cu +++ b/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(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 diff --git a/include/llama.h b/include/llama.h index 9fab69317006..7f494ba01a5e 100644 --- a/include/llama.h +++ b/include/llama.h @@ -330,6 +330,10 @@ extern "C" { // override key-value pairs of the model meta data const struct llama_model_kv_override * kv_overrides; + // MoE expert cache: keep the hottest routed experts per layer resident in GPU memory + const char * moe_cache_profile; // routing profile CSV from llama-moe-trace (NULL = disabled) + int32_t moe_cache_slots; // experts cached per layer (0 = disabled) + // Keep the booleans together to avoid misalignment during copy-by-value. bool vocab_only; // only load the vocabulary, no weights bool check_tensors; // validate model tensor data diff --git a/src/llama-graph.cpp b/src/llama-graph.cpp index 3a8f43909ed6..41ddf9673065 100644 --- a/src/llama-graph.cpp +++ b/src/llama-graph.cpp @@ -1771,7 +1771,8 @@ ggml_tensor * llm_graph_context::build_moe_ffn( ggml_tensor * up_exps_s, ggml_tensor * gate_exps_s, ggml_tensor * down_exps_s, - ggml_tensor * selected_experts_in) const { + ggml_tensor * selected_experts_in, + const llama_layer * moe_cache) const { return build_moe_ffn( cur, gate_inp, /* gate_inp_b */ nullptr, @@ -1793,7 +1794,8 @@ ggml_tensor * llm_graph_context::build_moe_ffn( gate_exps_s, down_exps_s, selected_experts_in - ); + , + moe_cache); } ggml_tensor * llm_graph_context::build_moe_ffn( @@ -1820,7 +1822,8 @@ ggml_tensor * llm_graph_context::build_moe_ffn( ggml_tensor * up_exps_s, ggml_tensor * gate_exps_s, ggml_tensor * down_exps_s, - ggml_tensor * selected_experts_in) const { + ggml_tensor * selected_experts_in, + const llama_layer * moe_cache) const { const int64_t n_embd = cur->ne[0]; const int64_t n_tokens = cur->ne[1]; const bool weight_before_ffn = arch == LLM_ARCH_LLAMA4; // for llama4, we apply the sigmoid-ed weights before the FFN @@ -1917,6 +1920,23 @@ ggml_tensor * llm_graph_context::build_moe_ffn( } cb(selected_experts, "ffn_moe_topk", il); + // MoE expert cache: split routed ids into hot-pack ids and cold ids. + // scope: plain fused-SILU gated FFN (no clamp, no expert biases/scales, + // no pre-FFN weighting) — the dual chains below reproduce exactly that + ggml_tensor * ids_hot = nullptr; + ggml_tensor * ids_cold = nullptr; + const bool use_moe_packs = moe_cache && moe_cache->moe_map_hot && !gate_up_exps && + !up_exps_s && !gate_exps_s && !down_exps_s && + type_op == LLM_FFN_SILU && gate_exps && !up_exps_b && !gate_exps_b && !weight_before_ffn && + (il < 0 || hparams.swiglu_clamp_exp[il] <= 1e-6f); + if (use_moe_packs) { + ggml_tensor * ids_flat = ggml_cont_1d(ctx0, selected_experts, n_expert_used*n_tokens); // topk ids are a strided view + ids_hot = ggml_reshape_2d(ctx0, ggml_get_rows(ctx0, moe_cache->moe_map_hot, ids_flat), n_expert_used, n_tokens); + ids_cold = ggml_reshape_2d(ctx0, ggml_get_rows(ctx0, moe_cache->moe_map_cold, ids_flat), n_expert_used, n_tokens); + cb(ids_hot, "ffn_moe_ids_hot", il); + cb(ids_cold, "ffn_moe_ids_cold", il); + } + if (arch == LLM_ARCH_GROVEMOE && n_expert != hparams.n_expert) { // TODO: Use scalar div instead when/if implemented ggml_tensor * f_sel = ggml_cast(ctx0, selected_experts, GGML_TYPE_F32); @@ -1972,7 +1992,32 @@ ggml_tensor * llm_graph_context::build_moe_ffn( ggml_tensor * up = nullptr; ggml_tensor * experts = nullptr; - if (gate_up_exps) { + if (use_moe_packs) { + // MoE expert cache: one complete FFN chain per residency side. Each + // chain is unbroken so the scheduler never splices CPU ops between GPU + // ops (that migrates the hot pack weights to CPU every layer). Skipped + // (-1) rows are zero and swiglu(0,0) = 0, so the two chain outputs are + // disjoint and one add reconstructs the exact single-tensor result. + auto build_pack_chain = [&](ggml_tensor * w_gate, ggml_tensor * w_up, ggml_tensor * w_down, ggml_tensor * ids) { + ggml_tensor * gate = ggml_mul_mat_id(ctx0, w_gate, cur, ids); + gate->op_params[0] = 1; // ids may contain -1 + ggml_tensor * up_p = ggml_mul_mat_id(ctx0, w_up, cur, ids); + up_p->op_params[0] = 1; + ggml_tensor * act = ggml_swiglu_split(ctx0, gate, up_p); + ggml_tensor * down = ggml_mul_mat_id(ctx0, w_down, act, ids); + down->op_params[0] = 1; + return down; + }; + + ggml_tensor * hot = build_pack_chain(moe_cache->ffn_gate_exps_hot, moe_cache->ffn_up_exps_hot, moe_cache->ffn_down_exps_hot, ids_hot); + cb(hot, "ffn_moe_down_hot", il); + + ggml_tensor * cold = build_pack_chain(gate_exps, up_exps, down_exps, ids_cold); + cb(cold, "ffn_moe_down_cold", il); + + experts = ggml_add(ctx0, hot, cold); + cb(experts, "ffn_moe_down", il); + } else if (gate_up_exps) { // merged gate_up path: one mul_mat_id, then split into gate and up views ggml_tensor * gate_up = build_lora_mm_id(gate_up_exps, cur, selected_experts, up_exps_s); // [n_ff*2, n_expert_used, n_tokens] cb(gate_up, "ffn_moe_gate_up", il); @@ -2024,82 +2069,84 @@ ggml_tensor * llm_graph_context::build_moe_ffn( const bool has_gate = gate_exps || gate_up_exps; - switch (type_op) { - case LLM_FFN_SILU: - if (gate_exps) { - if (il >= 0) { - const float limit = hparams.swiglu_clamp_exp[il]; - constexpr float eps = 1e-6f; - if (limit > eps) { - up = ggml_clamp(ctx0, up, -limit, limit); - cb(up, "ffn_moe_up_clamped", il); - - if (arch == LLM_ARCH_DEEPSEEK4) { - cur = ggml_clamp(ctx0, cur, -INFINITY, limit); - cb(cur, "ffn_moe_gate_clamped", il); - cur = ggml_swiglu_split(ctx0, cur, up); - } else { - ggml_tensor * gate_act = ggml_silu(ctx0, cur); - cb(gate_act, "ffn_moe_silu", il); - gate_act = ggml_clamp(ctx0, gate_act, -INFINITY, limit); - cb(gate_act, "ffn_moe_silu_clamped", il); - cur = ggml_mul(ctx0, gate_act, up); + if (!use_moe_packs) { + switch (type_op) { + case LLM_FFN_SILU: + if (gate_exps) { + if (il >= 0) { + const float limit = hparams.swiglu_clamp_exp[il]; + constexpr float eps = 1e-6f; + if (limit > eps) { + up = ggml_clamp(ctx0, up, -limit, limit); + cb(up, "ffn_moe_up_clamped", il); + + if (arch == LLM_ARCH_DEEPSEEK4) { + cur = ggml_clamp(ctx0, cur, -INFINITY, limit); + cb(cur, "ffn_moe_gate_clamped", il); + cur = ggml_swiglu_split(ctx0, cur, up); + } else { + ggml_tensor * gate_act = ggml_silu(ctx0, cur); + cb(gate_act, "ffn_moe_silu", il); + gate_act = ggml_clamp(ctx0, gate_act, -INFINITY, limit); + cb(gate_act, "ffn_moe_silu_clamped", il); + cur = ggml_mul(ctx0, gate_act, up); + } + cb(cur, "ffn_moe_swiglu_limited", il); + break; } - cb(cur, "ffn_moe_swiglu_limited", il); - break; } } - } - if (has_gate) { - cur = ggml_swiglu_split(ctx0, cur, up); - cb(cur, "ffn_moe_swiglu", il); - } else { - cur = ggml_silu(ctx0, cur); - cb(cur, "ffn_moe_silu", il); - } break; - case LLM_FFN_GELU: - if (has_gate) { - cur = ggml_geglu_split(ctx0, cur, up); - cb(cur, "ffn_moe_geglu", il); - } else { - cur = ggml_gelu(ctx0, cur); - cb(cur, "ffn_moe_gelu", il); - } break; - case LLM_FFN_SWIGLU_OAI_MOE: - { - // TODO: move to hparams? - constexpr float alpha = 1.702f; - constexpr float limit = 7.0f; - cur = ggml_swiglu_oai(ctx0, cur, up, alpha, limit); - cb(cur, "ffn_moe_swiglu_oai", il); - } break; - case LLM_FFN_RELU: - if (has_gate) { - cur = ggml_reglu_split(ctx0, cur, up); - cb(cur, "ffn_moe_reglu", il); - } else { - cur = ggml_relu(ctx0, cur); - cb(cur, "ffn_moe_relu", il); - } break; - case LLM_FFN_RELU_SQR: - if (has_gate) { - // TODO: add support for gated squared relu - GGML_ABORT("fatal error: gated squared relu not implemented"); - } else { - cur = ggml_relu(ctx0, cur); - cur = ggml_sqr(ctx0, cur); - cb(cur, "ffn_moe_relu_sqr", il); - } break; - default: - GGML_ABORT("fatal error"); - } + if (has_gate) { + cur = ggml_swiglu_split(ctx0, cur, up); + cb(cur, "ffn_moe_swiglu", il); + } else { + cur = ggml_silu(ctx0, cur); + cb(cur, "ffn_moe_silu", il); + } break; + case LLM_FFN_GELU: + if (has_gate) { + cur = ggml_geglu_split(ctx0, cur, up); + cb(cur, "ffn_moe_geglu", il); + } else { + cur = ggml_gelu(ctx0, cur); + cb(cur, "ffn_moe_gelu", il); + } break; + case LLM_FFN_SWIGLU_OAI_MOE: + { + // TODO: move to hparams? + constexpr float alpha = 1.702f; + constexpr float limit = 7.0f; + cur = ggml_swiglu_oai(ctx0, cur, up, alpha, limit); + cb(cur, "ffn_moe_swiglu_oai", il); + } break; + case LLM_FFN_RELU: + if (has_gate) { + cur = ggml_reglu_split(ctx0, cur, up); + cb(cur, "ffn_moe_reglu", il); + } else { + cur = ggml_relu(ctx0, cur); + cb(cur, "ffn_moe_relu", il); + } break; + case LLM_FFN_RELU_SQR: + if (has_gate) { + // TODO: add support for gated squared relu + GGML_ABORT("fatal error: gated squared relu not implemented"); + } else { + cur = ggml_relu(ctx0, cur); + cur = ggml_sqr(ctx0, cur); + cb(cur, "ffn_moe_relu_sqr", il); + } break; + default: + GGML_ABORT("fatal error"); + } - experts = build_lora_mm_id(down_exps, cur, selected_experts, down_exps_s); // [n_embd, n_expert_used, n_tokens] - cb(experts, "ffn_moe_down", il); + experts = build_lora_mm_id(down_exps, cur, selected_experts, down_exps_s); // [n_embd, n_expert_used, n_tokens] + cb(experts, "ffn_moe_down", il); - if (down_exps_s) { - cb(experts, "ffn_moe_down_scaled", il); + if (down_exps_s) { + cb(experts, "ffn_moe_down_scaled", il); + } } if (down_exps_b) { diff --git a/src/llama-graph.h b/src/llama-graph.h index 7ed490ce6728..e8d20f30b333 100644 --- a/src/llama-graph.h +++ b/src/llama-graph.h @@ -1017,7 +1017,8 @@ struct llm_graph_context { ggml_tensor * up_exps_s = nullptr, ggml_tensor * gate_exps_s = nullptr, ggml_tensor * down_exps_s = nullptr, - ggml_tensor * selected_experts_in = nullptr) const; + ggml_tensor * selected_experts_in = nullptr, + const llama_layer * moe_cache = nullptr) const; ggml_tensor * build_moe_ffn( ggml_tensor * cur, @@ -1043,7 +1044,8 @@ struct llm_graph_context { ggml_tensor * up_exps_s = nullptr, ggml_tensor * gate_exps_s = nullptr, ggml_tensor * down_exps_s = nullptr, - ggml_tensor * selected_experts_in = nullptr) const; + ggml_tensor * selected_experts_in = nullptr, + const llama_layer * moe_cache = nullptr) const; // // inputs diff --git a/src/llama-model.cpp b/src/llama-model.cpp index c9b290e99fb8..85e495e23897 100644 --- a/src/llama-model.cpp +++ b/src/llama-model.cpp @@ -1653,9 +1653,165 @@ bool llama_model_base::load_tensors(llama_model_loader & ml) { } } + init_moe_expert_cache(); + return true; } +void llama_model_base::init_moe_expert_cache() { + // flags take precedence; env vars kept as a fallback + const char * profile_path = params.moe_cache_profile; + int n_slots = params.moe_cache_slots; + if (profile_path == nullptr || profile_path[0] == '\0') { + profile_path = getenv("GGML_MOE_CACHE_PROFILE"); + } + if (n_slots <= 0) { + const char * slots_env = getenv("GGML_MOE_CACHE_SLOTS"); + n_slots = slots_env ? atoi(slots_env) : 0; + } + if (profile_path == nullptr || profile_path[0] == '\0' || n_slots <= 0) { + return; + } + + // routing profile: moe-trace CSV (pos,layer,id0,...), decode rows only + std::map> freq; // layer -> expert -> count + { + FILE * f = fopen(profile_path, "r"); + if (!f) { + LLAMA_LOG_WARN("%s: cannot open profile '%s' - expert cache disabled\n", __func__, profile_path); + return; + } + char line[4096]; + while (fgets(line, sizeof(line), f)) { + char * p = line; + const long pos = strtol(p, &p, 10); + if (*p != ',' || pos < 0) { continue; } + p++; + const long il = strtol(p, &p, 10); + while (*p == ',') { + p++; + const long e = strtol(p, &p, 10); + if (e >= 0) { + freq[(int) il][(int) e]++; + } + } + } + fclose(f); + } + if (freq.empty()) { + LLAMA_LOG_WARN("%s: profile '%s' has no decode rows - expert cache disabled\n", __func__, profile_path); + return; + } + + ggml_backend_dev_t dev = nullptr; + for (const auto & d : devices) { + if (!d.is_meta) { dev = d.dev; break; } + } + if (dev == nullptr || ggml_backend_dev_type(dev) != GGML_BACKEND_DEVICE_TYPE_GPU) { + LLAMA_LOG_WARN("%s: no GPU device - expert cache disabled\n", __func__); + return; + } + ggml_backend_buffer_type_t buft = ggml_backend_dev_buffer_type(dev); + + // candidate layers: routed experts resident in host memory + std::vector pack_layers; + for (int il = 0; il < (int) layers.size(); il++) { + const auto & l = layers[il]; + if (l.ffn_gate_exps && l.ffn_up_exps && l.ffn_down_exps && freq.count(il) && + l.ffn_gate_exps->buffer && ggml_backend_buft_is_host(ggml_backend_buffer_get_type(l.ffn_gate_exps->buffer))) { + pack_layers.push_back(il); + } + } + if (pack_layers.empty()) { + LLAMA_LOG_INFO("%s: no CPU-resident MoE layers - expert cache not built\n", __func__); + return; + } + + ggml_init_params ctx_params = { + /*.mem_size =*/ (5*pack_layers.size() + 1)*ggml_tensor_overhead(), + /*.mem_buffer =*/ nullptr, + /*.no_alloc =*/ true, + }; + ggml_context * ctx = ggml_init(ctx_params); + + for (int il : pack_layers) { + auto & l = layers[il]; + const ggml_tensor * g = l.ffn_gate_exps; + const ggml_tensor * u = l.ffn_up_exps; + const ggml_tensor * d = l.ffn_down_exps; + const int64_t n_expert = g->ne[2]; + const int64_t S = std::min(n_slots, n_expert); + l.ffn_gate_exps_hot = ggml_new_tensor_3d(ctx, g->type, g->ne[0], g->ne[1], S); + l.ffn_up_exps_hot = ggml_new_tensor_3d(ctx, u->type, u->ne[0], u->ne[1], S); + l.ffn_down_exps_hot = ggml_new_tensor_3d(ctx, d->type, d->ne[0], d->ne[1], S); + l.moe_map_hot = ggml_new_tensor_2d(ctx, GGML_TYPE_I32, 1, n_expert); + l.moe_map_cold = ggml_new_tensor_2d(ctx, GGML_TYPE_I32, 1, n_expert); + ggml_format_name(l.ffn_gate_exps_hot, "blk.%d.ffn_gate_exps_hot", il); + ggml_format_name(l.ffn_up_exps_hot, "blk.%d.ffn_up_exps_hot", il); + ggml_format_name(l.ffn_down_exps_hot, "blk.%d.ffn_down_exps_hot", il); + ggml_format_name(l.moe_map_hot, "blk.%d.moe_map_hot", il); + ggml_format_name(l.moe_map_cold, "blk.%d.moe_map_cold", il); + } + + ggml_backend_buffer_t buf = ggml_backend_alloc_ctx_tensors_from_buft(ctx, buft); + if (buf == nullptr) { + LLAMA_LOG_WARN("%s: pack allocation failed - expert cache disabled\n", __func__); + ggml_free(ctx); + for (int il : pack_layers) { + auto & l = layers[il]; + l.ffn_gate_exps_hot = l.ffn_up_exps_hot = l.ffn_down_exps_hot = nullptr; + l.moe_map_hot = l.moe_map_cold = nullptr; + } + return; + } + + // fill packs (expert dim is outermost: one contiguous slab per expert) + std::vector slab; + std::vector map_hot, map_cold; + size_t total_bytes = 0; + for (int il : pack_layers) { + auto & l = layers[il]; + const int64_t n_expert = l.ffn_gate_exps->ne[2]; + const int64_t S = l.ffn_gate_exps_hot->ne[2]; + + std::vector> ranked; // (-count, expert) + for (const auto & [e, c] : freq[il]) { + if (e < n_expert) { + ranked.push_back({-c, e}); + } + } + std::sort(ranked.begin(), ranked.end()); + + map_hot.assign(n_expert, -1); + map_cold.resize(n_expert); + for (int64_t e = 0; e < n_expert; e++) { + map_cold[e] = (int32_t) e; + } + for (int64_t s = 0; s < S && s < (int64_t) ranked.size(); s++) { + const int32_t e = ranked[s].second; + map_hot[e] = (int32_t) s; + map_cold[e] = -1; + const ggml_tensor * srcs[3] = { l.ffn_gate_exps, l.ffn_up_exps, l.ffn_down_exps }; + ggml_tensor * dsts[3] = { l.ffn_gate_exps_hot, l.ffn_up_exps_hot, l.ffn_down_exps_hot }; + for (int t = 0; t < 3; t++) { + const size_t nb = srcs[t]->nb[2]; + slab.resize(nb); + ggml_backend_tensor_get(srcs[t], slab.data(), e*nb, nb); + ggml_backend_tensor_set(dsts[t], slab.data(), s*nb, nb); + total_bytes += nb; + } + } + ggml_backend_tensor_set(l.moe_map_hot, map_hot.data(), 0, n_expert*sizeof(int32_t)); + ggml_backend_tensor_set(l.moe_map_cold, map_cold.data(), 0, n_expert*sizeof(int32_t)); + } + + pimpl->ctxs_bufs.emplace_back(ggml_context_ptr{ctx}, std::vector{}); + pimpl->ctxs_bufs.back().second.emplace_back(buf); + + LLAMA_LOG_INFO("%s: expert cache: %zu layers x %d slots, %.2f MiB uploaded to %s\n", + __func__, pack_layers.size(), n_slots, total_bytes/1024.0/1024.0, ggml_backend_buft_name(buft)); +} + ggml_tensor * llama_model_base::create_tensor(llama_model_loader & ml, const LLM_TN_IMPL & tn, const std::initializer_list & ne, int flags) { const buft_list_t * buft_list_layer = tn.bid == -1 ? nullptr : pimpl->dev_layer.at(tn.bid).buft_list; return ml.create_tensor( @@ -2325,6 +2481,8 @@ llama_model_params llama_model_default_params() { /*.progress_callback =*/ nullptr, /*.progress_callback_user_data =*/ nullptr, /*.kv_overrides =*/ nullptr, + /*.moe_cache_profile =*/ nullptr, + /*.moe_cache_slots =*/ 0, /*.vocab_only =*/ false, /*.check_tensors =*/ false, /*.use_extra_bufts =*/ true, diff --git a/src/llama-model.h b/src/llama-model.h index 45b054cedf1d..62ccb4db1129 100644 --- a/src/llama-model.h +++ b/src/llama-model.h @@ -303,6 +303,15 @@ struct llama_layer { struct ggml_tensor * ffn_down_exps = nullptr; struct ggml_tensor * ffn_up_exps = nullptr; struct ggml_tensor * ffn_gate_up_exps = nullptr; + + // MoE expert cache (hot/cold split): GPU-resident packs of the S most + // frequently routed experts of a CPU-offloaded layer + id remap tables. + // Cold side reuses the original ffn_*_exps tensors with hot ids masked to -1. + struct ggml_tensor * ffn_gate_exps_hot = nullptr; + struct ggml_tensor * ffn_down_exps_hot = nullptr; + struct ggml_tensor * ffn_up_exps_hot = nullptr; + struct ggml_tensor * moe_map_hot = nullptr; // i32[n_expert]: pack slot or -1 + struct ggml_tensor * moe_map_cold = nullptr; // i32[n_expert]: global id or -1 struct ggml_tensor * ffn_gate_inp_b = nullptr; struct ggml_tensor * ffn_gate_exps_b = nullptr; struct ggml_tensor * ffn_down_exps_b = nullptr; @@ -728,6 +737,10 @@ struct llama_model_base : public llama_model { void load_vocab (llama_model_loader & ml) override; bool load_tensors(llama_model_loader & ml) override; + // GGML_MOE_CACHE_PROFILE + GGML_MOE_CACHE_SLOTS: build GPU-resident hot + // expert packs for CPU-offloaded MoE layers (see llama_layer::*_exps_hot) + void init_moe_expert_cache(); + // model must define these void load_arch_hparams(llama_model_loader & ml) override = 0; void load_arch_tensors(llama_model_loader & ml) override = 0; diff --git a/src/models/deepseek2.cpp b/src/models/deepseek2.cpp index a9e8bc514036..93fd67f02e4f 100644 --- a/src/models/deepseek2.cpp +++ b/src/models/deepseek2.cpp @@ -396,7 +396,9 @@ llama_model_deepseek2::graph::graph(const llama_model & model, const llm_graph_p (llama_expert_gating_func_type) hparams.expert_gating_func, il, nullptr, - model.layers[il].ffn_gate_up_exps); + model.layers[il].ffn_gate_up_exps, + nullptr, nullptr, nullptr, + nullptr, &model.layers[il]); cb(moe_out, "ffn_moe_out", il); // FFN shared expert diff --git a/src/models/qwen35moe.cpp b/src/models/qwen35moe.cpp index 7b0876cbb04b..08d63e487768 100644 --- a/src/models/qwen35moe.cpp +++ b/src/models/qwen35moe.cpp @@ -511,7 +511,8 @@ ggml_tensor * llama_model_qwen35moe::graph::build_layer_ffn(ggml_tensor * cur, c nullptr, model.layers[il].ffn_gate_up_exps, model.layers[il].ffn_up_exps_s, model.layers[il].ffn_gate_exps_s, - model.layers[il].ffn_down_exps_s); + model.layers[il].ffn_down_exps_s, + nullptr, &model.layers[il]); cb(moe_out, "ffn_moe_out", il); // Add shared experts if present - following Qwen3Next reference implementation @@ -689,7 +690,8 @@ llama_model_qwen35moe::graph_mtp::graph_mtp(const llama_model & model, const llm nullptr, layer.ffn_gate_up_exps, layer.ffn_up_exps_s, layer.ffn_gate_exps_s, - layer.ffn_down_exps_s); + layer.ffn_down_exps_s, + nullptr, nullptr); cb(moe_out, "mtp_ffn_moe_out", il); if (layer.ffn_up_shexp != nullptr) { diff --git a/tests/test-backend-ops.cpp b/tests/test-backend-ops.cpp index 51463dc1b558..153fec655865 100644 --- a/tests/test-backend-ops.cpp +++ b/tests/test-backend-ops.cpp @@ -4390,7 +4390,7 @@ struct test_mul_mat_hadamard : public test_mul_mat { } }; -static void init_mul_mat_id_tensors(ggml_context * ctx, int n_mats) { +static void init_mul_mat_id_tensors(ggml_context * ctx, int n_mats, bool skip_ids = 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)) { @@ -4403,6 +4403,13 @@ static void init_mul_mat_id_tensors(ggml_context * ctx, int n_mats) { data[i] = i % n_mats; } std::shuffle(data.begin(), data.end(), rng); + if (skip_ids) { + // id == -1 marks "expert not owned by this pack" (hot/cold expert + // split); keep slot 0 valid so every row computes something + for (int i = 1; i < t->ne[0]; i += 2) { + data[i] = -1; + } + } ggml_backend_tensor_set(t, data.data(), r * t->nb[1], t->ne[0] * sizeof(int32_t)); } } else { @@ -4421,9 +4428,10 @@ struct test_mul_mat_id : public test_case { const int64_t m; const int64_t n; const int64_t k; + const bool skip_ids; // some ids are -1 (hot/cold expert-pack split) std::string vars() override { - return VARS_TO_STR8(type_a, type_b, n_mats, n_used, b, m, n, k); + return VARS_TO_STR9(type_a, type_b, n_mats, n_used, b, m, n, k, skip_ids); } double max_nmse_err() override { @@ -4445,9 +4453,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) + int64_t m = 32, int64_t n = 32, int64_t k = 32, bool skip_ids = 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) { + m(m), n(n), k(k), skip_ids(skip_ids) { GGML_ASSERT(n_used <= n_mats); } @@ -4468,12 +4476,17 @@ struct test_mul_mat_id : public test_case { ggml_tensor * out = ggml_mul_mat_id(ctx, as, b, ids); ggml_set_name(out, "out"); + if (skip_ids) { + // announce that ids may contain -1 so backends route around + // kernels without skip support (mirrors llama's pack nodes) + out->op_params[0] = 1; + } return out; } void initialize_tensors(ggml_context * ctx) override { - init_mul_mat_id_tensors(ctx, n_mats); + init_mul_mat_id_tensors(ctx, n_mats, skip_ids); } }; @@ -8807,6 +8820,14 @@ static std::vector> make_test_cases_eval() { test_cases.emplace_back(new test_mul_mat(GGML_TYPE_Q8_0, GGML_TYPE_F32, 8192, 1, 5120, {128, 1}, {1, 1})); test_cases.emplace_back(new test_mul_mat(GGML_TYPE_Q8_0, GGML_TYPE_F32, 8192, 512, 5120, {128, 1}, {1, 1})); #endif + // hot/cold expert-pack split: ids may be -1 ("expert not owned by this pack") and the + // op must emit zero rows for those slots — cover mmvq (n=1), mmq (n=64), mmf (f16) and + // the general fallback across quantized/float types + for (ggml_type ta : {GGML_TYPE_Q4_0, GGML_TYPE_Q8_0, GGML_TYPE_F16, GGML_TYPE_F32}) { + test_cases.emplace_back(new test_mul_mat_id(ta, GGML_TYPE_F32, 16, 8, false, 256, 1, 256, /*skip_ids=*/true)); + test_cases.emplace_back(new test_mul_mat_id(ta, GGML_TYPE_F32, 16, 8, false, 256, 4, 256, /*skip_ids=*/true)); + test_cases.emplace_back(new test_mul_mat_id(ta, GGML_TYPE_F32, 16, 8, false, 256, 64, 256, /*skip_ids=*/true)); + } for (ggml_type type_a : all_types) { for (int i = 1; i < 10; ++i) { diff --git a/tools/CMakeLists.txt b/tools/CMakeLists.txt index 780df3266132..fef34c77e6c7 100644 --- a/tools/CMakeLists.txt +++ b/tools/CMakeLists.txt @@ -17,6 +17,7 @@ else() add_subdirectory(batched-bench) add_subdirectory(gguf-split) add_subdirectory(imatrix) + add_subdirectory(moe-trace) add_subdirectory(llama-bench) add_subdirectory(completion) add_subdirectory(perplexity) diff --git a/tools/moe-trace/CMakeLists.txt b/tools/moe-trace/CMakeLists.txt new file mode 100644 index 000000000000..a842a822b220 --- /dev/null +++ b/tools/moe-trace/CMakeLists.txt @@ -0,0 +1,4 @@ +set(TARGET llama-moe-trace) +add_executable(${TARGET} moe-trace.cpp) +target_link_libraries(${TARGET} PRIVATE llama-common llama ${CMAKE_THREAD_LIBS_INIT}) +target_compile_features(${TARGET} PRIVATE cxx_std_17) diff --git a/tools/moe-trace/moe-trace.cpp b/tools/moe-trace/moe-trace.cpp new file mode 100644 index 000000000000..032d49d977f9 --- /dev/null +++ b/tools/moe-trace/moe-trace.cpp @@ -0,0 +1,151 @@ +// moe-trace: dump routed expert ids per layer per decode step. +// +// Captures the "ffn_moe_topk-" id tensors through the scheduler eval +// callback (same mechanism as imatrix), so it works regardless of which +// backend computed the node (CPU-offloaded experts included). +// +// Output CSV, one row per (position, layer): pos,layer,id0,id1,... +// Prompt rows are tagged with negative positions so the simulator can +// separate prefill routing from decode routing. +// +// Usage: +// MOE_TRACE_OUT=trace.csv llama-moe-trace -m model.gguf -ngl 99 -ncmoe 26 -fa on \ +// -p "prompt text" -n 512 + +#include "arg.h" +#include "common.h" +#include "log.h" +#include "llama.h" + +#include +#include +#include +#include + +struct trace_ctx { + FILE * out = nullptr; + int pos = 0; // current decode position (negative = prefill) + bool in_prompt = true; + std::vector buf; +}; + +static bool trace_cb(struct ggml_tensor * t, bool ask, void * user_data) { + trace_ctx * tc = (trace_ctx *) user_data; + + const bool is_topk = strncmp(t->name, "ffn_moe_topk-", 13) == 0; + if (ask) { + return is_topk; + } + if (!is_topk || t->type != GGML_TYPE_I32) { + return true; + } + + const int layer = atoi(t->name + 13); + const int n_used = (int) t->ne[0]; + const int n_tokens = (int) t->ne[1]; + + // topk is a non-contiguous view over the argsort rows: copy the full + // strided byte range, then index by nb[] — sizing by n_used*n_tokens + // would under-allocate and tensor_get would smash the heap. + const size_t nbytes = ggml_nbytes(t); + tc->buf.resize((nbytes + sizeof(int32_t) - 1) / sizeof(int32_t)); + ggml_backend_tensor_get(t, tc->buf.data(), 0, nbytes); + const char * base = (const char *) tc->buf.data(); + + for (int j = 0; j < n_tokens; j++) { + // prefill batches carry n_tokens > 1; decode steps carry 1 + const int pos = tc->in_prompt ? -(tc->pos + n_tokens - j) : tc->pos; + fprintf(tc->out, "%d,%d", pos, layer); + for (int i = 0; i < n_used; i++) { + const int32_t id = *(const int32_t *)(base + j*t->nb[1] + i*t->nb[0]); + fprintf(tc->out, ",%d", id); + } + fputc('\n', tc->out); + } + return true; +} + +int main(int argc, char ** argv) { + common_params params; + params.n_predict = 256; + + // reuse the standard arg parser; -o (out_file) holds the trace path + if (!common_params_parse(argc, argv, params, LLAMA_EXAMPLE_PERPLEXITY)) { + return 1; + } + // output path via env — the arg registry has no free slot for this example + const char * out_path = getenv("MOE_TRACE_OUT"); + if (!out_path) { + out_path = "moe-trace.csv"; + } + + common_init(); + + trace_ctx tc; + tc.out = fopen(out_path, "w"); + if (!tc.out) { + LOG_ERR("failed to open %s for writing\n", out_path); + return 1; + } + + params.cb_eval = trace_cb; + params.cb_eval_user_data = &tc; + params.warmup = false; + + llama_backend_init(); + llama_numa_init(params.numa); + + common_init_result_ptr llama_init = common_init_from_params(params); + llama_model * model = llama_init ? llama_init->model() : nullptr; + llama_context * lctx = llama_init ? llama_init->context() : nullptr; + if (model == nullptr || lctx == nullptr) { + LOG_ERR("failed to load model\n"); + return 1; + } + const llama_vocab * vocab = llama_model_get_vocab(model); + + std::vector tokens = common_tokenize(lctx, params.prompt, true); + if (tokens.empty()) { + LOG_ERR("empty prompt\n"); + return 1; + } + LOG_INF("prompt: %zu tokens, decoding %d\n", tokens.size(), params.n_predict); + + // prefill + tc.in_prompt = true; + tc.pos = 0; + for (size_t i = 0; i < tokens.size(); i += params.n_batch) { + const int n_eval = std::min((int) (tokens.size() - i), params.n_batch); + if (llama_decode(lctx, llama_batch_get_one(tokens.data() + i, n_eval))) { + LOG_ERR("prefill failed at %zu\n", i); + return 1; + } + tc.pos += n_eval; + } + + // greedy decode + tc.in_prompt = false; + llama_sampler * smpl = llama_sampler_init_greedy(); + llama_token tok = 0; + for (int i = 0; i < params.n_predict; i++) { + tok = llama_sampler_sample(smpl, lctx, -1); + if (llama_vocab_is_eog(vocab, tok)) { + break; + } + tc.pos = i; + if (llama_decode(lctx, llama_batch_get_one(&tok, 1))) { + LOG_ERR("decode failed at %d\n", i); + return 1; + } + if (i % 64 == 0) { + LOG_INF("decoded %d/%d\n", i, params.n_predict); + } + } + llama_sampler_free(smpl); + + fclose(tc.out); + LOG_INF("trace written to %s\n", out_path); + + llama_backend_free(); + return 0; +} diff --git a/tools/moe-trace/simulate.py b/tools/moe-trace/simulate.py new file mode 100644 index 000000000000..974f55da0ddd --- /dev/null +++ b/tools/moe-trace/simulate.py @@ -0,0 +1,197 @@ +#!/usr/bin/env python3 +"""Replay a moe-trace CSV against expert-cache policies. + +Each layer gets an independent cache of S expert slots (experts are +layer-specific). A "hit" means a routed expert was already resident in +VRAM when the router asked for it. Hit rate directly bounds how much +decode-time RAM traffic a cache can remove. + +Policies: + static-oracle top-S experts per layer by whole-trace frequency (upper bound) + static-prefill top-S by PREFILL routing only, fixed for decode (deployable v1) + lru evict least-recently-used on miss + lfu-decay score = EMA of use; evict lowest; insert always + reelect-K EMA counters, cache membership re-elected every K decode steps + (the Phase 2 design: batch eviction, no per-token churn) + +Usage: simulate.py trace.csv [--budgets 0.125,0.25,0.5] [--reelect 32] +""" +import argparse +import csv +import math +import sys +from collections import defaultdict + + +def load(path): + prefill, decode = defaultdict(list), defaultdict(list) # layer -> [ids per step] + n_expert_seen = 0 + with open(path) as f: + for row in csv.reader(f): + pos, layer, ids = int(row[0]), int(row[1]), [int(x) for x in row[2:]] + (prefill if pos < 0 else decode)[layer].append((pos, ids)) + n_expert_seen = max(n_expert_seen, max(ids) + 1) + for d in (prefill, decode): + for layer in d: + d[layer].sort(key=lambda t: t[0]) + d[layer] = [ids for _, ids in d[layer]] + return prefill, decode, n_expert_seen + + +def sim_static(decode, resident): + hits = total = 0 + for layer, steps in decode.items(): + r = resident.get(layer, set()) + for ids in steps: + for e in ids: + hits += e in r + total += 1 + return hits / max(total, 1) + + +def top_by_freq(steps_by_layer, S): + resident = {} + for layer, steps in steps_by_layer.items(): + freq = defaultdict(int) + for ids in steps: + for e in ids: + freq[e] += 1 + resident[layer] = set(sorted(freq, key=freq.get, reverse=True)[:S]) + return resident + + +def sim_lru(decode, S): + hits = total = ins = 0 + n_steps = 0 + for layer, steps in decode.items(): + cache, clock = {}, 0 + n_steps = max(n_steps, len(steps)) + for ids in steps: + for e in ids: + clock += 1 + if e in cache: + hits += 1 + else: + if len(cache) >= S: + cache.pop(min(cache, key=cache.get)) + ins += 1 + cache[e] = clock + total += 1 + return hits / max(total, 1), ins / max(n_steps, 1) + + +def sim_lfu_decay(decode, S, alpha=0.95): + hits = total = ins = 0 + n_steps = 0 + for layer, steps in decode.items(): + score, cache = defaultdict(float), set() + n_steps = max(n_steps, len(steps)) + for ids in steps: + for k in score: + score[k] *= alpha + for e in ids: + score[e] += 1.0 + if e in cache: + hits += 1 + else: + if len(cache) < S: + cache.add(e) + ins += 1 + else: + victim = min(cache, key=lambda k: score[k]) + if score[e] >= score[victim]: + cache.discard(victim) + cache.add(e) + ins += 1 + total += 1 + return hits / max(total, 1), ins / max(n_steps, 1) + + +def sim_reelect(decode, prefill, S, K, alpha=0.98): + hits = total = ins = 0 + n_steps = 0 + for layer, steps in decode.items(): + score = defaultdict(float) + n_steps = max(n_steps, len(steps)) + for ids in prefill.get(layer, []): # free warm-up from prompt routing + for e in ids: + score[e] += 1.0 + cache = set(sorted(score, key=score.get, reverse=True)[:S]) + for step, ids in enumerate(steps): + for e in ids: + hits += e in cache + total += 1 + score[e] += 1.0 + if step % K == K - 1: + for k in score: + score[k] *= alpha + new = set(sorted(score, key=score.get, reverse=True)[:S]) + ins += len(new - cache) + cache = new + return hits / max(total, 1), ins / max(n_steps, 1) + + +def main(): + ap = argparse.ArgumentParser() + ap.add_argument("trace") + ap.add_argument("--budgets", default="0.0625,0.125,0.25,0.5", + help="cache size as fraction of expert count") + ap.add_argument("--reelect", type=int, default=32) + ap.add_argument("--slab-mb", type=float, default=1.9, + help="MB per (layer,expert) gate+up+down slab") + args = ap.parse_args() + + prefill, decode, n_expert = load(args.trace) + n_layers = len(decode) + n_steps = max(len(s) for s in decode.values()) if decode else 0 + n_used = len(next(iter(decode.values()))[0]) if decode else 0 + print(f"trace: {n_layers} MoE layers, {n_expert} experts, " + f"{n_used} active/token, {n_steps} decode steps\n") + + # skew snapshot: what share of routed traffic hits the top-k% experts + freq = defaultdict(int) + for steps in decode.values(): + for ids in steps: + for e in ids: + freq[e] += 1 + ranked = sorted(freq.values(), reverse=True) + tot = sum(ranked) + for frac in (0.1, 0.25, 0.5): + k = max(1, int(len(ranked) * frac)) + print(f"top {frac:>4.0%} of (layer,expert) pairs carry " + f"{sum(ranked[:k])/tot:.1%} of decode routing") + print() + + # cost model: expert slab MB, RAM GB/s (CPU miss read), PCIe GB/s (upload) + SLAB_MB, RAM_GBS, PCIE_GBS = args.slab_mb, 45.0, 12.4 + reads_per_step = sum(len(s2[0]) for s2 in decode.values()) # experts touched/step + + def cost_ms(hit, ins_per_step): + miss_mb = reads_per_step * (1 - hit) * SLAB_MB + upload_mb = ins_per_step * SLAB_MB + return miss_mb / RAM_GBS, upload_mb / PCIE_GBS # ms if GB/s and MB + + hdr = f"{'budget':>7} {'slots':>6} | {'policy':>10} {'hit':>7} {'up-MB/tok':>9} " \ + f"{'miss-ms':>8} {'up-ms':>6}" + print(hdr) + print("-" * len(hdr)) + for frac in [float(x) for x in args.budgets.split(",")]: + S = max(1, int(n_expert * frac)) + rows = [ + ("st-oracle", sim_static(decode, top_by_freq(decode, S)), 0.0), + ("st-prefill", sim_static(decode, top_by_freq(prefill, S)) if prefill else float("nan"), 0.0), + ] + for name, fn in (("lru", sim_lru), ("lfu-decay", sim_lfu_decay)): + h, i = fn(decode, S) + rows.append((name, h, i)) + h, i = sim_reelect(decode, prefill, S, args.reelect) + rows.append((f"reelect-{args.reelect}", h, i)) + for name, h, i in rows: + miss_ms, up_ms = cost_ms(h, i) + print(f"{frac:>7.4f} {S:>6} | {name:>10} {h:>7.1%} {i*SLAB_MB:>9.2f} " + f"{miss_ms:>8.2f} {up_ms:>6.2f}") + print() + + +if __name__ == "__main__": + main()