diff --git a/conversion/__init__.py b/conversion/__init__.py index 06c2c50ad..b13483f9f 100644 --- a/conversion/__init__.py +++ b/conversion/__init__.py @@ -27,6 +27,8 @@ TEXT_MODEL_MAP: dict[str, str] = { "BaichuanForCausalLM": "baichuan", "BailingMoeForCausalLM": "bailingmoe", "BailingMoeV2ForCausalLM": "bailingmoe", + "BailingMoeV3ForCausalLM": "bailing_hybrid", + "BailingMoeV3Model": "bailing_hybrid", "BambaForCausalLM": "granite", "BertForMaskedLM": "bert", "BertForSequenceClassification": "bert", diff --git a/conversion/bailing_hybrid.py b/conversion/bailing_hybrid.py new file mode 100644 index 000000000..c2765c146 --- /dev/null +++ b/conversion/bailing_hybrid.py @@ -0,0 +1,221 @@ +from __future__ import annotations + +import math +from typing import Iterable, TYPE_CHECKING + +import torch + +if TYPE_CHECKING: + from torch import Tensor + +from .base import ModelBase, TextModel, gguf, logger + + +@ModelBase.register("BailingMoeV3Model", "BailingMoeV3ForCausalLM") +class BailingHybridModel(TextModel): + """Ling 3.0 flash (inclusionAI): hybrid KDA + gated MLA MoE, `bailing_hybrid`. + + NOT the `bailingmoe2` stack (Ling 2.0) despite the shared family name -- only + the MoE router carries over. The KDA block comes from Kimi Linear and the MLA + block from Kimi's no-Q-compression variant, so conversion mirrors + conversion/kimi_linear.py. The differences that matter here: + + * The attention module is `attention.`, not `self_attn.`. None of Kimi's + tensor mappings match; bailing-hybrid entries were added alongside them. + + * `attention.g_proj.weight` exists on BOTH layer types with different + shapes and different meanings: on KDA layers it is the full-rank output + gate {n_embd, d_inner}, on MLA layers the head-wise attention gate + {n_embd, n_head}. A name->enum table cannot express that, so the KDA one + is renamed to `g_full_proj` here. Without this the loader silently binds + the wrong tensor. + + * A_log is stored as +exp(A_log), NOT Kimi's -exp(A_log). config sets + kda_safe_gate=true, which changes the decay to + g = kda_lower_bound * sigmoid(exp(A_log) * (f(x) + dt_bias)) + so the sign lives in kda_lower_bound (-5.0), written as a KV below. + Verified against fla ops/kda/gate.py (naive ref and Triton kernel agree). + + * `no_kda_lora: true` -> full-rank f_proj / g_proj, so SSM_F / SSM_G + replace Kimi's SSM_F_{A,B} / SSM_G_{A,B} pairs. + + Config fields that look load-bearing and are NOT (verified by grepping + modeling_bailing_moe_v3.py): expert_swiglu_limit_list and + share_expert_swiglu_limit_list (populated with non-zero values for the last + few layers, yet BailingMoeV3MLP.forward is a plain SwiGLU), use_qk_norm, + linear_silu, group_norm_size, max_window_layers, mtp_use_kda, use_mla_nope, + use_nGPT, scale_router_input, seq_aux. partial_rotary_factor is overwritten + to 1.0 by the rotary module itself, so rotary_dim == qk_rope_head_dim. + """ + + model_arch = gguf.MODEL_ARCH.BAILING_HYBRID + + _experts: list[dict[str, Tensor]] | None = None + + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + # the MTP/nextn head is a real block in the checkpoint; include it unless + # --no-mtp, matching glm/command_r. llama.cpp marks it TENSOR_SKIP. + if (n_nextn := int(self.hparams.get("num_nextn_predict_layers", 0) or 0)) > 0 and not self.no_mtp: + self.block_count = self.hparams["num_hidden_layers"] + n_nextn + # tensor_map was built from the old block_count in super().__init__(), + # so it must be rebuilt or every layer-42 tensor fails to map. + self.tensor_map = gguf.get_tensor_name_map(self.model_arch, self.block_count) + + def _is_kda_layer(self, bid: int) -> bool: + """KDA everywhere except the last layer of each group, and the MTP head. + + Mirrors modeling_bailing_moe_v3.py:1006 -- + MLA if (layer_idx + 1) % layer_group_size == 0 + or layer_idx >= num_hidden_layers // group_size * group_size + The second clause is what puts the MTP head (layer 42) on MLA. + """ + group = self.hparams["layer_group_size"] + n_layer = self.hparams["num_hidden_layers"] + is_mla = ((bid + 1) % group == 0) or (bid >= n_layer // group * group) + return not is_mla + + def set_gguf_parameters(self): + hparams = self.hparams + + # MLA KV cache requires the attention be converted to MQA (1 KV group). + hparams["num_key_value_heads"] = 1 + + super().set_gguf_parameters() + self.gguf_writer.add_vocab_size(hparams["vocab_size"]) + + assert hparams.get("no_kda_lora"), \ + "no_kda_lora is false: this checkpoint uses low-rank KDA gates, which " \ + "map to SSM_F_A/SSM_F_B (kimi-linear), not SSM_F/SSM_G" + assert hparams.get("kda_safe_gate"), \ + "kda_safe_gate is false: llama.cpp's bailing-hybrid graph only implements " \ + "the safe-gate decay form" + + # Per-layer KV head count: 0 marks a KDA (recurrent) layer, which is how + # llama.cpp tells the two branches apart. + # NOTE: this array must be block_count long, NOT num_hidden_layers -- the + # loader validates it against n_layer_all and rejects the model outright + # if the MTP/nextn block has no entry. The MTP head is MLA, so it gets 1. + _num_kv_heads = [0 if self._is_kda_layer(il) else 1 for il in range(self.block_count)] + assert any(_num_kv_heads), "no MLA layers found -- layer_group_size indexing is wrong" + assert len(_num_kv_heads) == self.block_count + self.gguf_writer.add_head_count_kv(_num_kv_heads) + logger.info(f"bailing-hybrid: {sum(1 for x in _num_kv_heads if x)} MLA / " + f"{sum(1 for x in _num_kv_heads if not x)} KDA layers") + + # ---- KDA ---- + self.gguf_writer.add_ssm_conv_kernel(hparams["short_conv_kernel_size"]) + self.gguf_writer.add_kda_head_dim(hparams["head_dim"]) + self.gguf_writer.add_kda_lower_bound(float(hparams["kda_lower_bound"])) + + # ---- MLA ---- + # q_lora_rank is null (no Q compression), so add_q_lora_rank is skipped. + assert hparams.get("q_lora_rank") is None, \ + "q_lora_rank is set: the graph builds a single wide q_proj and has no q_a/q_b path" + kv_lora_rank = hparams["kv_lora_rank"] + qk_rope_head_dim = hparams["qk_rope_head_dim"] + qk_nope_head_dim = hparams["qk_nope_head_dim"] + self.gguf_writer.add_kv_lora_rank(kv_lora_rank) + self.gguf_writer.add_rope_dimension_count(qk_rope_head_dim) + self.gguf_writer.add_key_length(kv_lora_rank + qk_rope_head_dim) + self.gguf_writer.add_key_length_mla(qk_nope_head_dim + qk_rope_head_dim) + self.gguf_writer.add_value_length_mla(hparams["v_head_dim"]) + + # ---- MoE (noaux_tc grouped top-k, bit-exact bailingmoe2) ---- + self.gguf_writer.add_expert_feed_forward_length(hparams["moe_intermediate_size"]) + self.gguf_writer.add_expert_shared_count(hparams["num_shared_experts"]) + self.gguf_writer.add_leading_dense_block_count(hparams["first_k_dense_replace"]) + self.gguf_writer.add_expert_weights_scale(hparams["routed_scaling_factor"]) + self.gguf_writer.add_expert_weights_norm(hparams["norm_topk_prob"]) + self.gguf_writer.add_expert_group_count(hparams["n_group"]) + self.gguf_writer.add_expert_group_used_count(hparams["topk_group"]) + + score = hparams.get("score_function", hparams.get("scoring_func")) + assert score == "sigmoid", f"unexpected router score function {score!r}" + self.gguf_writer.add_expert_gating_func(gguf.ExpertGatingFuncType.SIGMOID) + + if (n_nextn := int(hparams.get("num_nextn_predict_layers", 0) or 0)) > 0 and not self.no_mtp: + self.gguf_writer.add_nextn_predict_layers(n_nextn) + + def prepare_tensors(self): + super().prepare_tensors() + if self._experts is not None: + experts = [k for d in self._experts for k in d.keys()] + if len(experts) > 0: + raise ValueError(f"Unprocessed experts: {experts}") + + def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iterable[tuple[str, Tensor]]: + # KDA conv1d: HF [d_inner, d_conv] -> numpy (1, d_inner, 1, d_conv), + # which GGUF reverses into ggml ne = [d_conv, 1, d_inner, 1]. Memory + # layout is preserved either way (d_conv changes fastest). + if name.endswith((".q_conv1d.weight", ".k_conv1d.weight", ".v_conv1d.weight")): + if data_torch.ndim == 2: + d_inner, d_conv = data_torch.shape + data_torch = data_torch.reshape(1, d_inner, 1, d_conv) + elif data_torch.ndim == 3: + d_inner, _, d_conv = data_torch.shape + data_torch = data_torch.reshape(1, d_inner, 1, d_conv) + + # A_log is 1-D [n_head] here (kimi's is [1,H,1,1], solar_open2's + # [1,1,64,1] -- third layout in three ports). Store +exp(A_log): the + # negation lives in kda_lower_bound, unlike kimi which bakes in -exp(). + if name.endswith(".A_log"): + data_torch = torch.exp(data_torch.float()) + data_torch = data_torch.reshape(1, 1, -1, 1) + + if name.endswith(".dt_bias"): + name = name.rpartition(".dt_bias")[0] + ".dt_proj.bias" + + # llama.cpp asks for `blk.N.exp_probs_b.bias`, but `mlp.gate.expert_bias` + # has no .weight/.bias suffix for map_tensor_name to strip, so it would be + # written as a bare `blk.N.exp_probs_b` and the load fails on a missing + # tensor. Same fix as bailingmoe/afmoe/grovemoe. + if name.endswith(".expert_bias"): + name = name.replace(".expert_bias", ".expert_bias.bias") + + # Disambiguate the two g_proj tensors (see the class docstring). + if name.endswith(".attention.g_proj.weight"): + assert bid is not None + if self._is_kda_layer(bid): + name = name.replace(".attention.g_proj.", ".attention.g_full_proj.") + + # merge the routed experts into one 3-D tensor per projection + if ".mlp.experts." in name: + n_experts = self.hparams["num_experts"] + assert bid is not None + + if self._experts is None: + self._experts = [{} for _ in range(self.block_count)] + + self._experts[bid][name] = data_torch + + if len(self._experts[bid]) >= n_experts * 3: + for wid, tname in [("gate_proj", gguf.MODEL_TENSOR.FFN_GATE_EXP), + ("down_proj", gguf.MODEL_TENSOR.FFN_DOWN_EXP), + ("up_proj", gguf.MODEL_TENSOR.FFN_UP_EXP)]: + datas: list[Tensor] = [] + for xid in range(n_experts): + ename = f"model.layers.{bid}.mlp.experts.{xid}.{wid}.weight" + datas.append(self._experts[bid][ename]) + del self._experts[bid][ename] + data_torch = torch.stack(datas, dim=0) + new_name = self.format_tensor_name(tname, bid) + yield from super().modify_tensors(data_torch, new_name, bid) + return + + # MLA absorption needs kv_b split, with k_b transposed + if name.endswith("kv_b_proj.weight"): + name_kb = name.replace("kv_b_proj", "k_b_proj") + name_vb = name.replace("kv_b_proj", "v_b_proj") + n_head_kv = self.hparams["num_attention_heads"] + v_head_dim = self.hparams["v_head_dim"] + qk_nope_head_dim = self.hparams["qk_nope_head_dim"] + assert data_torch.shape[0] == n_head_kv * (v_head_dim + qk_nope_head_dim) + kv_b = data_torch.view(n_head_kv, v_head_dim + qk_nope_head_dim, data_torch.shape[-1]) + k_b, v_b = torch.split(kv_b, [qk_nope_head_dim, v_head_dim], dim=1) + k_b = k_b.transpose(1, 2) + yield from super().modify_tensors(k_b, name_kb, bid) + yield from super().modify_tensors(v_b, name_vb, bid) + return + + yield from super().modify_tensors(data_torch, name, bid) diff --git a/gguf-py/gguf/constants.py b/gguf-py/gguf/constants.py index 8516222cc..fb6b8a038 100644 --- a/gguf-py/gguf/constants.py +++ b/gguf-py/gguf/constants.py @@ -245,7 +245,9 @@ class Keys: DT_B_C_RMS = "{arch}.ssm.dt_b_c_rms" class KDA: - HEAD_DIM = "{arch}.kda.head_dim" + HEAD_DIM = "{arch}.kda.head_dim" + # bailing-hybrid safe-gate: g = LOWER_BOUND * sigmoid(exp(A_log) * (f(x) + dt_bias)) + LOWER_BOUND = "{arch}.kda.lower_bound" class WKV: HEAD_SIZE = "{arch}.wkv.head_size" @@ -568,6 +570,7 @@ class MODEL_ARCH(IntEnum): LLAMA_EMBED = auto() MAINCODER = auto() KIMI_LINEAR = auto() + BAILING_HYBRID = auto() TALKIE = auto() MELLUM = auto() NANBEIGE = auto() @@ -685,6 +688,8 @@ class MODEL_TENSOR(IntEnum): SSM_BETA = auto() # Kimi Linear qwen3.5 SSM_G_A = auto() # Kimi Linear SSM_G_B = auto() # Kimi Linear + SSM_F = auto() # bailing-hybrid (full-rank forget gate) + SSM_G = auto() # bailing-hybrid (full-rank output gate) TIME_MIX_W0 = auto() TIME_MIX_W1 = auto() TIME_MIX_W2 = auto() @@ -1240,6 +1245,7 @@ MODEL_ARCH_NAMES: dict[MODEL_ARCH, str] = { MODEL_ARCH.LLAMA_EMBED: "llama-embed", MODEL_ARCH.MAINCODER: "maincoder", MODEL_ARCH.KIMI_LINEAR: "kimi-linear", + MODEL_ARCH.BAILING_HYBRID: "bailing-hybrid", MODEL_ARCH.TALKIE: "talkie", MODEL_ARCH.MELLUM: "mellum", MODEL_ARCH.NANBEIGE: "nanbeige", @@ -1355,6 +1361,8 @@ TENSOR_NAMES: dict[MODEL_TENSOR, str] = { MODEL_TENSOR.SSM_BETA: "blk.{bid}.ssm_beta", # Kimi Linear qwen3.5 MODEL_TENSOR.SSM_G_A: "blk.{bid}.ssm_g_a", # Kimi Linear MODEL_TENSOR.SSM_G_B: "blk.{bid}.ssm_g_b", # Kimi Linear + MODEL_TENSOR.SSM_F: "blk.{bid}.ssm_f", # bailing-hybrid + MODEL_TENSOR.SSM_G: "blk.{bid}.ssm_g", # bailing-hybrid MODEL_TENSOR.TIME_MIX_W0: "blk.{bid}.time_mix_w0", MODEL_TENSOR.TIME_MIX_W1: "blk.{bid}.time_mix_w1", MODEL_TENSOR.TIME_MIX_W2: "blk.{bid}.time_mix_w2", @@ -4749,6 +4757,47 @@ MODEL_TENSORS: dict[MODEL_ARCH, list[MODEL_TENSOR]] = { MODEL_TENSOR.FFN_DOWN, MODEL_TENSOR.FFN_UP, ], + MODEL_ARCH.BAILING_HYBRID: [ + MODEL_TENSOR.TOKEN_EMBD, + MODEL_TENSOR.OUTPUT_NORM, + MODEL_TENSOR.OUTPUT, + MODEL_TENSOR.ATTN_NORM, + MODEL_TENSOR.ATTN_Q, + MODEL_TENSOR.ATTN_K, + MODEL_TENSOR.ATTN_V, + MODEL_TENSOR.ATTN_OUT, + MODEL_TENSOR.ATTN_GATE, + MODEL_TENSOR.ATTN_KV_A_MQA, + MODEL_TENSOR.ATTN_KV_A_NORM, + MODEL_TENSOR.ATTN_KV_B, + MODEL_TENSOR.ATTN_K_B, + MODEL_TENSOR.ATTN_V_B, + MODEL_TENSOR.SSM_CONV1D_Q, + MODEL_TENSOR.SSM_CONV1D_K, + MODEL_TENSOR.SSM_CONV1D_V, + MODEL_TENSOR.SSM_F, + MODEL_TENSOR.SSM_G, + MODEL_TENSOR.SSM_BETA, + MODEL_TENSOR.SSM_A, + MODEL_TENSOR.SSM_DT, + MODEL_TENSOR.SSM_NORM, + MODEL_TENSOR.FFN_NORM, + MODEL_TENSOR.FFN_GATE, + MODEL_TENSOR.FFN_DOWN, + MODEL_TENSOR.FFN_UP, + MODEL_TENSOR.FFN_GATE_INP, + MODEL_TENSOR.FFN_GATE_EXP, + MODEL_TENSOR.FFN_DOWN_EXP, + MODEL_TENSOR.FFN_UP_EXP, + MODEL_TENSOR.FFN_EXP_PROBS_B, + MODEL_TENSOR.FFN_GATE_SHEXP, + MODEL_TENSOR.FFN_DOWN_SHEXP, + MODEL_TENSOR.FFN_UP_SHEXP, + MODEL_TENSOR.NEXTN_EH_PROJ, + MODEL_TENSOR.NEXTN_ENORM, + MODEL_TENSOR.NEXTN_HNORM, + MODEL_TENSOR.NEXTN_SHARED_HEAD_NORM, + ], MODEL_ARCH.KIMI_LINEAR: [ MODEL_TENSOR.TOKEN_EMBD, MODEL_TENSOR.OUTPUT_NORM, diff --git a/gguf-py/gguf/gguf_writer.py b/gguf-py/gguf/gguf_writer.py index 39da9f2c0..ff6960091 100644 --- a/gguf-py/gguf/gguf_writer.py +++ b/gguf-py/gguf/gguf_writer.py @@ -1091,6 +1091,9 @@ class GGUFWriter: def add_kda_head_dim(self, value: int) -> None: self.add_uint32(Keys.KDA.HEAD_DIM.format(arch=self.arch), value) + def add_kda_lower_bound(self, value: float) -> None: + self.add_float32(Keys.KDA.LOWER_BOUND.format(arch=self.arch), 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 7892342e4..5618ca1c1 100644 --- a/gguf-py/gguf/tensor_mapping.py +++ b/gguf-py/gguf/tensor_mapping.py @@ -269,6 +269,7 @@ class TensorNameMap: "layers.{bid}.self_attn.q_proj", # qwen3-embedding "backbone.layers.{bid}.mixer.q_proj", # nemotron-h "model.blocks.{bid}.attn.attn_query", # talkie + "model.layers.{bid}.attention.q_proj", # bailing-hybrid ), # Attention key @@ -290,6 +291,7 @@ class TensorNameMap: "layers.{bid}.self_attn.k_proj", # qwen3-embedding "backbone.layers.{bid}.mixer.k_proj", # nemotron-h "model.blocks.{bid}.attn.attn_key", # talkie + "model.layers.{bid}.attention.k_proj", # bailing-hybrid ), # Attention value @@ -310,6 +312,7 @@ class TensorNameMap: "layers.{bid}.self_attn.v_proj", # qwen3-embedding "backbone.layers.{bid}.mixer.v_proj", # nemotron-h "model.blocks.{bid}.attn.attn_value", # talkie + "model.layers.{bid}.attention.v_proj", # bailing-hybrid ), # Attention output @@ -349,6 +352,7 @@ class TensorNameMap: "backbone.layers.{bid}.mixer.o_proj", # nemotron-h "model.layers.{bid}.self_attn.language_expert_dense", # cogvlm "model.blocks.{bid}.attn.attn_resid", # talkie + "model.layers.{bid}.attention.o_proj", # bailing-hybrid ), # Attention output norm @@ -385,6 +389,7 @@ class TensorNameMap: "model.layers.{bid}.self_attn.gate_proj", # afmoe "model.layers.{bid}.linear_attn.in_proj_z", # qwen3.5 "model.layers.{bid}.self_attn.g_proj", # step3.5 head-wise attention gate + "model.layers.{bid}.attention.g_proj", # bailing-hybrid ), # Feed-forward norm @@ -832,6 +837,7 @@ class TensorNameMap: "model.layers.{bid}.linear_attn.dt_proj", # qwen3next "backbone.layers.{bid}.mixer.dt", # nemotron-h-moe "model.layers.{bid}.self_attn.dt_proj", # kimi + "model.layers.{bid}.attention.dt_proj", # bailing-hybrid ), MODEL_TENSOR.SSM_DT_NORM: ( @@ -846,6 +852,7 @@ class TensorNameMap: "model.layers.layers.{bid}.mixer.A_log", # plamo2 "model.layers.{bid}.linear_attn.A_log", # qwen3next "model.layers.{bid}.self_attn.A_log", # kimi + "model.layers.{bid}.attention.A_log", # bailing-hybrid ), MODEL_TENSOR.SSM_B_NORM: ( @@ -872,6 +879,7 @@ class TensorNameMap: "model.layers.{bid}.linear_attn.norm", # qwen3next "backbone.layers.{bid}.mixer.norm", # mamba2 "model.layers.{bid}.self_attn.o_norm", # kimi + "model.layers.{bid}.attention.o_norm", # bailing-hybrid ), MODEL_TENSOR.SSM_OUT: ( @@ -893,12 +901,15 @@ class TensorNameMap: # Kimi Linear KDA (using SSM_ prefix for consistency) MODEL_TENSOR.SSM_CONV1D_Q: ( "model.layers.{bid}.self_attn.q_conv1d", + "model.layers.{bid}.attention.q_conv1d", # bailing-hybrid ), MODEL_TENSOR.SSM_CONV1D_K: ( "model.layers.{bid}.self_attn.k_conv1d", + "model.layers.{bid}.attention.k_conv1d", # bailing-hybrid ), MODEL_TENSOR.SSM_CONV1D_V: ( "model.layers.{bid}.self_attn.v_conv1d", + "model.layers.{bid}.attention.v_conv1d", # bailing-hybrid ), MODEL_TENSOR.SSM_F_A: ( "model.layers.{bid}.self_attn.f_a_proj", @@ -909,6 +920,7 @@ class TensorNameMap: MODEL_TENSOR.SSM_BETA: ( "model.layers.{bid}.linear_attn.in_proj_b", # qwen3.5 "model.layers.{bid}.self_attn.b_proj", # Kimi Linear + "model.layers.{bid}.attention.b_proj", # bailing-hybrid ), MODEL_TENSOR.SSM_G_A: ( "model.layers.{bid}.self_attn.g_a_proj", @@ -916,6 +928,12 @@ class TensorNameMap: MODEL_TENSOR.SSM_G_B: ( "model.layers.{bid}.self_attn.g_b_proj", ), + MODEL_TENSOR.SSM_F: ( + "model.layers.{bid}.attention.f_proj", # bailing-hybrid (full-rank) + ), + MODEL_TENSOR.SSM_G: ( + "model.layers.{bid}.attention.g_full_proj", # bailing-hybrid (renamed by converter) + ), MODEL_TENSOR.TIME_MIX_W0: ( "model.layers.{bid}.attention.w0", # rwkv7 ), @@ -1099,20 +1117,24 @@ class TensorNameMap: MODEL_TENSOR.ATTN_KV_A_MQA: ( "model.layers.{bid}.self_attn.kv_a_proj_with_mqa", # deepseek2 "layers.{bid}.attention.wkv_a_with_mqa", # mistral-large + "model.layers.{bid}.attention.kv_a_proj_with_mqa", # bailing-hybrid ), MODEL_TENSOR.ATTN_KV_B: ( "model.layers.{bid}.self_attn.kv_b_proj", # deepseek2 + "model.layers.{bid}.attention.kv_b_proj", # bailing-hybrid ), MODEL_TENSOR.ATTN_K_B: ( "model.layers.{bid}.self_attn.k_b_proj", # deepseek2 "layers.{bid}.attention.k_b_proj", # mistral-large + "model.layers.{bid}.attention.k_b_proj", # bailing-hybrid ), MODEL_TENSOR.ATTN_V_B: ( "model.layers.{bid}.self_attn.v_b_proj", # deepseek2 "layers.{bid}.attention.v_b_proj", # mistral-large + "model.layers.{bid}.attention.v_b_proj", # bailing-hybrid ), MODEL_TENSOR.ATTN_Q_A_NORM: ( @@ -1123,6 +1145,7 @@ class TensorNameMap: MODEL_TENSOR.ATTN_KV_A_NORM: ( "model.layers.{bid}.self_attn.kv_a_layernorm", # deepseek2 "layers.{bid}.attention.kv_a_norm", # mistral-large + "model.layers.{bid}.attention.kv_a_layernorm", # bailing-hybrid ), MODEL_TENSOR.ATTN_SUB_NORM: ( @@ -2564,6 +2587,7 @@ class TensorNameMap: MODEL_TENSOR.NEXTN_SHARED_HEAD_NORM: ( "model.layers.{bid}.shared_head.norm", + "model.layers.{bid}.final_layernorm", # bailing-hybrid ), } diff --git a/src/llama-arch.cpp b/src/llama-arch.cpp index 836cfade2..a942d3bf0 100644 --- a/src/llama-arch.cpp +++ b/src/llama-arch.cpp @@ -141,6 +141,7 @@ static const std::map LLM_ARCH_NAMES = { { LLM_ARCH_LLAMA_EMBED, "llama-embed" }, { LLM_ARCH_MAINCODER, "maincoder" }, { LLM_ARCH_KIMI_LINEAR, "kimi-linear" }, + { LLM_ARCH_BAILING_HYBRID, "bailing-hybrid" }, { LLM_ARCH_TALKIE, "talkie" }, { LLM_ARCH_MELLUM, "mellum" }, { LLM_ARCH_NANBEIGE, "nanbeige" }, @@ -303,7 +304,8 @@ static const std::map LLM_KV_NAMES = { { LLM_KV_SSM_GROUP_COUNT, "%s.ssm.group_count" }, { LLM_KV_SSM_DT_B_C_RMS, "%s.ssm.dt_b_c_rms" }, - { LLM_KV_KDA_HEAD_DIM, "%s.kda.head_dim" }, + { LLM_KV_KDA_HEAD_DIM, "%s.kda.head_dim" }, + { LLM_KV_KDA_LOWER_BOUND, "%s.kda.lower_bound" }, { LLM_KV_WKV_HEAD_SIZE, "%s.wkv.head_size" }, @@ -455,6 +457,8 @@ static const std::map LLM_TENSOR_NAMES = { { LLM_TENSOR_SSM_BETA, "blk.%d.ssm_beta" }, { LLM_TENSOR_SSM_G_A, "blk.%d.ssm_g_a" }, { LLM_TENSOR_SSM_G_B, "blk.%d.ssm_g_b" }, + { LLM_TENSOR_SSM_F, "blk.%d.ssm_f" }, + { LLM_TENSOR_SSM_G, "blk.%d.ssm_g" }, { LLM_TENSOR_SSM_NORM, "blk.%d.ssm_norm" }, { LLM_TENSOR_ATTN_Q_A_NORM, "blk.%d.attn_q_a_norm" }, { LLM_TENSOR_ATTN_KV_A_NORM, "blk.%d.attn_kv_a_norm" }, @@ -748,6 +752,8 @@ static const std::map LLM_TENSOR_INFOS = { {LLM_TENSOR_SSM_BETA, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL_MAT}}, {LLM_TENSOR_SSM_G_A, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL_MAT}}, {LLM_TENSOR_SSM_G_B, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL_MAT}}, + {LLM_TENSOR_SSM_F, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL_MAT}}, + {LLM_TENSOR_SSM_G, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL_MAT}}, {LLM_TENSOR_TIME_MIX_LERP_X, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL}}, {LLM_TENSOR_TIME_MIX_LN, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL}}, {LLM_TENSOR_CHANNEL_MIX_LERP_K, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL}}, @@ -967,6 +973,7 @@ bool llm_arch_is_hybrid(const llm_arch & arch) { case LLM_ARCH_NEMOTRON_H_MOE: case LLM_ARCH_QWEN3NEXT: case LLM_ARCH_KIMI_LINEAR: + case LLM_ARCH_BAILING_HYBRID: case LLM_ARCH_QWEN35: case LLM_ARCH_QWEN35MOE: case LLM_ARCH_DEEPSEEK4: @@ -1027,6 +1034,7 @@ bool llm_arch_supports_sm_tensor(const llm_arch & arch) { case LLM_ARCH_MINIMAX_M3: case LLM_ARCH_MISTRAL4: case LLM_ARCH_KIMI_LINEAR: + case LLM_ARCH_BAILING_HYBRID: case LLM_ARCH_QWEN3TTS: return false; default: diff --git a/src/llama-arch.h b/src/llama-arch.h index 49c2a6ac3..e16e16a8c 100644 --- a/src/llama-arch.h +++ b/src/llama-arch.h @@ -143,6 +143,7 @@ enum llm_arch { LLM_ARCH_LLAMA_EMBED, LLM_ARCH_MAINCODER, LLM_ARCH_KIMI_LINEAR, + LLM_ARCH_BAILING_HYBRID, LLM_ARCH_TALKIE, LLM_ARCH_MELLUM, LLM_ARCH_EAGLE3, @@ -309,6 +310,7 @@ enum llm_kv { LLM_KV_SSM_DT_B_C_RMS, LLM_KV_KDA_HEAD_DIM, + LLM_KV_KDA_LOWER_BOUND, LLM_KV_WKV_HEAD_SIZE, @@ -483,6 +485,8 @@ enum llm_tensor { LLM_TENSOR_SSM_BETA, // kimi: beta mixing coefficient and qwen3.5 LLM_TENSOR_SSM_G_A, // kimi: output gate projection A LLM_TENSOR_SSM_G_B, // kimi: output gate projection B + LLM_TENSOR_SSM_F, // bailing-hybrid: full-rank forget gate (no_kda_lora) + LLM_TENSOR_SSM_G, // bailing-hybrid: full-rank output gate (no_kda_lora) LLM_TENSOR_TIME_MIX_W0, LLM_TENSOR_TIME_MIX_W1, LLM_TENSOR_TIME_MIX_W2, diff --git a/src/llama-hparams.h b/src/llama-hparams.h index 6e8336c98..ac32aaa27 100644 --- a/src/llama-hparams.h +++ b/src/llama-hparams.h @@ -163,6 +163,10 @@ struct llama_hparams { // for Kimi Linear KDA uint32_t n_embd_head_kda = 0; + // bailing-hybrid KDA safe gate. 0.0f means "not a safe-gate model", i.e. use + // the kimi form g = -exp(A_log)*softplus(.) instead. + float f_kda_lower_bound = 0.0f; + bool ssm_dt_b_c_rms = false; float f_clamp_kqv = 0.0f; diff --git a/src/llama-model.cpp b/src/llama-model.cpp index dda311c47..02ed42bcc 100644 --- a/src/llama-model.cpp +++ b/src/llama-model.cpp @@ -312,6 +312,8 @@ static llama_model * llama_model_mapping(llm_arch arch, const llama_model_params return new llama_model_mimo2(params); case LLM_ARCH_KIMI_LINEAR: return new llama_model_kimi_linear(params); + case LLM_ARCH_BAILING_HYBRID: + return new llama_model_bailing_hybrid(params); case LLM_ARCH_STEP35: return new llama_model_step35(params); default: @@ -2572,6 +2574,11 @@ llama_rope_type llama_model_rope_type(const llama_model * model) { return LLAMA_ROPE_TYPE_NONE; // use what we call a normal RoPE, operating on pairs of consecutive head values + // bailing-hybrid: config rope_interleave=true. The reference de-interleaves + // (view(d/2,2).transpose) before a rotate_half, which is exactly pairwise + // rotation in stored order -- i.e. NORM, not the NEOX that DeepSeek-style + // MLA normally uses. The non-interleaved branch upstream is a literal 1/0. + case LLM_ARCH_BAILING_HYBRID: case LLM_ARCH_LLAMA: case LLM_ARCH_LLADA: case LLM_ARCH_LLAMA4: diff --git a/src/llama-model.h b/src/llama-model.h index 6b9e94a0a..8c50adb2c 100644 --- a/src/llama-model.h +++ b/src/llama-model.h @@ -510,6 +510,11 @@ struct llama_layer { struct ggml_tensor * ssm_g_b = nullptr; struct ggml_tensor * ssm_o_norm = nullptr; + // full-rank KDA forget/output gates (bailing-hybrid, no_kda_lora=true): + // single matmuls that replace the ssm_{f,g}_{a,b} low-rank pairs above + struct ggml_tensor * ssm_f = nullptr; + struct ggml_tensor * ssm_g = nullptr; + // DSA (deepseek sparse attention) struct ggml_tensor * indexer_k_norm = nullptr; struct ggml_tensor * indexer_k_norm_b = nullptr; diff --git a/src/models/bailing-hybrid.cpp b/src/models/bailing-hybrid.cpp new file mode 100644 index 000000000..a625412f0 --- /dev/null +++ b/src/models/bailing-hybrid.cpp @@ -0,0 +1,575 @@ +#include "models.h" +#include "llama-memory-recurrent.h" + +// Ling 3.0 flash (inclusionAI/Ling-3.0-flash) -- model_type "bailing_hybrid", +// BailingMoeV3ForCausalLM. 127.5B total / 5.1B active. +// +// 42 layers: 35 KDA (Kimi Delta Attention) + 7 gated MLA, MLA at every layer +// where (il + 1) % layer_group_size == 0 with layer_group_size 6, i.e. layers +// 5/11/17/23/29/35/41 -- MLA is LAST in each group (solar_open2 is the +// opposite, softmax-first). Layer 42 is an MTP/nextn head and is skipped. +// +// Derived from src/models/kimi-linear.cpp, which already has both halves: the +// KDA block, and MLA without Q compression at exactly this geometry +// (qk_rope 64 / qk_nope 128 / qk_head 192). The MoE router is bit-exact +// bailingmoe2 (noaux_tc grouped top-k + sigmoid + expert_bias), which +// build_moe_ffn handles from hparams with no code here. +// +// Deltas against kimi-linear, all verified against modeling_bailing_moe_v3.py +// and the fla kernels rather than inferred: +// +// 1. SAFE GATE. config kda_safe_gate=true, kda_lower_bound=-5.0 replaces +// g = -exp(A_log) * softplus(f(x) + dt_bias) [kimi] +// with +// g = lower_bound * sigmoid(exp(A_log) * (f(x) + dt_bias)) +// Confirmed identical in fla's naive reference (ops/kda/gate.py:57-69) and +// its Triton kernel (ops/kda/gate.py:116-119). Because g is built here and +// handed to GGML_OP_GATED_DELTA_NET as an input, the kernel is untouched -- +// same shape of fix as solar_open2's beta = 2*sigmoid(.). Note the +// converter must store +exp(A_log), NOT kimi's -exp(A_log): the sign now +// lives in lower_bound. Getting this wrong is a silent quality +// regression, never a crash. +// +// 2. no_kda_lora=true -> f_proj / g_proj are FULL-RANK {n_embd, d_inner} +// single matmuls, not kimi's low-rank f_a/f_b, g_a/g_b pairs. +// +// 3. A_log is 1-D [n_head]. Kimi's is [1,H,1,1], solar_open2's is [1,1,64,1] +// -- third layout in three ports. The converter reshapes it. +// +// 4. MLA carries a HEAD-WISE sigmoid output gate: g_proj {n_embd, n_head}, +// one scalar per head broadcast across v_head_dim, applied to the SDPA +// result before dense/o_proj. solar_open2's gate is elementwise and +// full-width -- do not copy that broadcast. +// +// 5. MLA USES RoPE, unlike kimi (rotary_emb=None there). rope_interleave=true +// resolves to llama.cpp's NORM rope: the reference de-interleaves with +// view(d/2,2).transpose before a rotate_half, which is pairwise rotation +// in stored order. theta 6e6 over the 64-dim rope slice only. +// +// Vestigial config fields, verified unreferenced by grepping the reference: +// expert_swiglu_limit_list / share_expert_swiglu_limit_list (BailingMoeV3MLP is +// a plain SwiGLU -- these are populated with non-zero values for the last few +// layers and are still dead), use_qk_norm, linear_silu, max_window_layers, +// mtp_use_kda, use_mla_nope, use_nGPT, scale_router_input, seq_aux. +// partial_rotary_factor is overwritten to 1.0 by the rotary module itself. + +void llama_model_bailing_hybrid::load_arch_hparams(llama_model_loader & ml) { + ml.get_key(LLM_KV_ATTENTION_LAYERNORM_RMS_EPS, hparams.f_norm_rms_eps); + ml.get_key(LLM_KV_ATTENTION_KEY_LENGTH_MLA, hparams.n_embd_head_k_mla_impl); + ml.get_key(LLM_KV_ATTENTION_VALUE_LENGTH_MLA, hparams.n_embd_head_v_mla_impl); + ml.get_key(LLM_KV_ATTENTION_KV_LORA_RANK, hparams.n_lora_kv); + ml.get_key(LLM_KV_SSM_CONV_KERNEL, hparams.ssm_d_conv); + ml.get_key(LLM_KV_KDA_HEAD_DIM, hparams.n_embd_head_kda); + ml.get_key(LLM_KV_KDA_LOWER_BOUND, hparams.f_kda_lower_bound, false); + + // KDA layers are marked with n_head_kv == 0 (same convention as Kimi Linear, + // solar_open2 and Jamba); MLA layers carry the real KV head count, which the + // converter forces to 1 so the MLA KV cache can be used. + for (uint32_t i = 0; i < hparams.n_layer(); ++i) { + hparams.is_recr_impl[i] = hparams.n_head_kv(i) == 0; + } + + ml.get_key(LLM_KV_EXPERT_FEED_FORWARD_LENGTH, hparams.n_ff_exp); + ml.get_key(LLM_KV_EXPERT_SHARED_COUNT, hparams.n_expert_shared); + ml.get_key(LLM_KV_LEADING_DENSE_BLOCK_COUNT, hparams.n_layer_dense_lead, false); + ml.get_key(LLM_KV_EXPERT_WEIGHTS_SCALE, hparams.expert_weights_scale, false); + ml.get_key(LLM_KV_EXPERT_WEIGHTS_NORM, hparams.expert_weights_norm, false); + ml.get_key(LLM_KV_EXPERT_GATING_FUNC, hparams.expert_gating_func); + ml.get_key(LLM_KV_NEXTN_PREDICT_LAYERS, hparams.n_layer_nextn, false); + + GGML_ASSERT(hparams.n_layer_nextn < hparams.n_layer_all && "n_layer_nextn must be < n_layer_all"); + + // The safe gate is what this arch is; a GGUF without it was converted by + // something that did not understand the model. + GGML_ASSERT(hparams.f_kda_lower_bound < 0.0f && + "bailing-hybrid requires a negative kda.lower_bound (safe gate); re-convert this model"); + + switch (hparams.n_layer()) { + case 42: type = LLM_TYPE_A13B; break; // Ling-3.0-flash 127.5B-A5.1B + default: type = LLM_TYPE_UNKNOWN; + } +} + +void llama_model_bailing_hybrid::load_arch_tensors(llama_model_loader &) { + LLAMA_LOAD_LOCALS; + + tok_embd = create_tensor(tn(LLM_TENSOR_TOKEN_EMBD, "weight"), {n_embd, n_vocab}, 0); + + output_norm = create_tensor(tn(LLM_TENSOR_OUTPUT_NORM, "weight"), {n_embd}, 0); + output = create_tensor(tn(LLM_TENSOR_OUTPUT, "weight"), {n_embd, n_vocab}, 0); + + const int64_t head_dim_kda = hparams.n_embd_head_kda; // 128 + const int64_t ssm_d_conv = hparams.ssm_d_conv; // 4 + const int64_t d_inner = head_dim_kda * n_head; // 32 * 128 = 4096 + + for (int i = 0; i < n_layer_all; ++i) { + // The MTP/nextn head (layer 42) ships in the checkpoint but is not part + // of the main forward pass -- allocate nothing for it. + const int flags = (i >= n_layer) ? TENSOR_SKIP : 0; + + auto & layer = layers[i]; + + layer.attn_norm = create_tensor(tn(LLM_TENSOR_ATTN_NORM, "weight", i), {n_embd}, flags); + + if (i < n_layer && hparams.is_recr(i)) { + // ---- KDA linear-attention layer ---- + // conv1d weights are 4D in the GGUF but quantisation may drop the + // trailing 1, so accept 3D too (same dance as kimi-linear.cpp). + layer.ssm_q_conv = create_tensor(tn(LLM_TENSOR_SSM_CONV1D_Q, "weight", i), {ssm_d_conv, 1, d_inner, 1}, TENSOR_NOT_REQUIRED); + if (!layer.ssm_q_conv) { + layer.ssm_q_conv = create_tensor(tn(LLM_TENSOR_SSM_CONV1D_Q, "weight", i), {ssm_d_conv, 1, d_inner}, 0); + } + layer.ssm_k_conv = create_tensor(tn(LLM_TENSOR_SSM_CONV1D_K, "weight", i), {ssm_d_conv, 1, d_inner, 1}, TENSOR_NOT_REQUIRED); + if (!layer.ssm_k_conv) { + layer.ssm_k_conv = create_tensor(tn(LLM_TENSOR_SSM_CONV1D_K, "weight", i), {ssm_d_conv, 1, d_inner}, 0); + } + layer.ssm_v_conv = create_tensor(tn(LLM_TENSOR_SSM_CONV1D_V, "weight", i), {ssm_d_conv, 1, d_inner, 1}, TENSOR_NOT_REQUIRED); + if (!layer.ssm_v_conv) { + layer.ssm_v_conv = create_tensor(tn(LLM_TENSOR_SSM_CONV1D_V, "weight", i), {ssm_d_conv, 1, d_inner}, 0); + } + + // num_kv_heads_for_linear_attn = 0 => K is full width, like Q/V + create_tensor_qkv(layer, i, n_embd, d_inner, d_inner, d_inner, 0); + + // full-rank forget/output gates (no_kda_lora = true) + layer.ssm_f = create_tensor(tn(LLM_TENSOR_SSM_F, "weight", i), {n_embd, d_inner}, 0); + layer.ssm_g = create_tensor(tn(LLM_TENSOR_SSM_G, "weight", i), {n_embd, d_inner}, 0); + + layer.ssm_beta = create_tensor(tn(LLM_TENSOR_SSM_BETA, "weight", i), {n_embd, n_head}, 0); + + // stored as +exp(A_log) by the converter; the negation lives in + // kda.lower_bound. Converter emits ggml ne = [1, n_head, 1, 1]. + layer.ssm_a = create_tensor(tn(LLM_TENSOR_SSM_A, i), {1, n_head, 1, 1}, TENSOR_NOT_REQUIRED); + if (!layer.ssm_a) { + layer.ssm_a = create_tensor(tn(LLM_TENSOR_SSM_A, i), {1, n_head}, 0); + } + + layer.ssm_dt_b = create_tensor(tn(LLM_TENSOR_SSM_DT, "bias", i), {d_inner}, 0); + + layer.ssm_o_norm = create_tensor(tn(LLM_TENSOR_SSM_NORM, "weight", i), {head_dim_kda}, 0); + + layer.wo = create_tensor(tn(LLM_TENSOR_ATTN_OUT, "weight", i), {d_inner, n_embd}, 0); + } else { + // ---- gated MLA layer (also the shape of the skipped MTP head) ---- + const int64_t kv_lora_rank = hparams.n_lora_kv; + const int64_t n_embd_head_k_mla = hparams.n_embd_head_k_mla(); // 192 + const int64_t n_embd_head_v_mla = hparams.n_embd_head_v_mla(); // 128 + const int64_t qk_rope_head_dim = hparams.n_rot(); // 64 + + // q_lora_rank is null in config => no Q compression, one wide q_proj + layer.wq = create_tensor(tn(LLM_TENSOR_ATTN_Q, "weight", i), {n_embd, n_head * n_embd_head_k_mla}, flags); + + layer.attn_kv_a_norm = create_tensor(tn(LLM_TENSOR_ATTN_KV_A_NORM, "weight", i), {kv_lora_rank}, flags); + layer.wkv_a_mqa = create_tensor(tn(LLM_TENSOR_ATTN_KV_A_MQA, "weight", i), {n_embd, kv_lora_rank + qk_rope_head_dim}, flags); + + // legacy GGUFs keep kv_b fused (MLA KV cache disabled) + layer.wkv_b = create_tensor(tn(LLM_TENSOR_ATTN_KV_B, "weight", i), + {kv_lora_rank, n_head * (n_embd_head_k_mla - qk_rope_head_dim + n_embd_head_v_mla)}, + flags | TENSOR_NOT_REQUIRED | TENSOR_SKIP_IF_VIRTUAL); + if (!layer.wkv_b) { + layer.wk_b = create_tensor(tn(LLM_TENSOR_ATTN_K_B, "weight", i), {n_embd_head_k_mla - qk_rope_head_dim, kv_lora_rank, n_head}, flags); + layer.wv_b = create_tensor(tn(LLM_TENSOR_ATTN_V_B, "weight", i), {kv_lora_rank, n_embd_head_v_mla, n_head}, flags); + } + + // head-wise gate: ONE scalar per head, not per output element + layer.wqkv_gate = create_tensor(tn(LLM_TENSOR_ATTN_GATE, "weight", i), {n_embd, n_head}, flags); + + layer.wo = create_tensor(tn(LLM_TENSOR_ATTN_OUT, "weight", i), {n_head * n_embd_head_v_mla, n_embd}, flags); + } + + layer.ffn_norm = create_tensor(tn(LLM_TENSOR_FFN_NORM, "weight", i), {n_embd}, flags); + + const int64_t n_ff_exp = hparams.n_ff_exp; + + if ((uint32_t) i < hparams.n_layer_dense_lead) { + // first_k_dense_replace = 2 -> layers 0 and 1 are dense + layer.ffn_gate = create_tensor(tn(LLM_TENSOR_FFN_GATE, "weight", i), {n_embd, n_ff}, flags); + layer.ffn_down = create_tensor(tn(LLM_TENSOR_FFN_DOWN, "weight", i), {n_ff, n_embd}, flags); + layer.ffn_up = create_tensor(tn(LLM_TENSOR_FFN_UP, "weight", i), {n_embd, n_ff}, flags); + } else { + layer.ffn_gate_inp = create_tensor(tn(LLM_TENSOR_FFN_GATE_INP, "weight", i), {n_embd, n_expert}, flags); + layer.ffn_gate_exps = create_tensor(tn(LLM_TENSOR_FFN_GATE_EXPS, "weight", i), {n_embd, n_ff_exp, n_expert}, flags); + layer.ffn_down_exps = create_tensor(tn(LLM_TENSOR_FFN_DOWN_EXPS, "weight", i), {n_ff_exp, n_embd, n_expert}, flags); + layer.ffn_up_exps = create_tensor(tn(LLM_TENSOR_FFN_UP_EXPS, "weight", i), {n_embd, n_ff_exp, n_expert}, flags); + + const int64_t n_ff_shexp = n_ff_exp * (hparams.n_expert_shared > 0 ? hparams.n_expert_shared : 1); + layer.ffn_gate_shexp = create_tensor(tn(LLM_TENSOR_FFN_GATE_SHEXP, "weight", i), {n_embd, n_ff_shexp}, flags); + layer.ffn_down_shexp = create_tensor(tn(LLM_TENSOR_FFN_DOWN_SHEXP, "weight", i), {n_ff_shexp, n_embd}, flags); + layer.ffn_up_shexp = create_tensor(tn(LLM_TENSOR_FFN_UP_SHEXP, "weight", i), {n_embd, n_ff_shexp}, flags); + + layer.ffn_exp_probs_b = create_tensor(tn(LLM_TENSOR_FFN_EXP_PROBS_B, "bias", i), {n_expert}, flags); + } + + // MTP/nextn head: preserved but unused. These MUST be created even though + // nothing reads them -- the loader throws if n_created < n_tensors, so an + // unclaimed tensor in the GGUF fails the load outright. Ling has no nextn + // embed_tokens / shared_head.head (it borrows the main model's), and names + // its final norm `final_layernorm` -> NEXTN_SHARED_HEAD_NORM. + if (i >= n_layer) { + layer.nextn.eh_proj = create_tensor(tn(LLM_TENSOR_NEXTN_EH_PROJ, "weight", i), {2 * n_embd, n_embd}, flags); + layer.nextn.enorm = create_tensor(tn(LLM_TENSOR_NEXTN_ENORM, "weight", i), {n_embd}, flags); + layer.nextn.hnorm = create_tensor(tn(LLM_TENSOR_NEXTN_HNORM, "weight", i), {n_embd}, flags); + layer.nextn.shared_head_norm = create_tensor(tn(LLM_TENSOR_NEXTN_SHARED_HEAD_NORM, "weight", i), {n_embd}, flags); + } + } +} + +std::unique_ptr llama_model_bailing_hybrid::build_arch_graph(const llm_graph_params & params) const { + return std::make_unique(*this, params); +} + +// Causal conv1d over Q/K/V. Copied from kimi-linear.cpp -- qkv selects which of +// the three conv states to read/write (0=Q, 1=K, 2=V). +static ggml_tensor * causal_conv1d(ggml_cgraph * gf, ggml_context * ctx0, ggml_tensor * conv_states_all, + ggml_tensor * conv_state_all, int64_t qkv, ggml_tensor * x, ggml_tensor * proj_w, ggml_tensor * conv_w, + int64_t d_conv, int64_t head_dim, int64_t n_head, int64_t n_seq_tokens, int64_t n_seqs, + int64_t n_tokens, int64_t kv_head) { + const int64_t d_inner = head_dim * n_head; + const int64_t conv_state_size = (d_conv - 1) * d_inner; + const int64_t n_embd_r_total = 3 * conv_state_size; // Q + K + V + + ggml_tensor * conv_state_x = ggml_view_3d(ctx0, conv_state_all, d_conv - 1, d_inner, n_seqs, + (d_conv - 1) * ggml_element_size(conv_state_all), + n_embd_r_total * ggml_element_size(conv_state_all), + qkv * conv_state_size * ggml_element_size(conv_state_all)); + + ggml_tensor * x_proj = ggml_mul_mat(ctx0, proj_w, x); + ggml_tensor * x_3d = ggml_reshape_3d(ctx0, x_proj, d_inner, n_seq_tokens, n_seqs); + + ggml_tensor * conv_x = ggml_concat(ctx0, conv_state_x, ggml_transpose(ctx0, x_3d), 0); + + ggml_tensor * last_conv_x = ggml_view_3d(ctx0, conv_x, d_conv - 1, d_inner, n_seqs, + conv_x->nb[1], conv_x->nb[2], n_seq_tokens * conv_x->nb[0]); + ggml_build_forward_expand(gf, + ggml_cpy(ctx0, last_conv_x, + ggml_view_3d(ctx0, conv_states_all, d_conv - 1, d_inner, n_seqs, + (d_conv - 1) * ggml_element_size(conv_states_all), + n_embd_r_total * ggml_element_size(conv_states_all), + (kv_head * n_embd_r_total + qkv * conv_state_size) * ggml_element_size(conv_states_all)))); + + ggml_tensor * conv_weight = ggml_reshape_2d(ctx0, conv_w, d_conv, d_inner); + + ggml_tensor * Xcur = ggml_ssm_conv(ctx0, conv_x, conv_weight); + Xcur = ggml_reshape_2d(ctx0, Xcur, d_inner, n_tokens); + Xcur = ggml_silu(ctx0, Xcur); + + return ggml_reshape_4d(ctx0, Xcur, head_dim, n_head, n_seq_tokens, n_seqs); +} + +llama_model_bailing_hybrid::graph::graph(const llama_model & model, const llm_graph_params & params) : + llm_build_delta_net_base(params), model(model) { + ggml_tensor * cur; + ggml_tensor * inpL; + + inpL = build_inp_embd(model.tok_embd); + cb(inpL, "model.embed_tokens", -1); + + // MLA layers are RoPE'd (unlike kimi-linear), so positions are needed. + ggml_tensor * inp_pos = build_inp_pos(); + + auto * inp_kv = !hparams.is_mla() ? build_inp_mem_hybrid() : nullptr; + auto * inp_k = hparams.is_mla() ? build_inp_mem_hybrid_k() : nullptr; + auto * inp_rs = hparams.is_mla() ? inp_k->get_recr() : inp_kv->get_recr(); + auto * inp_attn_kv = !hparams.is_mla() ? inp_kv->get_attn() : nullptr; + auto * inp_attn_k = hparams.is_mla() ? inp_k->get_attn() : nullptr; + + ggml_tensor * inp_out_ids = build_inp_out_ids(); + + const int64_t n_head = hparams.n_head(); + const int64_t head_dim = hparams.n_embd_head_kda; + const int64_t d_conv = hparams.ssm_d_conv; + const int64_t d_inner = n_head * head_dim; + const int64_t n_seqs = ubatch.n_seqs; + const int64_t n_seq_tokens = ubatch.n_seq_tokens; + + GGML_ASSERT(n_seqs != 0); + GGML_ASSERT(ubatch.equal_seqs()); + GGML_ASSERT(ubatch.n_tokens == n_seq_tokens * n_seqs); + + const int64_t n_embd_head_k_mla = hparams.n_embd_head_k_mla(); // 192 + const int64_t n_embd_head_v_mla = hparams.n_embd_head_v_mla(); // 128 + const int64_t kv_lora_rank = hparams.n_lora_kv; // 512 + const int64_t n_embd_head_qk_rope = hparams.n_rot(); // 64 + const int64_t n_embd_head_qk_nope = n_embd_head_k_mla - n_embd_head_qk_rope; // 128 + + // scaling = qk_head_dim ** -0.5 over the FULL 192, not the nope part + const float kq_scale_mla = 1.0f / sqrtf((float) n_embd_head_k_mla); + + const float kda_lower_bound = hparams.f_kda_lower_bound; + + // NORM rope over the 64-dim rope slice only -- see the header comment. + const int rope_type = LLAMA_ROPE_TYPE_NORM; + const int n_rot = n_embd_head_qk_rope; + const float freq_base = hparams.rope_freq_base_train; + const float freq_scale = hparams.rope_freq_scale_train; + const float ext_factor = cparams.yarn_ext_factor; + const float attn_factor = cparams.yarn_attn_factor; + const float beta_fast = cparams.yarn_beta_fast; + const float beta_slow = cparams.yarn_beta_slow; + const int n_ctx_orig = cparams.n_ctx_orig_yarn; + + for (int il = 0; il < n_layer; ++il) { + const auto & layer = model.layers[il]; + ggml_tensor * inpSA = inpL; + + cur = build_norm(inpL, layer.attn_norm, NULL, LLM_NORM_RMS, il); + cb(cur, "attn_norm", il); + + ggml_build_forward_expand(gf, cur); + + if (hparams.is_recr(il)) { + // ================= KDA linear-attention layer ================= + const auto * mctx_cur = inp_rs->mctx; + const auto kv_head = mctx_cur->get_head(); + + ggml_tensor * conv_states_all = mctx_cur->get_r_l(il); + cb(conv_states_all, "conv_states_all", il); + ggml_tensor * conv_state_all = build_rs(inp_rs, conv_states_all, hparams.n_embd_r(), n_seqs); + + ggml_tensor * Qcur = causal_conv1d(gf, ctx0, conv_states_all, conv_state_all, 0, cur, layer.wq, layer.ssm_q_conv, d_conv, head_dim, n_head, n_seq_tokens, n_seqs, n_tokens, kv_head); + ggml_tensor * Kcur = causal_conv1d(gf, ctx0, conv_states_all, conv_state_all, 1, cur, layer.wk, layer.ssm_k_conv, d_conv, head_dim, n_head, n_seq_tokens, n_seqs, n_tokens, kv_head); + ggml_tensor * Vcur = causal_conv1d(gf, ctx0, conv_states_all, conv_state_all, 2, cur, layer.wv, layer.ssm_v_conv, d_conv, head_dim, n_head, n_seq_tokens, n_seqs, n_tokens, kv_head); + + // *** delta 1+2 vs kimi-linear *** + // full-rank f_proj (one matmul, not f_b(f_a(x))), then the safe gate + // g = lower_bound * sigmoid(exp(A_log) * (f(x) + dt_bias)) + // ssm_a already holds +exp(A_log). dt_bias is added BEFORE the + // per-head A scaling and before the sigmoid -- fla adds the bias to + // the raw projection, then multiplies inside the sigmoid. + ggml_tensor * g1 = ggml_mul_mat(ctx0, layer.ssm_f, cur); + g1 = ggml_add(ctx0, g1, layer.ssm_dt_b); + g1 = ggml_reshape_3d(ctx0, g1, head_dim, n_head, n_tokens); + + // A is per-head: [1, n_head, 1] broadcast over head_dim and tokens + ggml_tensor * A = ggml_reshape_3d(ctx0, layer.ssm_a, 1, n_head, 1); + g1 = ggml_mul(ctx0, g1, A); + g1 = ggml_sigmoid(ctx0, g1); + g1 = ggml_scale(ctx0, g1, kda_lower_bound); + cb(g1, "kda_g1_safe_gate", il); + + g1 = ggml_reshape_4d(ctx0, g1, head_dim, n_head, n_seq_tokens, n_seqs); + + // allow_neg_eigval is off here: plain sigmoid, no 2x (that is + // solar_open2's delta, not this model's). + ggml_tensor * beta = ggml_mul_mat(ctx0, layer.ssm_beta, cur); + beta = ggml_reshape_4d(ctx0, beta, 1, n_head, n_seq_tokens, n_seqs); + beta = ggml_sigmoid(ctx0, beta); + cb(beta, "kda_beta", il); + + cur = ggml_reshape_3d(ctx0, cur, cur->ne[0], n_seq_tokens, n_seqs); + + ggml_tensor * ssm_states_all = mctx_cur->get_s_l(il); + ggml_tensor * state = build_rs(inp_rs, ssm_states_all, hparams.n_embd_s(), n_seqs); + state = ggml_reshape_4d(ctx0, state, head_dim, head_dim, n_head, n_seqs); + + const float eps_norm = hparams.f_norm_rms_eps; + Qcur = ggml_l2_norm(ctx0, Qcur, eps_norm); + Kcur = ggml_l2_norm(ctx0, Kcur, eps_norm); + + auto attn_out = build_delta_net(Qcur, Kcur, Vcur, g1, beta, state, il); + + ggml_tensor * output = ggml_cont(ctx0, attn_out.first); + ggml_tensor * new_state = attn_out.second; + + ggml_build_forward_expand(gf, + ggml_cpy(ctx0, new_state, + ggml_view_1d(ctx0, ssm_states_all, hparams.n_embd_s() * n_seqs, + kv_head * hparams.n_embd_s() * ggml_element_size(ssm_states_all)))); + + // full-rank output gate, then RMSNorm(x) * sigmoid(g) + ggml_tensor * cur_2d = ggml_reshape_2d(ctx0, cur, cur->ne[0], n_seq_tokens * n_seqs); + ggml_tensor * g2 = ggml_mul_mat(ctx0, layer.ssm_g, cur_2d); + g2 = ggml_reshape_3d(ctx0, g2, head_dim, n_head, n_seq_tokens * n_seqs); + + ggml_tensor * attn_out_final = ggml_reshape_3d(ctx0, output, head_dim, n_head, n_seq_tokens * n_seqs); + ggml_tensor * normed = build_norm(attn_out_final, layer.ssm_o_norm, nullptr, LLM_NORM_RMS, il); + ggml_tensor * gated = ggml_mul(ctx0, normed, ggml_sigmoid(ctx0, g2)); + + gated = ggml_cont_2d(ctx0, gated, d_inner, n_tokens); + cur = ggml_mul_mat(ctx0, layer.wo, gated); + cb(cur, "kda_out", il); + } else { + // ================= gated MLA layer ================= + // q_proj is one wide matmul (q_lora_rank is null). Per head the + // layout is [nope(128) | rope(64)], matching the reference's + // split(q, [qk_nope_head_dim, qk_rope_head_dim], dim=-1). + ggml_tensor * Qcur = ggml_mul_mat(ctx0, layer.wq, cur); + + ggml_tensor * kv_cmpr_pe = ggml_mul_mat(ctx0, layer.wkv_a_mqa, cur); + + ggml_tensor * kv_cmpr = ggml_view_2d(ctx0, kv_cmpr_pe, kv_lora_rank, n_tokens, + ggml_row_size(kv_cmpr_pe->type, kv_lora_rank + n_embd_head_qk_rope), 0); + ggml_tensor * k_pe = ggml_view_3d(ctx0, kv_cmpr_pe, n_embd_head_qk_rope, 1, n_tokens, + ggml_row_size(kv_cmpr_pe->type, kv_lora_rank + n_embd_head_qk_rope), + ggml_row_size(kv_cmpr_pe->type, kv_lora_rank + n_embd_head_qk_rope), + ggml_row_size(kv_cmpr_pe->type, kv_lora_rank)); + + // *** delta 5: kimi applies no RoPE here; this model does *** + k_pe = ggml_rope_ext(ctx0, ggml_cont(ctx0, k_pe), inp_pos, nullptr, + n_rot, rope_type, n_ctx_orig, freq_base, freq_scale, + ext_factor, attn_factor, beta_fast, beta_slow); + cb(k_pe, "k_pe", il); + + kv_cmpr = build_norm(kv_cmpr, layer.attn_kv_a_norm, nullptr, LLM_NORM_RMS, il); + + ggml_tensor * attn_out = nullptr; + + if (layer.wk_b && layer.wv_b) { // MLA KV cache enabled + ggml_tensor * q_nope = + ggml_view_3d(ctx0, Qcur, n_embd_head_qk_nope, n_head, n_tokens, + ggml_row_size(Qcur->type, n_embd_head_k_mla), + ggml_row_size(Qcur->type, n_embd_head_k_mla) * n_head, 0); + + ggml_tensor * q_pe = ggml_view_3d( + ctx0, Qcur, n_embd_head_qk_rope, n_head, n_tokens, + ggml_row_size(Qcur->type, n_embd_head_k_mla), + ggml_row_size(Qcur->type, n_embd_head_k_mla) * n_head, + ggml_row_size(Qcur->type, n_embd_head_qk_nope)); + + q_pe = ggml_rope_ext(ctx0, ggml_cont(ctx0, q_pe), inp_pos, nullptr, + n_rot, rope_type, n_ctx_orig, freq_base, freq_scale, + ext_factor, attn_factor, beta_fast, beta_slow); + cb(q_pe, "q_pe", il); + + // {n_embd_head_qk_nope, n_tokens, n_head} + q_nope = ggml_permute(ctx0, q_nope, 0, 2, 1, 3); + + ggml_tensor * q_nope_absorbed = ggml_mul_mat(ctx0, layer.wk_b, q_nope); + q_nope_absorbed = ggml_permute(ctx0, q_nope_absorbed, 0, 2, 1, 3); + + // note: rope must go first for in-place context shifting + Qcur = ggml_concat(ctx0, q_nope_absorbed, q_pe, 0); + cb(Qcur, "Qcur", il); + + kv_cmpr = ggml_reshape_3d(ctx0, kv_cmpr, kv_lora_rank, 1, n_tokens); + + ggml_tensor * Kcur = ggml_concat(ctx0, kv_cmpr, k_pe, 0); + ggml_tensor * Vcur = kv_cmpr; + + // wo is applied after the head-wise gate, so pass null here + attn_out = build_attn(inp_attn_k, nullptr, NULL, layer.wo_s, + Qcur, Kcur, Vcur, nullptr, nullptr, layer.wv_b, kq_scale_mla, il); + } else { // MLA KV cache disabled -- fall back to MHA + ggml_tensor * q_pe = ggml_view_3d( + ctx0, Qcur, n_embd_head_qk_rope, n_head, n_tokens, + ggml_row_size(Qcur->type, n_embd_head_k_mla), + ggml_row_size(Qcur->type, n_embd_head_k_mla) * n_head, + ggml_row_size(Qcur->type, n_embd_head_qk_nope)); + q_pe = ggml_rope_ext(ctx0, ggml_cont(ctx0, q_pe), inp_pos, nullptr, + n_rot, rope_type, n_ctx_orig, freq_base, freq_scale, + ext_factor, attn_factor, beta_fast, beta_slow); + + ggml_tensor * q_nope = + ggml_view_3d(ctx0, Qcur, n_embd_head_qk_nope, n_head, n_tokens, + ggml_row_size(Qcur->type, n_embd_head_k_mla), + ggml_row_size(Qcur->type, n_embd_head_k_mla) * n_head, 0); + + // rebuild Q as [nope | rope] to match the K layout below + Qcur = ggml_concat(ctx0, ggml_cont(ctx0, q_nope), q_pe, 0); + + ggml_tensor * kv = ggml_mul_mat(ctx0, layer.wkv_b, kv_cmpr); + const int64_t kv_per_head = n_embd_head_qk_nope + n_embd_head_v_mla; + + ggml_tensor * k_nope = ggml_view_3d(ctx0, kv, n_embd_head_qk_nope, n_head, n_tokens, + ggml_row_size(kv->type, kv_per_head), + ggml_row_size(kv->type, kv_per_head * n_head), 0); + ggml_tensor * Vcur = ggml_view_3d(ctx0, kv, n_embd_head_v_mla, n_head, n_tokens, + ggml_row_size(kv->type, kv_per_head), + ggml_row_size(kv->type, kv_per_head * n_head), + ggml_row_size(kv->type, n_embd_head_qk_nope)); + Vcur = ggml_cont(ctx0, Vcur); + + // k_pe is shared across heads (MQA) -> broadcast before concat + ggml_tensor * k_pe_target = ggml_new_tensor_3d(ctx0, k_pe->type, n_embd_head_qk_rope, n_head, n_tokens); + ggml_tensor * k_pe_repeated = ggml_repeat(ctx0, k_pe, k_pe_target); + ggml_tensor * Kcur = ggml_concat(ctx0, ggml_cont(ctx0, k_nope), k_pe_repeated, 0); + + attn_out = build_attn(inp_attn_kv, nullptr, NULL, layer.wo_s, + Qcur, Kcur, Vcur, nullptr, nullptr, nullptr, kq_scale_mla, il); + } + cb(attn_out, "attn_out", il); + + // *** delta 4: HEAD-WISE sigmoid gate *** + // g_proj is {n_embd, n_head}: one scalar per head, broadcast across + // v_head_dim. Reshaping attn_out to [head_dim, n_head, n_tokens] and + // the gate to [1, n_head, n_tokens] makes ggml_mul do that broadcast. + ggml_tensor * gate = ggml_mul_mat(ctx0, layer.wqkv_gate, cur); + gate = ggml_sigmoid(ctx0, gate); + gate = ggml_reshape_3d(ctx0, gate, 1, n_head, n_tokens); + cb(gate, "attn_gate_headwise", il); + + attn_out = ggml_reshape_3d(ctx0, attn_out, n_embd_head_v_mla, n_head, n_tokens); + attn_out = ggml_mul(ctx0, attn_out, gate); + attn_out = ggml_cont_2d(ctx0, attn_out, n_embd_head_v_mla * n_head, n_tokens); + cb(attn_out, "attn_gated", il); + + cur = ggml_mul_mat(ctx0, layer.wo, attn_out); + cb(cur, "mla_out", il); + } + + if (il == n_layer - 1 && inp_out_ids) { + cur = ggml_get_rows(ctx0, cur, inp_out_ids); + inpSA = ggml_get_rows(ctx0, inpSA, inp_out_ids); + } + + ggml_tensor * ffn_inp = ggml_add(ctx0, cur, inpSA); + cb(ffn_inp, "ffn_inp", il); + + cur = build_norm(ffn_inp, layer.ffn_norm, NULL, LLM_NORM_RMS, il); + cb(cur, "ffn_norm", il); + + if ((uint32_t) il < hparams.n_layer_dense_lead) { + cur = build_ffn(cur, + layer.ffn_up, NULL, NULL, + layer.ffn_gate, NULL, NULL, + layer.ffn_down, NULL, NULL, + NULL, LLM_FFN_SILU, LLM_FFN_PAR, il); + cb(cur, "ffn_out", il); + } else { + // noaux_tc grouped top-k + sigmoid + expert_bias: build_moe_ffn reads + // n_expert_groups / n_group_used straight from hparams. + ggml_tensor * moe_out = build_moe_ffn(cur, + layer.ffn_gate_inp, + layer.ffn_up_exps, + layer.ffn_gate_exps, + layer.ffn_down_exps, + layer.ffn_exp_probs_b, + hparams.n_expert, + hparams.n_expert_used, + LLM_FFN_SILU, hparams.expert_weights_norm, + hparams.expert_weights_scale, + (llama_expert_gating_func_type) hparams.expert_gating_func, + il); + cb(moe_out, "ffn_moe_out", il); + + ggml_tensor * ffn_shexp = build_ffn(cur, + layer.ffn_up_shexp, NULL, NULL, + layer.ffn_gate_shexp, NULL, NULL, + layer.ffn_down_shexp, NULL, NULL, + NULL, LLM_FFN_SILU, LLM_FFN_PAR, il); + cb(ffn_shexp, "ffn_shexp", il); + + cur = ggml_add(ctx0, moe_out, ffn_shexp); + cb(cur, "ffn_out", il); + } + + cur = ggml_add(ctx0, cur, ffn_inp); + + cur = build_cvec(cur, il); + cb(cur, "l_out", il); + + inpL = cur; + } + + cur = inpL; + + cur = build_norm(cur, model.output_norm, NULL, LLM_NORM_RMS, -1); + cb(cur, "result_norm", -1); + res->t_embd = cur; + + cur = ggml_mul_mat(ctx0, model.output, cur); + cb(cur, "result_output", -1); + res->t_logits = cur; + + ggml_build_forward_expand(gf, cur); +} diff --git a/src/models/models.h b/src/models/models.h index ad3dadaf3..ba0496edc 100644 --- a/src/models/models.h +++ b/src/models/models.h @@ -2211,6 +2211,24 @@ struct llama_model_kimi_linear : public llama_model_base { std::unique_ptr build_arch_graph(const llm_graph_params & params) const override; }; +// Ling 3.0 flash (inclusionAI/Ling-3.0-flash), model_type "bailing_hybrid". +// Hybrid KDA + gated MLA MoE. Despite the family name this is NOT the +// bailingmoe2 stack (Ling 2.0) -- only the MoE router carries over. See +// src/models/bailing-hybrid.cpp for the deltas against kimi-linear. +struct llama_model_bailing_hybrid : public llama_model_base { + llama_model_bailing_hybrid(const struct llama_model_params & params) : llama_model_base(params) {} + void load_arch_hparams(llama_model_loader & ml) override; + void load_arch_tensors(llama_model_loader & ml) override; + + struct graph : public llm_build_delta_net_base { + graph(const llama_model & model, const llm_graph_params & params); + + const llama_model & model; + }; + + std::unique_ptr build_arch_graph(const llm_graph_params & params) const override; +}; + struct llama_model_step35 : public llama_model_base { llama_model_step35(const struct llama_model_params & params) : llama_model_base(params) {}