spec: add eagle3-v3 support for gpt-oss model (#25794)

This commit is contained in:
Ruixiang Wang
2026-07-28 09:58:16 +03:00
committed by GitHub
parent c6292cfb8e
commit f87067841b
9 changed files with 59 additions and 6 deletions
+13 -3
View File
@@ -437,6 +437,7 @@ struct common_speculative_impl_draft_eagle3 : public common_speculative_impl {
int32_t n_embd_dec = 0; // draft hidden size
int32_t n_embd_enc = 0; // target_layer_ids_n * target_hidden_size
int32_t n_embd_tgt = 0; // target model hidden size
int32_t n_layer_tgt = 0; // target model layer count
const int32_t * target_layer_ids = nullptr; // model_dft's extract layer indices
uint32_t target_layer_ids_n = 0;
@@ -478,6 +479,7 @@ struct common_speculative_impl_draft_eagle3 : public common_speculative_impl {
n_embd_tgt = llama_model_n_embd(model_tgt);
n_embd_dec = llama_model_n_embd(model_dft);
n_embd_enc = (int32_t) target_layer_ids_n * n_embd_tgt;
n_layer_tgt = llama_model_n_layer(model_tgt);
const int32_t n_b = (int32_t) llama_n_batch(ctx_dft);
batch = llama_batch_init(/*n_tokens=*/ n_b, /*embd=*/ n_embd_dec, /*n_seq_max=*/ 1);
@@ -510,9 +512,15 @@ struct common_speculative_impl_draft_eagle3 : public common_speculative_impl {
}
}
// turn on extraction of the target layers' input embeddings
// turn on extraction of the target layers' hidden states
for (uint32_t k = 0; k < target_layer_ids_n; ++k) {
llama_set_embeddings_layer_inp(ctx_tgt, (uint32_t) target_layer_ids[k], true);
if (target_layer_ids[k] < n_layer_tgt) {
llama_set_embeddings_layer_inp(ctx_tgt, (uint32_t) target_layer_ids[k], true);
} else if (target_layer_ids[k] == n_layer_tgt) {
llama_set_embeddings_nextn(ctx_tgt, true, /*masked*/ false);
} else {
GGML_ABORT("EAGLE3: target layer id %d exceeds target n_layer %d", target_layer_ids[k], n_layer_tgt);
}
}
// turn on extraction of the draft model's pre-norm hidden state
@@ -600,7 +608,9 @@ struct common_speculative_impl_draft_eagle3 : public common_speculative_impl {
features_buf.resize((size_t) n_tokens * n_embd_enc, 0.0f);
for (uint32_t k = 0; k < target_layer_ids_n; ++k) {
const float * layer = llama_get_embeddings_layer_inp(ctx_tgt, (uint32_t) target_layer_ids[k]);
const float * layer = target_layer_ids[k] < n_layer_tgt
? llama_get_embeddings_layer_inp(ctx_tgt, (uint32_t) target_layer_ids[k])
: llama_get_embeddings_nextn(ctx_tgt);
if (!layer) {
GGML_ABORT("EAGLE3: target layer %d input not extracted.", target_layer_ids[k]);
}
+16 -2
View File
@@ -69,9 +69,14 @@ class LlamaModel(TextModel):
target_config = {**target_config, **target_config["text_config"]}
self.target_vocab_size = target_config["vocab_size"]
# target_layers: derived from target model layer count (low/mid/high)
# target_layers: use the eagle3 config's explicit aux hidden-state layer ids
# if present, else derive from the target layer count.
target_num_layers = target_config["num_hidden_layers"]
target_layers = [2, target_num_layers // 2, target_num_layers - 3]
aux_layer_ids = eagle3_raw_config.get("eagle_aux_hidden_state_layer_ids")
if aux_layer_ids:
target_layers = aux_layer_ids
else:
target_layers = [2, target_num_layers // 2, target_num_layers - 3]
logger.info(f"EAGLE-3: target_layers = {target_layers} (target model has {target_num_layers} layers)")
self.gguf_writer.add_target_layers(target_layers)
@@ -90,6 +95,12 @@ class LlamaModel(TextModel):
logger.info(f"EAGLE-3: norm_before_residual = {norm_before_residual}")
self.gguf_writer.add_norm_before_residual(norm_before_residual)
# norm_before_fc: RMSNorm applied to the fused target features before the
# fc projection (e.g. nvidia/gpt-oss-120b-Eagle3-v3)
norm_before_fc = eagle3_raw_config.get("norm_before_fc", False)
logger.info(f"EAGLE-3: norm_before_fc = {norm_before_fc}")
self.gguf_writer.add_norm_before_fc(norm_before_fc)
def set_vocab(self):
# eagle3: use tokenizer from target model if provided
original_dir_model = None
@@ -222,6 +233,9 @@ class LlamaModel(TextModel):
if name == "fc.weight":
yield (name, data_torch)
return
if name == "input_norm.weight":
yield (self.format_tensor_name(gguf.MODEL_TENSOR.ENC_OUTPUT_NORM), data_torch)
return
if name == "d2t":
# store for manual int64 handling in prepare_tensors (avoid F32 conversion)
if not hasattr(self, '_eagle3_int_tensors'):
+2
View File
@@ -161,6 +161,7 @@ class Keys:
TARGET_HIDDEN_SIZE = "{arch}.target_hidden_size"
BLOCK_SIZE = "{arch}.block_size"
NORM_BEFORE_RESIDUAL = "{arch}.norm_before_residual"
NORM_BEFORE_FC = "{arch}.norm_before_fc"
class Attention:
HEAD_COUNT = "{arch}.attention.head_count"
@@ -4343,6 +4344,7 @@ MODEL_TENSORS: dict[MODEL_ARCH, list[MODEL_TENSOR]] = {
MODEL_TENSOR.FFN_DOWN,
MODEL_TENSOR.FFN_UP,
MODEL_TENSOR.FC,
MODEL_TENSOR.ENC_OUTPUT_NORM,
MODEL_TENSOR.D2T,
],
MODEL_ARCH.DFLASH: [
+3
View File
@@ -971,6 +971,9 @@ class GGUFWriter:
def add_norm_before_residual(self, value: bool) -> None:
self.add_bool(Keys.LLM.NORM_BEFORE_RESIDUAL.format(arch=self.arch), value)
def add_norm_before_fc(self, value: bool) -> None:
self.add_bool(Keys.LLM.NORM_BEFORE_FC.format(arch=self.arch), value)
def add_attention_output_group_count(self, count: int) -> None:
self.add_uint32(Keys.Attention.OUTPUT_GROUP_COUNT.format(arch=self.arch), count)
+1
View File
@@ -317,6 +317,7 @@ static const std::map<llm_kv, const char *> LLM_KV_NAMES = {
{ LLM_KV_TARGET_LAYERS, "%s.target_layers" },
{ LLM_KV_TARGET_HIDDEN_SIZE, "%s.target_hidden_size" },
{ LLM_KV_NORM_BEFORE_RESIDUAL, "%s.norm_before_residual" },
{ LLM_KV_NORM_BEFORE_FC, "%s.norm_before_fc" },
{ LLM_KV_SHORTCONV_L_CACHE, "%s.shortconv.l_cache" },
// sentence-transformers dense modules feature dims
+1
View File
@@ -363,6 +363,7 @@ enum llm_kv {
LLM_KV_TARGET_LAYERS,
LLM_KV_TARGET_HIDDEN_SIZE,
LLM_KV_NORM_BEFORE_RESIDUAL,
LLM_KV_NORM_BEFORE_FC,
LLM_KV_SHORTCONV_L_CACHE,
+1
View File
@@ -47,6 +47,7 @@ struct llama_hparams {
bool use_par_res;
bool swin_norm;
bool norm_before_residual = false;
bool norm_before_fc = false;
uint32_t n_ctx_train; // context size the model was trained on
uint32_t n_embd;
+15
View File
@@ -28,6 +28,10 @@ void llama_model_eagle3::load_arch_hparams(llama_model_loader & ml) {
LLAMA_LOG_INFO("%s: EAGLE3gnorm_before_residual = true\n", __func__);
}
// eagle3 norm_before_fc (optional, default false)
// compatible with eagle3.1 (e.g. nvidia/gpt-oss-120b-Eagle3-v3)
ml.get_key(LLM_KV_NORM_BEFORE_FC, hparams.norm_before_fc, false);
type = LLM_TYPE_UNKNOWN;
}
@@ -53,6 +57,11 @@ void llama_model_eagle3::load_arch_tensors(llama_model_loader &) {
// Feature fusion layer: projects 3 target layers to draft hidden size
fc = create_tensor(tn(LLM_TENSOR_FC, "weight"), {n_embd_inp, n_embd}, 0);
// RMSNorm on the fused target features (input to fc), only when norm_before_fc is set.
if (hparams.norm_before_fc) {
output_norm_enc = create_tensor(tn(LLM_TENSOR_ENC_OUTPUT_NORM, "weight"), {n_embd_inp}, 0);
}
// Output layer (uses draft vocab size)
output_norm = create_tensor(tn(LLM_TENSOR_OUTPUT_NORM, "weight"), {n_embd}, 0);
output = create_tensor(tn(LLM_TENSOR_OUTPUT, "weight"), {n_embd, n_draft_vocab}, TENSOR_NOT_REQUIRED);
@@ -130,6 +139,12 @@ llama_model_eagle3::graph<true>::graph(const llama_model & model, const llm_grap
cur = build_inp_embd_enc();
// RMSNorm on the fused target features before fc
if (hparams.norm_before_fc) {
cur = build_norm(cur, model.output_norm_enc, NULL, LLM_NORM_RMS, -1);
cb(cur, "enc_input_norm", -1);
}
// Feature fusion layer
cur = build_lora_mm(model.fc, cur);
cb(cur, "fc_out", -1);
+7 -1
View File
@@ -116,7 +116,7 @@ llama_model_openai_moe::graph::graph(const llama_model & model, const llm_graph_
cb(cur, "attn_out", il);
}
if (il == n_layer - 1) {
if (il == n_layer - 1 && inp_out_ids && cparams.embeddings_nextn_masked) {
// skip computing output for unused tokens
cur = ggml_get_rows(ctx0, cur, inp_out_ids);
inpSA = ggml_get_rows(ctx0, inpSA, inp_out_ids);
@@ -154,6 +154,12 @@ llama_model_openai_moe::graph::graph(const llama_model & model, const llm_graph_
}
cur = inpL;
res->t_h_nextn = cur;
if (!cparams.embeddings_nextn_masked && inp_out_ids) {
cur = ggml_get_rows(ctx0, cur, inp_out_ids);
}
cur = build_norm(cur,
model.output_norm, NULL,
LLM_NORM_RMS, -1);