From 843e7dcc15098dd6b3aaf33b4847052293d8ef1a Mon Sep 17 00:00:00 2001 From: Kartik Gulia Date: Thu, 15 Oct 2026 09:19:25 +0530 Subject: [PATCH] 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 * typo fix --------- Co-authored-by: Georgi Gerganov --- common/speculative.cpp | 59 +++++-- conversion/__init__.py | 1 + conversion/nemotron.py | 101 +++++++++++- gguf-py/gguf/constants.py | 49 ++++++ gguf-py/gguf/gguf_writer.py | 24 +++ gguf-py/gguf/tensor_mapping.py | 47 ++++++ include/llama.h | 3 +- src/llama-arch.cpp | 36 +++++ src/llama-arch.h | 21 +++ src/llama-batch.cpp | 4 + src/llama-hparams.h | 6 + src/llama-model.h | 44 ++++++ src/models/dflash.cpp | 274 +++++++++++++++++++++++++++++++-- 13 files changed, 648 insertions(+), 21 deletions(-) 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 lilicorr_embd; + std::vector 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_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_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_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 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(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::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::graph(const llama_model & model, const llm_gra tok_embd = model_other->tok_embd; } - auto inp = std::make_unique(n_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; - 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(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(n_embd); + inp->tokens = inp_tokens; + res->add_input(std::move(inp)); + } - 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::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); } }