diff --git a/src/models/models.h b/src/models/models.h index 1cac1d4328f..8ec39093dba 100644 --- a/src/models/models.h +++ b/src/models/models.h @@ -1608,7 +1608,14 @@ struct llama_model_openai_moe : public llama_model_base { void load_arch_tensors(llama_model_loader & ml) override; struct graph : public llm_graph_context { + const llama_model & model; + graph(const llama_model & model, const llm_graph_params & params); + + ggml_tensor * build_layer_moe_hot( + ggml_tensor * cur, + ggml_tensor * logits, + int il); }; std::unique_ptr build_arch_graph(const llm_graph_params & params) const override; diff --git a/src/models/openai-moe.cpp b/src/models/openai-moe.cpp index 3ab15d61f08..13b018409a6 100644 --- a/src/models/openai-moe.cpp +++ b/src/models/openai-moe.cpp @@ -1,4 +1,5 @@ #include "models.h" +#include "moe-hot-cache/llama-moe-hot-cache.h" void llama_model_openai_moe::load_arch_hparams(llama_model_loader & ml) { ml.get_key(LLM_KV_ATTENTION_LAYERNORM_RMS_EPS, hparams.f_norm_rms_eps); @@ -61,7 +62,9 @@ std::unique_ptr llama_model_openai_moe::build_arch_graph(cons return std::make_unique(*this, params); } -llama_model_openai_moe::graph::graph(const llama_model & model, const llm_graph_params & params) : llm_graph_context(params) { +llama_model_openai_moe::graph::graph(const llama_model & model, const llm_graph_params & params) : + llm_graph_context(params), + model(model) { ggml_tensor * cur; ggml_tensor * inpL; @@ -129,17 +132,29 @@ llama_model_openai_moe::graph::graph(const llama_model & model, const llm_graph_ cb(cur, "attn_post_norm", il); // MoE branch - cur = build_moe_ffn(cur, - model.layers[il].ffn_gate_inp, model.layers[il].ffn_gate_inp_b, - model.layers[il].ffn_up_exps, model.layers[il].ffn_up_exps_b, - model.layers[il].ffn_gate_exps, model.layers[il].ffn_gate_exps_b, - model.layers[il].ffn_down_exps, model.layers[il].ffn_down_exps_b, - nullptr, - n_expert, n_expert_used, - LLM_FFN_SWIGLU_OAI_MOE, false, - hparams.expert_weights_scale, - LLAMA_EXPERT_GATING_FUNC_TYPE_SOFTMAX_WEIGHT, - il); + if (llama_moe_hot_cache_layer_active_for_graph(model, il, llama_moe_hot_cache_graph_kind::logits)) { + ggml_tensor * logits = build_lora_mm(model.layers[il].ffn_gate_inp, cur); + cb(logits, "ffn_moe_logits", il); + + if (model.layers[il].ffn_gate_inp_b) { + logits = ggml_add(ctx0, logits, model.layers[il].ffn_gate_inp_b); + cb(logits, "ffn_moe_logits_biased", il); + } + + cur = build_layer_moe_hot(cur, logits, il); + } else { + cur = build_moe_ffn(cur, + model.layers[il].ffn_gate_inp, model.layers[il].ffn_gate_inp_b, + model.layers[il].ffn_up_exps, model.layers[il].ffn_up_exps_b, + model.layers[il].ffn_gate_exps, model.layers[il].ffn_gate_exps_b, + model.layers[il].ffn_down_exps, model.layers[il].ffn_down_exps_b, + nullptr, + n_expert, n_expert_used, + LLM_FFN_SWIGLU_OAI_MOE, false, + hparams.expert_weights_scale, + LLAMA_EXPERT_GATING_FUNC_TYPE_SOFTMAX_WEIGHT, + il); + } cb(cur, "ffn_moe_out", il); cur = ggml_add(ctx0, cur, ffn_inp); diff --git a/src/moe-hot-cache/llama-moe-hot-cache-adapter.cpp b/src/moe-hot-cache/llama-moe-hot-cache-adapter.cpp index 366bf9f76ce..05973a3a629 100644 --- a/src/moe-hot-cache/llama-moe-hot-cache-adapter.cpp +++ b/src/moe-hot-cache/llama-moe-hot-cache-adapter.cpp @@ -67,11 +67,18 @@ static llama_moe_hot_cache_graph_profile mellum_profile() { return qwen35_profile(); } +static llama_moe_hot_cache_graph_profile openai_moe_profile() { + // GPT-OSS uses logits-based top-k routing with OpenAI SwiGLU experts. + // Reuse the conservative logits profile until GPT-OSS-specific PP shortcuts are measured. + return qwen35_profile(); +} + static const llama_moe_hot_cache_model_adapter ADAPTERS[] = { { LLM_ARCH_QWEN35MOE, "qwen35moe", llama_moe_hot_cache_graph_kind::qwen35_ffn, LLM_FFN_SILU }, { LLM_ARCH_QWEN3NEXT, "qwen3next", llama_moe_hot_cache_graph_kind::logits, LLM_FFN_SILU }, { LLM_ARCH_GEMMA4, "gemma4", llama_moe_hot_cache_graph_kind::logits, LLM_FFN_GELU }, { LLM_ARCH_MELLUM, "mellum", llama_moe_hot_cache_graph_kind::logits, LLM_FFN_SILU }, + { LLM_ARCH_OPENAI_MOE, "gpt-oss", llama_moe_hot_cache_graph_kind::logits, LLM_FFN_SWIGLU_OAI_MOE }, }; } // namespace @@ -197,6 +204,8 @@ llama_moe_hot_cache_graph_profile llama_moe_hot_cache_model_adapter::profile() c return gemma4_profile(); case LLM_ARCH_MELLUM: return mellum_profile(); + case LLM_ARCH_OPENAI_MOE: + return openai_moe_profile(); default: return {}; } diff --git a/src/moe-hot-cache/llama-moe-hot-cache-builder.cpp b/src/moe-hot-cache/llama-moe-hot-cache-builder.cpp index 2e039f999c6..cf3f0d1e568 100644 --- a/src/moe-hot-cache/llama-moe-hot-cache-builder.cpp +++ b/src/moe-hot-cache/llama-moe-hot-cache-builder.cpp @@ -39,6 +39,20 @@ static ggml_tensor * new_tensor_like_scale( return dst; } +static ggml_tensor * new_tensor_like_expert_bias( + ggml_context * ctx, + const ggml_tensor * src, + int64_t n_cache, + const char * name) { + if (src == nullptr) { + return nullptr; + } + + ggml_tensor * dst = ggml_new_tensor_2d(ctx, src->type, src->ne[0], n_cache); + ggml_set_name(dst, name); + return dst; +} + static void zero_tensor(ggml_tensor * t) { std::vector zeros(ggml_nbytes(t), 0); ggml_backend_tensor_set(t, zeros.data(), 0, zeros.size()); @@ -111,6 +125,21 @@ void llama_moe_hot_cache_copy_scale_slice( ggml_backend_tensor_set(dst, buf.data(), dst->nb[0]*dst_expert, bytes); } +void llama_moe_hot_cache_copy_bias_slice( + const ggml_tensor * src, + ggml_tensor * dst, + uint32_t src_expert, + uint32_t dst_expert) { + if (src == nullptr || dst == nullptr) { + return; + } + + const size_t bytes = ggml_nbytes(src) / size_t(src->ne[1]); + std::vector buf(bytes); + ggml_backend_tensor_get(src, buf.data(), src->nb[1]*src_expert, bytes); + ggml_backend_tensor_set(dst, buf.data(), dst->nb[1]*dst_expert, bytes); +} + void llama_moe_hot_cache_set_tensor_i32_1d(ggml_tensor * t, uint32_t index, int32_t value) { ggml_backend_tensor_set(t, &value, t->nb[1]*index, sizeof(value)); } @@ -148,6 +177,10 @@ std::unique_ptr llama_moe_hot_cache_build( n_tensors += 4; // map + hot mask + cold mask + down n_tensors += layer.ffn_gate_up_exps != nullptr ? 1 : 2; + n_tensors += layer.ffn_gate_up_exps_b != nullptr ? 1 : 0; + n_tensors += layer.ffn_gate_exps_b != nullptr ? 1 : 0; + n_tensors += layer.ffn_up_exps_b != nullptr ? 1 : 0; + n_tensors += layer.ffn_down_exps_b != nullptr ? 1 : 0; n_tensors += layer.ffn_gate_exps_s != nullptr ? 1 : 0; n_tensors += layer.ffn_up_exps_s != nullptr ? 1 : 0; n_tensors += layer.ffn_down_exps_s != nullptr ? 1 : 0; @@ -190,6 +223,10 @@ std::unique_ptr llama_moe_hot_cache_build( dst.ffn_gate_exps = new_tensor_like_experts(ctx.get(), src.ffn_gate_exps, n_cache, format("blk.%u.ffn_gate_exps.hot_cache", il).c_str()); dst.ffn_up_exps = new_tensor_like_experts(ctx.get(), src.ffn_up_exps, n_cache, format("blk.%u.ffn_up_exps.hot_cache", il).c_str()); dst.ffn_down_exps = new_tensor_like_experts(ctx.get(), src.ffn_down_exps, n_cache, format("blk.%u.ffn_down_exps.hot_cache", il).c_str()); + dst.ffn_gate_up_exps_b = new_tensor_like_expert_bias(ctx.get(), src.ffn_gate_up_exps_b, n_cache, format("blk.%u.ffn_gate_up_exps_b.hot_cache", il).c_str()); + dst.ffn_gate_exps_b = new_tensor_like_expert_bias(ctx.get(), src.ffn_gate_exps_b, n_cache, format("blk.%u.ffn_gate_exps_b.hot_cache", il).c_str()); + dst.ffn_up_exps_b = new_tensor_like_expert_bias(ctx.get(), src.ffn_up_exps_b, n_cache, format("blk.%u.ffn_up_exps_b.hot_cache", il).c_str()); + dst.ffn_down_exps_b = new_tensor_like_expert_bias(ctx.get(), src.ffn_down_exps_b, n_cache, format("blk.%u.ffn_down_exps_b.hot_cache", il).c_str()); dst.ffn_gate_exps_s = new_tensor_like_scale (ctx.get(), src.ffn_gate_exps_s, n_cache, format("blk.%u.ffn_gate_exps_s.hot_cache", il).c_str()); dst.ffn_up_exps_s = new_tensor_like_scale (ctx.get(), src.ffn_up_exps_s, n_cache, format("blk.%u.ffn_up_exps_s.hot_cache", il).c_str()); dst.ffn_down_exps_s = new_tensor_like_scale (ctx.get(), src.ffn_down_exps_s, n_cache, format("blk.%u.ffn_down_exps_s.hot_cache", il).c_str()); @@ -224,6 +261,10 @@ std::unique_ptr llama_moe_hot_cache_build( if (dst.ffn_gate_exps) { zero_tensor(dst.ffn_gate_exps); } if (dst.ffn_up_exps) { zero_tensor(dst.ffn_up_exps); } if (dst.ffn_down_exps) { zero_tensor(dst.ffn_down_exps); } + if (dst.ffn_gate_up_exps_b) { zero_tensor(dst.ffn_gate_up_exps_b); } + if (dst.ffn_gate_exps_b) { zero_tensor(dst.ffn_gate_exps_b); } + if (dst.ffn_up_exps_b) { zero_tensor(dst.ffn_up_exps_b); } + if (dst.ffn_down_exps_b) { zero_tensor(dst.ffn_down_exps_b); } if (dst.ffn_gate_exps_s) { zero_tensor(dst.ffn_gate_exps_s); } if (dst.ffn_up_exps_s) { zero_tensor(dst.ffn_up_exps_s); } if (dst.ffn_down_exps_s) { zero_tensor(dst.ffn_down_exps_s); } @@ -248,6 +289,10 @@ std::unique_ptr llama_moe_hot_cache_build( llama_moe_hot_cache_copy_expert_slice(src.ffn_gate_exps, dst.ffn_gate_exps, expert, cache_id); llama_moe_hot_cache_copy_expert_slice(src.ffn_up_exps, dst.ffn_up_exps, expert, cache_id); llama_moe_hot_cache_copy_expert_slice(src.ffn_down_exps, dst.ffn_down_exps, expert, cache_id); + llama_moe_hot_cache_copy_bias_slice(src.ffn_gate_up_exps_b, dst.ffn_gate_up_exps_b, expert, cache_id); + llama_moe_hot_cache_copy_bias_slice(src.ffn_gate_exps_b, dst.ffn_gate_exps_b, expert, cache_id); + llama_moe_hot_cache_copy_bias_slice(src.ffn_up_exps_b, dst.ffn_up_exps_b, expert, cache_id); + llama_moe_hot_cache_copy_bias_slice(src.ffn_down_exps_b, dst.ffn_down_exps_b, expert, cache_id); llama_moe_hot_cache_copy_scale_slice(src.ffn_gate_exps_s, dst.ffn_gate_exps_s, expert, cache_id); llama_moe_hot_cache_copy_scale_slice(src.ffn_up_exps_s, dst.ffn_up_exps_s, expert, cache_id); llama_moe_hot_cache_copy_scale_slice(src.ffn_down_exps_s, dst.ffn_down_exps_s, expert, cache_id); @@ -315,6 +360,10 @@ std::unique_ptr llama_moe_hot_cache_build_multi( n_tensors += 4; // map + hot mask + cold mask + down n_tensors += src.ffn_gate_up_exps != nullptr ? 1 : 2; + n_tensors += src.ffn_gate_up_exps_b != nullptr ? 1 : 0; + n_tensors += src.ffn_gate_exps_b != nullptr ? 1 : 0; + n_tensors += src.ffn_up_exps_b != nullptr ? 1 : 0; + n_tensors += src.ffn_down_exps_b != nullptr ? 1 : 0; n_tensors += src.ffn_gate_exps_s != nullptr ? 1 : 0; n_tensors += src.ffn_up_exps_s != nullptr ? 1 : 0; n_tensors += src.ffn_down_exps_s != nullptr ? 1 : 0; @@ -365,6 +414,10 @@ std::unique_ptr llama_moe_hot_cache_build_multi( dst_lane.ffn_gate_exps = new_tensor_like_experts(ctx.get(), src.ffn_gate_exps, n_cache, format("blk.%u.ffn_gate_exps.hot_cache.%zu", il, lane_index).c_str()); dst_lane.ffn_up_exps = new_tensor_like_experts(ctx.get(), src.ffn_up_exps, n_cache, format("blk.%u.ffn_up_exps.hot_cache.%zu", il, lane_index).c_str()); dst_lane.ffn_down_exps = new_tensor_like_experts(ctx.get(), src.ffn_down_exps, n_cache, format("blk.%u.ffn_down_exps.hot_cache.%zu", il, lane_index).c_str()); + dst_lane.ffn_gate_up_exps_b = new_tensor_like_expert_bias(ctx.get(), src.ffn_gate_up_exps_b, n_cache, format("blk.%u.ffn_gate_up_exps_b.hot_cache.%zu", il, lane_index).c_str()); + dst_lane.ffn_gate_exps_b = new_tensor_like_expert_bias(ctx.get(), src.ffn_gate_exps_b, n_cache, format("blk.%u.ffn_gate_exps_b.hot_cache.%zu", il, lane_index).c_str()); + dst_lane.ffn_up_exps_b = new_tensor_like_expert_bias(ctx.get(), src.ffn_up_exps_b, n_cache, format("blk.%u.ffn_up_exps_b.hot_cache.%zu", il, lane_index).c_str()); + dst_lane.ffn_down_exps_b = new_tensor_like_expert_bias(ctx.get(), src.ffn_down_exps_b, n_cache, format("blk.%u.ffn_down_exps_b.hot_cache.%zu", il, lane_index).c_str()); dst_lane.ffn_gate_exps_s = new_tensor_like_scale (ctx.get(), src.ffn_gate_exps_s, n_cache, format("blk.%u.ffn_gate_exps_s.hot_cache.%zu", il, lane_index).c_str()); dst_lane.ffn_up_exps_s = new_tensor_like_scale (ctx.get(), src.ffn_up_exps_s, n_cache, format("blk.%u.ffn_up_exps_s.hot_cache.%zu", il, lane_index).c_str()); dst_lane.ffn_down_exps_s = new_tensor_like_scale (ctx.get(), src.ffn_down_exps_s, n_cache, format("blk.%u.ffn_down_exps_s.hot_cache.%zu", il, lane_index).c_str()); @@ -400,6 +453,10 @@ std::unique_ptr llama_moe_hot_cache_build_multi( if (dst_lane.ffn_gate_exps) { zero_tensor(dst_lane.ffn_gate_exps); } if (dst_lane.ffn_up_exps) { zero_tensor(dst_lane.ffn_up_exps); } if (dst_lane.ffn_down_exps) { zero_tensor(dst_lane.ffn_down_exps); } + if (dst_lane.ffn_gate_up_exps_b) { zero_tensor(dst_lane.ffn_gate_up_exps_b); } + if (dst_lane.ffn_gate_exps_b) { zero_tensor(dst_lane.ffn_gate_exps_b); } + if (dst_lane.ffn_up_exps_b) { zero_tensor(dst_lane.ffn_up_exps_b); } + if (dst_lane.ffn_down_exps_b) { zero_tensor(dst_lane.ffn_down_exps_b); } if (dst_lane.ffn_gate_exps_s) { zero_tensor(dst_lane.ffn_gate_exps_s); } if (dst_lane.ffn_up_exps_s) { zero_tensor(dst_lane.ffn_up_exps_s); } if (dst_lane.ffn_down_exps_s) { zero_tensor(dst_lane.ffn_down_exps_s); } @@ -425,6 +482,10 @@ std::unique_ptr llama_moe_hot_cache_build_multi( llama_moe_hot_cache_copy_expert_slice(src.ffn_gate_exps, dst_lane.ffn_gate_exps, expert, cache_id); llama_moe_hot_cache_copy_expert_slice(src.ffn_up_exps, dst_lane.ffn_up_exps, expert, cache_id); llama_moe_hot_cache_copy_expert_slice(src.ffn_down_exps, dst_lane.ffn_down_exps, expert, cache_id); + llama_moe_hot_cache_copy_bias_slice(src.ffn_gate_up_exps_b, dst_lane.ffn_gate_up_exps_b, expert, cache_id); + llama_moe_hot_cache_copy_bias_slice(src.ffn_gate_exps_b, dst_lane.ffn_gate_exps_b, expert, cache_id); + llama_moe_hot_cache_copy_bias_slice(src.ffn_up_exps_b, dst_lane.ffn_up_exps_b, expert, cache_id); + llama_moe_hot_cache_copy_bias_slice(src.ffn_down_exps_b, dst_lane.ffn_down_exps_b, expert, cache_id); llama_moe_hot_cache_copy_scale_slice(src.ffn_gate_exps_s, dst_lane.ffn_gate_exps_s, expert, cache_id); llama_moe_hot_cache_copy_scale_slice(src.ffn_up_exps_s, dst_lane.ffn_up_exps_s, expert, cache_id); llama_moe_hot_cache_copy_scale_slice(src.ffn_down_exps_s, dst_lane.ffn_down_exps_s, expert, cache_id); diff --git a/src/moe-hot-cache/llama-moe-hot-cache-builder.h b/src/moe-hot-cache/llama-moe-hot-cache-builder.h index 1d80532f375..c46f6ac011d 100644 --- a/src/moe-hot-cache/llama-moe-hot-cache-builder.h +++ b/src/moe-hot-cache/llama-moe-hot-cache-builder.h @@ -39,6 +39,12 @@ void llama_moe_hot_cache_copy_scale_slice( uint32_t src_expert, uint32_t dst_expert); +void llama_moe_hot_cache_copy_bias_slice( + const ggml_tensor * src, + ggml_tensor * dst, + uint32_t src_expert, + uint32_t dst_expert); + void llama_moe_hot_cache_set_tensor_i32_1d(ggml_tensor * t, uint32_t index, int32_t value); void llama_moe_hot_cache_set_tensor_f32_1d(ggml_tensor * t, uint32_t index, float value); diff --git a/src/moe-hot-cache/llama-moe-hot-cache-graph.cpp b/src/moe-hot-cache/llama-moe-hot-cache-graph.cpp index 8c50e4952a6..a4d65a1a2aa 100644 --- a/src/moe-hot-cache/llama-moe-hot-cache-graph.cpp +++ b/src/moe-hot-cache/llama-moe-hot-cache-graph.cpp @@ -267,13 +267,17 @@ static ggml_tensor * llama_moe_hot_cache_build_moe_ffn_with_ids( ggml_tensor * selected_experts, ggml_tensor * weights, ggml_tensor * up_exps, + ggml_tensor * up_exps_b, ggml_tensor * gate_exps, + ggml_tensor * gate_exps_b, ggml_tensor * down_exps, + ggml_tensor * down_exps_b, int64_t n_expert, int64_t n_expert_used, llm_ffn_op_type type_op, int il, ggml_tensor * gate_up_exps, + ggml_tensor * gate_up_exps_b, ggml_tensor * up_exps_s, ggml_tensor * gate_exps_s, ggml_tensor * down_exps_s, @@ -291,7 +295,7 @@ static ggml_tensor * llama_moe_hot_cache_build_moe_ffn_with_ids( 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; - ggml_tensor * selected_experts_scale_ids = selected_experts; + ggml_tensor * selected_experts_safe_ids = selected_experts; const auto cb_moe = [&](ggml_tensor * t, const char * name) { if (branch_backend != nullptr) { ggml_backend_sched_set_tensor_backend(sched, t, branch_backend); @@ -319,12 +323,20 @@ static ggml_tensor * llama_moe_hot_cache_build_moe_ffn_with_ids( ggml_build_forward_expand(gf, weights); } - if ((flags & LLAMA_MOE_HOT_CACHE_MUL_MAT_ID_FLAG_ALLOW_NEGATIVE_IDS) && - (up_exps_s != nullptr || gate_exps_s != nullptr || down_exps_s != nullptr)) { + const bool needs_safe_ids = + up_exps_b != nullptr || + gate_exps_b != nullptr || + down_exps_b != nullptr || + gate_up_exps_b != nullptr || + up_exps_s != nullptr || + gate_exps_s != nullptr || + down_exps_s != nullptr; + + if ((flags & LLAMA_MOE_HOT_CACHE_MUL_MAT_ID_FLAG_ALLOW_NEGATIVE_IDS) && needs_safe_ids) { ggml_tensor * scale_ids_f32 = ggml_cast(ctx0, selected_experts, GGML_TYPE_F32); scale_ids_f32 = ggml_clamp(ctx0, scale_ids_f32, 0.0f, float(n_expert - 1)); - selected_experts_scale_ids = ggml_cast(ctx0, scale_ids_f32, GGML_TYPE_I32); - cb_moe(selected_experts_scale_ids, "ffn_moe_scale_ids"); + selected_experts_safe_ids = ggml_cast(ctx0, scale_ids_f32, GGML_TYPE_I32); + cb_moe(selected_experts_safe_ids, "ffn_moe_scale_ids"); } // Negative-ID rows in gate/up are ignored by the final down projection, so avoid @@ -353,10 +365,15 @@ static ggml_tensor * llama_moe_hot_cache_build_moe_ffn_with_ids( ggml_tensor * gate_up = llama_moe_hot_cache_build_lora_mm_id(graph, gate_up_exps, cur, selected_experts, intermediate_flags); cb_moe(gate_up, "ffn_moe_gate_up"); + if (gate_up_exps_b) { + gate_up = ggml_add_id(ctx0, gate_up, gate_up_exps_b, selected_experts_safe_ids); + cb_moe(gate_up, "ffn_moe_gate_up_biased"); + } + if (up_exps_s) { ggml_tensor * s = ggml_reshape_3d(ctx0, up_exps_s, 1, n_expert, 1); s = ggml_repeat_4d(ctx0, s, 1, n_expert, n_tokens, 1); - s = ggml_get_rows(ctx0, s, selected_experts_scale_ids); + s = ggml_get_rows(ctx0, s, selected_experts_safe_ids); gate_up = ggml_mul(ctx0, gate_up, s); cb_moe(gate_up, "ffn_moe_gate_up_scaled"); } @@ -370,10 +387,15 @@ static ggml_tensor * llama_moe_hot_cache_build_moe_ffn_with_ids( up = llama_moe_hot_cache_build_lora_mm_id(graph, up_exps, cur, selected_experts, intermediate_flags); cb_moe(up, "ffn_moe_up"); + if (up_exps_b) { + up = ggml_add_id(ctx0, up, up_exps_b, selected_experts_safe_ids); + cb_moe(up, "ffn_moe_up_biased"); + } + if (up_exps_s) { ggml_tensor * s = ggml_reshape_3d(ctx0, up_exps_s, 1, n_expert, 1); s = ggml_repeat_4d(ctx0, s, 1, n_expert, n_tokens, 1); - s = ggml_get_rows(ctx0, s, selected_experts_scale_ids); + s = ggml_get_rows(ctx0, s, selected_experts_safe_ids); up = ggml_mul(ctx0, up, s); cb_moe(up, "ffn_moe_up_scaled"); } @@ -385,10 +407,15 @@ static ggml_tensor * llama_moe_hot_cache_build_moe_ffn_with_ids( cur = up; } + if (gate_exps_b) { + cur = ggml_add_id(ctx0, cur, gate_exps_b, selected_experts_safe_ids); + cb_moe(cur, "ffn_moe_gate_biased"); + } + if (gate_exps_s) { ggml_tensor * s = ggml_reshape_3d(ctx0, gate_exps_s, 1, n_expert, 1); s = ggml_repeat_4d(ctx0, s, 1, n_expert, n_tokens, 1); - s = ggml_get_rows(ctx0, s, selected_experts_scale_ids); + s = ggml_get_rows(ctx0, s, selected_experts_safe_ids); cur = ggml_mul(ctx0, cur, s); cb_moe(cur, "ffn_moe_gate_scaled"); } @@ -413,6 +440,13 @@ static ggml_tensor * llama_moe_hot_cache_build_moe_ffn_with_ids( cur = ggml_gelu(ctx0, cur); cb_moe(cur, "ffn_moe_gelu"); } break; + case LLM_FFN_SWIGLU_OAI_MOE: + { + constexpr float alpha = 1.702f; + constexpr float limit = 7.0f; + cur = ggml_swiglu_oai(ctx0, cur, up, alpha, limit); + cb_moe(cur, "ffn_moe_swiglu_oai"); + } break; case LLM_FFN_RELU: if (has_gate) { cur = ggml_reglu_split(ctx0, cur, up); @@ -428,10 +462,15 @@ static ggml_tensor * llama_moe_hot_cache_build_moe_ffn_with_ids( experts = llama_moe_hot_cache_build_lora_mm_id(graph, down_exps, cur, selected_experts, output_flags); cb_moe(experts, "ffn_moe_down"); + if (down_exps_b) { + experts = ggml_add_id(ctx0, experts, down_exps_b, selected_experts_safe_ids); + cb_moe(experts, "ffn_moe_down_biased"); + } + if (down_exps_s) { ggml_tensor * s = ggml_reshape_3d(ctx0, down_exps_s, 1, n_expert, 1); s = ggml_repeat_4d(ctx0, s, 1, n_expert, n_tokens, 1); - s = ggml_get_rows(ctx0, s, selected_experts_scale_ids); + s = ggml_get_rows(ctx0, s, selected_experts_safe_ids); experts = ggml_mul(ctx0, experts, s); cb_moe(experts, "ffn_moe_down_scaled"); } @@ -696,13 +735,17 @@ static ggml_tensor * llama_moe_hot_cache_build_moe_hot_multi_from_logits( lane_ids, lane_weights, lane.ffn_up_exps, + lane.ffn_up_exps_b, lane.ffn_gate_exps, + lane.ffn_gate_exps_b, lane.ffn_down_exps, + lane.ffn_down_exps_b, lane.n_hot + 1, 1, adapter.ffn_op, il, lane.ffn_gate_up_exps, + lane.ffn_gate_up_exps_b, lane.ffn_up_exps_s, lane.ffn_gate_exps_s, lane.ffn_down_exps_s, @@ -747,13 +790,17 @@ static ggml_tensor * llama_moe_hot_cache_build_moe_hot_multi_from_logits( cold_ids, cold_weights, layer.ffn_up_exps, + layer.ffn_up_exps_b, layer.ffn_gate_exps, + layer.ffn_gate_exps_b, layer.ffn_down_exps, + layer.ffn_down_exps_b, n_expert, 1, adapter.ffn_op, il, layer.ffn_gate_up_exps, + layer.ffn_gate_up_exps_b, layer.ffn_up_exps_s, layer.ffn_gate_exps_s, layer.ffn_down_exps_s, @@ -1096,13 +1143,17 @@ static ggml_tensor * llama_moe_hot_cache_build_moe_hot_from_logits( hot_ids, hot_weights, cache.ffn_up_exps, + cache.ffn_up_exps_b, cache.ffn_gate_exps, + cache.ffn_gate_exps_b, cache.ffn_down_exps, + cache.ffn_down_exps_b, cache.n_hot + 1, 1, adapter.ffn_op, il, cache.ffn_gate_up_exps, + cache.ffn_gate_up_exps_b, cache.ffn_up_exps_s, cache.ffn_gate_exps_s, cache.ffn_down_exps_s, @@ -1164,13 +1215,17 @@ static ggml_tensor * llama_moe_hot_cache_build_moe_hot_from_logits( cold_ids, cold_weights, layer.ffn_up_exps, + layer.ffn_up_exps_b, layer.ffn_gate_exps, + layer.ffn_gate_exps_b, layer.ffn_down_exps, + layer.ffn_down_exps_b, n_expert, 1, adapter.ffn_op, il, layer.ffn_gate_up_exps, + layer.ffn_gate_up_exps_b, layer.ffn_up_exps_s, layer.ffn_gate_exps_s, layer.ffn_down_exps_s, @@ -1598,13 +1653,17 @@ ggml_tensor * llama_model_qwen35moe::graph::build_layer_ffn_hot(ggml_tensor * cu hot_ids, hot_weights, cache.ffn_up_exps, + nullptr, cache.ffn_gate_exps, + nullptr, cache.ffn_down_exps, + nullptr, cache.n_hot + 1, 1, LLM_FFN_SILU, il, cache.ffn_gate_up_exps, + nullptr, cache.ffn_up_exps_s, cache.ffn_gate_exps_s, cache.ffn_down_exps_s, @@ -1668,13 +1727,17 @@ ggml_tensor * llama_model_qwen35moe::graph::build_layer_ffn_hot(ggml_tensor * cu cold_ids, cold_weights, layer.ffn_up_exps, + nullptr, layer.ffn_gate_exps, + nullptr, layer.ffn_down_exps, + nullptr, n_expert, 1, LLM_FFN_SILU, il, layer.ffn_gate_up_exps, + nullptr, layer.ffn_up_exps_s, layer.ffn_gate_exps_s, layer.ffn_down_exps_s, @@ -1809,6 +1872,12 @@ ggml_tensor * llama_model_qwen3next::graph::build_layer_moe_hot(ggml_tensor * cu return llama_moe_hot_cache_build_moe_hot_from_logits(*this, model, cur, logits, il, adapter); } +ggml_tensor * llama_model_openai_moe::graph::build_layer_moe_hot(ggml_tensor * cur, ggml_tensor * logits, const int il) { + const llama_moe_hot_cache_model_adapter & adapter = + llama_moe_hot_cache_require_model_adapter(model.arch, llama_moe_hot_cache_graph_kind::logits); + return llama_moe_hot_cache_build_moe_hot_from_logits(*this, model, cur, logits, il, adapter); +} + template ggml_tensor * llama_model_mellum::graph::build_layer_moe_hot(ggml_tensor * cur, ggml_tensor * logits, const int il) { const llama_moe_hot_cache_model_adapter & adapter = diff --git a/src/moe-hot-cache/llama-moe-hot-cache-planner.cpp b/src/moe-hot-cache/llama-moe-hot-cache-planner.cpp index 6e10884176c..f297052960e 100644 --- a/src/moe-hot-cache/llama-moe-hot-cache-planner.cpp +++ b/src/moe-hot-cache/llama-moe-hot-cache-planner.cpp @@ -131,6 +131,16 @@ size_t llama_moe_hot_cache_tensor_expert_bytes(const ggml_tensor * t) { return ggml_nbytes(t) / size_t(t->ne[2]); } +size_t llama_moe_hot_cache_tensor_expert_bias_bytes(const ggml_tensor * t) { + if (t == nullptr) { + return 0; + } + if (t->ne[1] <= 0) { + throw std::runtime_error("MoE expert bias tensor has invalid expert dimension"); + } + return ggml_nbytes(t) / size_t(t->ne[1]); +} + std::vector llama_moe_hot_cache_collect_expert_sizes( const llama_model & model) { std::vector result; @@ -148,10 +158,14 @@ std::vector llama_moe_hot_cache_collect_expert_ if (layer.ffn_gate_up_exps != nullptr) { bytes += llama_moe_hot_cache_tensor_expert_bytes(layer.ffn_gate_up_exps); + bytes += llama_moe_hot_cache_tensor_expert_bias_bytes(layer.ffn_gate_up_exps_b); } else { bytes += llama_moe_hot_cache_tensor_expert_bytes(layer.ffn_gate_exps); bytes += llama_moe_hot_cache_tensor_expert_bytes(layer.ffn_up_exps); + bytes += llama_moe_hot_cache_tensor_expert_bias_bytes(layer.ffn_gate_exps_b); + bytes += llama_moe_hot_cache_tensor_expert_bias_bytes(layer.ffn_up_exps_b); } + bytes += llama_moe_hot_cache_tensor_expert_bias_bytes(layer.ffn_down_exps_b); if (bytes > 0) { result.push_back({ il, uint32_t(ex), bytes }); diff --git a/src/moe-hot-cache/llama-moe-hot-cache-planner.h b/src/moe-hot-cache/llama-moe-hot-cache-planner.h index 83fe7333a30..172cf23033b 100644 --- a/src/moe-hot-cache/llama-moe-hot-cache-planner.h +++ b/src/moe-hot-cache/llama-moe-hot-cache-planner.h @@ -5,6 +5,7 @@ #include size_t llama_moe_hot_cache_tensor_expert_bytes(const ggml_tensor * t); +size_t llama_moe_hot_cache_tensor_expert_bias_bytes(const ggml_tensor * t); std::vector llama_moe_hot_cache_collect_expert_sizes( const llama_model & model); diff --git a/src/moe-hot-cache/llama-moe-hot-cache-updater.cpp b/src/moe-hot-cache/llama-moe-hot-cache-updater.cpp index bfe291770f6..15cfedcedff 100644 --- a/src/moe-hot-cache/llama-moe-hot-cache-updater.cpp +++ b/src/moe-hot-cache/llama-moe-hot-cache-updater.cpp @@ -247,6 +247,10 @@ llama_moe_hot_cache_update_stats llama_moe_hot_cache_update_from_scored_observat llama_moe_hot_cache_copy_expert_slice(src.ffn_gate_exps, dst.ffn_gate_exps, candidate.add_expert, candidate.cache_id); llama_moe_hot_cache_copy_expert_slice(src.ffn_up_exps, dst.ffn_up_exps, candidate.add_expert, candidate.cache_id); llama_moe_hot_cache_copy_expert_slice(src.ffn_down_exps, dst.ffn_down_exps, candidate.add_expert, candidate.cache_id); + llama_moe_hot_cache_copy_bias_slice(src.ffn_gate_up_exps_b, dst.ffn_gate_up_exps_b, candidate.add_expert, candidate.cache_id); + llama_moe_hot_cache_copy_bias_slice(src.ffn_gate_exps_b, dst.ffn_gate_exps_b, candidate.add_expert, candidate.cache_id); + llama_moe_hot_cache_copy_bias_slice(src.ffn_up_exps_b, dst.ffn_up_exps_b, candidate.add_expert, candidate.cache_id); + llama_moe_hot_cache_copy_bias_slice(src.ffn_down_exps_b, dst.ffn_down_exps_b, candidate.add_expert, candidate.cache_id); llama_moe_hot_cache_copy_scale_slice(src.ffn_gate_exps_s, dst.ffn_gate_exps_s, candidate.add_expert, candidate.cache_id); llama_moe_hot_cache_copy_scale_slice(src.ffn_up_exps_s, dst.ffn_up_exps_s, candidate.add_expert, candidate.cache_id); llama_moe_hot_cache_copy_scale_slice(src.ffn_down_exps_s, dst.ffn_down_exps_s, candidate.add_expert, candidate.cache_id); diff --git a/src/moe-hot-cache/llama-moe-hot-cache.cpp b/src/moe-hot-cache/llama-moe-hot-cache.cpp index 7ed288b9b5e..cf27cb71d05 100644 --- a/src/moe-hot-cache/llama-moe-hot-cache.cpp +++ b/src/moe-hot-cache/llama-moe-hot-cache.cpp @@ -251,10 +251,11 @@ void llama_moe_hot_cache_init(llama_model & model, const llama_model_params & pa const uint32_t n_expert_per_layer = model.hparams.n_expert; if (n_expert_per_layer > 0) { - const double cpu_moe_layer_equiv = (double) plan.selected_count() / (double) n_expert_per_layer; - LLAMA_LOG_WARN("%s: selected %zu/%zu observed experts for hot-cache across %zu lanes (n-cpu-moe equivalent = %.1f layers @ %u experts/layer, %zu/%zu MiB)\n", + const double hot_cache_layer_equiv = (double) plan.selected_count() / (double) n_expert_per_layer; + const double gpu_moe_layer_equiv = std::max(0.0, (double) model.hparams.n_layer() - hot_cache_layer_equiv); + LLAMA_LOG_WARN("%s: selected %zu/%zu observed experts for hot-cache across %zu lanes (n-gpu-moe equivalent = %.1f layers @ %u experts/layer, %zu/%zu MiB)\n", __func__, plan.selected_count(), plan.observed.size(), plan.lanes.size(), - cpu_moe_layer_equiv, n_expert_per_layer, + gpu_moe_layer_equiv, n_expert_per_layer, plan.used_bytes()/LLAMA_MOE_HOT_CACHE_MIB, plan.budget_bytes()/LLAMA_MOE_HOT_CACHE_MIB); } else { LLAMA_LOG_WARN("%s: selected %zu/%zu observed experts for hot-cache across %zu lanes (%zu/%zu MiB)\n", diff --git a/src/moe-hot-cache/llama-moe-hot-cache.h b/src/moe-hot-cache/llama-moe-hot-cache.h index 84bde410299..7eff4effc4f 100644 --- a/src/moe-hot-cache/llama-moe-hot-cache.h +++ b/src/moe-hot-cache/llama-moe-hot-cache.h @@ -102,6 +102,10 @@ struct llama_moe_hot_cache_layer_lane { ggml_tensor * ffn_gate_exps = nullptr; ggml_tensor * ffn_up_exps = nullptr; ggml_tensor * ffn_down_exps = nullptr; + ggml_tensor * ffn_gate_up_exps_b = nullptr; + ggml_tensor * ffn_gate_exps_b = nullptr; + ggml_tensor * ffn_up_exps_b = nullptr; + ggml_tensor * ffn_down_exps_b = nullptr; ggml_tensor * ffn_gate_exps_s = nullptr; ggml_tensor * ffn_up_exps_s = nullptr; ggml_tensor * ffn_down_exps_s = nullptr; diff --git a/tests/test-moe-hot-cache-adapter.cpp b/tests/test-moe-hot-cache-adapter.cpp index 495bb2327ce..24f6ac7598d 100644 --- a/tests/test-moe-hot-cache-adapter.cpp +++ b/tests/test-moe-hot-cache-adapter.cpp @@ -65,6 +65,12 @@ static void test_find_supported_adapters() { require(std::string(mellum->name) == "mellum"); require(mellum->graph_kind == llama_moe_hot_cache_graph_kind::logits); require(mellum->ffn_op == LLM_FFN_SILU); + + const auto * openai_moe = llama_moe_hot_cache_find_model_adapter(LLM_ARCH_OPENAI_MOE); + require(openai_moe != nullptr); + require(std::string(openai_moe->name) == "gpt-oss"); + require(openai_moe->graph_kind == llama_moe_hot_cache_graph_kind::logits); + require(openai_moe->ffn_op == LLM_FFN_SWIGLU_OAI_MOE); } static void test_rejects_unsupported_arch() { @@ -90,6 +96,9 @@ static void test_graph_kind_capability_checks() { require(llama_moe_hot_cache_adapter_supports_graph_kind(LLM_ARCH_MELLUM, llama_moe_hot_cache_graph_kind::logits)); require(!llama_moe_hot_cache_adapter_supports_graph_kind(LLM_ARCH_MELLUM, llama_moe_hot_cache_graph_kind::qwen35_ffn)); + require(llama_moe_hot_cache_adapter_supports_graph_kind(LLM_ARCH_OPENAI_MOE, llama_moe_hot_cache_graph_kind::logits)); + require(!llama_moe_hot_cache_adapter_supports_graph_kind(LLM_ARCH_OPENAI_MOE, llama_moe_hot_cache_graph_kind::qwen35_ffn)); + const auto * qwen_any = llama_moe_hot_cache_find_model_adapter(LLM_ARCH_QWEN35MOE, llama_moe_hot_cache_graph_kind::none); require(qwen_any != nullptr); require(qwen_any->graph_kind == llama_moe_hot_cache_graph_kind::qwen35_ffn); @@ -126,6 +135,13 @@ static void test_profile_defaults_are_arch_specific() { require(mellum.merge_sum_rows); require(!mellum.branch_reduce_merge); require(mellum.cpu_decode_routing_max_tokens == 1); + + const auto openai_moe = llama_moe_hot_cache_graph_profile_for_arch(LLM_ARCH_OPENAI_MOE); + require(openai_moe.cpu_decode_routing); + require(openai_moe.decode_direct_merge); + require(openai_moe.merge_sum_rows); + require(!openai_moe.branch_reduce_merge); + require(openai_moe.cpu_decode_routing_max_tokens == 1); } static void test_parallel_mode_is_runtime_switchable() { diff --git a/tests/test-moe-hot-cache-planner.cpp b/tests/test-moe-hot-cache-planner.cpp index 91ce4339fbf..1e54100f8ec 100644 --- a/tests/test-moe-hot-cache-planner.cpp +++ b/tests/test-moe-hot-cache-planner.cpp @@ -28,9 +28,12 @@ static void test_tensor_expert_bytes_splits_by_expert_dimension() { require(ctx != nullptr); ggml_tensor * tensor = ggml_new_tensor_3d(ctx.get(), GGML_TYPE_F32, 4, 5, 3); + ggml_tensor * bias = ggml_new_tensor_2d(ctx.get(), GGML_TYPE_F32, 4, 3); require(llama_moe_hot_cache_tensor_expert_bytes(nullptr) == 0); require(llama_moe_hot_cache_tensor_expert_bytes(tensor) == ggml_nbytes(tensor)/3); + require(llama_moe_hot_cache_tensor_expert_bias_bytes(nullptr) == 0); + require(llama_moe_hot_cache_tensor_expert_bias_bytes(bias) == ggml_nbytes(bias)/3); } static void test_select_accounts_for_one_dummy_expert_per_active_layer() {