Commit 843e7dcc1 for llama.cpp
commit 843e7dcc15098dd6b3aaf33b4847052293d8ef1a
Author: Kartik Gulia <kgulia@nvidia.com>
Date: Thu Oct 15 09:19:25 2026 +0530
spec : add lilicorr support (#30506)
* spec : add lilicorr support
* convert : use convention for hparam assignment
* add todo for hparam use in llama-batch.cpp
Co-authored-by: Georgi Gerganov <ggerganov@gmail.com>
* typo fix
---------
Co-authored-by: Georgi Gerganov <ggerganov@gmail.com>
diff --git a/common/speculative.cpp b/common/speculative.cpp
index 50647dbb2..701093308 100644
--- a/common/speculative.cpp
+++ b/common/speculative.cpp
@@ -999,6 +999,12 @@ struct common_speculative_impl_draft_dflash : public common_speculative_impl {
int32_t block_size = 0;
llama_token mask_token_id = 0;
+ bool is_lilicorr = false;
+ int32_t lilicorr_top_k = 0;
+ size_t lilicorr_rows = 0;
+ std::vector<float> lilicorr_embd;
+ std::vector<llama_pos> lilicorr_pos;
+
bool is_dflash2 = false;
bool is_mrope = false;
int32_t selector_top_k = 0;
@@ -1053,6 +1059,16 @@ struct common_speculative_impl_draft_dflash : public common_speculative_impl {
selector_top_k = llama_model_dflash_selector_top_k(model_dft);
is_dflash2 = selector_top_k > 0;
+ char lilicorr_meta[32] = {};
+ lilicorr_top_k = llama_model_meta_val_str(model_dft, "lilicorr.candidate_topk", lilicorr_meta, sizeof(lilicorr_meta)) >= 0
+ ? std::atoi(lilicorr_meta) : 0;
+ is_lilicorr = lilicorr_top_k > 0;
+ if (is_lilicorr) {
+ lilicorr_rows = (size_t) block_size + 1;
+ const size_t lilicorr_slots = (size_t) n_seq * lilicorr_rows;
+ lilicorr_embd.resize(lilicorr_slots * (size_t) n_embd_dec);
+ lilicorr_pos.assign(lilicorr_slots, -1);
+ }
mask_token_id = llama_vocab_mask(llama_model_get_vocab(model_dft));
if (is_dspark && this->params.p_min > 0.0f) {
@@ -1098,7 +1114,7 @@ struct common_speculative_impl_draft_dflash : public common_speculative_impl {
// offload draft sampling to the backend
backend_chains.assign(n_seq, nullptr);
- if (this->params.backend_sampling && !is_dflash2) {
+ if (this->params.backend_sampling && !is_dflash2 && !is_lilicorr) {
for (llama_seq_id seq_id = 0; seq_id < (llama_seq_id) n_seq; ++seq_id) {
llama_sampler * chain = llama_sampler_chain_init(llama_sampler_chain_default_params());
llama_sampler_chain_add(chain, llama_sampler_init_top_k(10));
@@ -1118,7 +1134,7 @@ struct common_speculative_impl_draft_dflash : public common_speculative_impl {
}
// DFlash2 reads its selector lattice from h_nextn and never consumes raw logits.
- llama_set_embeddings_nextn(ctx_dft, true, /*masked*/ !is_dflash2);
+ llama_set_embeddings_nextn(ctx_dft, true, /*masked*/ !is_dflash2 && !is_lilicorr);
llama_set_causal_attn(ctx_dft, causal_attn); // DFlash needs non-causal attention unless the model says otherwise
}
@@ -1233,6 +1249,24 @@ struct common_speculative_impl_draft_dflash : public common_speculative_impl {
__func__, rc, (int) n_chunk, (int) offset);
return false;
}
+ if (is_lilicorr) {
+ // Ordinary DFlash can write to the KV cache, but LiliCorr needs to store embeddings across sequences.
+ const float * projected = llama_get_embeddings_nextn(ctx_dft);
+ if (projected) {
+ const size_t seq_offset = (size_t) seq_id * lilicorr_rows;
+ const llama_pos pos_min = std::max(0, batch_inject.tokens.back().pos[0] - block_size);
+ for (int32_t i = 0; i < n_chunk; ++i) {
+ const llama_pos pos = batch_inject.tokens[i].pos[0];
+ if (pos < pos_min) {
+ continue;
+ }
+ const size_t slot = seq_offset + (size_t) pos % lilicorr_rows;
+ std::memcpy(lilicorr_embd.data() + slot * (size_t) n_embd_dec,
+ projected + (size_t) i * n_embd_dec, (size_t) n_embd_dec * sizeof(float));
+ lilicorr_pos[slot] = pos;
+ }
+ }
+ }
}
}
@@ -1260,12 +1294,17 @@ struct common_speculative_impl_draft_dflash : public common_speculative_impl {
const int32_t n = (int32_t) dp.pos0;
const int32_t n_draft = params.n_max;
-
- const int32_t n_block_tokens = n_draft + (is_dspark && sample_from_anchor ? 0 : 1);
+ const int32_t n_block_tokens = is_lilicorr ? block_size : n_draft + (is_dspark && sample_from_anchor ? 0 : 1);
i_block_beg[seq_id] = batch.size();
n_block [seq_id] = n_block_tokens;
for (int32_t i = 0; i < n_block_tokens; ++i) {
- batch.add(i == 0 ? dp.id_last : mask_token_id, n + i, seq_id, !is_dflash2);
+ const int32_t idx = batch.add(i == 0 ? dp.id_last : mask_token_id, n + i, seq_id, !is_dflash2 && !is_lilicorr);
+ if (is_lilicorr) {
+ // For LiliCorr drafting batch, set state embeddings for use in the associated transformer head.
+ const llama_pos pos = dp.pos0 - 1;
+ const size_t slot = (size_t) seq_id * lilicorr_rows + (size_t) pos % lilicorr_rows;
+ batch.set_embd_state(idx, { lilicorr_embd.data() + slot * (size_t) n_embd_dec, 1, (size_t) n_embd_dec });
+ }
}
}
@@ -1293,21 +1332,23 @@ struct common_speculative_impl_draft_dflash : public common_speculative_impl {
auto & result = *dp.result;
- if (is_dflash2) {
+ if (is_dflash2 || is_lilicorr) {
+ // Both DFlash and LiliCorr greedily walk the candidate lattice.
const float * lattice = llama_get_embeddings_nextn(ctx_dft);
GGML_ASSERT(lattice && "DFlash2 selector produced no lattice");
int32_t predecessor = 0;
for (int32_t i = 1; i < n_block_tokens; ++i) {
const float * row = lattice + (size_t) (beg + i) * n_embd_dec;
- const float * scores = row + selector_top_k + (size_t) predecessor * selector_top_k;
+ const int32_t top_k = is_dflash2 ? selector_top_k : lilicorr_top_k;
+ const float * scores = row + top_k + (size_t) predecessor * top_k;
predecessor = (int32_t) std::distance(scores,
- std::max_element(scores, scores + selector_top_k));
+ std::max_element(scores, scores + top_k));
if (params.p_min > 0.0f) {
// softmax(scores) at the argmax, i.e. 1 / sum(exp(s_k - s_max))
float sum = 0.0f;
- for (int32_t k = 0; k < selector_top_k; ++k) {
+ for (int32_t k = 0; k < top_k; ++k) {
sum += std::exp(scores[k] - scores[predecessor]);
}
if (1.0f / sum < params.p_min) {
diff --git a/conversion/__init__.py b/conversion/__init__.py
index 7f32a87f0..d5b480a9b 100644
--- a/conversion/__init__.py
+++ b/conversion/__init__.py
@@ -59,6 +59,7 @@ TEXT_MODEL_MAP: dict[str, str] = {
"DeepseekV32ForCausalLM": "deepseek",
"DFlashDraftModel": "qwen",
"DFlash2DraftModel": "qwen",
+ "LiLiCorrDraftModel": "nemotron",
"Qwen3DSparkModel": "qwen",
"DSparkDraftModel": "qwen",
"DSparkSpeculator": "qwen",
diff --git a/conversion/nemotron.py b/conversion/nemotron.py
index f7eefcf8d..581f41a02 100644
--- a/conversion/nemotron.py
+++ b/conversion/nemotron.py
@@ -2,15 +2,114 @@ from __future__ import annotations
from typing import Any, Callable, Iterable, TYPE_CHECKING
+import numpy as np
import torch
if TYPE_CHECKING:
from pathlib import Path
from torch import Tensor
-from .base import MmprojModel, ModelBase, TextModel, gguf, logger
+from .base import MmprojModel, ModelBase, ModelType, TextModel, get_model_architecture, gguf, logger
from .granite import GraniteHybridModel
+from .qwen import DFlashModel
+
+
+@ModelBase.register("LiLiCorrDraftModel")
+class LiLiCorrDraftModel(DFlashModel):
+ model_arch = gguf.MODEL_ARCH.DFLASH
+
+ def __init__(self, *args, **kwargs):
+ super().__init__(*args, **kwargs)
+ prefix = "model.lilicorr."
+ if not any(name.startswith(prefix) for name in self.model_tensors):
+ raise ValueError("no lilicorr.* tensors found")
+ dflash_config = self.hparams["dflash_config"]
+ self.lilicorr_geometry = {
+ "hidden_size": int(dflash_config["lilicorr_hidden_size"]),
+ "candidate_topk": int(dflash_config["lilicorr_candidate_topk"]),
+ "num_heads": int(dflash_config["lilicorr_num_heads"]),
+ "num_layers": int(dflash_config["lilicorr_num_layers"]),
+ "factor_dim": int(dflash_config["lilicorr_factor_dim"]),
+ "mlp_hidden_size": self.model_tensors[prefix + "layers.0.mlp.up_proj.weight"]().shape[0],
+ }
+ self.tensor_map = gguf.get_tensor_name_map(self.model_arch, self.block_count + self.lilicorr_geometry["num_layers"] + 1)
+ block = int(dflash_config.get("block_size", self.hparams.get("block_size", 16)))
+ topk = self.lilicorr_geometry["candidate_topk"]
+ heads = self.lilicorr_geometry["num_heads"]
+ nodes = (block - 1) * topk
+ relative = gguf.LazyBase.to_eager(self.model_tensors.pop(prefix + "relative_slot_bias")()).float().cpu().numpy()
+ same = gguf.LazyBase.to_eager(self.model_tensors.pop(prefix + "same_slot_bias")()).float().cpu().numpy()
+ bias = np.empty((heads, nodes, nodes), dtype=np.float32)
+ for q in range(nodes):
+ q_slot = q // topk
+ for k in range(nodes):
+ k_slot = k // topk
+ bias[:, q, k] = relative[:, q_slot - k_slot + block - 1]
+ if q_slot == k_slot:
+ bias[:, q, k] += same
+ self.model_tensors[prefix + "attention_bias.bias"] = lambda bias=bias: torch.from_numpy(bias)
+ for name in ("slot_embedding", "rank_embedding"):
+ gen = self.model_tensors.pop(prefix + name)
+ self.model_tensors[prefix + name + ".weight"] = lambda gen=gen: gen().reshape(-1, self.lilicorr_geometry["hidden_size"])
+ for i in range(self.lilicorr_geometry["num_layers"]):
+ layer_prefix = prefix + f"layers.{i}.attn.in_proj"
+ for suffix in ("weight", "bias"):
+ self.model_tensors[layer_prefix + "." + suffix] = self.model_tensors.pop(layer_prefix + "_" + suffix)
+ dflash_config["selector_rank"] = 0
+ dflash_config["selector_top_k"] = 0
+
+ def set_vocab(self):
+ if self.target_model_dir is None:
+ raise ValueError("LiLiCorr draft requires --target-model-dir")
+ from . import get_model_class
+ target_text_config = ModelBase.load_hparams(self.target_model_dir, False)
+ target_text_config = {**target_text_config, **target_text_config.get("text_config", {})}
+ target_arch = get_model_architecture(target_text_config, ModelType.TEXT)
+ target_cls = get_model_class(target_arch)
+ original_dir = self.dir_model
+ self.dir_model = self.target_model_dir
+ try:
+ if target_arch == "NemotronHForCausalLM":
+ self.is_moe = "num_experts_per_tok" in target_text_config
+ target_cls.set_vocab(self) # ty: ignore[unresolved-attribute]
+ finally:
+ self.dir_model = original_dir
+ self.gguf_writer.add_mask_token_id(self.hparams["dflash_config"]["mask_token_id"])
+
+ def set_gguf_parameters(self):
+ super().set_gguf_parameters()
+ dflash_config = self.hparams["dflash_config"]
+ self.gguf_writer.add_lilicorr_hidden_size(int(self.lilicorr_geometry["hidden_size"]))
+ self.gguf_writer.add_lilicorr_candidate_topk(int(self.lilicorr_geometry["candidate_topk"]))
+ self.gguf_writer.add_lilicorr_num_heads(int(self.lilicorr_geometry["num_heads"]))
+ self.gguf_writer.add_lilicorr_num_layers(int(self.lilicorr_geometry["num_layers"]))
+ self.gguf_writer.add_lilicorr_factor_dim(int(self.lilicorr_geometry["factor_dim"]))
+ self.gguf_writer.add_lilicorr_mlp_hidden_size(int(self.lilicorr_geometry["mlp_hidden_size"]))
+ self.gguf_writer.add_lilicorr_vector_eps(float(dflash_config["lilicorr_vector_eps"]))
+ self.gguf_writer.add_lilicorr_logit_scale(float(dflash_config["lilicorr_logit_scale"]))
+
+ def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iterable[tuple[str, Tensor]]:
+ if name.startswith("model.lilicorr."):
+ if name.startswith("model.lilicorr.layers."):
+ parts = name.split(".", 4)
+ if not parts[4].startswith("attn.in_proj."):
+ parts[3] = str(self.block_count + int(parts[3]))
+ name = ".".join(parts)
+ elif name.startswith("model.lilicorr.feature_norm."):
+ feature_bid = self.block_count + self.lilicorr_geometry["num_layers"]
+ name = name.replace("model.lilicorr.feature_norm.", f"model.lilicorr.feature_mlp.{feature_bid}.norm.")
+ elif name.startswith("model.lilicorr.feature_mlp."):
+ feature_bid = self.block_count + self.lilicorr_geometry["num_layers"]
+ name = name.replace("model.lilicorr.feature_mlp.", f"model.lilicorr.feature_mlp.{feature_bid}.")
+ yield self.map_tensor_name(name), data_torch
+ return
+ yield from super().modify_tensors(data_torch, name, bid)
+
+ def tensor_force_quant(self, name: str, new_name: str, bid: int | None, n_dims: int) -> gguf.GGMLQuantizationType | bool:
+ if name.startswith("model.lilicorr."):
+ return gguf.GGMLQuantizationType.F32
+ return super().tensor_force_quant(name, new_name, bid, n_dims)
@ModelBase.register(
diff --git a/gguf-py/gguf/constants.py b/gguf-py/gguf/constants.py
index c31ab4a42..f9baf687e 100644
--- a/gguf-py/gguf/constants.py
+++ b/gguf-py/gguf/constants.py
@@ -309,6 +309,16 @@ class Keys:
class WKV:
HEAD_SIZE = "{arch}.wkv.head_size"
+ class LiliCorr:
+ HIDDEN_SIZE = "lilicorr.hidden_size"
+ CANDIDATE_TOPK = "lilicorr.candidate_topk"
+ NUM_HEADS = "lilicorr.num_heads"
+ NUM_LAYERS = "lilicorr.num_layers"
+ FACTOR_DIM = "lilicorr.factor_dim"
+ MLP_HIDDEN_SIZE = "lilicorr.mlp_hidden_size"
+ VECTOR_EPS = "lilicorr.vector_eps"
+ LOGIT_SCALE = "lilicorr.logit_scale"
+
class PosNet:
EMBEDDING_LENGTH = "{arch}.posnet.embedding_length"
BLOCK_COUNT = "{arch}.posnet.block_count"
@@ -1279,6 +1289,19 @@ class MODEL_TENSOR(IntEnum):
DFLASH_SELECTOR_PREV = auto()
DFLASH_SELECTOR_NEXT = auto()
DFLASH_SELECTOR_HIDDEN = auto()
+ LILICORR_TOKEN_PROJ = auto()
+ LILICORR_PASS_HIDDEN_PROJ = auto()
+ LILICORR_SLOT_EMBD = auto()
+ LILICORR_RANK_EMBD = auto()
+ LILICORR_ATTN_BIAS = auto()
+ LILICORR_CONTEXT_PROJ = auto()
+ LILICORR_ATTN_IN_PROJ = auto()
+ LILICORR_OUTPUT_NORM = auto()
+ LILICORR_ANCHOR_NORM = auto()
+ LILICORR_FACTOR_INPUT_PROJ = auto()
+ LILICORR_OUT_HEAD = auto()
+ LILICORR_IN_HEAD = auto()
+ LILICORR_ANCHOR_OUT_HEAD = auto()
# lfm2 audio
A_ENC_NORM_CONV = auto()
A_ENC_LINEAR_POS = auto()
@@ -2093,6 +2116,19 @@ TENSOR_NAMES: dict[MODEL_TENSOR, str] = {
MODEL_TENSOR.DFLASH_SELECTOR_PREV: "selector_predecessor",
MODEL_TENSOR.DFLASH_SELECTOR_NEXT: "selector_successor",
MODEL_TENSOR.DFLASH_SELECTOR_HIDDEN: "selector_hidden",
+ MODEL_TENSOR.LILICORR_TOKEN_PROJ: "lilicorr.token_proj",
+ MODEL_TENSOR.LILICORR_PASS_HIDDEN_PROJ: "lilicorr.pass_hidden_proj",
+ MODEL_TENSOR.LILICORR_SLOT_EMBD: "lilicorr.slot_embedding",
+ MODEL_TENSOR.LILICORR_RANK_EMBD: "lilicorr.rank_embedding",
+ MODEL_TENSOR.LILICORR_ATTN_BIAS: "lilicorr.attention_bias",
+ MODEL_TENSOR.LILICORR_CONTEXT_PROJ: "lilicorr.context_proj",
+ MODEL_TENSOR.LILICORR_ATTN_IN_PROJ: "lilicorr.layers.{bid}.attn.in_proj",
+ MODEL_TENSOR.LILICORR_OUTPUT_NORM: "lilicorr.output_norm",
+ MODEL_TENSOR.LILICORR_ANCHOR_NORM: "lilicorr.anchor_norm",
+ MODEL_TENSOR.LILICORR_FACTOR_INPUT_PROJ: "lilicorr.factor_input_proj",
+ MODEL_TENSOR.LILICORR_OUT_HEAD: "lilicorr.out_head",
+ MODEL_TENSOR.LILICORR_IN_HEAD: "lilicorr.in_head",
+ MODEL_TENSOR.LILICORR_ANCHOR_OUT_HEAD: "lilicorr.anchor_out_head",
MODEL_TENSOR.D2T: "d2t",
MODEL_TENSOR.ATTN_V_GATE: "blk.{bid}.attn_v_gate",
MODEL_TENSOR.ATTN_V_EXP: "blk.{bid}.attn_v_exps",
@@ -5548,6 +5584,19 @@ MODEL_TENSORS: dict[MODEL_ARCH, list[MODEL_TENSOR]] = {
MODEL_TENSOR.DFLASH_SELECTOR_PREV,
MODEL_TENSOR.DFLASH_SELECTOR_NEXT,
MODEL_TENSOR.DFLASH_SELECTOR_HIDDEN,
+ MODEL_TENSOR.LILICORR_TOKEN_PROJ,
+ MODEL_TENSOR.LILICORR_PASS_HIDDEN_PROJ,
+ MODEL_TENSOR.LILICORR_SLOT_EMBD,
+ MODEL_TENSOR.LILICORR_RANK_EMBD,
+ MODEL_TENSOR.LILICORR_ATTN_BIAS,
+ MODEL_TENSOR.LILICORR_CONTEXT_PROJ,
+ MODEL_TENSOR.LILICORR_ATTN_IN_PROJ,
+ MODEL_TENSOR.LILICORR_OUTPUT_NORM,
+ MODEL_TENSOR.LILICORR_ANCHOR_NORM,
+ MODEL_TENSOR.LILICORR_FACTOR_INPUT_PROJ,
+ MODEL_TENSOR.LILICORR_OUT_HEAD,
+ MODEL_TENSOR.LILICORR_IN_HEAD,
+ MODEL_TENSOR.LILICORR_ANCHOR_OUT_HEAD,
],
MODEL_ARCH.MISTRAL4: [
MODEL_TENSOR.TOKEN_EMBD,
diff --git a/gguf-py/gguf/gguf_writer.py b/gguf-py/gguf/gguf_writer.py
index 02ac3e823..cc2bd8e7b 100644
--- a/gguf-py/gguf/gguf_writer.py
+++ b/gguf-py/gguf/gguf_writer.py
@@ -1236,6 +1236,30 @@ class GGUFWriter:
def add_kda_gate_lower_bound(self, value: float) -> None:
self.add_float32(Keys.KDA.GATE_LOWER_BOUND.format(arch=self.arch), value)
+ def add_lilicorr_hidden_size(self, value: int) -> None:
+ self.add_uint32(Keys.LiliCorr.HIDDEN_SIZE, value)
+
+ def add_lilicorr_candidate_topk(self, value: int) -> None:
+ self.add_uint32(Keys.LiliCorr.CANDIDATE_TOPK, value)
+
+ def add_lilicorr_num_heads(self, value: int) -> None:
+ self.add_uint32(Keys.LiliCorr.NUM_HEADS, value)
+
+ def add_lilicorr_num_layers(self, value: int) -> None:
+ self.add_uint32(Keys.LiliCorr.NUM_LAYERS, value)
+
+ def add_lilicorr_factor_dim(self, value: int) -> None:
+ self.add_uint32(Keys.LiliCorr.FACTOR_DIM, value)
+
+ def add_lilicorr_mlp_hidden_size(self, value: int) -> None:
+ self.add_uint32(Keys.LiliCorr.MLP_HIDDEN_SIZE, value)
+
+ def add_lilicorr_vector_eps(self, value: float) -> None:
+ self.add_float32(Keys.LiliCorr.VECTOR_EPS, value)
+
+ def add_lilicorr_logit_scale(self, value: float) -> None:
+ self.add_float32(Keys.LiliCorr.LOGIT_SCALE, value)
+
def add_tokenizer_model(self, model: str) -> None:
self.add_string(Keys.Tokenizer.MODEL, model)
diff --git a/gguf-py/gguf/tensor_mapping.py b/gguf-py/gguf/tensor_mapping.py
index 6445784a3..5eab5b1fd 100644
--- a/gguf-py/gguf/tensor_mapping.py
+++ b/gguf-py/gguf/tensor_mapping.py
@@ -182,11 +182,48 @@ class TensorNameMap:
MODEL_TENSOR.A_QF_PROJ_LINEAR: (
"projector.linear",
),
+ MODEL_TENSOR.LILICORR_TOKEN_PROJ: (
+ "model.lilicorr.token_proj",
+ ),
+ MODEL_TENSOR.LILICORR_PASS_HIDDEN_PROJ: (
+ "model.lilicorr.pass_hidden_proj",
+ ),
+ MODEL_TENSOR.LILICORR_SLOT_EMBD: (
+ "model.lilicorr.slot_embedding",
+ ),
+ MODEL_TENSOR.LILICORR_RANK_EMBD: (
+ "model.lilicorr.rank_embedding",
+ ),
+ MODEL_TENSOR.LILICORR_ATTN_BIAS: (
+ "model.lilicorr.attention_bias",
+ ),
+ MODEL_TENSOR.LILICORR_CONTEXT_PROJ: (
+ "model.lilicorr.context_proj",
+ ),
+ MODEL_TENSOR.LILICORR_OUTPUT_NORM: (
+ "model.lilicorr.output_norm",
+ ),
+ MODEL_TENSOR.LILICORR_ANCHOR_NORM: (
+ "model.lilicorr.anchor_norm",
+ ),
+ MODEL_TENSOR.LILICORR_FACTOR_INPUT_PROJ: (
+ "model.lilicorr.factor_input_proj",
+ ),
+ MODEL_TENSOR.LILICORR_OUT_HEAD: (
+ "model.lilicorr.out_head",
+ ),
+ MODEL_TENSOR.LILICORR_IN_HEAD: (
+ "model.lilicorr.in_head",
+ ),
+ MODEL_TENSOR.LILICORR_ANCHOR_OUT_HEAD: (
+ "model.lilicorr.anchor_out_head",
+ ),
}
block_mappings_cfg: dict[MODEL_TENSOR, tuple[str, ...]] = {
# Attention norm
MODEL_TENSOR.ATTN_NORM: (
+ "model.lilicorr.layers.{bid}.attn_norm", # lilicorr
"gpt_neox.layers.{bid}.input_layernorm", # gptneox
"transformer.h.{bid}.ln_1", # gpt2 gpt-j refact qwen jais exaone
"transformer.blocks.{bid}.norm_1", # mpt
@@ -324,6 +361,7 @@ class TensorNameMap:
# Attention output
MODEL_TENSOR.ATTN_OUT: (
+ "model.lilicorr.layers.{bid}.attn.out_proj", # lilicorr
"gpt_neox.layers.{bid}.attention.dense", # gptneox
"transformer.h.{bid}.attn.c_proj", # gpt2 refact qwen jais
"transformer.blocks.{bid}.attn.out_proj", # mpt
@@ -407,6 +445,8 @@ class TensorNameMap:
# Feed-forward norm
MODEL_TENSOR.FFN_NORM: (
+ "model.lilicorr.layers.{bid}.mlp_norm", # lilicorr
+ "model.lilicorr.feature_mlp.{bid}.norm", # lilicorr
"gpt_neox.layers.{bid}.post_attention_layernorm", # gptneox
"transformer.h.{bid}.ln_2", # gpt2 refact qwen jais exaone
"h.{bid}.post_attention_layernorm", # bloom
@@ -504,6 +544,8 @@ class TensorNameMap:
# Feed-forward up
MODEL_TENSOR.FFN_UP: (
+ "model.lilicorr.layers.{bid}.mlp.up_proj", # lilicorr
+ "model.lilicorr.feature_mlp.{bid}.up_proj", # lilicorr
"gpt_neox.layers.{bid}.mlp.dense_h_to_4h", # gptneox
"transformer.h.{bid}.mlp.c_fc", # gpt2 jais
"transformer.blocks.{bid}.ffn.up_proj", # mpt
@@ -640,6 +682,8 @@ class TensorNameMap:
# Feed-forward down
MODEL_TENSOR.FFN_DOWN: (
+ "model.lilicorr.layers.{bid}.mlp.down_proj", # lilicorr
+ "model.lilicorr.feature_mlp.{bid}.down_proj", # lilicorr
"gpt_neox.layers.{bid}.mlp.dense_4h_to_h", # gptneox
"transformer.h.{bid}.mlp.c_proj", # gpt2 refact qwen jais
"transformer.blocks.{bid}.ffn.down_proj", # mpt
@@ -1554,6 +1598,9 @@ class TensorNameMap:
MODEL_TENSOR.DFLASH_SELECTOR_HIDDEN: (
"model.candidate_selector.hidden_projection",
),
+ MODEL_TENSOR.LILICORR_ATTN_IN_PROJ: (
+ "model.lilicorr.layers.{bid}.attn.in_proj",
+ ),
MODEL_TENSOR.CLS: (
"classifier", # jina
diff --git a/include/llama.h b/include/llama.h
index f15ca9f75..3af6e2223 100644
--- a/include/llama.h
+++ b/include/llama.h
@@ -1056,7 +1056,8 @@ extern "C" {
// "state" here means extra hidden state carried over from a previous stage, e.g.:
// - MTP: state from N layers of the target model
// - Qwen3 VL (deepstack): state from N layers of the vision encoder
- // Returns false if the context does not take a state embedding (currently only MTP contexts do)
+ // - LiliCorr: context embeddings for the transformers
+ // Returns false if the context does not take a state embedding (MTP and LiliCorr contexts do)
LLAMA_API bool llama_batch_ext_set_embd_state(
struct llama_batch_ext * batch,
int32_t idx,
diff --git a/src/llama-arch.cpp b/src/llama-arch.cpp
index 1acda79d8..5744c5fb8 100644
--- a/src/llama-arch.cpp
+++ b/src/llama-arch.cpp
@@ -389,6 +389,15 @@ static const std::map<llm_kv, const char *> LLM_KV_NAMES = {
{ LLM_KV_DFLASH_SELECTOR_RANK, "%s.selector_rank" },
{ LLM_KV_DFLASH_SELECTOR_TOP_K, "%s.selector_top_k" },
+ { LLM_KV_LILICORR_HIDDEN_SIZE, "lilicorr.hidden_size" },
+ { LLM_KV_LILICORR_CANDIDATE_TOPK, "lilicorr.candidate_topk" },
+ { LLM_KV_LILICORR_NUM_HEADS, "lilicorr.num_heads" },
+ { LLM_KV_LILICORR_NUM_LAYERS, "lilicorr.num_layers" },
+ { LLM_KV_LILICORR_VECTOR_EPS, "lilicorr.vector_eps" },
+ { LLM_KV_LILICORR_LOGIT_SCALE, "lilicorr.logit_scale" },
+ { LLM_KV_LILICORR_FACTOR_DIM, "lilicorr.factor_dim" },
+ { LLM_KV_LILICORR_MLP_HIDDEN_SIZE, "lilicorr.mlp_hidden_size" },
+
{ LLM_KV_SHORTCONV_L_CACHE, "%s.shortconv.l_cache" },
// sentence-transformers dense modules feature dims
{ LLM_KV_DENSE_2_FEAT_IN, "%s.dense_2_feat_in" },
@@ -759,6 +768,19 @@ static const std::map<llm_tensor, const char *> LLM_TENSOR_NAMES = {
{ LLM_TENSOR_DFLASH_SELECTOR_PREV, "selector_predecessor" },
{ LLM_TENSOR_DFLASH_SELECTOR_NEXT, "selector_successor" },
{ LLM_TENSOR_DFLASH_SELECTOR_HIDDEN, "selector_hidden" },
+ { LLM_TENSOR_LILICORR_TOKEN_PROJ, "lilicorr.token_proj" },
+ { LLM_TENSOR_LILICORR_PASS_HIDDEN_PROJ, "lilicorr.pass_hidden_proj" },
+ { LLM_TENSOR_LILICORR_SLOT_EMBD, "lilicorr.slot_embedding" },
+ { LLM_TENSOR_LILICORR_RANK_EMBD, "lilicorr.rank_embedding" },
+ { LLM_TENSOR_LILICORR_ATTN_BIAS, "lilicorr.attention_bias" },
+ { LLM_TENSOR_LILICORR_CONTEXT_PROJ, "lilicorr.context_proj" },
+ { LLM_TENSOR_LILICORR_ATTN_IN_PROJ, "lilicorr.layers.%d.attn.in_proj" },
+ { LLM_TENSOR_LILICORR_OUTPUT_NORM, "lilicorr.output_norm" },
+ { LLM_TENSOR_LILICORR_ANCHOR_NORM, "lilicorr.anchor_norm" },
+ { LLM_TENSOR_LILICORR_FACTOR_INPUT_PROJ, "lilicorr.factor_input_proj" },
+ { LLM_TENSOR_LILICORR_OUT_HEAD, "lilicorr.out_head" },
+ { LLM_TENSOR_LILICORR_IN_HEAD, "lilicorr.in_head" },
+ { LLM_TENSOR_LILICORR_ANCHOR_OUT_HEAD, "lilicorr.anchor_out_head" },
{ LLM_TENSOR_ATTN_V_GATE, "blk.%d.attn_v_gate" },
{ LLM_TENSOR_ATTN_V_EXPS, "blk.%d.attn_v_exps" },
};
@@ -1076,6 +1098,20 @@ static const std::map<llm_tensor, llm_tensor_info> LLM_TENSOR_INFOS = {
{LLM_TENSOR_DFLASH_SELECTOR_PREV, {LLM_TENSOR_LAYER_OUTPUT, GGML_OP_GET_ROWS}},
{LLM_TENSOR_DFLASH_SELECTOR_NEXT, {LLM_TENSOR_LAYER_OUTPUT, GGML_OP_GET_ROWS}},
{LLM_TENSOR_DFLASH_SELECTOR_HIDDEN, {LLM_TENSOR_LAYER_OUTPUT, GGML_OP_MUL_MAT}},
+ {LLM_TENSOR_LILICORR_TOKEN_PROJ, {LLM_TENSOR_LAYER_OUTPUT, GGML_OP_MUL_MAT}},
+ {LLM_TENSOR_LILICORR_PASS_HIDDEN_PROJ, {LLM_TENSOR_LAYER_OUTPUT, GGML_OP_MUL_MAT}},
+ {LLM_TENSOR_LILICORR_SLOT_EMBD, {LLM_TENSOR_LAYER_OUTPUT, GGML_OP_ADD}},
+ {LLM_TENSOR_LILICORR_RANK_EMBD, {LLM_TENSOR_LAYER_OUTPUT, GGML_OP_ADD}},
+ {LLM_TENSOR_LILICORR_ATTN_BIAS, {LLM_TENSOR_LAYER_OUTPUT, GGML_OP_ADD}},
+ {LLM_TENSOR_LILICORR_CONTEXT_PROJ, {LLM_TENSOR_LAYER_OUTPUT, GGML_OP_MUL_MAT}},
+ {LLM_TENSOR_LILICORR_ATTN_IN_PROJ, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL_MAT}},
+ {LLM_TENSOR_LILICORR_OUTPUT_NORM, {LLM_TENSOR_LAYER_OUTPUT, GGML_OP_MUL}},
+ {LLM_TENSOR_LILICORR_ANCHOR_NORM, {LLM_TENSOR_LAYER_OUTPUT, GGML_OP_MUL}},
+ {LLM_TENSOR_LILICORR_FACTOR_INPUT_PROJ, {LLM_TENSOR_LAYER_OUTPUT, GGML_OP_MUL_MAT}},
+ {LLM_TENSOR_LILICORR_OUT_HEAD, {LLM_TENSOR_LAYER_OUTPUT, GGML_OP_MUL_MAT}},
+ {LLM_TENSOR_LILICORR_IN_HEAD, {LLM_TENSOR_LAYER_OUTPUT, GGML_OP_MUL_MAT}},
+ {LLM_TENSOR_LILICORR_ANCHOR_OUT_HEAD, {LLM_TENSOR_LAYER_OUTPUT, GGML_OP_MUL_MAT}},
+
{LLM_TENSOR_ATTN_V_GATE, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL_MAT}},
{LLM_TENSOR_ATTN_V_EXPS, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL_MAT_ID}},
};
diff --git a/src/llama-arch.h b/src/llama-arch.h
index e6e9bd4b9..94771c2a1 100644
--- a/src/llama-arch.h
+++ b/src/llama-arch.h
@@ -432,6 +432,14 @@ enum llm_kv {
LLM_KV_DFLASH_CONV_GROUP_SIZE,
LLM_KV_DFLASH_SELECTOR_RANK,
LLM_KV_DFLASH_SELECTOR_TOP_K,
+ LLM_KV_LILICORR_HIDDEN_SIZE,
+ LLM_KV_LILICORR_CANDIDATE_TOPK,
+ LLM_KV_LILICORR_NUM_HEADS,
+ LLM_KV_LILICORR_NUM_LAYERS,
+ LLM_KV_LILICORR_VECTOR_EPS,
+ LLM_KV_LILICORR_LOGIT_SCALE,
+ LLM_KV_LILICORR_FACTOR_DIM,
+ LLM_KV_LILICORR_MLP_HIDDEN_SIZE,
LLM_KV_NORM_BEFORE_RESIDUAL,
LLM_KV_NORM_BEFORE_FC,
@@ -766,6 +774,19 @@ enum llm_tensor {
LLM_TENSOR_DFLASH_SELECTOR_PREV,
LLM_TENSOR_DFLASH_SELECTOR_NEXT,
LLM_TENSOR_DFLASH_SELECTOR_HIDDEN,
+ LLM_TENSOR_LILICORR_TOKEN_PROJ,
+ LLM_TENSOR_LILICORR_PASS_HIDDEN_PROJ,
+ LLM_TENSOR_LILICORR_SLOT_EMBD,
+ LLM_TENSOR_LILICORR_RANK_EMBD,
+ LLM_TENSOR_LILICORR_ATTN_BIAS,
+ LLM_TENSOR_LILICORR_CONTEXT_PROJ,
+ LLM_TENSOR_LILICORR_ATTN_IN_PROJ,
+ LLM_TENSOR_LILICORR_OUTPUT_NORM,
+ LLM_TENSOR_LILICORR_ANCHOR_NORM,
+ LLM_TENSOR_LILICORR_FACTOR_INPUT_PROJ,
+ LLM_TENSOR_LILICORR_OUT_HEAD,
+ LLM_TENSOR_LILICORR_IN_HEAD,
+ LLM_TENSOR_LILICORR_ANCHOR_OUT_HEAD,
LLM_TENSOR_ATTN_V_GATE,
LLM_TENSOR_ATTN_V_EXPS,
};
diff --git a/src/llama-batch.cpp b/src/llama-batch.cpp
index 3c4b52008..321aa9d23 100644
--- a/src/llama-batch.cpp
+++ b/src/llama-batch.cpp
@@ -1147,6 +1147,10 @@ size_t llama_batch_ext_select_n_embd_state(llama_context_type ctx_type, const ll
if (ctx_type == LLAMA_CONTEXT_TYPE_MTP) {
return hparams.n_embd_out();
}
+ // TODO: this probably should become a function of `llm_arch`
+ if (hparams.lilicorr_candidate_topk > 0) {
+ return hparams.n_embd_inp();
+ }
return 0;
}
diff --git a/src/llama-hparams.h b/src/llama-hparams.h
index 064de68f3..922164a54 100644
--- a/src/llama-hparams.h
+++ b/src/llama-hparams.h
@@ -264,6 +264,12 @@ struct llama_hparams {
uint32_t dflash_conv_group_size = 0;
uint32_t dflash_selector_rank = 0;
uint32_t dflash_selector_top_k = 0;
+ uint32_t lilicorr_hidden_size = 0;
+ uint32_t lilicorr_candidate_topk = 0;
+ uint32_t lilicorr_num_heads = 0;
+ uint32_t lilicorr_layers = 0;
+ float lilicorr_vector_eps = 0.0f;
+ float lilicorr_logit_scale = 0.0f;
// llama4 smallthinker
uint32_t n_moe_layer_step = 0;
diff --git a/src/llama-model.h b/src/llama-model.h
index bc640e0dd..bcfbf1e2c 100644
--- a/src/llama-model.h
+++ b/src/llama-model.h
@@ -386,6 +386,20 @@ struct llama_layer {
struct ggml_tensor * dflash_ffn_conv_base = nullptr;
struct ggml_tensor * dflash_ffn_conv_proj = nullptr;
+ struct ggml_tensor * lilicorr_attn_norm = nullptr;
+ struct ggml_tensor * lilicorr_mlp_norm = nullptr;
+ struct ggml_tensor * lilicorr_attn_in_proj = nullptr;
+ struct ggml_tensor * lilicorr_attn_in_proj_b = nullptr;
+ struct ggml_tensor * lilicorr_attn_out_proj = nullptr;
+ struct ggml_tensor * lilicorr_attn_out_proj_b = nullptr;
+ struct ggml_tensor * lilicorr_attn_out_proj_s = nullptr;
+ struct ggml_tensor * lilicorr_mlp_0 = nullptr;
+ struct ggml_tensor * lilicorr_mlp_0_b = nullptr;
+ struct ggml_tensor * lilicorr_mlp_0_s = nullptr;
+ struct ggml_tensor * lilicorr_mlp_2 = nullptr;
+ struct ggml_tensor * lilicorr_mlp_2_b = nullptr;
+ struct ggml_tensor * lilicorr_mlp_2_s = nullptr;
+
// mamba proj
struct ggml_tensor * ssm_in = nullptr;
struct ggml_tensor * ssm_x = nullptr;
@@ -764,6 +778,36 @@ struct llama_model {
struct ggml_tensor * dflash_selector_next = nullptr;
struct ggml_tensor * dflash_selector_hidden = nullptr;
+ struct ggml_tensor * lilicorr_token_proj = nullptr;
+ struct ggml_tensor * lilicorr_token_proj_b = nullptr;
+ struct ggml_tensor * lilicorr_token_proj_s = nullptr;
+ struct ggml_tensor * lilicorr_pass_hidden_proj = nullptr;
+ struct ggml_tensor * lilicorr_pass_hidden_proj_b = nullptr;
+ struct ggml_tensor * lilicorr_pass_hidden_proj_s = nullptr;
+ struct ggml_tensor * lilicorr_feature_norm = nullptr;
+ struct ggml_tensor * lilicorr_feature_norm_b = nullptr;
+ struct ggml_tensor * lilicorr_feature_mlp_1 = nullptr;
+ struct ggml_tensor * lilicorr_feature_mlp_1_b = nullptr;
+ struct ggml_tensor * lilicorr_feature_mlp_3 = nullptr;
+ struct ggml_tensor * lilicorr_feature_mlp_3_b = nullptr;
+ struct ggml_tensor * lilicorr_slot_embedding = nullptr;
+ struct ggml_tensor * lilicorr_rank_embedding = nullptr;
+ struct ggml_tensor * lilicorr_attention_bias = nullptr;
+ struct ggml_tensor * lilicorr_context_proj = nullptr;
+ struct ggml_tensor * lilicorr_context_proj_b = nullptr;
+ struct ggml_tensor * lilicorr_context_proj_s = nullptr;
+ struct ggml_tensor * lilicorr_output_norm = nullptr;
+ struct ggml_tensor * lilicorr_anchor_norm = nullptr;
+ struct ggml_tensor * lilicorr_factor_input_proj = nullptr;
+ struct ggml_tensor * lilicorr_factor_input_proj_b = nullptr;
+ struct ggml_tensor * lilicorr_out_head = nullptr;
+ struct ggml_tensor * lilicorr_out_head_b = nullptr;
+ struct ggml_tensor * lilicorr_in_head = nullptr;
+ struct ggml_tensor * lilicorr_in_head_b = nullptr;
+ struct ggml_tensor * lilicorr_anchor_out_head = nullptr;
+ struct ggml_tensor * lilicorr_anchor_out_head_b = nullptr;
+
+
// unified vector to store target-model extracted layer ids in eagle3, dflash, etc.
std::vector<int32_t> target_layer_ids;
diff --git a/src/models/dflash.cpp b/src/models/dflash.cpp
index aeeadbbd1..982351dc8 100644
--- a/src/models/dflash.cpp
+++ b/src/models/dflash.cpp
@@ -33,6 +33,19 @@ void llama_model_dflash::load_arch_hparams(llama_model_loader & ml) {
ml.get_key(LLM_KV_DFLASH_CONV_GROUP_SIZE, hparams.dflash_conv_group_size, false);
ml.get_key(LLM_KV_DFLASH_SELECTOR_RANK, hparams.dflash_selector_rank, false);
ml.get_key(LLM_KV_DFLASH_SELECTOR_TOP_K, hparams.dflash_selector_top_k, false);
+ const bool has_lilicorr = ml.get_key(LLM_KV_LILICORR_HIDDEN_SIZE, hparams.lilicorr_hidden_size, false);
+ ml.get_key(LLM_KV_LILICORR_CANDIDATE_TOPK, hparams.lilicorr_candidate_topk, has_lilicorr);
+ ml.get_key(LLM_KV_LILICORR_NUM_HEADS, hparams.lilicorr_num_heads, has_lilicorr);
+ ml.get_key(LLM_KV_LILICORR_NUM_LAYERS, hparams.lilicorr_layers, has_lilicorr);
+ ml.get_key(LLM_KV_LILICORR_VECTOR_EPS, hparams.lilicorr_vector_eps, has_lilicorr);
+ ml.get_key(LLM_KV_LILICORR_LOGIT_SCALE, hparams.lilicorr_logit_scale, has_lilicorr);
+ if (has_lilicorr && (hparams.dflash_block_size < 2 || hparams.lilicorr_hidden_size == 0 ||
+ hparams.lilicorr_num_heads == 0 || hparams.lilicorr_hidden_size % hparams.lilicorr_num_heads ||
+ hparams.lilicorr_candidate_topk == 0 || hparams.lilicorr_layers == 0 ||
+ hparams.dflash_selector_top_k > 0 ||
+ (uint64_t) hparams.lilicorr_candidate_topk * (hparams.lilicorr_candidate_topk + 1) > hparams.n_embd)) {
+ throw std::runtime_error("invalid LiLiCorr DFlash metadata");
+ }
if (!ml.get_arr(LLM_KV_TARGET_LAYERS, target_layer_ids, false)) {
throw std::runtime_error("DFlash model requires 'target_layers' in GGUF metadata");
@@ -137,6 +150,10 @@ void llama_model_dflash::load_arch_tensors(llama_model_loader &) {
}
const struct ggml_tensor * selector_meta = ml->get_tensor_meta("selector_hidden.weight");
+ const bool conv_meta = hparams.dflash_conv_kernel_size > 0 && hparams.dflash_conv_group_size > 0;
+ if (conv_meta && n_embd % hparams.dflash_conv_group_size != 0) {
+ throw std::runtime_error("DFlash hidden size must be divisible by conv_group_size");
+ }
if (selector_meta) {
const int64_t rank = hparams.dflash_selector_rank;
if (rank <= 0 || hparams.dflash_block_size <= 0 || hparams.dflash_selector_top_k <= 0 ||
@@ -159,6 +176,66 @@ void llama_model_dflash::load_arch_tensors(llama_model_loader &) {
hparams.dflash_selector_rank, hparams.dflash_selector_top_k);
}
+ if (hparams.lilicorr_layers > 0) {
+ const int64_t hidden = hparams.lilicorr_hidden_size;
+ const int64_t top_k = hparams.lilicorr_candidate_topk;
+ const int64_t heads = hparams.lilicorr_num_heads;
+ uint32_t factor = 0, mlp = 0;
+ ml->get_key(LLM_KV_LILICORR_FACTOR_DIM, factor);
+ ml->get_key(LLM_KV_LILICORR_MLP_HIDDEN_SIZE, mlp);
+ const int64_t slots = hparams.dflash_block_size - 1;
+ const int feature_bid = n_layer_all + hparams.lilicorr_layers;
+ if (factor == 0 || mlp == 0 || hparams.lilicorr_layers > layers.size()) {
+ throw std::runtime_error("invalid LiLiCorr head dimensions");
+ }
+ if (hidden != n_embd) {
+ lilicorr_token_proj = create_tensor(tn(LLM_TENSOR_LILICORR_TOKEN_PROJ, "weight"), { n_embd, hidden }, 0);
+ lilicorr_token_proj_b = create_tensor(tn(LLM_TENSOR_LILICORR_TOKEN_PROJ, "bias"), { hidden }, 0);
+ lilicorr_token_proj_s = create_tensor(tn(LLM_TENSOR_LILICORR_TOKEN_PROJ, "scale"), { 1 }, TENSOR_NOT_REQUIRED);
+ }
+ lilicorr_pass_hidden_proj = create_tensor(tn(LLM_TENSOR_LILICORR_PASS_HIDDEN_PROJ, "weight"), { hparams.n_embd, hidden }, 0);
+ lilicorr_pass_hidden_proj_b = create_tensor(tn(LLM_TENSOR_LILICORR_PASS_HIDDEN_PROJ, "bias"), { hidden }, 0);
+ lilicorr_pass_hidden_proj_s = create_tensor(tn(LLM_TENSOR_LILICORR_PASS_HIDDEN_PROJ, "scale"), { 1 }, TENSOR_NOT_REQUIRED);
+ lilicorr_feature_norm = create_tensor(tn(LLM_TENSOR_FFN_NORM, "weight", feature_bid), { 5 }, 0);
+ lilicorr_feature_norm_b = create_tensor(tn(LLM_TENSOR_FFN_NORM, "bias", feature_bid), { 5 }, 0);
+ lilicorr_feature_mlp_1 = create_tensor(tn(LLM_TENSOR_FFN_UP, "weight", feature_bid), { 5, hidden }, 0);
+ lilicorr_feature_mlp_1_b = create_tensor(tn(LLM_TENSOR_FFN_UP, "bias", feature_bid), { hidden }, 0);
+ lilicorr_feature_mlp_3 = create_tensor(tn(LLM_TENSOR_FFN_DOWN, "weight", feature_bid), { hidden, hidden }, 0);
+ lilicorr_feature_mlp_3_b = create_tensor(tn(LLM_TENSOR_FFN_DOWN, "bias", feature_bid), { hidden }, 0);
+ lilicorr_slot_embedding = create_tensor(tn(LLM_TENSOR_LILICORR_SLOT_EMBD, "weight"), { hidden, slots }, 0);
+ lilicorr_rank_embedding = create_tensor(tn(LLM_TENSOR_LILICORR_RANK_EMBD, "weight"), { hidden, top_k }, 0);
+ lilicorr_attention_bias = create_tensor(tn(LLM_TENSOR_LILICORR_ATTN_BIAS, "bias"), { slots * top_k, slots * top_k, heads }, 0);
+ lilicorr_context_proj = create_tensor(tn(LLM_TENSOR_LILICORR_CONTEXT_PROJ, "weight"), { n_embd, hidden }, 0);
+ lilicorr_context_proj_b = create_tensor(tn(LLM_TENSOR_LILICORR_CONTEXT_PROJ, "bias"), { hidden }, 0);
+ lilicorr_context_proj_s = create_tensor(tn(LLM_TENSOR_LILICORR_CONTEXT_PROJ, "scale"), { 1 }, TENSOR_NOT_REQUIRED);
+ for (uint32_t i = 0; i < hparams.lilicorr_layers; ++i) {
+ const int bid = n_layer_all + i;
+ layers[i].lilicorr_attn_norm = create_tensor(tn(LLM_TENSOR_ATTN_NORM, "weight", bid), { hidden }, 0);
+ layers[i].lilicorr_mlp_norm = create_tensor(tn(LLM_TENSOR_FFN_NORM, "weight", bid), { hidden }, 0);
+ layers[i].lilicorr_attn_in_proj = create_tensor(tn(LLM_TENSOR_LILICORR_ATTN_IN_PROJ, "weight", i), { hidden, 3 * hidden }, 0);
+ layers[i].lilicorr_attn_in_proj_b = create_tensor(tn(LLM_TENSOR_LILICORR_ATTN_IN_PROJ, "bias", i), { 3 * hidden }, 0);
+ layers[i].lilicorr_attn_out_proj = create_tensor(tn(LLM_TENSOR_ATTN_OUT, "weight", bid), { hidden, hidden }, 0);
+ layers[i].lilicorr_attn_out_proj_b = create_tensor(tn(LLM_TENSOR_ATTN_OUT, "bias", bid), { hidden }, 0);
+ layers[i].lilicorr_attn_out_proj_s = create_tensor(tn(LLM_TENSOR_ATTN_OUT, "scale", bid), { 1 }, TENSOR_NOT_REQUIRED);
+ layers[i].lilicorr_mlp_0 = create_tensor(tn(LLM_TENSOR_FFN_UP, "weight", bid), { hidden, mlp }, 0);
+ layers[i].lilicorr_mlp_0_b = create_tensor(tn(LLM_TENSOR_FFN_UP, "bias", bid), { mlp }, 0);
+ layers[i].lilicorr_mlp_0_s = create_tensor(tn(LLM_TENSOR_FFN_UP, "scale", bid), { 1 }, TENSOR_NOT_REQUIRED);
+ layers[i].lilicorr_mlp_2 = create_tensor(tn(LLM_TENSOR_FFN_DOWN, "weight", bid), { mlp, hidden }, 0);
+ layers[i].lilicorr_mlp_2_b = create_tensor(tn(LLM_TENSOR_FFN_DOWN, "bias", bid), { hidden }, 0);
+ layers[i].lilicorr_mlp_2_s = create_tensor(tn(LLM_TENSOR_FFN_DOWN, "scale", bid), { 1 }, TENSOR_NOT_REQUIRED);
+ }
+ lilicorr_output_norm = create_tensor(tn(LLM_TENSOR_LILICORR_OUTPUT_NORM, "weight"), { hidden }, 0);
+ lilicorr_anchor_norm = create_tensor(tn(LLM_TENSOR_LILICORR_ANCHOR_NORM, "weight"), { hidden }, 0);
+ lilicorr_factor_input_proj = create_tensor(tn(LLM_TENSOR_LILICORR_FACTOR_INPUT_PROJ, "weight"), { 3 * hidden, hidden }, 0);
+ lilicorr_factor_input_proj_b = create_tensor(tn(LLM_TENSOR_LILICORR_FACTOR_INPUT_PROJ, "bias"), { hidden }, 0);
+ lilicorr_out_head = create_tensor(tn(LLM_TENSOR_LILICORR_OUT_HEAD, "weight"), { hidden, factor }, 0);
+ lilicorr_out_head_b = create_tensor(tn(LLM_TENSOR_LILICORR_OUT_HEAD, "bias"), { factor }, 0);
+ lilicorr_in_head = create_tensor(tn(LLM_TENSOR_LILICORR_IN_HEAD, "weight"), { hidden, factor }, 0);
+ lilicorr_in_head_b = create_tensor(tn(LLM_TENSOR_LILICORR_IN_HEAD, "bias"), { factor }, 0);
+ lilicorr_anchor_out_head = create_tensor(tn(LLM_TENSOR_LILICORR_ANCHOR_OUT_HEAD, "weight"), { hidden, factor }, 0);
+ lilicorr_anchor_out_head_b = create_tensor(tn(LLM_TENSOR_LILICORR_ANCHOR_OUT_HEAD, "bias"), { factor }, 0);
+ }
+
fc = create_tensor(tn(LLM_TENSOR_FC, "weight"), { n_embd_inp, n_embd }, 0);
fc_s = create_tensor(tn(LLM_TENSOR_FC, "scale"), { 1 }, TENSOR_NOT_REQUIRED);
output_norm_enc = create_tensor(tn(LLM_TENSOR_ENC_OUTPUT_NORM, "weight"), { n_embd }, 0); // encoder hidden_norm (after fc)
@@ -244,7 +321,7 @@ void llama_model_dflash::load_arch_tensors(llama_model_loader &) {
layer.ffn_down = create_tensor(tn(LLM_TENSOR_FFN_DOWN, "weight", i), { n_ff, n_embd }, 0);
layer.ffn_up = create_tensor(tn(LLM_TENSOR_FFN_UP, "weight", i), { n_embd, n_ff }, 0);
- if (selector_meta) {
+ if (conv_meta) {
const int64_t kernel = hparams.dflash_conv_kernel_size;
const int64_t groups = n_embd / hparams.dflash_conv_group_size;
const int64_t projected = 2 * kernel * groups;
@@ -561,6 +638,168 @@ static void build_dflash2_selector(llm_graph_context & g, const llama_model & mo
ggml_build_forward_expand(g.gf, packed);
}
+static void build_lilicorr_scores(llm_graph_context & g, const llama_model & model,
+ ggml_tensor * tok_embd, ggml_tensor * anchors) {
+ ggml_context * ctx0 = g.ctx0;
+ auto & res = g.res;
+
+ const auto & hparams = g.hparams;
+ const int64_t n_tokens = g.n_tokens;
+ const int64_t n_embd = g.n_embd;
+
+ const int64_t hidden = hparams.lilicorr_hidden_size;
+ const int64_t top_k = hparams.lilicorr_candidate_topk;
+ const int64_t heads = hparams.lilicorr_num_heads;
+ const int64_t n_blocks = g.ubatch.n_seqs_unq;
+ GGML_ASSERT(n_blocks > 0 && n_tokens % n_blocks == 0);
+ GGML_ASSERT(res->t_logits->ne[1] == n_tokens);
+ GGML_ASSERT(anchors && tok_embd);
+
+ const int64_t tokens_per_block = n_tokens / n_blocks;
+ const int64_t block_size = hparams.dflash_block_size;
+ const int64_t row_used = top_k + top_k * top_k;
+ const int64_t slots = block_size - 1;
+ const int64_t nodes = slots * top_k;
+ const int64_t factor_dim = model.lilicorr_out_head->ne[1];
+ GGML_ASSERT(tokens_per_block == block_size);
+
+ ggml_tensor * candidates = ggml_top_k(ctx0, res->t_logits, top_k);
+ ggml_tensor * logits_rows = ggml_reshape_3d(ctx0, res->t_logits, 1, res->t_logits->ne[0], n_tokens);
+ ggml_tensor * candidate_logits = ggml_reshape_2d(ctx0,
+ ggml_get_rows(ctx0, logits_rows, candidates), top_k, n_tokens);
+ // Rank features and embeddings require candidates in descending logit order.
+ ggml_tensor * order = ggml_argsort(ctx0, candidate_logits, GGML_SORT_ORDER_DESC);
+ ggml_tensor * candidate_rows = ggml_reshape_3d(ctx0, candidates, 1, top_k, n_tokens);
+ candidates = ggml_reshape_2d(ctx0, ggml_get_rows(ctx0, candidate_rows, order), top_k, n_tokens);
+ ggml_tensor * prob_rows = ggml_reshape_3d(ctx0, ggml_soft_max(ctx0, res->t_logits), 1, res->t_logits->ne[0], n_tokens);
+ ggml_tensor * probs = ggml_reshape_2d(ctx0,
+ ggml_get_rows(ctx0, prob_rows, candidates), top_k, n_tokens);
+
+ ggml_tensor * cand_run = ggml_reshape_3d(ctx0, candidates, top_k, n_tokens / n_blocks, n_blocks);
+ ggml_tensor * probs_run = ggml_reshape_3d(ctx0, probs, top_k, n_tokens / n_blocks, n_blocks);
+ cand_run = ggml_cont(ctx0, ggml_view_3d(ctx0, cand_run, top_k, slots, n_blocks,
+ cand_run->nb[1], cand_run->nb[2], cand_run->nb[1]));
+ probs_run = ggml_cont(ctx0, ggml_view_3d(ctx0, probs_run, top_k, slots, n_blocks,
+ probs_run->nb[1], probs_run->nb[2], probs_run->nb[1]));
+
+ ggml_tensor * lp = ggml_log(ctx0, probs_run);
+ ggml_tensor * lp_blk = ggml_reshape_3d(ctx0, lp, top_k, slots, n_blocks);
+ ggml_tensor * top_lp = ggml_view_3d(ctx0, lp_blk, 1, slots, n_blocks,
+ lp_blk->nb[1], lp_blk->nb[2], 0);
+ ggml_tensor * relative = ggml_sub(ctx0, lp_blk, top_lp);
+ ggml_tensor * feature = ggml_reshape_4d(ctx0, lp_blk, 1, top_k, slots, n_blocks);
+
+ ggml_tensor * rank = ggml_arange(ctx0, 0, top_k, 1);
+ rank = ggml_repeat(ctx0, ggml_reshape_3d(ctx0, rank, top_k, 1, 1), lp_blk);
+ ggml_tensor * first = ggml_scale_bias(ctx0, ggml_step(ctx0, rank), -1.0f, 1.0f);
+
+ for (ggml_tensor * item : { probs_run, relative,
+ ggml_scale(ctx0, rank, top_k == 1 ? 0.0f : 1.0f / (top_k - 1)), first }) {
+ feature = ggml_concat(ctx0, feature, ggml_reshape_4d(ctx0, item, 1, top_k, slots, n_blocks), 0);
+ }
+ feature = ggml_reshape_2d(ctx0, feature, 5, nodes * n_blocks);
+
+ ggml_tensor * embeds = ggml_get_rows(ctx0, tok_embd, ggml_reshape_1d(ctx0, cand_run, nodes * n_blocks));
+ ggml_tensor * token_states = n_embd == hidden ? embeds : ggml_add(ctx0, g.build_lora_mm(model.lilicorr_token_proj, embeds, model.lilicorr_token_proj_s), model.lilicorr_token_proj_b);
+
+ ggml_tensor * pass_blk = ggml_reshape_3d(ctx0, res->t_embd, n_embd, block_size, n_blocks);
+ ggml_tensor * passes = ggml_reshape_2d(ctx0, ggml_cont(ctx0, ggml_view_3d(ctx0, pass_blk, n_embd, slots, n_blocks,
+ pass_blk->nb[1], pass_blk->nb[2], pass_blk->nb[1])),
+ n_embd, slots * n_blocks);
+
+ ggml_tensor * pass_states = ggml_add(ctx0, g.build_lora_mm(model.lilicorr_pass_hidden_proj, passes, model.lilicorr_pass_hidden_proj_s), model.lilicorr_pass_hidden_proj_b);
+
+ pass_states = ggml_reshape_4d(ctx0, pass_states, hidden, 1, slots, n_blocks);
+ pass_states = ggml_reshape_2d(ctx0, ggml_repeat_4d(ctx0, pass_states, hidden, top_k, slots, n_blocks), hidden, nodes * n_blocks);
+
+ ggml_tensor * feature_states = ggml_norm(ctx0, feature, 1.0e-5f);
+ feature_states = ggml_add(ctx0, ggml_mul(ctx0, feature_states, model.lilicorr_feature_norm), model.lilicorr_feature_norm_b);
+ feature_states = ggml_silu(ctx0, ggml_add(ctx0, g.build_lora_mm(model.lilicorr_feature_mlp_1, feature_states), model.lilicorr_feature_mlp_1_b));
+ feature_states = ggml_add(ctx0, g.build_lora_mm(model.lilicorr_feature_mlp_3, feature_states), model.lilicorr_feature_mlp_3_b);
+
+ ggml_tensor * slot = ggml_reshape_4d(ctx0, model.lilicorr_slot_embedding, hidden, 1, slots, 1);
+ ggml_tensor * rank_embd = ggml_reshape_4d(ctx0, model.lilicorr_rank_embedding, hidden, top_k, 1, 1);
+ ggml_tensor * pos = ggml_add(ctx0, ggml_repeat_4d(ctx0, slot, hidden, top_k, slots, n_blocks), ggml_repeat_4d(ctx0, rank_embd, hidden, top_k, slots, n_blocks));
+ pos = ggml_reshape_2d(ctx0, pos, hidden, nodes * n_blocks);
+
+ ggml_tensor * x = ggml_add(ctx0, ggml_add(ctx0, token_states, pass_states), feature_states);
+ x = ggml_add(ctx0, x, pos);
+ ggml_tensor * anchor_blk = ggml_reshape_3d(ctx0, anchors, n_embd, block_size, n_blocks);
+ ggml_tensor * anchor = ggml_reshape_2d(ctx0,
+ ggml_cont(ctx0, ggml_view_3d(ctx0, anchor_blk, n_embd, 1, n_blocks,
+ anchor_blk->nb[1], anchor_blk->nb[2], 0)), n_embd, n_blocks);
+ ggml_tensor * anchor_state = ggml_add(ctx0, g.build_lora_mm(model.lilicorr_context_proj, anchor, model.lilicorr_context_proj_s), model.lilicorr_context_proj_b);
+
+ const int64_t head_dim = hidden / heads;
+ for (uint32_t layer = 0; layer < hparams.lilicorr_layers; ++layer) {
+ const auto & weights = model.layers[layer];
+ ggml_tensor * attn_input = ggml_mul(ctx0, ggml_rms_norm(ctx0, x, hparams.f_norm_rms_eps), weights.lilicorr_attn_norm);
+ ggml_tensor * qkv = ggml_add(ctx0, g.build_lora_mm(weights.lilicorr_attn_in_proj, attn_input), weights.lilicorr_attn_in_proj_b);
+ qkv = ggml_reshape_3d(ctx0, qkv, 3 * hidden, nodes, n_blocks);
+ ggml_tensor * q = ggml_cont(ctx0, ggml_view_4d(ctx0, qkv, head_dim, nodes, heads, n_blocks,
+ qkv->nb[1], head_dim * qkv->nb[0], qkv->nb[2], 0));
+ ggml_tensor * k = ggml_cont(ctx0, ggml_view_4d(ctx0, qkv, head_dim, nodes, heads, n_blocks,
+ qkv->nb[1], head_dim * qkv->nb[0], qkv->nb[2], hidden * qkv->nb[0]));
+ ggml_tensor * v = ggml_view_4d(ctx0, qkv, head_dim, nodes, heads, n_blocks,
+ qkv->nb[1], head_dim * qkv->nb[0], qkv->nb[2], 2 * hidden * qkv->nb[0]);
+ ggml_tensor * attn = ggml_soft_max_ext(ctx0, ggml_mul_mat(ctx0, k, q), model.lilicorr_attention_bias,
+ 1.0f / std::sqrt(static_cast<float>(head_dim)), 0.0f);
+ ggml_tensor * joined = ggml_out_prod(ctx0, v, ggml_transpose(ctx0, attn));
+ joined = ggml_permute(ctx0, joined, 0, 2, 1, 3);
+ ggml_tensor * attn_output = ggml_cont_2d(ctx0, joined, hidden, nodes * n_blocks);
+ attn_output = ggml_add(ctx0, g.build_lora_mm(weights.lilicorr_attn_out_proj, attn_output, weights.lilicorr_attn_out_proj_s), weights.lilicorr_attn_out_proj_b);
+ x = ggml_add(ctx0, x, attn_output);
+ ggml_tensor * mlp = ggml_mul(ctx0, ggml_rms_norm(ctx0, x, hparams.f_norm_rms_eps), weights.lilicorr_mlp_norm);
+ mlp = ggml_silu(ctx0, ggml_add(ctx0, g.build_lora_mm(weights.lilicorr_mlp_0, mlp, weights.lilicorr_mlp_0_s), weights.lilicorr_mlp_0_b));
+ mlp = ggml_add(ctx0, g.build_lora_mm(weights.lilicorr_mlp_2, mlp, weights.lilicorr_mlp_2_s), weights.lilicorr_mlp_2_b);
+ x = ggml_add(ctx0, x, mlp);
+ }
+
+ x = ggml_mul(ctx0, ggml_rms_norm(ctx0, x, hparams.f_norm_rms_eps), model.lilicorr_output_norm);
+ anchor_state = ggml_mul(ctx0, ggml_rms_norm(ctx0, anchor_state, hparams.f_norm_rms_eps), model.lilicorr_anchor_norm);
+ ggml_tensor * anchor_nodes = ggml_reshape_3d(ctx0, anchor_state, hidden, 1, n_blocks);
+ anchor_nodes = ggml_reshape_2d(ctx0, ggml_repeat_4d(ctx0, anchor_nodes, hidden, nodes, n_blocks, 1), hidden, nodes * n_blocks);
+ ggml_tensor * factor_in = ggml_concat(ctx0, x, anchor_nodes, 0);
+ factor_in = ggml_concat(ctx0, factor_in, ggml_mul(ctx0, x, anchor_nodes), 0);
+ ggml_tensor * factors = ggml_add(ctx0, g.build_lora_mm(model.lilicorr_factor_input_proj, factor_in), model.lilicorr_factor_input_proj_b);
+ factors = ggml_silu(ctx0, factors);
+
+ ggml_tensor * out = ggml_add(ctx0, g.build_lora_mm(model.lilicorr_out_head, factors), model.lilicorr_out_head_b);
+ out = ggml_reshape_4d(ctx0, ggml_l2_norm(ctx0, out, hparams.lilicorr_vector_eps), factor_dim, top_k, slots, n_blocks);
+
+ ggml_tensor * in = ggml_add(ctx0, g.build_lora_mm(model.lilicorr_in_head, factors), model.lilicorr_in_head_b);
+ in = ggml_reshape_4d(ctx0, ggml_l2_norm(ctx0, in, hparams.lilicorr_vector_eps), factor_dim, top_k, slots, n_blocks);
+
+ ggml_tensor * anchor_out = ggml_add(ctx0, g.build_lora_mm(model.lilicorr_anchor_out_head, anchor_state), model.lilicorr_anchor_out_head_b);
+ anchor_out = ggml_l2_norm(ctx0, anchor_out, hparams.lilicorr_vector_eps);
+ anchor_out = ggml_reshape_3d(ctx0, anchor_out, factor_dim, 1, n_blocks);
+
+ ggml_tensor * first_in = ggml_cont(ctx0, ggml_view_4d(ctx0, in, factor_dim, top_k, 1, n_blocks,
+ in->nb[1], in->nb[2], in->nb[3], 0));
+ ggml_tensor * first_score = ggml_scale(ctx0,
+ ggml_mul_mat(ctx0, first_in, ggml_reshape_4d(ctx0, anchor_out, factor_dim, 1, 1, n_blocks)),
+ hparams.lilicorr_logit_scale);
+ ggml_tensor * scores = ggml_pad(ctx0,
+ ggml_reshape_3d(ctx0, first_score, top_k, 1, n_blocks), top_k * top_k - top_k, 0, 0, 0);
+ if (slots > 1) {
+ ggml_tensor * next = ggml_cont(ctx0, ggml_view_4d(ctx0, in, factor_dim, top_k, slots - 1, n_blocks,
+ in->nb[1], in->nb[2], in->nb[3], in->nb[2]));
+ ggml_tensor * previous = ggml_cont(ctx0, ggml_view_4d(ctx0, out, factor_dim, top_k, slots - 1, n_blocks,
+ out->nb[1], out->nb[2], out->nb[3], 0));
+ ggml_tensor * score = ggml_scale(ctx0, ggml_mul_mat(ctx0, next, previous), hparams.lilicorr_logit_scale);
+ scores = ggml_concat(ctx0, scores,
+ ggml_reshape_3d(ctx0, score, top_k * top_k, slots - 1, n_blocks), 1);
+ }
+
+ ggml_tensor * row = ggml_concat(ctx0, ggml_cast(ctx0, ggml_reshape_3d(ctx0, cand_run, top_k, slots, n_blocks),
+ GGML_TYPE_F32), scores, 0);
+ ggml_tensor * packed = ggml_pad_ext(ctx0, row, 0, n_embd - row_used, 1, 0, 0, 0, 0, 0);
+ packed = ggml_reshape_2d(ctx0, packed, n_embd, n_tokens);
+ g.cb(packed, "lilicorr_lattice", -1);
+ res->t_h_nextn = packed;
+ ggml_build_forward_expand(g.gf, packed);
+}
+
// DFlash decoder, dual-mode by batch type:
// * embd batch -> fused target features: project + inject K/V into the cache.
// * token batch -> noise-block diffusion: attend over [committed, MASK...] to generate draft tokens
@@ -667,6 +906,10 @@ llama_model_dflash::graph<false>::graph(const llama_model & model, const llm_gra
}
res->t_embd = inp_g;
+ if (model.lilicorr_pass_hidden_proj) {
+ // Expose context features for retrieval by LiliCorr
+ res->t_h_nextn = inp_g;
+ }
ggml_build_forward_expand(gf, inp_g);
return;
@@ -682,22 +925,31 @@ llama_model_dflash::graph<false>::graph(const llama_model & model, const llm_gra
tok_embd = model_other->tok_embd;
}
- auto inp = std::make_unique<llm_graph_input_embd>(n_embd);
-
- inp->tokens = ggml_new_tensor_1d(ctx0, GGML_TYPE_I32, n_tokens);
- ggml_set_input(inp->tokens);
- res->t_inp_tokens = inp->tokens;
-
- ggml_tensor * inp_tokens = inp->tokens;
+ ggml_tensor * inp_tokens = ggml_new_tensor_1d(ctx0, GGML_TYPE_I32, n_tokens);
+ ggml_set_input(inp_tokens);
+ res->t_inp_tokens = inp_tokens;
+
+ ggml_tensor * inp_anchor = nullptr;
+ if (model.lilicorr_pass_hidden_proj && ubatch.embd_state) {
+ // LiliCorr sets inp_anchor to pass the context features to the lattice scorer
+ auto inp = std::make_unique<llm_graph_input_embd_h>(n_embd, n_embd);
+ inp->tokens = inp_tokens;
+ inp->h = ggml_new_tensor_2d(ctx0, GGML_TYPE_F32, n_embd, n_tokens);
+ ggml_set_input(inp->h);
+ inp_anchor = inp->h;
+ res->add_input(std::move(inp));
+ } else {
+ auto inp = std::make_unique<llm_graph_input_embd>(n_embd);
+ inp->tokens = inp_tokens;
+ res->add_input(std::move(inp));
+ }
- ggml_tensor * inpL = ggml_get_rows(ctx0, tok_embd, inp->tokens);
+ ggml_tensor * inpL = ggml_get_rows(ctx0, tok_embd, inp_tokens);
if (hparams.f_embedding_scale != 0.0f) {
inpL = ggml_scale(ctx0, inpL, hparams.f_embedding_scale);
}
cb(inpL, "inp_noise_embd", -1);
- res->add_input(std::move(inp));
-
for (int il = 0; il < n_layer; ++il) {
const auto & layer = model.layers[il];
@@ -848,6 +1100,8 @@ llama_model_dflash::graph<false>::graph(const llama_model & model, const llm_gra
if (model.dflash_selector_hidden) {
build_dflash2_selector(*this, model, inp_tokens);
+ } else if (model.lilicorr_pass_hidden_proj && ubatch.embd_state) {
+ build_lilicorr_scores(*this, model, tok_embd, inp_anchor);
}
}