darthrubik/llama-cpp
1
0
Fork 0
Code Issues Pull Requests 1 Packages Projects Releases Wiki Activity

spec : add lilicorr support (#30506)
Copilot Setup Steps / copilot-setup-steps (push) Has been cancelled
Check Pre-Tokenizer Hashes / pre-tokenizer-hashes (push) Has been cancelled
Python check requirements.txt / check-requirements (push) Has been cancelled
Python Type-Check / python type-check (push) Has been cancelled
Update Operations Documentation / update-ops-docs (push) Has been cancelled

Browse Source 
* 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>
This commit is contained in:
Kartik Gulia
2026-10-15 09:19:25 +05:30
committed by GitHub
parent 3e390af18d
commit 843e7dcc15
13 changed files with 648 additions and 21 deletions
Download Patch File Download Diff File Expand all files Collapse all files
common/speculative.cpp
+50 -9
View File
@@ -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) {
conversion/__init__.py
+1
View File
@@ -59,6 +59,7 @@ TEXT_MODEL_MAP: dict[str, str] = {
"DeepseekV32ForCausalLM": "deepseek",
"DFlashDraftModel": "qwen",
"DFlash2DraftModel": "qwen",
"LiLiCorrDraftModel": "nemotron",
"Qwen3DSparkModel": "qwen",
"DSparkDraftModel": "qwen",
"DSparkSpeculator": "qwen",
conversion/nemotron.py
+100 -1
View File
@@ -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(
gguf-py/gguf/constants.py
+49
View File
@@ -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,
gguf-py/gguf/gguf_writer.py
+24
View File
@@ -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)
gguf-py/gguf/tensor_mapping.py
+47
View File
@@ -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
include/llama.h
+2 -1
View File
@@ -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,
src/llama-arch.cpp
+36
View File
@@ -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}},
};
src/llama-arch.h
+21
View File
@@ -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,
};
src/llama-batch.cpp
+4
View File
@@ -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;
}
src/llama-hparams.h
+6
View File
@@ -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;
src/llama-model.h
+44
View File
@@ -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;
src/models/dflash.cpp
+264 -10
View File
@@ -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;
}
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));
}
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 * 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);
}
}