Skip to content

add gpt-oss support #7

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
adrianhoehne merged 1 commit into from
Jun 10, 2026
Merged

add gpt-oss support #7

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
7 changes: 7 additions & 0 deletions src/models/models.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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<llm_graph_context> build_arch_graph(const llm_graph_params & params) const override;
Expand Down
39 changes: 27 additions & 12 deletions src/models/openai-moe.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -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);
Expand Down Expand Up @@ -61,7 +62,9 @@ std::unique_ptr<llm_graph_context> llama_model_openai_moe::build_arch_graph(cons
return std::make_unique<graph>(*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;

Expand Down Expand Up @@ -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);
Expand Down
9 changes: 9 additions & 0 deletions src/moe-hot-cache/llama-moe-hot-cache-adapter.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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 {};
}
Expand Down
61 changes: 61 additions & 0 deletions src/moe-hot-cache/llama-moe-hot-cache-builder.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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<uint8_t> zeros(ggml_nbytes(t), 0);
ggml_backend_tensor_set(t, zeros.data(), 0, zeros.size());
Expand Down Expand Up @@ -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<uint8_t> 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));
}
Expand Down Expand Up @@ -148,6 +177,10 @@ std::unique_ptr<llama_moe_hot_cache> 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;
Expand Down Expand Up @@ -190,6 +223,10 @@ std::unique_ptr<llama_moe_hot_cache> 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());
Expand Down Expand Up @@ -224,6 +261,10 @@ std::unique_ptr<llama_moe_hot_cache> 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); }
Expand All @@ -248,6 +289,10 @@ std::unique_ptr<llama_moe_hot_cache> 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);
Expand Down Expand Up @@ -315,6 +360,10 @@ std::unique_ptr<llama_moe_hot_cache> 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;
Expand Down Expand Up @@ -365,6 +414,10 @@ std::unique_ptr<llama_moe_hot_cache> 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());
Expand Down Expand Up @@ -400,6 +453,10 @@ std::unique_ptr<llama_moe_hot_cache> 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); }
Expand All @@ -425,6 +482,10 @@ std::unique_ptr<llama_moe_hot_cache> 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);
Expand Down
6 changes: 6 additions & 0 deletions src/moe-hot-cache/llama-moe-hot-cache-builder.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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);

Expand Down
Loading
Loading