From 8922349d03b3269f84ea1b41e9bb4f8f97b330a0 Mon Sep 17 00:00:00 2001 From: Bar Haim Date: Thu, 18 Jun 2026 12:43:43 -0700 Subject: [PATCH 01/40] granite-switch: add llama.cpp backend (POC, CPU) New "granite-switch" architecture: a dense, all-attention Granite-4.1 model with N embedded LoRA adapters selected per-token by control tokens. - gguf-py schema (arch, KV keys, stacked LoRA tensor names) + writer helpers - conversion/granite.py: GraniteSwitchModel converter (stacks N adapters + zero base slot into per-projection A/B tensors; emits switch metadata) - C++ arch registration (llama-arch.{h,cpp}, llama-model.{h,cpp}) - src/models/granite_switch.cpp: load + per-token switched-LoRA graph via ggml_mul_mat_id over stacked tensors; sticky per-token index + control-token substitution in llm_graph_input_switch::set_input - llm_graph_input_switch in src/models/models.h Runs end-to-end on CPU: convert 3b checkpoint (842 tensors, stacked dim 13) and generate on both base and control-token paths. Sticky switch state is single-sequence (POC); full multi-sequence machinery is a follow-up. --- conversion/__init__.py | 1 + conversion/granite.py | 230 +++++++++++++++++++ gguf-py/gguf/constants.py | 65 ++++++ gguf-py/gguf/gguf_writer.py | 15 ++ src/llama-arch.cpp | 36 +++ src/llama-arch.h | 21 ++ src/llama-model.cpp | 4 + src/llama-model.h | 26 +++ src/models/granite_switch.cpp | 407 ++++++++++++++++++++++++++++++++++ src/models/models.h | 84 +++++++ 10 files changed, 889 insertions(+) create mode 100644 src/models/granite_switch.cpp diff --git a/conversion/__init__.py b/conversion/__init__.py index 02ea6385208a..86a292bf1258 100644 --- a/conversion/__init__.py +++ b/conversion/__init__.py @@ -98,6 +98,7 @@ "GraniteMoeForCausalLM": "granite", "GraniteMoeHybridForCausalLM": "granite", "GraniteMoeSharedForCausalLM": "granite", + "GraniteSwitchForCausalLM": "granite", "GraniteSpeechForConditionalGeneration": "granite", "GraniteSpeechPlusForConditionalGeneration": "granite", "Grok1ForCausalLM": "grok", diff --git a/conversion/granite.py b/conversion/granite.py index 8367ed225da6..736f35a9e231 100644 --- a/conversion/granite.py +++ b/conversion/granite.py @@ -123,6 +123,236 @@ def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iter yield from super().modify_tensors(data_torch, name, bid) +@ModelBase.register("GraniteSwitchForCausalLM") +class GraniteSwitchModel(GraniteMoeModel): + """Conversion for IBM's GraniteSwitchForCausalLM. + + Granite Switch is a dense, all-attention Granite model with N embedded LoRA + adapters selected per-token by control tokens. The on-disk checkpoint: + + - Uses a fused ``self_attn.qkv_proj`` and fused ``shared_mlp.input_linear`` + (gate+up), like GraniteMoe's shared expert, plus ``self_attn.o_proj`` and + ``shared_mlp.output_linear``. Base weights live under ``.base_layer.weight``. + - Reserves one cache slot for the (weightless) switch, so the decoder has + ``num_hidden_layers - 1`` blocks. The decoder layers are an ``nn.ModuleList`` + saved 0-based as ``model.layers.0 .. model.layers.{N-2}`` (the switch holds no + weights and is not in the list), so ``blk.{bid}`` maps straight from ``bid``. + We emit ``block_count = num_hidden_layers - 1``. + - Stacks each LoRA over the adapter dim: A ``[n_adapters, 1, max_rank, in]``, + B ``[n_adapters, 1, out_slice, max_rank]`` (lower ranks zero-padded). B is + pre-scaled by ``alpha/rank`` at compose time, so the runtime LoRA scale is + 1.0 and A is left unscaled. + + To make per-token selection branch-free with ``ggml_mul_mat_id`` (which has no + "skip/base" id), we materialize a zero adapter at slot 0 so the stacked first + dim is ``N = num_adapters + 1``: index 0 (base tokens) adds an exact-zero delta. + """ + model_arch = gguf.MODEL_ARCH.GRANITE_SWITCH + + # ggml's granite arch uses LLAMA_ROPE_TYPE_NORM, so q/k weights must be stored + # in the interleaved layout ggml expects. HF Granite uses transformers' + # rotate_half (split-half) layout, so we apply LlamaModel's permute to the q and + # k row-blocks of the fused qkv base AND to the q/k LoRA-B output rows. We do it + # explicitly per-slice below, so disable the parent's automatic + # (separate-projection) permute. + undo_permute = False + + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + # The (weightless) switch reserves one cache slot; the decoder has one fewer + # block than num_hidden_layers. The decoder layers are a 0-based ModuleList, + # so block ids map straight through (no index shift). Fix block_count and + # rebuild the tensor name map for the reduced count. + self.block_count = self.block_count - 1 + self.tensor_map = gguf.get_tensor_name_map(self.model_arch, self.block_count) + + self._n_adapters = int(self.hparams["num_adapters"]) + self._max_lora_rank = int(self.hparams["max_lora_rank"]) + # Stacked adapter dim: +1 for the materialized zero slot at index 0. + self._n_slots = self._n_adapters + 1 + + n_head = int(self.hparams["num_attention_heads"]) + n_kv_head = int(self.hparams["num_key_value_heads"]) + # The qkv projection emits projection_head_dim-wide vectors (matches the HF + # model's GraniteLoRAEmbeddedAttention.head_dim). Fall back to head_dim, then + # to hidden/heads. + head_dim = ( + self.hparams.get("projection_head_dim") + or self.hparams.get("head_dim") + or (self.hparams["hidden_size"] // n_head) + ) + self._n_head = n_head + self._n_kv_head = n_kv_head + self._head_dim = int(head_dim) + self._q_size = n_head * self._head_dim + self._kv_size = n_kv_head * self._head_dim + + def set_gguf_parameters(self): + # Emits the four Granite multipliers, shared FFN length, head counts, rope, + # vocab, etc. add_block_count uses self.block_count (already decremented). + super().set_gguf_parameters() + + # Granite Switch is dense (num_local_experts == 0). The llama parent may + # have emitted a nonzero expert_used_count (e.g. 2) from num_experts_per_tok; + # force the dense invariant so the loader's n_expert_used <= n_expert holds. + self.gguf_writer.add_expert_count(0) + self.gguf_writer.add_expert_used_count(0) + + self.gguf_writer.add_num_adapters(self._n_adapters) + self.gguf_writer.add_max_lora_rank(self._max_lora_rank) + self.gguf_writer.add_adapter_token_ids(self.hparams["adapter_token_ids"]) + self.gguf_writer.add_adapter_substitute_token_ids(self.hparams["adapter_substitute_token_ids"]) + self.gguf_writer.add_adapter_ranks(self.hparams["adapter_ranks"]) + logger.info( + "gguf: (granite-switch) num_adapters=%s max_lora_rank=%s n_slots=%s", + self._n_adapters, self._max_lora_rank, self._n_slots, + ) + # NOTE: control_token_gain is intentionally not emitted — runtime adapter + # selection is exact control-token-id matching, not the attention trick. + + def _permute_qk(self, w: Tensor, n_head: int) -> Tensor: + # The same interleave permute LlamaModel applies to q/k weights to produce + # ggml's NORM-rope layout, operating on a [out, ...] slice where + # `out = n_head * head_dim`. + return LlamaModel.permute(w, n_head, n_head) + + def _lora_a(self, data: Tensor) -> Tensor: + # on-disk A: [n_adapters, 1, max_rank, in] -> [n_adapters+1, max_rank, in] + a = data.squeeze(1) + zero = torch.zeros_like(a[:1]) + return torch.cat([zero, a], dim=0).contiguous() + + def _lora_b(self, data: Tensor, permute_n_head: int | None = None) -> Tensor: + # on-disk B: [n_adapters, 1, out, max_rank] -> [n_adapters+1, out, max_rank] + b = data.squeeze(1) + if permute_n_head is not None: + # Permute the output (row) dimension of each adapter's B slice to match + # the NEOX layout used for the (permuted) q/k base weights. + b = torch.stack([self._permute_qk(b[i], permute_n_head) for i in range(b.shape[0])], dim=0) + zero = torch.zeros_like(b[:1]) + return torch.cat([zero, b], dim=0).contiguous() + + def tensor_force_quant(self, name, new_name, bid, n_dims): + # Keep all stacked LoRA tensors in F16 (small; selected per-token via mul_mat_id). + if ".lora_a" in new_name or ".lora_b" in new_name: + return gguf.GGMLQuantizationType.F16 + return super().tensor_force_quant(name, new_name, bid, n_dims) + + def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iterable[tuple[str, Tensor]]: + T = gguf.MODEL_TENSOR + + # ---- Skip the (weightless) switch + control-token buffers ---- + # The switch is a hand-constructed cumsum-attention with no learned weights; + # control-token selection/substitution is reconstructed at load time from the + # adapter id arrays we emit in set_gguf_parameters. These persistent buffers + # (model.switch.*, model.adapter_token_ids) carry no parameters to convert. + bare = name.split(".")[-1] + if ( + name.startswith("model.switch.") or name.startswith("switch.") + or bare in ("adapter_token_ids", "control_to_substitute_lut") + ): + return + + # ---- Attention: fused QKV ---- + if "self_attn.qkv_proj" in name: + obid = bid + if name.endswith("base_layer.weight"): + # Fused [q|k|v] rows. Permute the q and k row-blocks (split-half rope). + q, k, v = data_torch.split([self._q_size, self._kv_size, self._kv_size], dim=0) + q = self._permute_qk(q, self._n_head) + k = self._permute_qk(k, self._n_kv_head) + fused = torch.cat([q, k, v], dim=0) + yield (self.format_tensor_name(T.ATTN_QKV, obid), fused) + return + if "lora_A_slices." in name: + slot = int(name.rsplit(".", 1)[1]) + key = {0: T.ATTN_QKV_LORA_A_Q, 1: T.ATTN_QKV_LORA_A_K, 2: T.ATTN_QKV_LORA_A_V}[slot] + yield (self.format_tensor_name(key, obid, suffix=""), self._lora_a(data_torch)) + return + if "lora_B_slices." in name: + slot = int(name.rsplit(".", 1)[1]) + key, ph = { + 0: (T.ATTN_QKV_LORA_B_Q, self._n_head), + 1: (T.ATTN_QKV_LORA_B_K, self._n_kv_head), + 2: (T.ATTN_QKV_LORA_B_V, None), + }[slot] + yield (self.format_tensor_name(key, obid, suffix=""), self._lora_b(data_torch, ph)) + return + raise ValueError(f"Unexpected qkv_proj tensor: {name}") + + # ---- Attention: output projection ---- + if "self_attn.o_proj" in name: + obid = bid + if name.endswith("base_layer.weight"): + yield (self.format_tensor_name(T.ATTN_OUT, obid), data_torch) + return + if name.endswith("lora_A"): + yield (self.format_tensor_name(T.ATTN_OUT_LORA_A, obid, suffix=""), self._lora_a(data_torch)) + return + if name.endswith("lora_B"): + yield (self.format_tensor_name(T.ATTN_OUT_LORA_B, obid, suffix=""), self._lora_b(data_torch)) + return + raise ValueError(f"Unexpected o_proj tensor: {name}") + + # ---- MLP: fused gate/up (shared_mlp.input_linear) ---- + if "shared_mlp.input_linear" in name: + obid = bid + ffn = self.hparams["shared_intermediate_size"] + if name.endswith("base_layer.weight"): + gate, up = data_torch.split([ffn, ffn], dim=0) + yield (self.format_tensor_name(T.FFN_GATE, obid), gate) + yield (self.format_tensor_name(T.FFN_UP, obid), up) + return + if "lora_A_slices." in name: + slot = int(name.rsplit(".", 1)[1]) + key = {0: T.FFN_GATE_LORA_A, 1: T.FFN_UP_LORA_A}[slot] + yield (self.format_tensor_name(key, obid, suffix=""), self._lora_a(data_torch)) + return + if "lora_B_slices." in name: + slot = int(name.rsplit(".", 1)[1]) + key = {0: T.FFN_GATE_LORA_B, 1: T.FFN_UP_LORA_B}[slot] + yield (self.format_tensor_name(key, obid, suffix=""), self._lora_b(data_torch)) + return + raise ValueError(f"Unexpected shared_mlp.input_linear tensor: {name}") + + # ---- MLP: down (shared_mlp.output_linear) ---- + if "shared_mlp.output_linear" in name: + obid = bid + if name.endswith("base_layer.weight"): + yield (self.format_tensor_name(T.FFN_DOWN, obid), data_torch) + return + if name.endswith("lora_A"): + yield (self.format_tensor_name(T.FFN_DOWN_LORA_A, obid, suffix=""), self._lora_a(data_torch)) + return + if name.endswith("lora_B"): + yield (self.format_tensor_name(T.FFN_DOWN_LORA_B, obid, suffix=""), self._lora_b(data_torch)) + return + raise ValueError(f"Unexpected shared_mlp.output_linear tensor: {name}") + + # ---- Per-layer norms (input_layernorm / post_attention_layernorm) ---- + if bid is not None and ".layers." in name and ( + "input_layernorm" in name or "post_attention_layernorm" in name + ): + obid = bid + key = T.ATTN_NORM if "input_layernorm" in name else T.FFN_NORM + yield (self.format_tensor_name(key, obid), data_torch) + return + + # ---- Everything else: embeddings, final norm (no layer index) ---- + # embed_tokens -> token_embd, model.norm -> output_norm. lm_head is tied. + if name in ("model.embed_tokens.weight", "embed_tokens.weight"): + yield (self.format_tensor_name(T.TOKEN_EMBD), data_torch) + return + if name in ("model.norm.weight", "norm.weight"): + yield (self.format_tensor_name(T.OUTPUT_NORM), data_torch) + return + if name == "lm_head.weight": + # Tied to token embeddings; the runtime ties output to token_embd. + return + + raise ValueError(f"granite-switch: unhandled tensor {name!r} (bid={bid})") + + @ModelBase.register("GraniteMoeHybridForCausalLM", "BambaForCausalLM") class GraniteHybridModel(Mamba2Model, GraniteMoeModel): """GraniteHybrid is a hybrid SSM + Attention model that uses Mamba2 SSM diff --git a/gguf-py/gguf/constants.py b/gguf-py/gguf/constants.py index 869e436acd5c..d124f50796f2 100644 --- a/gguf-py/gguf/constants.py +++ b/gguf-py/gguf/constants.py @@ -159,6 +159,12 @@ class LLM: TARGET_HIDDEN_SIZE = "{arch}.target_hidden_size" BLOCK_SIZE = "{arch}.block_size" NORM_BEFORE_RESIDUAL = "{arch}.norm_before_residual" + # Granite Switch: per-token embedded LoRA adapters + NUM_ADAPTERS = "{arch}.num_adapters" + ADAPTER_TOKEN_IDS = "{arch}.adapter_token_ids" + ADAPTER_SUBSTITUTE_TOKEN_IDS = "{arch}.adapter_substitute_token_ids" + ADAPTER_RANKS = "{arch}.adapter_ranks" + MAX_LORA_RANK = "{arch}.max_lora_rank" class Attention: HEAD_COUNT = "{arch}.attention.head_count" @@ -499,6 +505,7 @@ class MODEL_ARCH(IntEnum): GRANITE = auto() GRANITE_MOE = auto() GRANITE_HYBRID = auto() + GRANITE_SWITCH = auto() CHAMELEON = auto() WAVTOKENIZER_DEC = auto() PLM = auto() @@ -983,6 +990,21 @@ class MODEL_TENSOR(IntEnum): A_QF_FFN_UP = auto() A_QF_FFN_DOWN = auto() A_QF_FFN_NORM = auto() + # Granite Switch: per-token embedded LoRA adapters (stacked over N = num_adapters + 1) + ATTN_QKV_LORA_A_Q = auto() + ATTN_QKV_LORA_B_Q = auto() + ATTN_QKV_LORA_A_K = auto() + ATTN_QKV_LORA_B_K = auto() + ATTN_QKV_LORA_A_V = auto() + ATTN_QKV_LORA_B_V = auto() + ATTN_OUT_LORA_A = auto() + ATTN_OUT_LORA_B = auto() + FFN_GATE_LORA_A = auto() + FFN_GATE_LORA_B = auto() + FFN_UP_LORA_A = auto() + FFN_UP_LORA_B = auto() + FFN_DOWN_LORA_A = auto() + FFN_DOWN_LORA_B = auto() MODEL_ARCH_NAMES: dict[MODEL_ARCH, str] = { @@ -1079,6 +1101,7 @@ class MODEL_TENSOR(IntEnum): MODEL_ARCH.GRANITE: "granite", MODEL_ARCH.GRANITE_MOE: "granitemoe", MODEL_ARCH.GRANITE_HYBRID: "granitehybrid", + MODEL_ARCH.GRANITE_SWITCH: "granite-switch", MODEL_ARCH.CHAMELEON: "chameleon", MODEL_ARCH.WAVTOKENIZER_DEC: "wavtokenizer-dec", MODEL_ARCH.PLM: "plm", @@ -1560,6 +1583,21 @@ class MODEL_TENSOR(IntEnum): MODEL_TENSOR.NEXTN_SHARED_HEAD_NORM: "blk.{bid}.nextn.shared_head_norm", MODEL_TENSOR.FC: "fc", MODEL_TENSOR.D2T: "d2t", + # Granite Switch: per-token embedded LoRA adapters + MODEL_TENSOR.ATTN_QKV_LORA_A_Q: "blk.{bid}.attn_qkv.lora_a_q", + MODEL_TENSOR.ATTN_QKV_LORA_B_Q: "blk.{bid}.attn_qkv.lora_b_q", + MODEL_TENSOR.ATTN_QKV_LORA_A_K: "blk.{bid}.attn_qkv.lora_a_k", + MODEL_TENSOR.ATTN_QKV_LORA_B_K: "blk.{bid}.attn_qkv.lora_b_k", + MODEL_TENSOR.ATTN_QKV_LORA_A_V: "blk.{bid}.attn_qkv.lora_a_v", + MODEL_TENSOR.ATTN_QKV_LORA_B_V: "blk.{bid}.attn_qkv.lora_b_v", + MODEL_TENSOR.ATTN_OUT_LORA_A: "blk.{bid}.attn_output.lora_a", + MODEL_TENSOR.ATTN_OUT_LORA_B: "blk.{bid}.attn_output.lora_b", + MODEL_TENSOR.FFN_GATE_LORA_A: "blk.{bid}.ffn_gate.lora_a", + MODEL_TENSOR.FFN_GATE_LORA_B: "blk.{bid}.ffn_gate.lora_b", + MODEL_TENSOR.FFN_UP_LORA_A: "blk.{bid}.ffn_up.lora_a", + MODEL_TENSOR.FFN_UP_LORA_B: "blk.{bid}.ffn_up.lora_b", + MODEL_TENSOR.FFN_DOWN_LORA_A: "blk.{bid}.ffn_down.lora_a", + MODEL_TENSOR.FFN_DOWN_LORA_B: "blk.{bid}.ffn_down.lora_b", } MODEL_TENSORS: dict[MODEL_ARCH, list[MODEL_TENSOR]] = { @@ -3669,6 +3707,33 @@ class MODEL_TENSOR(IntEnum): MODEL_TENSOR.FFN_DOWN, MODEL_TENSOR.FFN_UP, ], + MODEL_ARCH.GRANITE_SWITCH: [ + MODEL_TENSOR.TOKEN_EMBD, + MODEL_TENSOR.OUTPUT_NORM, + MODEL_TENSOR.OUTPUT, + MODEL_TENSOR.ATTN_NORM, + MODEL_TENSOR.ATTN_QKV, + MODEL_TENSOR.ATTN_OUT, + MODEL_TENSOR.FFN_NORM, + MODEL_TENSOR.FFN_GATE, + MODEL_TENSOR.FFN_DOWN, + MODEL_TENSOR.FFN_UP, + # Per-token embedded LoRA adapters + MODEL_TENSOR.ATTN_QKV_LORA_A_Q, + MODEL_TENSOR.ATTN_QKV_LORA_B_Q, + MODEL_TENSOR.ATTN_QKV_LORA_A_K, + MODEL_TENSOR.ATTN_QKV_LORA_B_K, + MODEL_TENSOR.ATTN_QKV_LORA_A_V, + MODEL_TENSOR.ATTN_QKV_LORA_B_V, + MODEL_TENSOR.ATTN_OUT_LORA_A, + MODEL_TENSOR.ATTN_OUT_LORA_B, + MODEL_TENSOR.FFN_GATE_LORA_A, + MODEL_TENSOR.FFN_GATE_LORA_B, + MODEL_TENSOR.FFN_UP_LORA_A, + MODEL_TENSOR.FFN_UP_LORA_B, + MODEL_TENSOR.FFN_DOWN_LORA_A, + MODEL_TENSOR.FFN_DOWN_LORA_B, + ], MODEL_ARCH.CHAMELEON: [ 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 1e277f0687c5..2f5d6adc1dec 100644 --- a/gguf-py/gguf/gguf_writer.py +++ b/gguf-py/gguf/gguf_writer.py @@ -892,6 +892,21 @@ def add_residual_scale(self, value: float) -> None: def add_embedding_scale(self, value: float) -> None: self.add_float32(Keys.LLM.EMBEDDING_SCALE.format(arch=self.arch), value) + def add_num_adapters(self, count: int) -> None: + self.add_uint32(Keys.LLM.NUM_ADAPTERS.format(arch=self.arch), count) + + def add_adapter_token_ids(self, ids: Sequence[int]) -> None: + self.add_array(Keys.LLM.ADAPTER_TOKEN_IDS.format(arch=self.arch), ids) + + def add_adapter_substitute_token_ids(self, ids: Sequence[int]) -> None: + self.add_array(Keys.LLM.ADAPTER_SUBSTITUTE_TOKEN_IDS.format(arch=self.arch), ids) + + def add_adapter_ranks(self, ranks: Sequence[int]) -> None: + self.add_array(Keys.LLM.ADAPTER_RANKS.format(arch=self.arch), ranks) + + def add_max_lora_rank(self, rank: int) -> None: + self.add_uint32(Keys.LLM.MAX_LORA_RANK.format(arch=self.arch), rank) + def add_wkv_head_size(self, size: int) -> None: self.add_uint32(Keys.WKV.HEAD_SIZE.format(arch=self.arch), size) diff --git a/src/llama-arch.cpp b/src/llama-arch.cpp index b890e66fcf6e..6f4218f04a5b 100644 --- a/src/llama-arch.cpp +++ b/src/llama-arch.cpp @@ -100,6 +100,7 @@ static const std::map LLM_ARCH_NAMES = { { LLM_ARCH_GRANITE, "granite" }, { LLM_ARCH_GRANITE_MOE, "granitemoe" }, { LLM_ARCH_GRANITE_HYBRID, "granitehybrid" }, + { LLM_ARCH_GRANITE_SWITCH, "granite-switch" }, { LLM_ARCH_CHAMELEON, "chameleon" }, { LLM_ARCH_WAVTOKENIZER_DEC, "wavtokenizer-dec" }, { LLM_ARCH_PLM, "plm" }, @@ -215,6 +216,11 @@ static const std::map LLM_KV_NAMES = { { LLM_KV_TIME_DECAY_EXTRA_DIM, "%s.time_decay_extra_dim" }, { LLM_KV_RESIDUAL_SCALE, "%s.residual_scale" }, { LLM_KV_EMBEDDING_SCALE, "%s.embedding_scale" }, + { LLM_KV_NUM_ADAPTERS, "%s.num_adapters" }, + { LLM_KV_ADAPTER_TOKEN_IDS, "%s.adapter_token_ids" }, + { LLM_KV_ADAPTER_SUBSTITUTE_TOKEN_IDS, "%s.adapter_substitute_token_ids" }, + { LLM_KV_ADAPTER_RANKS, "%s.adapter_ranks" }, + { LLM_KV_MAX_LORA_RANK, "%s.max_lora_rank" }, { LLM_KV_TOKEN_SHIFT_COUNT, "%s.token_shift_count" }, { LLM_KV_INTERLEAVE_MOE_LAYER_STEP, "%s.interleave_moe_layer_step" }, { LLM_KV_FULL_ATTENTION_INTERVAL, "%s.full_attention_interval" }, @@ -603,6 +609,21 @@ static const std::map LLM_TENSOR_NAMES = { { LLM_TENSOR_MASKED_EMBD_ORDERING, "masked_embd_ordering" }, { LLM_TENSOR_FC, "fc" }, { LLM_TENSOR_D2T, "d2t" }, + // Granite Switch: per-token embedded LoRA adapters (no ".weight" suffix on disk) + { LLM_TENSOR_ATTN_QKV_LORA_A_Q, "blk.%d.attn_qkv.lora_a_q" }, + { LLM_TENSOR_ATTN_QKV_LORA_B_Q, "blk.%d.attn_qkv.lora_b_q" }, + { LLM_TENSOR_ATTN_QKV_LORA_A_K, "blk.%d.attn_qkv.lora_a_k" }, + { LLM_TENSOR_ATTN_QKV_LORA_B_K, "blk.%d.attn_qkv.lora_b_k" }, + { LLM_TENSOR_ATTN_QKV_LORA_A_V, "blk.%d.attn_qkv.lora_a_v" }, + { LLM_TENSOR_ATTN_QKV_LORA_B_V, "blk.%d.attn_qkv.lora_b_v" }, + { LLM_TENSOR_ATTN_OUT_LORA_A, "blk.%d.attn_output.lora_a" }, + { LLM_TENSOR_ATTN_OUT_LORA_B, "blk.%d.attn_output.lora_b" }, + { LLM_TENSOR_FFN_GATE_LORA_A, "blk.%d.ffn_gate.lora_a" }, + { LLM_TENSOR_FFN_GATE_LORA_B, "blk.%d.ffn_gate.lora_b" }, + { LLM_TENSOR_FFN_UP_LORA_A, "blk.%d.ffn_up.lora_a" }, + { LLM_TENSOR_FFN_UP_LORA_B, "blk.%d.ffn_up.lora_b" }, + { LLM_TENSOR_FFN_DOWN_LORA_A, "blk.%d.ffn_down.lora_a" }, + { LLM_TENSOR_FFN_DOWN_LORA_B, "blk.%d.ffn_down.lora_b" }, }; // declare information about the model weight tensors: @@ -854,6 +875,21 @@ static const std::map LLM_TENSOR_INFOS = { // eagle3 {LLM_TENSOR_FC, {LLM_TENSOR_LAYER_OUTPUT, GGML_OP_MUL_MAT}}, {LLM_TENSOR_D2T, {LLM_TENSOR_LAYER_OUTPUT, GGML_OP_GET_ROWS}}, + // granite-switch: per-token embedded LoRA adapters (selected via mul_mat_id) + {LLM_TENSOR_ATTN_QKV_LORA_A_Q, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL_MAT_ID}}, + {LLM_TENSOR_ATTN_QKV_LORA_B_Q, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL_MAT_ID}}, + {LLM_TENSOR_ATTN_QKV_LORA_A_K, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL_MAT_ID}}, + {LLM_TENSOR_ATTN_QKV_LORA_B_K, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL_MAT_ID}}, + {LLM_TENSOR_ATTN_QKV_LORA_A_V, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL_MAT_ID}}, + {LLM_TENSOR_ATTN_QKV_LORA_B_V, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL_MAT_ID}}, + {LLM_TENSOR_ATTN_OUT_LORA_A, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL_MAT_ID}}, + {LLM_TENSOR_ATTN_OUT_LORA_B, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL_MAT_ID}}, + {LLM_TENSOR_FFN_GATE_LORA_A, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL_MAT_ID}}, + {LLM_TENSOR_FFN_GATE_LORA_B, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL_MAT_ID}}, + {LLM_TENSOR_FFN_UP_LORA_A, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL_MAT_ID}}, + {LLM_TENSOR_FFN_UP_LORA_B, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL_MAT_ID}}, + {LLM_TENSOR_FFN_DOWN_LORA_A, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL_MAT_ID}}, + {LLM_TENSOR_FFN_DOWN_LORA_B, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL_MAT_ID}}, }; LLM_KV::LLM_KV(llm_arch arch, const char * suffix) : arch(arch), suffix(suffix) {} diff --git a/src/llama-arch.h b/src/llama-arch.h index a4f5091e7170..fad3d33d7ed2 100644 --- a/src/llama-arch.h +++ b/src/llama-arch.h @@ -105,6 +105,7 @@ enum llm_arch { LLM_ARCH_GRANITE, LLM_ARCH_GRANITE_MOE, LLM_ARCH_GRANITE_HYBRID, + LLM_ARCH_GRANITE_SWITCH, LLM_ARCH_CHAMELEON, LLM_ARCH_WAVTOKENIZER_DEC, LLM_ARCH_PLM, @@ -220,6 +221,11 @@ enum llm_kv { LLM_KV_TIME_DECAY_EXTRA_DIM, LLM_KV_RESIDUAL_SCALE, LLM_KV_EMBEDDING_SCALE, + LLM_KV_NUM_ADAPTERS, + LLM_KV_ADAPTER_TOKEN_IDS, + LLM_KV_ADAPTER_SUBSTITUTE_TOKEN_IDS, + LLM_KV_ADAPTER_RANKS, + LLM_KV_MAX_LORA_RANK, LLM_KV_TOKEN_SHIFT_COUNT, LLM_KV_INTERLEAVE_MOE_LAYER_STEP, LLM_KV_FULL_ATTENTION_INTERVAL, @@ -611,6 +617,21 @@ enum llm_tensor { LLM_TENSOR_MASKED_EMBD_ORDERING, LLM_TENSOR_FC, LLM_TENSOR_D2T, + // Granite Switch: per-token embedded LoRA adapters (stacked over N = num_adapters + 1) + LLM_TENSOR_ATTN_QKV_LORA_A_Q, + LLM_TENSOR_ATTN_QKV_LORA_B_Q, + LLM_TENSOR_ATTN_QKV_LORA_A_K, + LLM_TENSOR_ATTN_QKV_LORA_B_K, + LLM_TENSOR_ATTN_QKV_LORA_A_V, + LLM_TENSOR_ATTN_QKV_LORA_B_V, + LLM_TENSOR_ATTN_OUT_LORA_A, + LLM_TENSOR_ATTN_OUT_LORA_B, + LLM_TENSOR_FFN_GATE_LORA_A, + LLM_TENSOR_FFN_GATE_LORA_B, + LLM_TENSOR_FFN_UP_LORA_A, + LLM_TENSOR_FFN_UP_LORA_B, + LLM_TENSOR_FFN_DOWN_LORA_A, + LLM_TENSOR_FFN_DOWN_LORA_B, }; diff --git a/src/llama-model.cpp b/src/llama-model.cpp index d87481381e46..46650ccff237 100644 --- a/src/llama-model.cpp +++ b/src/llama-model.cpp @@ -228,6 +228,8 @@ static llama_model * llama_model_mapping(llm_arch arch, const llama_model_params return new llama_model_granite(params); case LLM_ARCH_GRANITE_MOE: return new llama_model_granite_moe(params); + case LLM_ARCH_GRANITE_SWITCH: + return new llama_model_granite_switch(params); case LLM_ARCH_MINICPM: return new llama_model_minicpm(params); case LLM_ARCH_GRANITE_HYBRID: @@ -1885,6 +1887,7 @@ void llama_model::print_info() const { arch == LLM_ARCH_GRANITE || arch == LLM_ARCH_GRANITE_MOE || arch == LLM_ARCH_GRANITE_HYBRID || + arch == LLM_ARCH_GRANITE_SWITCH || arch == LLM_ARCH_NEMOTRON_H_MOE) { LLAMA_LOG_INFO("%s: f_embedding_scale = %f\n", __func__, hparams.f_embedding_scale); LLAMA_LOG_INFO("%s: f_residual_scale = %f\n", __func__, hparams.f_residual_scale); @@ -2458,6 +2461,7 @@ llama_rope_type llama_model_rope_type(const llama_model * model) { case LLM_ARCH_GRANITE: case LLM_ARCH_GRANITE_MOE: case LLM_ARCH_GRANITE_HYBRID: + case LLM_ARCH_GRANITE_SWITCH: case LLM_ARCH_CHAMELEON: case LLM_ARCH_BAILINGMOE: case LLM_ARCH_NEO_BERT: diff --git a/src/llama-model.h b/src/llama-model.h index 45b054cedf1d..871fa96ed50b 100644 --- a/src/llama-model.h +++ b/src/llama-model.h @@ -221,6 +221,30 @@ struct llama_layer_nextn { struct ggml_tensor * shared_head_norm = nullptr; }; +// granite-switch: per-token embedded LoRA adapters, stacked across N = num_adapters + 1 +// slots in dim 2 (slot 0 is zero-filled so base tokens add an exact-zero delta). +// A: {n_embd_in, max_lora_rank, N} B: {max_lora_rank, n_embd_out, N} +// 7 injection sites (q, k, v, o, gate, up, down), each with an A and a B = 14 tensors. +struct llama_layer_switch_lora { + // attention: q/k/v are separate slices of the fused qkv projection (different out dims) + struct ggml_tensor * a_q = nullptr; + struct ggml_tensor * b_q = nullptr; + struct ggml_tensor * a_k = nullptr; + struct ggml_tensor * b_k = nullptr; + struct ggml_tensor * a_v = nullptr; + struct ggml_tensor * b_v = nullptr; + struct ggml_tensor * a_o = nullptr; + struct ggml_tensor * b_o = nullptr; + + // feed-forward + struct ggml_tensor * a_gate = nullptr; + struct ggml_tensor * b_gate = nullptr; + struct ggml_tensor * a_up = nullptr; + struct ggml_tensor * b_up = nullptr; + struct ggml_tensor * a_down = nullptr; + struct ggml_tensor * b_down = nullptr; +}; + struct llama_layer { // normalization struct ggml_tensor * attn_norm = nullptr; @@ -525,6 +549,8 @@ struct llama_layer { struct llama_layer_shortconv shortconv; struct llama_layer_nextn nextn; + + struct llama_layer_switch_lora switch_lora; }; struct llama_device { diff --git a/src/models/granite_switch.cpp b/src/models/granite_switch.cpp new file mode 100644 index 000000000000..3f500e65bcf0 --- /dev/null +++ b/src/models/granite_switch.cpp @@ -0,0 +1,407 @@ +#include "models.h" + +#include + +// ============================================================================ +// Granite Switch: a dense, all-attention Granite-4.1 model with N embedded LoRA +// adapters selected per-token by control tokens. +// +// Two cheap CPU steps (no router weights) implement the switch, done in +// llm_graph_input_switch::set_input: +// 1. sticky per-token index — each position takes the index of the most recent +// control token at-or-before it (0 = base, i+1 = adapter i; cannot revert). +// 2. token-exchange — rewrite each control token id to its substitute id before +// embedding (read index first, then swap). +// +// Then a standard Granite decoder, with LoRA added per-token on qkv / o / gate / +// up / down via ggml_mul_mat_id over stacked tensors. The stacked dim is +// N = num_adapters + 1; slot 0 is zero-filled so base tokens add an exact-zero +// delta. B is pre-scaled by alpha/rank at compose time, so the runtime LoRA +// scale is 1.0. +// ============================================================================ + +void llama_model_granite_switch::load_arch_hparams(llama_model_loader & ml) { + // Granite multipliers / eps / rope (mirrors llama_model_granite). + ml.get_key(LLM_KV_ATTENTION_LAYERNORM_RMS_EPS, hparams.f_norm_rms_eps); + ml.get_key(LLM_KV_LOGIT_SCALE, hparams.f_logit_scale); + ml.get_key(LLM_KV_RESIDUAL_SCALE, hparams.f_residual_scale, false); + ml.get_key(LLM_KV_EMBEDDING_SCALE, hparams.f_embedding_scale, false); + ml.get_key(LLM_KV_ATTENTION_SCALE, hparams.f_attention_scale, false); + + // Granite uses rope_finetuned as a switch for rope, so default to true. + bool rope_finetuned = true; + ml.get_key(LLM_KV_ROPE_SCALING_FINETUNED, rope_finetuned, false); + hparams.rope_finetuned = rope_finetuned; + + switch (hparams.n_layer()) { + case 40: type = LLM_TYPE_3B; break; + default: type = LLM_TYPE_UNKNOWN; + } + + // Shared FFN length (dense path uses n_ff; this is informational/parity). + ml.get_key(LLM_KV_EXPERT_SHARED_FEED_FORWARD_LENGTH, hparams.n_ff_shexp, /* required */ false); + + // The per-token LoRA tensors use GGML_OP_MUL_MAT_ID with n_expert_used = 1. + // The model is dense (n_expert == 0, so the generic loader checks require + // n_expert_used == 0), but the buffer-type probe for MUL_MAT_ID tensors in + // create_tensor builds a synthetic op sized by hparams.n_expert_used and + // asserts it is > 0. Set it to 1 here (after the generic checks have run) so + // the probe matches how we actually invoke mul_mat_id. + hparams.n_expert_used = 1; + + // --- switch metadata --- + ml.get_key(LLM_KV_NUM_ADAPTERS, n_adapters); + ml.get_key(LLM_KV_MAX_LORA_RANK, max_lora_rank); + n_slots = n_adapters + 1; + + std::vector token_ids; + std::vector substitute_ids; + ml.get_arr(LLM_KV_ADAPTER_TOKEN_IDS, token_ids); + ml.get_arr(LLM_KV_ADAPTER_SUBSTITUTE_TOKEN_IDS, substitute_ids); + + GGML_ASSERT(token_ids.size() == n_adapters && + "granite-switch: adapter_token_ids length must equal num_adapters"); + GGML_ASSERT(substitute_ids.size() == n_adapters && + "granite-switch: adapter_substitute_token_ids length must equal num_adapters"); + + control_token_to_index.clear(); + control_token_to_substitute.clear(); + for (uint32_t i = 0; i < n_adapters; ++i) { + // adapter i -> stacked slot i+1 (slot 0 is the base/zero delta). + control_token_to_index[(llama_token) token_ids[i]] = (int32_t) (i + 1); + control_token_to_substitute[(llama_token) token_ids[i]] = (llama_token) substitute_ids[i]; + } +} + +void llama_model_granite_switch::load_arch_tensors(llama_model_loader &) { + LLAMA_LOAD_LOCALS; + + const int64_t n_slots_i64 = (int64_t) n_slots; + const int64_t n_rank = (int64_t) max_lora_rank; + const int64_t n_embd_q = n_embd_head_k * n_head; // fused Q out dim + const int64_t n_embd_kv = n_embd_k_gqa; // K (== V) out dim per slice + + tok_embd = create_tensor(tn(LLM_TENSOR_TOKEN_EMBD, "weight"), {n_embd, n_vocab}, 0); + + // output + 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}, TENSOR_NOT_REQUIRED); + if (output == NULL) { + // tied to the input embeddings + output = create_tensor(tn(LLM_TENSOR_TOKEN_EMBD, "weight"), {n_embd, n_vocab}, TENSOR_DUPLICATED); + } + + for (int i = 0; i < n_layer; ++i) { + auto & layer = layers[i]; + + layer.attn_norm = create_tensor(tn(LLM_TENSOR_ATTN_NORM, "weight", i), {n_embd}, 0); + + // fused qkv base + separate o + layer.wqkv = create_tensor(tn(LLM_TENSOR_ATTN_QKV, "weight", i), {n_embd, n_embd_q + 2*n_embd_kv}, 0); + layer.wo = create_tensor(tn(LLM_TENSOR_ATTN_OUT, "weight", i), {n_embd_q, n_embd}, 0); + + layer.ffn_norm = create_tensor(tn(LLM_TENSOR_FFN_NORM, "weight", i), {n_embd}, 0); + + // dense (non-MoE) FFN + layer.ffn_gate = create_tensor(tn(LLM_TENSOR_FFN_GATE, "weight", i), {n_embd, n_ff}, 0); + 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); + + // ---- stacked per-token LoRA tensors (no .weight suffix) ---- + // A: {n_in, max_rank, N} B: {max_rank, n_out, N} + auto & sl = layer.switch_lora; + + sl.a_q = create_tensor(tn(LLM_TENSOR_ATTN_QKV_LORA_A_Q, i), {n_embd, n_rank, n_slots_i64}, 0); + sl.b_q = create_tensor(tn(LLM_TENSOR_ATTN_QKV_LORA_B_Q, i), {n_rank, n_embd_q, n_slots_i64}, 0); + sl.a_k = create_tensor(tn(LLM_TENSOR_ATTN_QKV_LORA_A_K, i), {n_embd, n_rank, n_slots_i64}, 0); + sl.b_k = create_tensor(tn(LLM_TENSOR_ATTN_QKV_LORA_B_K, i), {n_rank, n_embd_kv, n_slots_i64}, 0); + sl.a_v = create_tensor(tn(LLM_TENSOR_ATTN_QKV_LORA_A_V, i), {n_embd, n_rank, n_slots_i64}, 0); + sl.b_v = create_tensor(tn(LLM_TENSOR_ATTN_QKV_LORA_B_V, i), {n_rank, n_embd_kv, n_slots_i64}, 0); + + sl.a_o = create_tensor(tn(LLM_TENSOR_ATTN_OUT_LORA_A, i), {n_embd_q, n_rank, n_slots_i64}, 0); + sl.b_o = create_tensor(tn(LLM_TENSOR_ATTN_OUT_LORA_B, i), {n_rank, n_embd, n_slots_i64}, 0); + + sl.a_gate = create_tensor(tn(LLM_TENSOR_FFN_GATE_LORA_A, i), {n_embd, n_rank, n_slots_i64}, 0); + sl.b_gate = create_tensor(tn(LLM_TENSOR_FFN_GATE_LORA_B, i), {n_rank, n_ff, n_slots_i64}, 0); + sl.a_up = create_tensor(tn(LLM_TENSOR_FFN_UP_LORA_A, i), {n_embd, n_rank, n_slots_i64}, 0); + sl.b_up = create_tensor(tn(LLM_TENSOR_FFN_UP_LORA_B, i), {n_rank, n_ff, n_slots_i64}, 0); + sl.a_down = create_tensor(tn(LLM_TENSOR_FFN_DOWN_LORA_A, i), { n_ff, n_rank, n_slots_i64}, 0); + sl.b_down = create_tensor(tn(LLM_TENSOR_FFN_DOWN_LORA_B, i), {n_rank, n_embd, n_slots_i64}, 0); + } +} + +// ---------------------------------------------------------------------------- +// switch input: compute sticky adapter index + token-exchanged ids per token. +// ---------------------------------------------------------------------------- +void llm_graph_input_switch::set_input(const llama_ubatch * ubatch) { + if (!ubatch->token) { + return; // embedding input path is not used by this arch + } + + const int64_t n_tokens = ubatch->n_tokens; + + std::vector ids(n_tokens); + std::vector sub(n_tokens); + + // POC single-sequence sticky state: carry the last index across ubatches, + // resetting when this ubatch starts a fresh sequence (contains pos 0). + bool has_pos0 = false; + if (ubatch->pos) { + for (int64_t i = 0; i < n_tokens; ++i) { + if (ubatch->pos[i] == 0) { has_pos0 = true; break; } + } + } + if (has_pos0) { + smodel.poc_sticky_index = 0; + } + + int32_t cur = smodel.poc_sticky_index; + for (int64_t i = 0; i < n_tokens; ++i) { + const llama_token tok = ubatch->token[i]; + + // 1. read index first (sticky): update on a control token. + auto it = smodel.control_token_to_index.find(tok); + if (it != smodel.control_token_to_index.end()) { + cur = it->second; + } + ids[i] = cur; + + // 2. then token-exchange: rewrite control token to its substitute. + auto sit = smodel.control_token_to_substitute.find(tok); + sub[i] = (sit != smodel.control_token_to_substitute.end()) + ? (int32_t) sit->second + : (int32_t) tok; + } + smodel.poc_sticky_index = cur; + + ggml_backend_tensor_set(adapter_ids, ids.data(), 0, n_tokens*ggml_element_size(adapter_ids)); + ggml_backend_tensor_set(sub_tokens, sub.data(), 0, n_tokens*ggml_element_size(sub_tokens)); +} + +// ---------------------------------------------------------------------------- +// graph +// ---------------------------------------------------------------------------- +std::unique_ptr llama_model_granite_switch::build_arch_graph(const llm_graph_params & params) const { + return std::make_unique(*this, params); +} + +// per-token switched LoRA delta only: B_a·(A_a·x), selected per token. +// cur: {n_in, n_tokens}, ids: {n_tokens} -> returns {n_out, n_tokens}. +ggml_tensor * llama_model_granite_switch::graph::build_switched_lora_delta( + ggml_tensor * lora_a, + ggml_tensor * lora_b, + ggml_tensor * cur, + ggml_tensor * ids) { + const int64_t n_in = cur->ne[0]; + const int64_t n_tokens = cur->ne[1]; + + // mul_mat_id wants the activation as {n_in, n_expert_used=1, n_tokens} + // and ids as {n_expert_used=1, n_tokens}. + ggml_tensor * x = ggml_reshape_3d(ctx0, cur, n_in, 1, n_tokens); + ggml_tensor * ids2 = ggml_reshape_2d(ctx0, ids, 1, n_tokens); + + ggml_tensor * a = ggml_mul_mat_id(ctx0, lora_a, x, ids2); // {max_rank, 1, n_tokens} + ggml_tensor * d = ggml_mul_mat_id(ctx0, lora_b, a, ids2); // {n_out, 1, n_tokens} + + return ggml_reshape_2d(ctx0, d, d->ne[0], n_tokens); // {n_out, n_tokens} +} + +// per-token switched matmul: base (plain 2D weight) + selected LoRA delta. +ggml_tensor * llama_model_granite_switch::graph::build_switched_lora_mm( + ggml_tensor * w, + ggml_tensor * lora_a, + ggml_tensor * lora_b, + ggml_tensor * cur, // {n_in, n_tokens} + ggml_tensor * ids) { // {n_tokens} + ggml_tensor * base = ggml_mul_mat(ctx0, w, cur); // {n_out, n_tokens} + ggml_tensor * delta = build_switched_lora_delta(lora_a, lora_b, cur, ids); + return ggml_add(ctx0, base, delta); +} + +llama_model_granite_switch::graph::graph( + const llama_model & model, + const llm_graph_params & params) + : llm_graph_context(params) { + + const auto & smodel = static_cast(model); + + const int64_t n_embd_head = hparams.n_embd_head_v(); + GGML_ASSERT(n_embd_head == hparams.n_embd_head_k()); + GGML_ASSERT(n_embd_head == n_rot); + + // --- switch input: substituted ids + per-token adapter indices --- + auto inp_switch = std::make_unique(smodel); + inp_switch->sub_tokens = ggml_new_tensor_1d(ctx0, GGML_TYPE_I32, n_tokens); + inp_switch->adapter_ids = ggml_new_tensor_1d(ctx0, GGML_TYPE_I32, n_tokens); + ggml_set_input(inp_switch->sub_tokens); + ggml_set_input(inp_switch->adapter_ids); + ggml_tensor * sub_tokens = inp_switch->sub_tokens; + ggml_tensor * adapter_ids = inp_switch->adapter_ids; + res->add_input(std::move(inp_switch)); + + // embed the token-exchanged ids ourselves (we cannot use build_inp_embd + // because it embeds the raw ubatch tokens). Apply Granite embedding scale. + ggml_tensor * inpL = ggml_get_rows(ctx0, model.tok_embd, sub_tokens); + if (hparams.f_embedding_scale != 0.0f) { + inpL = ggml_scale(ctx0, inpL, hparams.f_embedding_scale); + } + cb(inpL, "inp_embd", -1); + + ggml_tensor * inp_pos = nullptr; + if (hparams.rope_finetuned) { + inp_pos = build_inp_pos(); + } + auto * inp_attn = build_attn_inp_kv(); + + ggml_tensor * inp_out_ids = build_inp_out_ids(); + + ggml_tensor * cur; + + for (int il = 0; il < n_layer; ++il) { + ggml_tensor * inpSA = inpL; + + // attn norm + cur = build_norm(inpL, model.layers[il].attn_norm, NULL, LLM_NORM_RMS, il); + cb(cur, "attn_norm", il); + + // self-attention (with per-token switched LoRA) + cur = build_attention_layer(cur, inp_pos, adapter_ids, inp_attn, model, n_embd_head, 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); + // adapter_ids is 1D {n_tokens}; ggml_get_rows on a 1D tensor would + // gather whole rows of width n_tokens. Reshape to {1, n_tokens} so we + // select per-token columns, then flatten back to 1D {n_out}. + const int64_t n_out = inp_out_ids->ne[0]; + adapter_ids = ggml_get_rows(ctx0, + ggml_reshape_2d(ctx0, adapter_ids, 1, adapter_ids->ne[0]), inp_out_ids); + adapter_ids = ggml_reshape_1d(ctx0, adapter_ids, n_out); + } + + // ffn (with per-token switched LoRA) + cur = build_layer_ffn(cur, inpSA, adapter_ids, model, 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; + + // lm_head + cur = build_lora_mm(model.output, cur, model.output_s); + + // Granite logit scaling + cur = ggml_scale(ctx0, cur, 1.0f / hparams.f_logit_scale); + cb(cur, "result_output", -1); + res->t_logits = cur; + + ggml_build_forward_expand(gf, cur); +} + +ggml_tensor * llama_model_granite_switch::graph::build_attention_layer( + ggml_tensor * cur, + ggml_tensor * inp_pos, + ggml_tensor * adapter_ids, + llm_graph_input_attn_kv * inp_attn, + const llama_model & model, + const int64_t n_embd_head, + const int il) { + + const auto & layer = model.layers[il]; + const auto & sl = layer.switch_lora; + + const int64_t n_head = hparams.n_head(il); + const int64_t n_head_kv = hparams.n_head_kv(il); + + // fused qkv base + per-slice switched LoRA deltas. + ggml_tensor * qkv = ggml_mul_mat(ctx0, layer.wqkv, cur); + cb(qkv, "wqkv", il); + + const int64_t n_embd_q = n_embd_head * n_head; + const int64_t n_embd_kv = n_embd_head * n_head_kv; + + // slice base qkv into Q/K/V (contiguous copies so we can add LoRA deltas). + ggml_tensor * Qcur = ggml_cont(ctx0, ggml_view_2d(ctx0, qkv, n_embd_q, qkv->ne[1], qkv->nb[1], 0)); + ggml_tensor * Kcur = ggml_cont(ctx0, ggml_view_2d(ctx0, qkv, n_embd_kv, qkv->ne[1], qkv->nb[1], n_embd_q*ggml_element_size(qkv))); + ggml_tensor * Vcur = ggml_cont(ctx0, ggml_view_2d(ctx0, qkv, n_embd_kv, qkv->ne[1], qkv->nb[1], (n_embd_q + n_embd_kv)*ggml_element_size(qkv))); + + // add per-token switched LoRA deltas to each base slice (delta only — the + // fused base qkv is already computed above). + Qcur = ggml_add(ctx0, Qcur, build_switched_lora_delta(sl.a_q, sl.b_q, cur, adapter_ids)); + Kcur = ggml_add(ctx0, Kcur, build_switched_lora_delta(sl.a_k, sl.b_k, cur, adapter_ids)); + Vcur = ggml_add(ctx0, Vcur, build_switched_lora_delta(sl.a_v, sl.b_v, cur, adapter_ids)); + + Qcur = ggml_reshape_3d(ctx0, Qcur, n_embd_head, n_head, n_tokens); + Kcur = ggml_reshape_3d(ctx0, Kcur, n_embd_head, n_head_kv, n_tokens); + Vcur = ggml_reshape_3d(ctx0, Vcur, n_embd_head, n_head_kv, n_tokens); + + if (hparams.rope_finetuned) { + ggml_tensor * rope_factors = model.get_rope_factors(cparams, il); + Qcur = ggml_rope_ext(ctx0, Qcur, inp_pos, rope_factors, + n_rot, rope_type, n_ctx_orig, freq_base, freq_scale, + ext_factor, attn_factor, beta_fast, beta_slow); + Kcur = ggml_rope_ext(ctx0, Kcur, inp_pos, rope_factors, + n_rot, rope_type, n_ctx_orig, freq_base, freq_scale, + ext_factor, attn_factor, beta_fast, beta_slow); + } + cb(Qcur, "Qcur", il); + cb(Kcur, "Kcur", il); + cb(Vcur, "Vcur", il); + + const float kq_scale = hparams.f_attention_scale == 0.0f + ? 1.0f/sqrtf(float(n_embd_head)) : hparams.f_attention_scale; + + // wo = nullptr: build_attn returns the concatenated heads (pre-o-proj), + // then we apply the switched o-proj manually. + ggml_tensor * attn = build_attn(inp_attn, + nullptr, nullptr, nullptr, + Qcur, Kcur, Vcur, nullptr, nullptr, nullptr, kq_scale, il); + cb(attn, "attn_pre_o", il); + + cur = build_switched_lora_mm(layer.wo, sl.a_o, sl.b_o, attn, adapter_ids); + cb(cur, "attn_out", il); + return cur; +} + +ggml_tensor * llama_model_granite_switch::graph::build_layer_ffn( + ggml_tensor * cur, + ggml_tensor * inpSA, + ggml_tensor * adapter_ids, + const llama_model & model, + const int il) { + + const auto & layer = model.layers[il]; + const auto & sl = layer.switch_lora; + + // Granite residual scale + if (hparams.f_residual_scale) { + cur = ggml_scale(ctx0, cur, hparams.f_residual_scale); + } + 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); + + // gated SILU FFN with per-token switched LoRA on gate / up / down. + ggml_tensor * g = build_switched_lora_mm(layer.ffn_gate, sl.a_gate, sl.b_gate, cur, adapter_ids); + ggml_tensor * u = build_switched_lora_mm(layer.ffn_up, sl.a_up, sl.b_up, cur, adapter_ids); + g = ggml_silu(ctx0, g); + ggml_tensor * gu = ggml_mul(ctx0, g, u); + cur = build_switched_lora_mm(layer.ffn_down, sl.a_down, sl.b_down, gu, adapter_ids); + cb(cur, "ffn_out", il); + + if (hparams.f_residual_scale) { + cur = ggml_scale(ctx0, cur, hparams.f_residual_scale); + } + cur = ggml_add(ctx0, cur, ffn_inp); + + cur = build_cvec(cur, il); + cb(cur, "l_out", il); + + return cur; +} diff --git a/src/models/models.h b/src/models/models.h index 7a52e7bc1ab7..fe026d3b5b3f 100644 --- a/src/models/models.h +++ b/src/models/models.h @@ -1536,6 +1536,90 @@ struct llama_model_granite_moe : public llama_model_base { }; +struct llama_model_granite_switch : public llama_model_base { + llama_model_granite_switch(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; + + // --- per-token switch metadata (read from GGUF in load_arch_hparams) --- + uint32_t n_adapters = 0; // number of real adapters (excludes the zero slot) + uint32_t n_slots = 0; // n_adapters + 1 (slot 0 = base/zero delta) + uint32_t max_lora_rank = 0; + + // token id -> sticky adapter index (1 + adapter, so slot 0 stays "base") + std::unordered_map control_token_to_index; + // control token id -> substitute token id (token-exchange before embedding) + std::unordered_map control_token_to_substitute; + + // POC: single-sequence sticky adapter state. The most recent control index + // seen, carried across ubatches within one decode. Reset to 0 when a ubatch + // contains sequence position 0 (start of a fresh prompt). This is sufficient + // for llama-cli / the smoke + parity tests; a full multi-sequence + // (reset-on-seq_rm) implementation is deferred — see the [[llamacpp-poc-scope]] + // note. `mutable` because set_input runs from a const-model graph context. + mutable int32_t poc_sticky_index = 0; + + struct graph : public llm_graph_context { + graph(const llama_model & model, const llm_graph_params & params); + + private: + // per-token switched LoRA delta only: B_a·(A_a·x), selected per token. + // Returns {n_out, n_tokens}. Used for the fused-qkv slices where the base + // term is already computed. + ggml_tensor * build_switched_lora_delta( + ggml_tensor * lora_a, // {n_in, max_lora_rank, N} + ggml_tensor * lora_b, // {max_lora_rank, n_out, N} + ggml_tensor * cur, // {n_in, n_tokens} + ggml_tensor * ids); // {n_tokens} I32 adapter index per token + + // per-token switched LoRA: y = W·x + B_a·(A_a·x), selected per token via mul_mat_id + ggml_tensor * build_switched_lora_mm( + ggml_tensor * w, // base weight (plain 2D) + ggml_tensor * lora_a, // {n_in, max_lora_rank, N} + ggml_tensor * lora_b, // {max_lora_rank, n_out, N} + ggml_tensor * cur, // {n_in, n_tokens} + ggml_tensor * ids); // {n_tokens} I32 adapter index per token + + ggml_tensor * build_attention_layer( + ggml_tensor * cur, + ggml_tensor * inp_pos, + ggml_tensor * adapter_ids, + llm_graph_input_attn_kv * inp_attn, + const llama_model & model, + const int64_t n_embd_head, + const int il); + + ggml_tensor * build_layer_ffn( + ggml_tensor * cur, + ggml_tensor * inpSA, + ggml_tensor * adapter_ids, + const llama_model & model, + const int il); + }; + + std::unique_ptr build_arch_graph(const llm_graph_params & params) const override; +}; + + +// Custom graph input for the per-token switch. In set_input it scans the ubatch, +// computes the sticky adapter index per token (writing adapter_ids), and rewrites +// control-token ids to their substitute ids (writing sub_tokens) — the +// token-exchange. Both tensors are I32 [n_tokens]. can_reuse() returns false so +// these are recomputed every ubatch. +class llm_graph_input_switch : public llm_graph_input_i { +public: + llm_graph_input_switch(const llama_model_granite_switch & smodel) : smodel(smodel) {} + virtual ~llm_graph_input_switch() = default; + + void set_input(const llama_ubatch * ubatch) override; + + ggml_tensor * sub_tokens = nullptr; // I32 [n_tokens] control-substituted token ids + ggml_tensor * adapter_ids = nullptr; // I32 [n_tokens] sticky per-token adapter index + + const llama_model_granite_switch & smodel; +}; + + struct llama_model_minicpm : public llama_model_base { llama_model_minicpm(const struct llama_model_params & params) : llama_model_base(params) {} void load_arch_hparams(llama_model_loader & ml) override; From 7633e95493b94c1442c989f0a650212d4be5b6b0 Mon Sep 17 00:00:00 2001 From: Bar Haim Date: Thu, 18 Jun 2026 13:29:00 -0700 Subject: [PATCH 02/40] granite-switch: add Mac (Metal) build + mid-sequence switch demo script Self-contained script to build llama.cpp on Apple Silicon (Metal), convert the composed 3b checkpoint, and run the crisp mid-sequence adapter-switch demos verified on Vela: - answerability: <|answerability|> mid-seq -> "unanswerable" - query_rewrite: <|query_rewrite|> mid-seq -> {"rewritten_question": ...} Each demo runs the same prompt twice, differing only by a control token placed before the assistant turn, so the per-token switch is visible. --- granite-switch-mac-demo.sh | 101 +++++++++++++++++++++++++++++++++++++ 1 file changed, 101 insertions(+) create mode 100755 granite-switch-mac-demo.sh diff --git a/granite-switch-mac-demo.sh b/granite-switch-mac-demo.sh new file mode 100755 index 000000000000..61c20d3b22ed --- /dev/null +++ b/granite-switch-mac-demo.sh @@ -0,0 +1,101 @@ +#!/usr/bin/env bash +# granite-switch-mac-demo.sh +# +# End-to-end Granite Switch demo on an Apple-Silicon Mac (Metal GPU backend). +# Builds llama.cpp from this branch, converts the composed checkpoint to GGUF, +# and runs the crisp mid-sequence adapter-switch demos verified on Vela: +# - answerability : <|answerability|> -> "unanswerable" (base would just answer) +# - query_rewrite : <|query_rewrite|> -> {"rewritten_question": ...} (base would answer) +# +# Each demo runs the SAME prompt twice, differing only by a control token placed +# mid-sequence (right before the assistant turn). The outputs differ structurally, +# which is the switch firing per-token. +# +# Usage: +# ./granite-switch-mac-demo.sh # build + convert (first run) then demo +# ./granite-switch-mac-demo.sh demo # skip build/convert, just run the demos +# GGUF=/path/to/gs-f16.gguf ./granite-switch-mac-demo.sh demo # use an existing GGUF +# +# Requirements: Xcode CLT (`xcode-select --install`), cmake (`brew install cmake`), +# python3, and ~16GB+ unified memory for the f16 model (8.4 GB on disk). + +set -euo pipefail + +REPO_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +cd "$REPO_DIR" + +SRC_MODEL="${SRC_MODEL:-ibm-granite/granite-switch-4.1-3b-preview}" +GGUF="${GGUF:-$REPO_DIR/gs-f16.gguf}" +BIN="$REPO_DIR/build/bin/llama-completion" +NGL="${NGL:-99}" # offload all layers to the Metal GPU; set NGL=0 to force CPU + +# --------------------------------------------------------------------------- +build() { + echo "=== Build llama.cpp (Metal is on by default on macOS) ===" + if [ ! -x "$BIN" ]; then + cmake -S . -B build -DLLAMA_CURL=OFF + cmake --build build -j --target llama llama-completion llama-tokenize + else + echo " (build/bin/llama-completion already present — skipping)" + fi +} + +convert() { + if [ -f "$GGUF" ]; then + echo "=== GGUF already present at $GGUF — skipping convert ===" + return + fi + echo "=== Set up python deps for the converter ===" + python3 -m venv .venv-convert + # shellcheck disable=SC1091 + . .venv-convert/bin/activate + pip install --upgrade pip + pip install numpy torch transformers safetensors sentencepiece "huggingface_hub[cli]" + + echo "=== Download $SRC_MODEL (~8 GB) ===" + SRC_DIR="$(python -c "from huggingface_hub import snapshot_download; print(snapshot_download('$SRC_MODEL'))")" + + echo "=== Convert HF -> GGUF (f16) ===" + PYTHONPATH=gguf-py python convert_hf_to_gguf.py "$SRC_DIR" --outfile "$GGUF" --outtype f16 + deactivate +} + +# run NAME PROMPT N +run() { + local name="$1" prompt="$2" n="${3:-16}" + echo "=== $name ===" + "$BIN" -m "$GGUF" -ngl "$NGL" --temp 0 -n "$n" -p "$prompt" 2>/dev/null + echo; echo "-----" +} + +demo() { + [ -x "$BIN" ] || { echo "FATAL: $BIN not built — run without 'demo' first"; exit 1; } + [ -f "$GGUF" ] || { echo "FATAL: $GGUF missing — run without 'demo' first"; exit 1; } + + echo + echo "############ DEMO 1: answerability (<|answerability|>, id 100356) ############" + echo "# Document is about the Eiffel Tower; question asks Australia's capital." + echo "# Base ANSWERS the question; adapter judges it UNANSWERABLE from the doc." + local ANS_CTX='<|start_of_role|>user<|end_of_role|>Document: The Eiffel Tower is in Paris. Question: What is the capital of Australia?<|end_of_role|><|end_of_text|><|start_of_role|>assistant<|end_of_role|>' + run "answerability OFF (base, no control)" "${ANS_CTX}" 12 + run "answerability ON (<|answerability|> mid-seq)" "${ANS_CTX}<|answerability|>" 12 + + echo + echo "############ DEMO 2: query_rewrite (<|query_rewrite|>, id 100353) ############" + echo "# Follow-up 'What other movies has HE made?' — base answers it;" + echo "# adapter REWRITES it into a standalone query (resolves 'he')." + local QR_CTX='<|start_of_role|>user<|end_of_role|>Who directed Inception?<|end_of_text|><|start_of_role|>assistant<|end_of_role|>Christopher Nolan directed Inception.<|end_of_text|><|start_of_role|>user<|end_of_role|>What other movies has he made?<|end_of_text|><|start_of_role|>assistant<|end_of_role|>' + run "query_rewrite OFF (base, no control)" "${QR_CTX}" 24 + run "query_rewrite ON (<|query_rewrite|> mid-seq)" "${QR_CTX}<|query_rewrite|>" 24 + + echo "=== DONE. In each pair, only the mid-sequence control token differs. ===" +} + +# --------------------------------------------------------------------------- +case "${1:-all}" in + demo) demo ;; + build) build ;; + convert) convert ;; + all) build; convert; demo ;; + *) echo "usage: $0 [all|build|convert|demo]"; exit 1 ;; +esac From 80a470cfd8f99813e8b81a5983f6d6ab2196b772 Mon Sep 17 00:00:00 2001 From: Bar Haim Date: Thu, 18 Jun 2026 13:35:58 -0700 Subject: [PATCH 03/40] granite-switch mac demo: add -no-cnv so each run is one-shot The composed model ships a chat template, so llama-completion auto-enables interactive conversation mode and halts at a `>` prompt after generating, stalling the script. -no-cnv disables conversation mode: generate once from the raw prompt and exit (also prints special tokens, making the switch visible). --- granite-switch-mac-demo.sh | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/granite-switch-mac-demo.sh b/granite-switch-mac-demo.sh index 61c20d3b22ed..d14dda945a91 100755 --- a/granite-switch-mac-demo.sh +++ b/granite-switch-mac-demo.sh @@ -61,10 +61,14 @@ convert() { } # run NAME PROMPT N +# -no-cnv: this model ships a chat template, so llama-completion AUTO-ENABLES +# interactive conversation mode and stops at a `>` prompt waiting for input. +# -no-cnv disables that: generate once from the raw prompt and exit (and it +# also prints special tokens, which makes the adapter switch easy to see). run() { local name="$1" prompt="$2" n="${3:-16}" echo "=== $name ===" - "$BIN" -m "$GGUF" -ngl "$NGL" --temp 0 -n "$n" -p "$prompt" 2>/dev/null + "$BIN" -m "$GGUF" -ngl "$NGL" --temp 0 -n "$n" -no-cnv -p "$prompt" 2>/dev/null echo; echo "-----" } From 3096444ce1a902992cdb44842594534a7ef2baf0 Mon Sep 17 00:00:00 2001 From: Bar Haim Date: Mon, 22 Jun 2026 05:57:05 -0700 Subject: [PATCH 04/40] granite-switch: replace global sticky index with in-graph router attention MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The POC computed the per-token adapter index on the CPU and carried it across ubatches in ONE global `mutable int32_t poc_sticky_index`, reset only when a ubatch contained sequence position 0. That global had two problems: 1. Concurrency: with multiple sequences in a batch it was last-writer- wins — one sequence's adapter leaked into the others. 2. Multi-turn: an interactive `ollama run` chat continues one KV cache, so turn 2 never saw position 0 and the index never reset — the adapter stayed stuck on across turns. Port the vLLM/HF backend mechanism faithfully: a single-head causal "router" attention recovers the adapter index in-graph. Per token, only dim 0 carries signal — Q[0]=1, K[0]=+gain for a control token / -gain otherwise, V[0]=adapter slot / 0 — and the causal softmax over the single visible control token recovers that adapter's slot (readback = clamp(round(V[0]), 0, n_adapters)). gain=15 matches config.py and is F16-safe (no F32 cache). The router's K/V live in the model KV cache at an extra layer R == hparams.router_layer (== n_layer). We bump n_layer_all to n_real+1 so the cache allocator gives the router its own per-sequence slot, and set n_layer_nextn=1 so n_layer() stays n_real — the decoder loop and tensor loading are untouched and never reference layer R. The router K is exempted from the k-shift RoPE loop (its dim-0 value is a literal magnitude, not a rotation). Because the selection now lives in the per-sequence KV cache, CONCURRENT requests are isolated for free (problem 1 fixed; verified by scratch/concurrent_switch_test.cpp). set_input becomes stateless pure per-token maps; the global is gone. Single-switch contract / known limitation, identical to vLLM & HF: the gain is flat (no recency), so within one sequence there is no mechanism to revert to base mid-sequence — once an adapter fires it stays on until that sequence ends (problem 2 is therefore NOT fixed by a faithful copy; vLLM/HF avoid it only because each served request is a fresh sequence). A client continuing one KV cache across turns must start a fresh sequence per turn, or opt into a recency-biased router (a deliberate divergence, not done here). Documented in granite_switch.cpp and asserted by scratch/multiturn_leak_test.cpp. Verified (CPU): both demos unchanged (answerability -> "unanswerable", query_rewrite -> rewritten query); concurrent two-sequence isolation passes; multi-turn carry-over matches the vLLM/HF contract. --- scratch/concurrent_switch_test.cpp | 128 +++++++++++++++++++++++++ scratch/multiturn_leak_test.cpp | 134 ++++++++++++++++++++++++++ src/llama-hparams.h | 4 + src/llama-kv-cache.cpp | 7 ++ src/models/granite_switch.cpp | 147 ++++++++++++++++++++++------- src/models/models.h | 37 +++++--- 6 files changed, 410 insertions(+), 47 deletions(-) create mode 100644 scratch/concurrent_switch_test.cpp create mode 100644 scratch/multiturn_leak_test.cpp diff --git a/scratch/concurrent_switch_test.cpp b/scratch/concurrent_switch_test.cpp new file mode 100644 index 000000000000..8bb4062ace28 --- /dev/null +++ b/scratch/concurrent_switch_test.cpp @@ -0,0 +1,128 @@ +// Concurrent multi-sequence isolation test for granite-switch. +// +// Two sequences are decoded together in one batch / one context: +// seq 0: answerability ON (<|answerability|>) -> expect "unanswerable" +// seq 1: NO control token -> expect a normal answer +// +// The OLD global `poc_sticky_index` could not isolate these (last writer wins): +// seq 1 would inherit seq 0's adapter. The in-graph router stores each sequence's +// selection in its own KV-cache cells, so they must not cross-contaminate. +// +// PASS: seq0 contains "unanswerable" AND seq1 does NOT. + +#include "llama.h" +#include +#include +#include +#include + +static std::vector tokz(const llama_vocab * v, const std::string & s, bool bos) { + int n = -llama_tokenize(v, s.c_str(), s.size(), NULL, 0, bos, true); + std::vector out(n); + llama_tokenize(v, s.c_str(), s.size(), out.data(), out.size(), bos, true); + return out; +} + +int main(int argc, char ** argv) { + std::string model_path; + for (int i = 1; i < argc; ++i) + if (strcmp(argv[i], "-m") == 0 && i + 1 < argc) model_path = argv[++i]; + if (model_path.empty()) { fprintf(stderr, "usage: %s -m model.gguf\n", argv[0]); return 1; } + + ggml_backend_load_all(); + + llama_model_params mp = llama_model_default_params(); + mp.n_gpu_layers = 0; + llama_model * model = llama_model_load_from_file(model_path.c_str(), mp); + if (!model) { fprintf(stderr, "load failed\n"); return 1; } + const llama_vocab * vocab = llama_model_get_vocab(model); + + const std::string p0 = + "<|start_of_role|>user<|end_of_role|>Document: The Eiffel Tower is in Paris. " + "Question: What is the capital of Australia?<|end_of_role|><|end_of_text|>" + "<|start_of_role|>assistant<|end_of_role|><|answerability|>"; + const std::string p1 = + "<|start_of_role|>user<|end_of_role|>What is the capital of Australia?" + "<|end_of_role|><|end_of_text|><|start_of_role|>assistant<|end_of_role|>"; + + std::vector t0 = tokz(vocab, p0, true); + std::vector t1 = tokz(vocab, p1, true); + + llama_context_params cp = llama_context_default_params(); + cp.n_ctx = 2048; + cp.n_batch = 2048; + cp.n_seq_max = 2; + llama_context * ctx = llama_init_from_model(model, cp); + if (!ctx) { fprintf(stderr, "ctx failed\n"); return 1; } + + auto sp = llama_sampler_chain_default_params(); + llama_sampler * smpl = llama_sampler_chain_init(sp); + llama_sampler_chain_add(smpl, llama_sampler_init_greedy()); + + // build a batch holding BOTH prompts, seq 0 and seq 1 interleaved by sequence id. + const int n_gen = 14; + llama_batch batch = llama_batch_init((int) (t0.size() + t1.size()), 0, 2); + + auto add = [&](llama_token id, llama_pos pos, llama_seq_id seq, bool logits) { + int k = batch.n_tokens; + batch.token[k] = id; + batch.pos[k] = pos; + batch.n_seq_id[k] = 1; + batch.seq_id[k][0] = seq; + batch.logits[k] = logits; + batch.n_tokens++; + }; + + for (size_t i = 0; i < t0.size(); ++i) add(t0[i], (llama_pos) i, 0, i == t0.size() - 1); + for (size_t i = 0; i < t1.size(); ++i) add(t1[i], (llama_pos) i, 1, i == t1.size() - 1); + + if (llama_decode(ctx, batch)) { fprintf(stderr, "prompt decode failed\n"); return 1; } + + // remember the logits index for each sequence's last prompt token. + int idx0 = (int) t0.size() - 1; + int idx1 = (int) (t0.size() + t1.size()) - 1; + llama_pos pos0 = (llama_pos) t0.size(); + llama_pos pos1 = (llama_pos) t1.size(); + + std::string out0, out1; + for (int step = 0; step < n_gen; ++step) { + llama_token id0 = llama_sampler_sample(smpl, ctx, idx0); + llama_token id1 = llama_sampler_sample(smpl, ctx, idx1); + + auto piece = [&](llama_token id, std::string & dst) { + if (llama_vocab_is_eog(vocab, id)) return; + char buf[256]; + int n = llama_token_to_piece(vocab, id, buf, sizeof(buf), 0, true); + if (n > 0) dst.append(buf, n); + }; + piece(id0, out0); + piece(id1, out1); + + batch.n_tokens = 0; + bool e0 = llama_vocab_is_eog(vocab, id0); + bool e1 = llama_vocab_is_eog(vocab, id1); + if (e0 && e1) break; + if (!e0) { add(id0, pos0++, 0, true); idx0 = batch.n_tokens - 1; } + if (!e1) { add(id1, pos1++, 1, true); idx1 = batch.n_tokens - 1; } + if (llama_decode(ctx, batch)) { fprintf(stderr, "gen decode failed\n"); break; } + } + + printf("\n============ CONCURRENT ISOLATION TEST ============\n"); + printf("SEQ 0 (answerability ON): %s\n", out0.c_str()); + printf("SEQ 1 (no control token): %s\n", out1.c_str()); + auto has = [](const std::string & h, const char * n){ return h.find(n) != std::string::npos; }; + bool s0_ok = has(out0, "unanswerable"); + bool s1_clean = !has(out1, "unanswerable"); + printf("---------------------------------------------------\n"); + printf("seq0 fired its adapter: %s\n", s0_ok ? "YES" : "NO"); + printf("seq1 isolated from seq0: %s\n", s1_clean ? "clean" : "CONTAMINATED"); + bool pass = s0_ok && s1_clean; + printf("RESULT: %s\n", pass ? "PASS (isolated)" : "FAIL"); + printf("===================================================\n"); + + llama_batch_free(batch); + llama_sampler_free(smpl); + llama_free(ctx); + llama_model_free(model); + return pass ? 0 : 1; +} diff --git a/scratch/multiturn_leak_test.cpp b/scratch/multiturn_leak_test.cpp new file mode 100644 index 000000000000..6967f48d2559 --- /dev/null +++ b/scratch/multiturn_leak_test.cpp @@ -0,0 +1,134 @@ +// Multi-turn behavior probe for granite-switch (DOCUMENTS A KNOWN LIMITATION). +// +// Both turns run on the SAME llama_context — one persistent KV cache, no reset +// between turns, exactly like `ollama run` continuing a chat: +// turn 1: answerability ON (<|answerability|>) -> "unanswerable" +// turn 2: NO control token -> ??? (same cache) +// +// EXPECTED (documented) behavior: turn 2 STILL reads the adapter. This is the +// single-switch contract, identical to the vLLM/HF backends, whose own docstring +// states: "SingleSwitch has no mechanism to transition back to base +// mid-sequence." The in-graph router selection lives in the per-sequence KV +// cache; turn 1's control-token K is still causally visible to turn 2's tokens, +// and with flat (recency-free) gain it keeps winning the softmax. +// +// vLLM/HF never observe this because each served request is a FRESH sequence +// (turn 2's context does not contain turn 1's control token). A chat client that +// continues one cache across turns (ollama) must therefore start a fresh +// sequence per turn, OR opt into a recency-biased router (a deliberate +// divergence from vLLM — not done here). +// +// This probe asserts the DOCUMENTED behavior so it stays a faithful vLLM copy: +// turn 1 -> "unanswerable" +// turn 2 -> ALSO "unanswerable" (carry-over expected on a persisted cache) + +#include "llama.h" +#include +#include +#include +#include + +static std::string g_model; + +static std::vector tok(const llama_vocab * v, const std::string & s, bool bos) { + int n = -llama_tokenize(v, s.c_str(), s.size(), NULL, 0, bos, true); + std::vector out(n); + llama_tokenize(v, s.c_str(), s.size(), out.data(), out.size(), bos, true); + return out; +} + +// Decode `prompt_tokens` (appended at current position `pos`) then greedily +// generate up to n_gen tokens, advancing `pos`. Returns the generated text. +static std::string run_turn(llama_context * ctx, const llama_vocab * vocab, + llama_sampler * smpl, + std::vector prompt_tokens, + int & pos, int n_gen) { + std::string text; + + // decode the prompt chunk (positions assigned automatically by the batch helper + // continue from the cache, since we never clear it between turns). + llama_batch batch = llama_batch_get_one(prompt_tokens.data(), prompt_tokens.size()); + if (llama_decode(ctx, batch)) { fprintf(stderr, "decode (prompt) failed\n"); return text; } + pos += (int) prompt_tokens.size(); + + for (int i = 0; i < n_gen; ++i) { + llama_token id = llama_sampler_sample(smpl, ctx, -1); + if (llama_vocab_is_eog(vocab, id)) break; + char buf[256]; + int n = llama_token_to_piece(vocab, id, buf, sizeof(buf), 0, true); + if (n > 0) text.append(buf, n); + llama_batch nb = llama_batch_get_one(&id, 1); + if (llama_decode(ctx, nb)) { fprintf(stderr, "decode (gen) failed\n"); break; } + pos += 1; + } + return text; +} + +int main(int argc, char ** argv) { + for (int i = 1; i < argc; ++i) { + if (strcmp(argv[i], "-m") == 0 && i + 1 < argc) g_model = argv[++i]; + } + if (g_model.empty()) { fprintf(stderr, "usage: %s -m model.gguf\n", argv[0]); return 1; } + + ggml_backend_load_all(); + + llama_model_params mp = llama_model_default_params(); + mp.n_gpu_layers = 0; + llama_model * model = llama_model_load_from_file(g_model.c_str(), mp); + if (!model) { fprintf(stderr, "load failed\n"); return 1; } + const llama_vocab * vocab = llama_model_get_vocab(model); + + llama_context_params cp = llama_context_default_params(); + cp.n_ctx = 2048; + cp.n_batch = 2048; + llama_context * ctx = llama_init_from_model(model, cp); + if (!ctx) { fprintf(stderr, "ctx failed\n"); return 1; } + + auto sp = llama_sampler_chain_default_params(); + llama_sampler * smpl = llama_sampler_chain_init(sp); + llama_sampler_chain_add(smpl, llama_sampler_init_greedy()); + + int pos = 0; + + // --- TURN 1: answerability ON (control token fires the adapter) --- + // Doc is about the Eiffel Tower; question asks Australia's capital -> adapter + // should judge it "unanswerable". + std::string t1 = + "<|start_of_role|>user<|end_of_role|>Document: The Eiffel Tower is in Paris. " + "Question: What is the capital of Australia?<|end_of_role|><|end_of_text|>" + "<|start_of_role|>assistant<|end_of_role|><|answerability|>"; + std::string out1 = run_turn(ctx, vocab, smpl, tok(vocab, t1, /*bos=*/true), pos, 12); + + // --- TURN 2: NO control token, same context (cache persists) --- + // A normal question. If the adapter leaked from turn 1, this would also be + // judged "unanswerable". With the router fix it answers normally. + std::string t2 = + "<|end_of_text|><|start_of_role|>user<|end_of_role|>" + "What is the capital of Australia?<|end_of_role|><|end_of_text|>" + "<|start_of_role|>assistant<|end_of_role|>"; + std::string out2 = run_turn(ctx, vocab, smpl, tok(vocab, t2, /*bos=*/false), pos, 16); + + printf("\n========== MULTI-TURN CARRY-OVER PROBE (KNOWN LIMIT) ==========\n"); + printf("TURN 1 (answerability ON): %s\n", out1.c_str()); + printf("TURN 2 (no control token): %s\n", out2.c_str()); + + auto contains = [](const std::string & h, const char * n) { + return h.find(n) != std::string::npos; + }; + bool t1_ok = contains(out1, "unanswerable"); + bool t2_carried = contains(out2, "unanswerable"); + + printf("--------------------------------------------------------------\n"); + printf("turn1 fired adapter (expect unanswerable): %s\n", t1_ok ? "YES" : "NO"); + printf("turn2 carried adapter (DOCUMENTED: expect carry): %s\n", t2_carried ? "carried" : "reverted"); + // Faithful vLLM/HF contract: on a persisted cache, selection carries over. + bool pass = t1_ok && t2_carried; + printf("RESULT: %s\n", pass ? "PASS (matches vLLM/HF single-switch contract)" + : "UNEXPECTED (behavior differs from vLLM/HF)"); + printf("==============================================================\n"); + + llama_sampler_free(smpl); + llama_free(ctx); + llama_model_free(model); + return pass ? 0 : 1; +} diff --git a/src/llama-hparams.h b/src/llama-hparams.h index 8be5f28f39e6..7199f43f1a77 100644 --- a/src/llama-hparams.h +++ b/src/llama-hparams.h @@ -52,6 +52,10 @@ struct llama_hparams { uint32_t n_embd; uint32_t n_layer_all; uint32_t n_layer_nextn = 0; + + // granite-switch: index of the single-head "router" KV layer that encodes + // per-token adapter selection. -1 when the model has no such layer. + int32_t router_layer = -1; uint32_t n_expert = 0; uint32_t n_expert_used = 0; uint32_t n_rel_attn_bkts = 0; diff --git a/src/llama-kv-cache.cpp b/src/llama-kv-cache.cpp index e70583e64152..3606f831f391 100644 --- a/src/llama-kv-cache.cpp +++ b/src/llama-kv-cache.cpp @@ -1925,6 +1925,13 @@ ggml_cgraph * llama_kv_cache::build_graph_shift(llm_graph_result * res, llama_co for (const auto & layer : layers) { const uint32_t il = layer.il; + // granite-switch: the router layer's K encodes the adapter slot as a + // literal dim-0 magnitude (no positional rotation), so it must NOT be + // RoPE-shifted when the KV cache is defragmented / position-shifted. + if (hparams.router_layer >= 0 && (int) il == hparams.router_layer) { + continue; + } + const int64_t n_head_kv = hparams.n_head_kv(il); const int64_t n_embd_k_gqa = hparams.n_embd_k_gqa(il); diff --git a/src/models/granite_switch.cpp b/src/models/granite_switch.cpp index 3f500e65bcf0..6a1b3a8aab51 100644 --- a/src/models/granite_switch.cpp +++ b/src/models/granite_switch.cpp @@ -6,12 +6,29 @@ // Granite Switch: a dense, all-attention Granite-4.1 model with N embedded LoRA // adapters selected per-token by control tokens. // -// Two cheap CPU steps (no router weights) implement the switch, done in -// llm_graph_input_switch::set_input: -// 1. sticky per-token index — each position takes the index of the most recent -// control token at-or-before it (0 = base, i+1 = adapter i; cannot revert). -// 2. token-exchange — rewrite each control token id to its substitute id before -// embedding (read index first, then swap). +// The per-token adapter index is recovered IN-GRAPH by a single-head causal +// "router" attention (ported faithfully from the vLLM/HF backends), whose K/V +// live in the model KV cache at an extra layer R == hparams.router_layer: +// - set_input fills pure per-token signals (no history): router Q[0]=1, +// K[0]=±gain (+ for a control token, - otherwise), V[0]=adapter slot / 0, +// plus the token-exchange (control id -> substitute id before embedding). +// - the causal softmax over the single visible control token (single-switch +// contract) puts ~all weight on it, so attended V[0] ≈ that adapter's slot; +// readback = clamp(round(V[0]), 0, n_adapters) -> I32 adapter index. +// Because the selection state lives in the per-sequence KV cache (like vLLM's +// PagedAttention router), CONCURRENT requests are isolated for free — there is +// no global sticky variable to leak between sequences (the prior POC's bug). +// +// SINGLE-SWITCH CONTRACT / KNOWN LIMITATION (identical to the vLLM & HF +// backends): the gain is flat (no recency), so within one sequence there is "no +// mechanism to transition back to base mid-sequence" (vLLM single.py docstring). +// A control token's K stays causally visible to all later tokens in the SAME +// sequence, so once an adapter fires it stays on until that sequence ends. +// vLLM/HF never observe this because each served request is a fresh sequence; a +// client that continues ONE KV cache across chat turns (e.g. `ollama run`) must +// start a fresh sequence per turn to return to base. A recency-biased router +// would lift this limit but is a deliberate divergence from vLLM and is not done +// here. See scratch/multiturn_leak_test.cpp and scratch/concurrent_switch_test.cpp. // // Then a standard Granite decoder, with LoRA added per-token on qkv / o / gate / // up / down via ggml_mul_mat_id over stacked tensors. The stacked dim is @@ -71,6 +88,26 @@ void llama_model_granite_switch::load_arch_hparams(llama_model_loader & ml) { control_token_to_index[(llama_token) token_ids[i]] = (int32_t) (i + 1); control_token_to_substitute[(llama_token) token_ids[i]] = (llama_token) substitute_ids[i]; } + + // --- router KV layer (per-token adapter selection lives in the KV cache) --- + // The GGUF block_count covers only the real decoder layers, so at this point + // hparams.n_layer_all == n_real and hparams.n_layer() == n_real. Append ONE + // extra single-head attention layer at index R = n_real to hold the router + // K/V. We bump n_layer_all to n_real+1 so the KV-cache allocator (which loops + // [0, n_layer_all)) gives the router its own per-sequence slot, and set + // n_layer_nextn = 1 so n_layer() == n_layer_all - n_layer_nextn stays n_real + // — the decoder loop and tensor loading (both bounded by n_layer()) are left + // completely unchanged and never touch layer R. + const uint32_t n_real = hparams.n_layer(); // == n_layer_all here (nextn==0) + hparams.router_layer = (int32_t) n_real; + hparams.n_layer_all = n_real + 1; + hparams.n_layer_nextn = 1; + + // The router is a single causal head. n_embd_head_k/v are global (the router + // inherits the model head dim and uses only dim 0; the rest is zero-padded). + hparams.n_head_arr[n_real] = 1; + hparams.n_head_kv_arr[n_real] = 1; + hparams.n_ff_arr[n_real] = 0; // no FFN for the router layer } void llama_model_granite_switch::load_arch_tensors(llama_model_loader &) { @@ -130,8 +167,19 @@ void llama_model_granite_switch::load_arch_tensors(llama_model_loader &) { } } +// Router differential gain. Each token's router K dim-0 is +GAIN for a control +// token and -GAIN otherwise; Q dim-0 is 1.0. In the causal softmax a single +// visible control token then dominates (its logit beats a normal token's by +// 2*GAIN), so the readback recovers that adapter's integer slot. 15.0 matches the +// vLLM/HF backend default (control_token_gain in config.py) and is F16-safe +// (no F32 KV cache needed). +static constexpr float GRANITE_SWITCH_ROUTER_GAIN = 15.0f; + // ---------------------------------------------------------------------------- -// switch input: compute sticky adapter index + token-exchanged ids per token. +// switch input: pure per-token maps. The adapter index is NOT computed here — +// it is recovered in-graph by the router attention (see build_arch_graph). This +// is intentionally stateless: no cross-ubatch carry. Per-sequence isolation +// (no leak between concurrent sequences) comes from the KV cache, not from here. // ---------------------------------------------------------------------------- void llm_graph_input_switch::set_input(const llama_ubatch * ubatch) { if (!ubatch->token) { @@ -140,42 +188,35 @@ void llm_graph_input_switch::set_input(const llama_ubatch * ubatch) { const int64_t n_tokens = ubatch->n_tokens; - std::vector ids(n_tokens); - std::vector sub(n_tokens); + std::vector sub (n_tokens); + std::vector ksig(n_tokens); + std::vector vval(n_tokens); + std::vector q (n_tokens, 1.0f); // Q dim-0 is constant 1.0 - // POC single-sequence sticky state: carry the last index across ubatches, - // resetting when this ubatch starts a fresh sequence (contains pos 0). - bool has_pos0 = false; - if (ubatch->pos) { - for (int64_t i = 0; i < n_tokens; ++i) { - if (ubatch->pos[i] == 0) { has_pos0 = true; break; } - } - } - if (has_pos0) { - smodel.poc_sticky_index = 0; - } - - int32_t cur = smodel.poc_sticky_index; for (int64_t i = 0; i < n_tokens; ++i) { const llama_token tok = ubatch->token[i]; - // 1. read index first (sticky): update on a control token. + // router K/V signals: control token -> (+gain, slot); else -> (-gain, 0). auto it = smodel.control_token_to_index.find(tok); if (it != smodel.control_token_to_index.end()) { - cur = it->second; + ksig[i] = +GRANITE_SWITCH_ROUTER_GAIN; + vval[i] = (float) it->second; // adapter slot (1 + adapter) + } else { + ksig[i] = -GRANITE_SWITCH_ROUTER_GAIN; + vval[i] = 0.0f; // base } - ids[i] = cur; - // 2. then token-exchange: rewrite control token to its substitute. + // token-exchange: rewrite a control token to its substitute id. auto sit = smodel.control_token_to_substitute.find(tok); sub[i] = (sit != smodel.control_token_to_substitute.end()) ? (int32_t) sit->second : (int32_t) tok; } - smodel.poc_sticky_index = cur; - ggml_backend_tensor_set(adapter_ids, ids.data(), 0, n_tokens*ggml_element_size(adapter_ids)); - ggml_backend_tensor_set(sub_tokens, sub.data(), 0, n_tokens*ggml_element_size(sub_tokens)); + ggml_backend_tensor_set(sub_tokens, sub.data(), 0, n_tokens*ggml_element_size(sub_tokens)); + ggml_backend_tensor_set(router_ksig, ksig.data(), 0, n_tokens*ggml_element_size(router_ksig)); + ggml_backend_tensor_set(router_vval, vval.data(), 0, n_tokens*ggml_element_size(router_vval)); + ggml_backend_tensor_set(router_q, q.data(), 0, n_tokens*ggml_element_size(router_q)); } // ---------------------------------------------------------------------------- @@ -229,14 +270,20 @@ llama_model_granite_switch::graph::graph( GGML_ASSERT(n_embd_head == hparams.n_embd_head_k()); GGML_ASSERT(n_embd_head == n_rot); - // --- switch input: substituted ids + per-token adapter indices --- + // --- switch input: substituted ids + per-token router K/V/Q signals --- auto inp_switch = std::make_unique(smodel); inp_switch->sub_tokens = ggml_new_tensor_1d(ctx0, GGML_TYPE_I32, n_tokens); - inp_switch->adapter_ids = ggml_new_tensor_1d(ctx0, GGML_TYPE_I32, n_tokens); + inp_switch->router_ksig = ggml_new_tensor_1d(ctx0, GGML_TYPE_F32, n_tokens); + inp_switch->router_vval = ggml_new_tensor_1d(ctx0, GGML_TYPE_F32, n_tokens); + inp_switch->router_q = ggml_new_tensor_1d(ctx0, GGML_TYPE_F32, n_tokens); ggml_set_input(inp_switch->sub_tokens); - ggml_set_input(inp_switch->adapter_ids); + ggml_set_input(inp_switch->router_ksig); + ggml_set_input(inp_switch->router_vval); + ggml_set_input(inp_switch->router_q); ggml_tensor * sub_tokens = inp_switch->sub_tokens; - ggml_tensor * adapter_ids = inp_switch->adapter_ids; + ggml_tensor * router_ksig = inp_switch->router_ksig; + ggml_tensor * router_vval = inp_switch->router_vval; + ggml_tensor * router_q = inp_switch->router_q; res->add_input(std::move(inp_switch)); // embed the token-exchanged ids ourselves (we cannot use build_inp_embd @@ -253,6 +300,40 @@ llama_model_granite_switch::graph::graph( } auto * inp_attn = build_attn_inp_kv(); + // --- router attention: recover the per-token adapter index in-graph --- + // A single causal head whose K/V live in the KV cache at layer R == router_layer. + // Only dim 0 carries signal (rest zero-padded): Q[0]=1, K[0]=±gain, V[0]=slot/0. + // The causal softmax over a single visible control token (the single-switch + // contract) puts ~all weight on it, so attended V[0] ≈ that adapter's slot. + // State persists across decode steps and is isolated per-sequence because the + // control token's K/V stay in the per-sequence KV cache (vLLM/HF parity). No + // RoPE on the router Q/K — positions are irrelevant to the selection. + const int R = hparams.router_layer; + GGML_ASSERT(R >= 0 && "granite-switch: router_layer not configured"); + // build [n_embd_head, 1, n_tokens] with row 0 = signal, rows 1.. = 0. + auto router_lane = [&](ggml_tensor * sig1d) { + ggml_tensor * t = ggml_reshape_3d(ctx0, sig1d, 1, 1, n_tokens); + return ggml_pad(ctx0, t, (int) n_embd_head - 1, 0, 0, 0); // pad dim0 right with zeros + }; + ggml_tensor * Qr = router_lane(router_q); + ggml_tensor * Kr = router_lane(router_ksig); + ggml_tensor * Vr = router_lane(router_vval); + + ggml_tensor * router_out = build_attn(inp_attn, + nullptr, nullptr, nullptr, + Qr, Kr, Vr, nullptr, nullptr, nullptr, /*kq_scale=*/1.0f, /*il=*/R); + cb(router_out, "router_out", R); + + // readback: router_out is {n_embd_head, n_tokens}; row 0 is the attended slot. + // adapter_index = clamp(round(row0), 0, n_adapters), as I32 for mul_mat_id. + ggml_tensor * slot_f = ggml_cont(ctx0, + ggml_view_2d(ctx0, router_out, 1, n_tokens, router_out->nb[1], 0)); + slot_f = ggml_reshape_1d(ctx0, slot_f, n_tokens); + slot_f = ggml_clamp(ctx0, slot_f, 0.0f, (float) smodel.n_adapters); + slot_f = ggml_round(ctx0, slot_f); // round BEFORE cast (cast truncates) + ggml_tensor * adapter_ids = ggml_cast(ctx0, slot_f, GGML_TYPE_I32); + cb(adapter_ids, "adapter_ids", -1); + ggml_tensor * inp_out_ids = build_inp_out_ids(); ggml_tensor * cur; diff --git a/src/models/models.h b/src/models/models.h index fe026d3b5b3f..41b6cc7b0c88 100644 --- a/src/models/models.h +++ b/src/models/models.h @@ -1546,18 +1546,20 @@ struct llama_model_granite_switch : public llama_model_base { uint32_t n_slots = 0; // n_adapters + 1 (slot 0 = base/zero delta) uint32_t max_lora_rank = 0; - // token id -> sticky adapter index (1 + adapter, so slot 0 stays "base") + // token id -> adapter slot (1 + adapter, so slot 0 stays "base") std::unordered_map control_token_to_index; // control token id -> substitute token id (token-exchange before embedding) std::unordered_map control_token_to_substitute; - // POC: single-sequence sticky adapter state. The most recent control index - // seen, carried across ubatches within one decode. Reset to 0 when a ubatch - // contains sequence position 0 (start of a fresh prompt). This is sufficient - // for llama-cli / the smoke + parity tests; a full multi-sequence - // (reset-on-seq_rm) implementation is deferred — see the [[llamacpp-poc-scope]] - // note. `mutable` because set_input runs from a const-model graph context. - mutable int32_t poc_sticky_index = 0; + // Per-token adapter selection is computed in-graph by a single-head causal + // "router" attention whose K/V live in the model KV cache at layer index + // hparams.router_layer (== n_layer). Because that state is per-sequence (it + // lives in the cache, like vLLM's PagedAttention router), CONCURRENT requests + // are isolated for free — there is no global sticky state to leak between + // sequences (the prior POC's bug). Within one sequence the selection carries + // over once fired (flat gain, no recency) — the vLLM/HF single-switch + // contract; see the limitation note in granite_switch.cpp. Mirrors the vLLM + // backend (single.py) and HF SingleSwitch. struct graph : public llm_graph_context { graph(const llama_model & model, const llm_graph_params & params); @@ -1601,11 +1603,16 @@ struct llama_model_granite_switch : public llama_model_base { }; -// Custom graph input for the per-token switch. In set_input it scans the ubatch, -// computes the sticky adapter index per token (writing adapter_ids), and rewrites -// control-token ids to their substitute ids (writing sub_tokens) — the -// token-exchange. Both tensors are I32 [n_tokens]. can_reuse() returns false so -// these are recomputed every ubatch. +// Custom graph input for the per-token switch. In set_input it fills, per token: +// - sub_tokens : the token-exchange (control id -> substitute id, else the id) +// - router_ksig : +gain for a control token, -gain otherwise (router K, dim 0) +// - router_vval : float(adapter_slot) for a control token, 0 otherwise (router V, dim 0) +// - router_q : constant 1.0 (router Q, dim 0) +// The adapter index per token is NOT computed here — it is recovered in-graph by a +// single-head causal attention over (router_q, router_ksig, router_vval) whose K/V +// live in the KV cache, so the selection is per-sequence (no cross-sequence leak). +// set_input is pure per-token maps (no history / no cross-ubatch carry). +// can_reuse() returns false so these are recomputed every ubatch. class llm_graph_input_switch : public llm_graph_input_i { public: llm_graph_input_switch(const llama_model_granite_switch & smodel) : smodel(smodel) {} @@ -1614,7 +1621,9 @@ class llm_graph_input_switch : public llm_graph_input_i { void set_input(const llama_ubatch * ubatch) override; ggml_tensor * sub_tokens = nullptr; // I32 [n_tokens] control-substituted token ids - ggml_tensor * adapter_ids = nullptr; // I32 [n_tokens] sticky per-token adapter index + ggml_tensor * router_ksig = nullptr; // F32 [n_tokens] router K signal (±gain) + ggml_tensor * router_vval = nullptr; // F32 [n_tokens] router V value (adapter slot / 0) + ggml_tensor * router_q = nullptr; // F32 [n_tokens] router Q value (constant 1.0) const llama_model_granite_switch & smodel; }; From b4eb4043415947952ad9689e67225b2ced0036c3 Mon Sep 17 00:00:00 2001 From: Bar Haim Date: Wed, 24 Jun 2026 17:09:55 +0300 Subject: [PATCH 05/40] granite-switch: drop scratch tests and mac demo for upstream PR Remove the local-only development artifacts that should not ship in the upstream PR: - granite-switch-mac-demo.sh (local Metal build + demo driver) - scratch/concurrent_switch_test.cpp - scratch/multiturn_leak_test.cpp Also drop the now-dangling reference to the scratch tests from the granite_switch.cpp header comment. Leaves only the core architecture support (conversion, gguf constants, llama-arch/model/kv-cache, and the granite_switch graph). --- granite-switch-mac-demo.sh | 105 ---------------------- scratch/concurrent_switch_test.cpp | 128 --------------------------- scratch/multiturn_leak_test.cpp | 134 ----------------------------- src/models/granite_switch.cpp | 2 +- 4 files changed, 1 insertion(+), 368 deletions(-) delete mode 100755 granite-switch-mac-demo.sh delete mode 100644 scratch/concurrent_switch_test.cpp delete mode 100644 scratch/multiturn_leak_test.cpp diff --git a/granite-switch-mac-demo.sh b/granite-switch-mac-demo.sh deleted file mode 100755 index d14dda945a91..000000000000 --- a/granite-switch-mac-demo.sh +++ /dev/null @@ -1,105 +0,0 @@ -#!/usr/bin/env bash -# granite-switch-mac-demo.sh -# -# End-to-end Granite Switch demo on an Apple-Silicon Mac (Metal GPU backend). -# Builds llama.cpp from this branch, converts the composed checkpoint to GGUF, -# and runs the crisp mid-sequence adapter-switch demos verified on Vela: -# - answerability : <|answerability|> -> "unanswerable" (base would just answer) -# - query_rewrite : <|query_rewrite|> -> {"rewritten_question": ...} (base would answer) -# -# Each demo runs the SAME prompt twice, differing only by a control token placed -# mid-sequence (right before the assistant turn). The outputs differ structurally, -# which is the switch firing per-token. -# -# Usage: -# ./granite-switch-mac-demo.sh # build + convert (first run) then demo -# ./granite-switch-mac-demo.sh demo # skip build/convert, just run the demos -# GGUF=/path/to/gs-f16.gguf ./granite-switch-mac-demo.sh demo # use an existing GGUF -# -# Requirements: Xcode CLT (`xcode-select --install`), cmake (`brew install cmake`), -# python3, and ~16GB+ unified memory for the f16 model (8.4 GB on disk). - -set -euo pipefail - -REPO_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" -cd "$REPO_DIR" - -SRC_MODEL="${SRC_MODEL:-ibm-granite/granite-switch-4.1-3b-preview}" -GGUF="${GGUF:-$REPO_DIR/gs-f16.gguf}" -BIN="$REPO_DIR/build/bin/llama-completion" -NGL="${NGL:-99}" # offload all layers to the Metal GPU; set NGL=0 to force CPU - -# --------------------------------------------------------------------------- -build() { - echo "=== Build llama.cpp (Metal is on by default on macOS) ===" - if [ ! -x "$BIN" ]; then - cmake -S . -B build -DLLAMA_CURL=OFF - cmake --build build -j --target llama llama-completion llama-tokenize - else - echo " (build/bin/llama-completion already present — skipping)" - fi -} - -convert() { - if [ -f "$GGUF" ]; then - echo "=== GGUF already present at $GGUF — skipping convert ===" - return - fi - echo "=== Set up python deps for the converter ===" - python3 -m venv .venv-convert - # shellcheck disable=SC1091 - . .venv-convert/bin/activate - pip install --upgrade pip - pip install numpy torch transformers safetensors sentencepiece "huggingface_hub[cli]" - - echo "=== Download $SRC_MODEL (~8 GB) ===" - SRC_DIR="$(python -c "from huggingface_hub import snapshot_download; print(snapshot_download('$SRC_MODEL'))")" - - echo "=== Convert HF -> GGUF (f16) ===" - PYTHONPATH=gguf-py python convert_hf_to_gguf.py "$SRC_DIR" --outfile "$GGUF" --outtype f16 - deactivate -} - -# run NAME PROMPT N -# -no-cnv: this model ships a chat template, so llama-completion AUTO-ENABLES -# interactive conversation mode and stops at a `>` prompt waiting for input. -# -no-cnv disables that: generate once from the raw prompt and exit (and it -# also prints special tokens, which makes the adapter switch easy to see). -run() { - local name="$1" prompt="$2" n="${3:-16}" - echo "=== $name ===" - "$BIN" -m "$GGUF" -ngl "$NGL" --temp 0 -n "$n" -no-cnv -p "$prompt" 2>/dev/null - echo; echo "-----" -} - -demo() { - [ -x "$BIN" ] || { echo "FATAL: $BIN not built — run without 'demo' first"; exit 1; } - [ -f "$GGUF" ] || { echo "FATAL: $GGUF missing — run without 'demo' first"; exit 1; } - - echo - echo "############ DEMO 1: answerability (<|answerability|>, id 100356) ############" - echo "# Document is about the Eiffel Tower; question asks Australia's capital." - echo "# Base ANSWERS the question; adapter judges it UNANSWERABLE from the doc." - local ANS_CTX='<|start_of_role|>user<|end_of_role|>Document: The Eiffel Tower is in Paris. Question: What is the capital of Australia?<|end_of_role|><|end_of_text|><|start_of_role|>assistant<|end_of_role|>' - run "answerability OFF (base, no control)" "${ANS_CTX}" 12 - run "answerability ON (<|answerability|> mid-seq)" "${ANS_CTX}<|answerability|>" 12 - - echo - echo "############ DEMO 2: query_rewrite (<|query_rewrite|>, id 100353) ############" - echo "# Follow-up 'What other movies has HE made?' — base answers it;" - echo "# adapter REWRITES it into a standalone query (resolves 'he')." - local QR_CTX='<|start_of_role|>user<|end_of_role|>Who directed Inception?<|end_of_text|><|start_of_role|>assistant<|end_of_role|>Christopher Nolan directed Inception.<|end_of_text|><|start_of_role|>user<|end_of_role|>What other movies has he made?<|end_of_text|><|start_of_role|>assistant<|end_of_role|>' - run "query_rewrite OFF (base, no control)" "${QR_CTX}" 24 - run "query_rewrite ON (<|query_rewrite|> mid-seq)" "${QR_CTX}<|query_rewrite|>" 24 - - echo "=== DONE. In each pair, only the mid-sequence control token differs. ===" -} - -# --------------------------------------------------------------------------- -case "${1:-all}" in - demo) demo ;; - build) build ;; - convert) convert ;; - all) build; convert; demo ;; - *) echo "usage: $0 [all|build|convert|demo]"; exit 1 ;; -esac diff --git a/scratch/concurrent_switch_test.cpp b/scratch/concurrent_switch_test.cpp deleted file mode 100644 index 8bb4062ace28..000000000000 --- a/scratch/concurrent_switch_test.cpp +++ /dev/null @@ -1,128 +0,0 @@ -// Concurrent multi-sequence isolation test for granite-switch. -// -// Two sequences are decoded together in one batch / one context: -// seq 0: answerability ON (<|answerability|>) -> expect "unanswerable" -// seq 1: NO control token -> expect a normal answer -// -// The OLD global `poc_sticky_index` could not isolate these (last writer wins): -// seq 1 would inherit seq 0's adapter. The in-graph router stores each sequence's -// selection in its own KV-cache cells, so they must not cross-contaminate. -// -// PASS: seq0 contains "unanswerable" AND seq1 does NOT. - -#include "llama.h" -#include -#include -#include -#include - -static std::vector tokz(const llama_vocab * v, const std::string & s, bool bos) { - int n = -llama_tokenize(v, s.c_str(), s.size(), NULL, 0, bos, true); - std::vector out(n); - llama_tokenize(v, s.c_str(), s.size(), out.data(), out.size(), bos, true); - return out; -} - -int main(int argc, char ** argv) { - std::string model_path; - for (int i = 1; i < argc; ++i) - if (strcmp(argv[i], "-m") == 0 && i + 1 < argc) model_path = argv[++i]; - if (model_path.empty()) { fprintf(stderr, "usage: %s -m model.gguf\n", argv[0]); return 1; } - - ggml_backend_load_all(); - - llama_model_params mp = llama_model_default_params(); - mp.n_gpu_layers = 0; - llama_model * model = llama_model_load_from_file(model_path.c_str(), mp); - if (!model) { fprintf(stderr, "load failed\n"); return 1; } - const llama_vocab * vocab = llama_model_get_vocab(model); - - const std::string p0 = - "<|start_of_role|>user<|end_of_role|>Document: The Eiffel Tower is in Paris. " - "Question: What is the capital of Australia?<|end_of_role|><|end_of_text|>" - "<|start_of_role|>assistant<|end_of_role|><|answerability|>"; - const std::string p1 = - "<|start_of_role|>user<|end_of_role|>What is the capital of Australia?" - "<|end_of_role|><|end_of_text|><|start_of_role|>assistant<|end_of_role|>"; - - std::vector t0 = tokz(vocab, p0, true); - std::vector t1 = tokz(vocab, p1, true); - - llama_context_params cp = llama_context_default_params(); - cp.n_ctx = 2048; - cp.n_batch = 2048; - cp.n_seq_max = 2; - llama_context * ctx = llama_init_from_model(model, cp); - if (!ctx) { fprintf(stderr, "ctx failed\n"); return 1; } - - auto sp = llama_sampler_chain_default_params(); - llama_sampler * smpl = llama_sampler_chain_init(sp); - llama_sampler_chain_add(smpl, llama_sampler_init_greedy()); - - // build a batch holding BOTH prompts, seq 0 and seq 1 interleaved by sequence id. - const int n_gen = 14; - llama_batch batch = llama_batch_init((int) (t0.size() + t1.size()), 0, 2); - - auto add = [&](llama_token id, llama_pos pos, llama_seq_id seq, bool logits) { - int k = batch.n_tokens; - batch.token[k] = id; - batch.pos[k] = pos; - batch.n_seq_id[k] = 1; - batch.seq_id[k][0] = seq; - batch.logits[k] = logits; - batch.n_tokens++; - }; - - for (size_t i = 0; i < t0.size(); ++i) add(t0[i], (llama_pos) i, 0, i == t0.size() - 1); - for (size_t i = 0; i < t1.size(); ++i) add(t1[i], (llama_pos) i, 1, i == t1.size() - 1); - - if (llama_decode(ctx, batch)) { fprintf(stderr, "prompt decode failed\n"); return 1; } - - // remember the logits index for each sequence's last prompt token. - int idx0 = (int) t0.size() - 1; - int idx1 = (int) (t0.size() + t1.size()) - 1; - llama_pos pos0 = (llama_pos) t0.size(); - llama_pos pos1 = (llama_pos) t1.size(); - - std::string out0, out1; - for (int step = 0; step < n_gen; ++step) { - llama_token id0 = llama_sampler_sample(smpl, ctx, idx0); - llama_token id1 = llama_sampler_sample(smpl, ctx, idx1); - - auto piece = [&](llama_token id, std::string & dst) { - if (llama_vocab_is_eog(vocab, id)) return; - char buf[256]; - int n = llama_token_to_piece(vocab, id, buf, sizeof(buf), 0, true); - if (n > 0) dst.append(buf, n); - }; - piece(id0, out0); - piece(id1, out1); - - batch.n_tokens = 0; - bool e0 = llama_vocab_is_eog(vocab, id0); - bool e1 = llama_vocab_is_eog(vocab, id1); - if (e0 && e1) break; - if (!e0) { add(id0, pos0++, 0, true); idx0 = batch.n_tokens - 1; } - if (!e1) { add(id1, pos1++, 1, true); idx1 = batch.n_tokens - 1; } - if (llama_decode(ctx, batch)) { fprintf(stderr, "gen decode failed\n"); break; } - } - - printf("\n============ CONCURRENT ISOLATION TEST ============\n"); - printf("SEQ 0 (answerability ON): %s\n", out0.c_str()); - printf("SEQ 1 (no control token): %s\n", out1.c_str()); - auto has = [](const std::string & h, const char * n){ return h.find(n) != std::string::npos; }; - bool s0_ok = has(out0, "unanswerable"); - bool s1_clean = !has(out1, "unanswerable"); - printf("---------------------------------------------------\n"); - printf("seq0 fired its adapter: %s\n", s0_ok ? "YES" : "NO"); - printf("seq1 isolated from seq0: %s\n", s1_clean ? "clean" : "CONTAMINATED"); - bool pass = s0_ok && s1_clean; - printf("RESULT: %s\n", pass ? "PASS (isolated)" : "FAIL"); - printf("===================================================\n"); - - llama_batch_free(batch); - llama_sampler_free(smpl); - llama_free(ctx); - llama_model_free(model); - return pass ? 0 : 1; -} diff --git a/scratch/multiturn_leak_test.cpp b/scratch/multiturn_leak_test.cpp deleted file mode 100644 index 6967f48d2559..000000000000 --- a/scratch/multiturn_leak_test.cpp +++ /dev/null @@ -1,134 +0,0 @@ -// Multi-turn behavior probe for granite-switch (DOCUMENTS A KNOWN LIMITATION). -// -// Both turns run on the SAME llama_context — one persistent KV cache, no reset -// between turns, exactly like `ollama run` continuing a chat: -// turn 1: answerability ON (<|answerability|>) -> "unanswerable" -// turn 2: NO control token -> ??? (same cache) -// -// EXPECTED (documented) behavior: turn 2 STILL reads the adapter. This is the -// single-switch contract, identical to the vLLM/HF backends, whose own docstring -// states: "SingleSwitch has no mechanism to transition back to base -// mid-sequence." The in-graph router selection lives in the per-sequence KV -// cache; turn 1's control-token K is still causally visible to turn 2's tokens, -// and with flat (recency-free) gain it keeps winning the softmax. -// -// vLLM/HF never observe this because each served request is a FRESH sequence -// (turn 2's context does not contain turn 1's control token). A chat client that -// continues one cache across turns (ollama) must therefore start a fresh -// sequence per turn, OR opt into a recency-biased router (a deliberate -// divergence from vLLM — not done here). -// -// This probe asserts the DOCUMENTED behavior so it stays a faithful vLLM copy: -// turn 1 -> "unanswerable" -// turn 2 -> ALSO "unanswerable" (carry-over expected on a persisted cache) - -#include "llama.h" -#include -#include -#include -#include - -static std::string g_model; - -static std::vector tok(const llama_vocab * v, const std::string & s, bool bos) { - int n = -llama_tokenize(v, s.c_str(), s.size(), NULL, 0, bos, true); - std::vector out(n); - llama_tokenize(v, s.c_str(), s.size(), out.data(), out.size(), bos, true); - return out; -} - -// Decode `prompt_tokens` (appended at current position `pos`) then greedily -// generate up to n_gen tokens, advancing `pos`. Returns the generated text. -static std::string run_turn(llama_context * ctx, const llama_vocab * vocab, - llama_sampler * smpl, - std::vector prompt_tokens, - int & pos, int n_gen) { - std::string text; - - // decode the prompt chunk (positions assigned automatically by the batch helper - // continue from the cache, since we never clear it between turns). - llama_batch batch = llama_batch_get_one(prompt_tokens.data(), prompt_tokens.size()); - if (llama_decode(ctx, batch)) { fprintf(stderr, "decode (prompt) failed\n"); return text; } - pos += (int) prompt_tokens.size(); - - for (int i = 0; i < n_gen; ++i) { - llama_token id = llama_sampler_sample(smpl, ctx, -1); - if (llama_vocab_is_eog(vocab, id)) break; - char buf[256]; - int n = llama_token_to_piece(vocab, id, buf, sizeof(buf), 0, true); - if (n > 0) text.append(buf, n); - llama_batch nb = llama_batch_get_one(&id, 1); - if (llama_decode(ctx, nb)) { fprintf(stderr, "decode (gen) failed\n"); break; } - pos += 1; - } - return text; -} - -int main(int argc, char ** argv) { - for (int i = 1; i < argc; ++i) { - if (strcmp(argv[i], "-m") == 0 && i + 1 < argc) g_model = argv[++i]; - } - if (g_model.empty()) { fprintf(stderr, "usage: %s -m model.gguf\n", argv[0]); return 1; } - - ggml_backend_load_all(); - - llama_model_params mp = llama_model_default_params(); - mp.n_gpu_layers = 0; - llama_model * model = llama_model_load_from_file(g_model.c_str(), mp); - if (!model) { fprintf(stderr, "load failed\n"); return 1; } - const llama_vocab * vocab = llama_model_get_vocab(model); - - llama_context_params cp = llama_context_default_params(); - cp.n_ctx = 2048; - cp.n_batch = 2048; - llama_context * ctx = llama_init_from_model(model, cp); - if (!ctx) { fprintf(stderr, "ctx failed\n"); return 1; } - - auto sp = llama_sampler_chain_default_params(); - llama_sampler * smpl = llama_sampler_chain_init(sp); - llama_sampler_chain_add(smpl, llama_sampler_init_greedy()); - - int pos = 0; - - // --- TURN 1: answerability ON (control token fires the adapter) --- - // Doc is about the Eiffel Tower; question asks Australia's capital -> adapter - // should judge it "unanswerable". - std::string t1 = - "<|start_of_role|>user<|end_of_role|>Document: The Eiffel Tower is in Paris. " - "Question: What is the capital of Australia?<|end_of_role|><|end_of_text|>" - "<|start_of_role|>assistant<|end_of_role|><|answerability|>"; - std::string out1 = run_turn(ctx, vocab, smpl, tok(vocab, t1, /*bos=*/true), pos, 12); - - // --- TURN 2: NO control token, same context (cache persists) --- - // A normal question. If the adapter leaked from turn 1, this would also be - // judged "unanswerable". With the router fix it answers normally. - std::string t2 = - "<|end_of_text|><|start_of_role|>user<|end_of_role|>" - "What is the capital of Australia?<|end_of_role|><|end_of_text|>" - "<|start_of_role|>assistant<|end_of_role|>"; - std::string out2 = run_turn(ctx, vocab, smpl, tok(vocab, t2, /*bos=*/false), pos, 16); - - printf("\n========== MULTI-TURN CARRY-OVER PROBE (KNOWN LIMIT) ==========\n"); - printf("TURN 1 (answerability ON): %s\n", out1.c_str()); - printf("TURN 2 (no control token): %s\n", out2.c_str()); - - auto contains = [](const std::string & h, const char * n) { - return h.find(n) != std::string::npos; - }; - bool t1_ok = contains(out1, "unanswerable"); - bool t2_carried = contains(out2, "unanswerable"); - - printf("--------------------------------------------------------------\n"); - printf("turn1 fired adapter (expect unanswerable): %s\n", t1_ok ? "YES" : "NO"); - printf("turn2 carried adapter (DOCUMENTED: expect carry): %s\n", t2_carried ? "carried" : "reverted"); - // Faithful vLLM/HF contract: on a persisted cache, selection carries over. - bool pass = t1_ok && t2_carried; - printf("RESULT: %s\n", pass ? "PASS (matches vLLM/HF single-switch contract)" - : "UNEXPECTED (behavior differs from vLLM/HF)"); - printf("==============================================================\n"); - - llama_sampler_free(smpl); - llama_free(ctx); - llama_model_free(model); - return pass ? 0 : 1; -} diff --git a/src/models/granite_switch.cpp b/src/models/granite_switch.cpp index 6a1b3a8aab51..4aaef8f4532b 100644 --- a/src/models/granite_switch.cpp +++ b/src/models/granite_switch.cpp @@ -28,7 +28,7 @@ // client that continues ONE KV cache across chat turns (e.g. `ollama run`) must // start a fresh sequence per turn to return to base. A recency-biased router // would lift this limit but is a deliberate divergence from vLLM and is not done -// here. See scratch/multiturn_leak_test.cpp and scratch/concurrent_switch_test.cpp. +// here. // // Then a standard Granite decoder, with LoRA added per-token on qkv / o / gate / // up / down via ggml_mul_mat_id over stacked tensors. The stacked dim is From 5b4fd0e150cf05de56d84357b9ed2fdbf3393d1a Mon Sep 17 00:00:00 2001 From: Bar Haim Date: Wed, 24 Jun 2026 17:52:04 +0300 Subject: [PATCH 06/40] granite-switch: trim comments to match native llama.cpp style --- src/models/granite_switch.cpp | 190 ++++++++++------------------------ src/models/models.h | 31 ++---- 2 files changed, 63 insertions(+), 158 deletions(-) diff --git a/src/models/granite_switch.cpp b/src/models/granite_switch.cpp index 4aaef8f4532b..3d9fe7dd679d 100644 --- a/src/models/granite_switch.cpp +++ b/src/models/granite_switch.cpp @@ -2,50 +2,14 @@ #include -// ============================================================================ -// Granite Switch: a dense, all-attention Granite-4.1 model with N embedded LoRA -// adapters selected per-token by control tokens. -// -// The per-token adapter index is recovered IN-GRAPH by a single-head causal -// "router" attention (ported faithfully from the vLLM/HF backends), whose K/V -// live in the model KV cache at an extra layer R == hparams.router_layer: -// - set_input fills pure per-token signals (no history): router Q[0]=1, -// K[0]=±gain (+ for a control token, - otherwise), V[0]=adapter slot / 0, -// plus the token-exchange (control id -> substitute id before embedding). -// - the causal softmax over the single visible control token (single-switch -// contract) puts ~all weight on it, so attended V[0] ≈ that adapter's slot; -// readback = clamp(round(V[0]), 0, n_adapters) -> I32 adapter index. -// Because the selection state lives in the per-sequence KV cache (like vLLM's -// PagedAttention router), CONCURRENT requests are isolated for free — there is -// no global sticky variable to leak between sequences (the prior POC's bug). -// -// SINGLE-SWITCH CONTRACT / KNOWN LIMITATION (identical to the vLLM & HF -// backends): the gain is flat (no recency), so within one sequence there is "no -// mechanism to transition back to base mid-sequence" (vLLM single.py docstring). -// A control token's K stays causally visible to all later tokens in the SAME -// sequence, so once an adapter fires it stays on until that sequence ends. -// vLLM/HF never observe this because each served request is a fresh sequence; a -// client that continues ONE KV cache across chat turns (e.g. `ollama run`) must -// start a fresh sequence per turn to return to base. A recency-biased router -// would lift this limit but is a deliberate divergence from vLLM and is not done -// here. -// -// Then a standard Granite decoder, with LoRA added per-token on qkv / o / gate / -// up / down via ggml_mul_mat_id over stacked tensors. The stacked dim is -// N = num_adapters + 1; slot 0 is zero-filled so base tokens add an exact-zero -// delta. B is pre-scaled by alpha/rank at compose time, so the runtime LoRA -// scale is 1.0. -// ============================================================================ - void llama_model_granite_switch::load_arch_hparams(llama_model_loader & ml) { - // Granite multipliers / eps / rope (mirrors llama_model_granite). ml.get_key(LLM_KV_ATTENTION_LAYERNORM_RMS_EPS, hparams.f_norm_rms_eps); ml.get_key(LLM_KV_LOGIT_SCALE, hparams.f_logit_scale); ml.get_key(LLM_KV_RESIDUAL_SCALE, hparams.f_residual_scale, false); ml.get_key(LLM_KV_EMBEDDING_SCALE, hparams.f_embedding_scale, false); ml.get_key(LLM_KV_ATTENTION_SCALE, hparams.f_attention_scale, false); - // Granite uses rope_finetuned as a switch for rope, so default to true. + // Granite uses rope_finetuned as a switch for rope, so default to true bool rope_finetuned = true; ml.get_key(LLM_KV_ROPE_SCALING_FINETUNED, rope_finetuned, false); hparams.rope_finetuned = rope_finetuned; @@ -55,18 +19,12 @@ void llama_model_granite_switch::load_arch_hparams(llama_model_loader & ml) { default: type = LLM_TYPE_UNKNOWN; } - // Shared FFN length (dense path uses n_ff; this is informational/parity). ml.get_key(LLM_KV_EXPERT_SHARED_FEED_FORWARD_LENGTH, hparams.n_ff_shexp, /* required */ false); - // The per-token LoRA tensors use GGML_OP_MUL_MAT_ID with n_expert_used = 1. - // The model is dense (n_expert == 0, so the generic loader checks require - // n_expert_used == 0), but the buffer-type probe for MUL_MAT_ID tensors in - // create_tensor builds a synthetic op sized by hparams.n_expert_used and - // asserts it is > 0. Set it to 1 here (after the generic checks have run) so - // the probe matches how we actually invoke mul_mat_id. + // the per-token LoRA tensors use MUL_MAT_ID with n_expert_used = 1; the model + // is otherwise dense, so set it after the generic loader's n_expert checks hparams.n_expert_used = 1; - // --- switch metadata --- ml.get_key(LLM_KV_NUM_ADAPTERS, n_adapters); ml.get_key(LLM_KV_MAX_LORA_RANK, max_lora_rank); n_slots = n_adapters + 1; @@ -76,47 +34,37 @@ void llama_model_granite_switch::load_arch_hparams(llama_model_loader & ml) { ml.get_arr(LLM_KV_ADAPTER_TOKEN_IDS, token_ids); ml.get_arr(LLM_KV_ADAPTER_SUBSTITUTE_TOKEN_IDS, substitute_ids); - GGML_ASSERT(token_ids.size() == n_adapters && - "granite-switch: adapter_token_ids length must equal num_adapters"); - GGML_ASSERT(substitute_ids.size() == n_adapters && - "granite-switch: adapter_substitute_token_ids length must equal num_adapters"); + GGML_ASSERT(token_ids.size() == n_adapters); + GGML_ASSERT(substitute_ids.size() == n_adapters); control_token_to_index.clear(); control_token_to_substitute.clear(); for (uint32_t i = 0; i < n_adapters; ++i) { - // adapter i -> stacked slot i+1 (slot 0 is the base/zero delta). + // adapter i -> stacked slot i+1 (slot 0 is the base/zero delta) control_token_to_index[(llama_token) token_ids[i]] = (int32_t) (i + 1); control_token_to_substitute[(llama_token) token_ids[i]] = (llama_token) substitute_ids[i]; } - // --- router KV layer (per-token adapter selection lives in the KV cache) --- - // The GGUF block_count covers only the real decoder layers, so at this point - // hparams.n_layer_all == n_real and hparams.n_layer() == n_real. Append ONE - // extra single-head attention layer at index R = n_real to hold the router - // K/V. We bump n_layer_all to n_real+1 so the KV-cache allocator (which loops - // [0, n_layer_all)) gives the router its own per-sequence slot, and set - // n_layer_nextn = 1 so n_layer() == n_layer_all - n_layer_nextn stays n_real - // — the decoder loop and tensor loading (both bounded by n_layer()) are left - // completely unchanged and never touch layer R. - const uint32_t n_real = hparams.n_layer(); // == n_layer_all here (nextn==0) + // append one extra single-head attention layer at index R = n_real to hold + // the router K/V; bump n_layer_all so the KV cache allocates it, and set + // n_layer_nextn = 1 so n_layer() stays n_real (decoder loop unchanged) + const uint32_t n_real = hparams.n_layer(); hparams.router_layer = (int32_t) n_real; hparams.n_layer_all = n_real + 1; hparams.n_layer_nextn = 1; - // The router is a single causal head. n_embd_head_k/v are global (the router - // inherits the model head dim and uses only dim 0; the rest is zero-padded). hparams.n_head_arr[n_real] = 1; hparams.n_head_kv_arr[n_real] = 1; - hparams.n_ff_arr[n_real] = 0; // no FFN for the router layer + hparams.n_ff_arr[n_real] = 0; } void llama_model_granite_switch::load_arch_tensors(llama_model_loader &) { LLAMA_LOAD_LOCALS; const int64_t n_slots_i64 = (int64_t) n_slots; - const int64_t n_rank = (int64_t) max_lora_rank; - const int64_t n_embd_q = n_embd_head_k * n_head; // fused Q out dim - const int64_t n_embd_kv = n_embd_k_gqa; // K (== V) out dim per slice + const int64_t n_rank = (int64_t) max_lora_rank; + const int64_t n_embd_q = n_embd_head_k * n_head; + const int64_t n_embd_kv = n_embd_k_gqa; tok_embd = create_tensor(tn(LLM_TENSOR_TOKEN_EMBD, "weight"), {n_embd, n_vocab}, 0); @@ -124,7 +72,6 @@ void llama_model_granite_switch::load_arch_tensors(llama_model_loader &) { 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}, TENSOR_NOT_REQUIRED); if (output == NULL) { - // tied to the input embeddings output = create_tensor(tn(LLM_TENSOR_TOKEN_EMBD, "weight"), {n_embd, n_vocab}, TENSOR_DUPLICATED); } @@ -139,13 +86,11 @@ void llama_model_granite_switch::load_arch_tensors(llama_model_loader &) { layer.ffn_norm = create_tensor(tn(LLM_TENSOR_FFN_NORM, "weight", i), {n_embd}, 0); - // dense (non-MoE) FFN layer.ffn_gate = create_tensor(tn(LLM_TENSOR_FFN_GATE, "weight", i), {n_embd, n_ff}, 0); 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); - // ---- stacked per-token LoRA tensors (no .weight suffix) ---- - // A: {n_in, max_rank, N} B: {max_rank, n_out, N} + // stacked per-token LoRA tensors: A {n_in, max_rank, N}, B {max_rank, n_out, N} auto & sl = layer.switch_lora; sl.a_q = create_tensor(tn(LLM_TENSOR_ATTN_QKV_LORA_A_Q, i), {n_embd, n_rank, n_slots_i64}, 0); @@ -167,23 +112,14 @@ void llama_model_granite_switch::load_arch_tensors(llama_model_loader &) { } } -// Router differential gain. Each token's router K dim-0 is +GAIN for a control -// token and -GAIN otherwise; Q dim-0 is 1.0. In the causal softmax a single -// visible control token then dominates (its logit beats a normal token's by -// 2*GAIN), so the readback recovers that adapter's integer slot. 15.0 matches the -// vLLM/HF backend default (control_token_gain in config.py) and is F16-safe -// (no F32 KV cache needed). +// router K dim-0 is +GAIN for a control token and -GAIN otherwise; in the causal +// softmax a single visible control token then dominates, so the readback recovers +// its adapter slot. 15.0 matches the vLLM/HF default and is F16-safe. static constexpr float GRANITE_SWITCH_ROUTER_GAIN = 15.0f; -// ---------------------------------------------------------------------------- -// switch input: pure per-token maps. The adapter index is NOT computed here — -// it is recovered in-graph by the router attention (see build_arch_graph). This -// is intentionally stateless: no cross-ubatch carry. Per-sequence isolation -// (no leak between concurrent sequences) comes from the KV cache, not from here. -// ---------------------------------------------------------------------------- void llm_graph_input_switch::set_input(const llama_ubatch * ubatch) { if (!ubatch->token) { - return; // embedding input path is not used by this arch + return; } const int64_t n_tokens = ubatch->n_tokens; @@ -191,22 +127,22 @@ void llm_graph_input_switch::set_input(const llama_ubatch * ubatch) { std::vector sub (n_tokens); std::vector ksig(n_tokens); std::vector vval(n_tokens); - std::vector q (n_tokens, 1.0f); // Q dim-0 is constant 1.0 + std::vector q (n_tokens, 1.0f); for (int64_t i = 0; i < n_tokens; ++i) { const llama_token tok = ubatch->token[i]; - // router K/V signals: control token -> (+gain, slot); else -> (-gain, 0). + // router K/V signals: control token -> (+gain, slot); else -> (-gain, 0) auto it = smodel.control_token_to_index.find(tok); if (it != smodel.control_token_to_index.end()) { ksig[i] = +GRANITE_SWITCH_ROUTER_GAIN; - vval[i] = (float) it->second; // adapter slot (1 + adapter) + vval[i] = (float) it->second; } else { ksig[i] = -GRANITE_SWITCH_ROUTER_GAIN; - vval[i] = 0.0f; // base + vval[i] = 0.0f; } - // token-exchange: rewrite a control token to its substitute id. + // rewrite a control token to its substitute id before embedding auto sit = smodel.control_token_to_substitute.find(tok); sub[i] = (sit != smodel.control_token_to_substitute.end()) ? (int32_t) sit->second @@ -219,15 +155,12 @@ void llm_graph_input_switch::set_input(const llama_ubatch * ubatch) { ggml_backend_tensor_set(router_q, q.data(), 0, n_tokens*ggml_element_size(router_q)); } -// ---------------------------------------------------------------------------- -// graph -// ---------------------------------------------------------------------------- std::unique_ptr llama_model_granite_switch::build_arch_graph(const llm_graph_params & params) const { return std::make_unique(*this, params); } -// per-token switched LoRA delta only: B_a·(A_a·x), selected per token. -// cur: {n_in, n_tokens}, ids: {n_tokens} -> returns {n_out, n_tokens}. +// per-token switched LoRA delta: B_a·(A_a·x), adapter selected per token via ids. +// cur: {n_in, n_tokens}, ids: {n_tokens} -> {n_out, n_tokens} ggml_tensor * llama_model_granite_switch::graph::build_switched_lora_delta( ggml_tensor * lora_a, ggml_tensor * lora_b, @@ -236,25 +169,23 @@ ggml_tensor * llama_model_granite_switch::graph::build_switched_lora_delta( const int64_t n_in = cur->ne[0]; const int64_t n_tokens = cur->ne[1]; - // mul_mat_id wants the activation as {n_in, n_expert_used=1, n_tokens} - // and ids as {n_expert_used=1, n_tokens}. - ggml_tensor * x = ggml_reshape_3d(ctx0, cur, n_in, 1, n_tokens); + ggml_tensor * x = ggml_reshape_3d(ctx0, cur, n_in, 1, n_tokens); ggml_tensor * ids2 = ggml_reshape_2d(ctx0, ids, 1, n_tokens); ggml_tensor * a = ggml_mul_mat_id(ctx0, lora_a, x, ids2); // {max_rank, 1, n_tokens} ggml_tensor * d = ggml_mul_mat_id(ctx0, lora_b, a, ids2); // {n_out, 1, n_tokens} - return ggml_reshape_2d(ctx0, d, d->ne[0], n_tokens); // {n_out, n_tokens} + return ggml_reshape_2d(ctx0, d, d->ne[0], n_tokens); } -// per-token switched matmul: base (plain 2D weight) + selected LoRA delta. +// per-token switched matmul: base 2D weight + selected LoRA delta ggml_tensor * llama_model_granite_switch::graph::build_switched_lora_mm( ggml_tensor * w, ggml_tensor * lora_a, ggml_tensor * lora_b, - ggml_tensor * cur, // {n_in, n_tokens} - ggml_tensor * ids) { // {n_tokens} - ggml_tensor * base = ggml_mul_mat(ctx0, w, cur); // {n_out, n_tokens} + ggml_tensor * cur, + ggml_tensor * ids) { + ggml_tensor * base = ggml_mul_mat(ctx0, w, cur); ggml_tensor * delta = build_switched_lora_delta(lora_a, lora_b, cur, ids); return ggml_add(ctx0, base, delta); } @@ -270,7 +201,7 @@ llama_model_granite_switch::graph::graph( GGML_ASSERT(n_embd_head == hparams.n_embd_head_k()); GGML_ASSERT(n_embd_head == n_rot); - // --- switch input: substituted ids + per-token router K/V/Q signals --- + // switch input: substituted ids + per-token router K/V/Q signals auto inp_switch = std::make_unique(smodel); inp_switch->sub_tokens = ggml_new_tensor_1d(ctx0, GGML_TYPE_I32, n_tokens); inp_switch->router_ksig = ggml_new_tensor_1d(ctx0, GGML_TYPE_F32, n_tokens); @@ -286,8 +217,8 @@ llama_model_granite_switch::graph::graph( ggml_tensor * router_q = inp_switch->router_q; res->add_input(std::move(inp_switch)); - // embed the token-exchanged ids ourselves (we cannot use build_inp_embd - // because it embeds the raw ubatch tokens). Apply Granite embedding scale. + // embed the token-exchanged ids (cannot use build_inp_embd, which embeds the + // raw ubatch tokens), then apply Granite embedding scale ggml_tensor * inpL = ggml_get_rows(ctx0, model.tok_embd, sub_tokens); if (hparams.f_embedding_scale != 0.0f) { inpL = ggml_scale(ctx0, inpL, hparams.f_embedding_scale); @@ -300,20 +231,15 @@ llama_model_granite_switch::graph::graph( } auto * inp_attn = build_attn_inp_kv(); - // --- router attention: recover the per-token adapter index in-graph --- - // A single causal head whose K/V live in the KV cache at layer R == router_layer. - // Only dim 0 carries signal (rest zero-padded): Q[0]=1, K[0]=±gain, V[0]=slot/0. - // The causal softmax over a single visible control token (the single-switch - // contract) puts ~all weight on it, so attended V[0] ≈ that adapter's slot. - // State persists across decode steps and is isolated per-sequence because the - // control token's K/V stay in the per-sequence KV cache (vLLM/HF parity). No - // RoPE on the router Q/K — positions are irrelevant to the selection. - const int R = hparams.router_layer; - GGML_ASSERT(R >= 0 && "granite-switch: router_layer not configured"); - // build [n_embd_head, 1, n_tokens] with row 0 = signal, rows 1.. = 0. + // router attention: a single causal head whose K/V live in the KV cache at + // layer R, recovering the per-token adapter index in-graph. Only dim 0 carries + // signal (Q[0]=1, K[0]=±gain, V[0]=slot/0), rest zero-padded; no RoPE. State is + // isolated per-sequence because the K/V stay in the per-sequence KV cache. + const int R = hparams.router_layer; + GGML_ASSERT(R >= 0); auto router_lane = [&](ggml_tensor * sig1d) { ggml_tensor * t = ggml_reshape_3d(ctx0, sig1d, 1, 1, n_tokens); - return ggml_pad(ctx0, t, (int) n_embd_head - 1, 0, 0, 0); // pad dim0 right with zeros + return ggml_pad(ctx0, t, (int) n_embd_head - 1, 0, 0, 0); }; ggml_tensor * Qr = router_lane(router_q); ggml_tensor * Kr = router_lane(router_ksig); @@ -324,13 +250,13 @@ llama_model_granite_switch::graph::graph( Qr, Kr, Vr, nullptr, nullptr, nullptr, /*kq_scale=*/1.0f, /*il=*/R); cb(router_out, "router_out", R); - // readback: router_out is {n_embd_head, n_tokens}; row 0 is the attended slot. - // adapter_index = clamp(round(row0), 0, n_adapters), as I32 for mul_mat_id. + // readback: row 0 of router_out is the attended slot; adapter_index = + // clamp(round(row0), 0, n_adapters) as I32 for mul_mat_id (round before cast) ggml_tensor * slot_f = ggml_cont(ctx0, ggml_view_2d(ctx0, router_out, 1, n_tokens, router_out->nb[1], 0)); slot_f = ggml_reshape_1d(ctx0, slot_f, n_tokens); slot_f = ggml_clamp(ctx0, slot_f, 0.0f, (float) smodel.n_adapters); - slot_f = ggml_round(ctx0, slot_f); // round BEFORE cast (cast truncates) + slot_f = ggml_round(ctx0, slot_f); ggml_tensor * adapter_ids = ggml_cast(ctx0, slot_f, GGML_TYPE_I32); cb(adapter_ids, "adapter_ids", -1); @@ -341,26 +267,25 @@ llama_model_granite_switch::graph::graph( for (int il = 0; il < n_layer; ++il) { ggml_tensor * inpSA = inpL; - // attn norm + // norm cur = build_norm(inpL, model.layers[il].attn_norm, NULL, LLM_NORM_RMS, il); cb(cur, "attn_norm", il); - // self-attention (with per-token switched LoRA) + // self-attention cur = build_attention_layer(cur, inp_pos, adapter_ids, inp_attn, model, n_embd_head, 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); - // adapter_ids is 1D {n_tokens}; ggml_get_rows on a 1D tensor would - // gather whole rows of width n_tokens. Reshape to {1, n_tokens} so we - // select per-token columns, then flatten back to 1D {n_out}. + // reshape adapter_ids to {1, n_tokens} so get_rows selects per-token + // columns, then flatten back to 1D {n_out} const int64_t n_out = inp_out_ids->ne[0]; adapter_ids = ggml_get_rows(ctx0, ggml_reshape_2d(ctx0, adapter_ids, 1, adapter_ids->ne[0]), inp_out_ids); adapter_ids = ggml_reshape_1d(ctx0, adapter_ids, n_out); } - // ffn (with per-token switched LoRA) + // ffn cur = build_layer_ffn(cur, inpSA, adapter_ids, model, il); inpL = cur; @@ -375,7 +300,7 @@ llama_model_granite_switch::graph::graph( // lm_head cur = build_lora_mm(model.output, cur, model.output_s); - // Granite logit scaling + // For Granite architectures - scale logits cur = ggml_scale(ctx0, cur, 1.0f / hparams.f_logit_scale); cb(cur, "result_output", -1); res->t_logits = cur; @@ -398,20 +323,17 @@ ggml_tensor * llama_model_granite_switch::graph::build_attention_layer( const int64_t n_head = hparams.n_head(il); const int64_t n_head_kv = hparams.n_head_kv(il); - // fused qkv base + per-slice switched LoRA deltas. ggml_tensor * qkv = ggml_mul_mat(ctx0, layer.wqkv, cur); cb(qkv, "wqkv", il); const int64_t n_embd_q = n_embd_head * n_head; const int64_t n_embd_kv = n_embd_head * n_head_kv; - // slice base qkv into Q/K/V (contiguous copies so we can add LoRA deltas). + // slice fused qkv into Q/K/V (contiguous so we can add LoRA deltas) ggml_tensor * Qcur = ggml_cont(ctx0, ggml_view_2d(ctx0, qkv, n_embd_q, qkv->ne[1], qkv->nb[1], 0)); ggml_tensor * Kcur = ggml_cont(ctx0, ggml_view_2d(ctx0, qkv, n_embd_kv, qkv->ne[1], qkv->nb[1], n_embd_q*ggml_element_size(qkv))); ggml_tensor * Vcur = ggml_cont(ctx0, ggml_view_2d(ctx0, qkv, n_embd_kv, qkv->ne[1], qkv->nb[1], (n_embd_q + n_embd_kv)*ggml_element_size(qkv))); - // add per-token switched LoRA deltas to each base slice (delta only — the - // fused base qkv is already computed above). Qcur = ggml_add(ctx0, Qcur, build_switched_lora_delta(sl.a_q, sl.b_q, cur, adapter_ids)); Kcur = ggml_add(ctx0, Kcur, build_switched_lora_delta(sl.a_k, sl.b_k, cur, adapter_ids)); Vcur = ggml_add(ctx0, Vcur, build_switched_lora_delta(sl.a_v, sl.b_v, cur, adapter_ids)); @@ -436,8 +358,8 @@ ggml_tensor * llama_model_granite_switch::graph::build_attention_layer( const float kq_scale = hparams.f_attention_scale == 0.0f ? 1.0f/sqrtf(float(n_embd_head)) : hparams.f_attention_scale; - // wo = nullptr: build_attn returns the concatenated heads (pre-o-proj), - // then we apply the switched o-proj manually. + // wo = nullptr: build_attn returns the concatenated heads, then apply the + // switched o-proj manually ggml_tensor * attn = build_attn(inp_attn, nullptr, nullptr, nullptr, Qcur, Kcur, Vcur, nullptr, nullptr, nullptr, kq_scale, il); @@ -458,7 +380,7 @@ ggml_tensor * llama_model_granite_switch::graph::build_layer_ffn( const auto & layer = model.layers[il]; const auto & sl = layer.switch_lora; - // Granite residual scale + // For Granite architectures - scale residual if (hparams.f_residual_scale) { cur = ggml_scale(ctx0, cur, hparams.f_residual_scale); } @@ -468,7 +390,7 @@ ggml_tensor * llama_model_granite_switch::graph::build_layer_ffn( cur = build_norm(ffn_inp, layer.ffn_norm, NULL, LLM_NORM_RMS, il); cb(cur, "ffn_norm", il); - // gated SILU FFN with per-token switched LoRA on gate / up / down. + // gated SILU FFN with per-token switched LoRA on gate / up / down ggml_tensor * g = build_switched_lora_mm(layer.ffn_gate, sl.a_gate, sl.b_gate, cur, adapter_ids); ggml_tensor * u = build_switched_lora_mm(layer.ffn_up, sl.a_up, sl.b_up, cur, adapter_ids); g = ggml_silu(ctx0, g); diff --git a/src/models/models.h b/src/models/models.h index 41b6cc7b0c88..07627c04d4d0 100644 --- a/src/models/models.h +++ b/src/models/models.h @@ -1541,7 +1541,7 @@ struct llama_model_granite_switch : public llama_model_base { void load_arch_hparams(llama_model_loader & ml) override; void load_arch_tensors(llama_model_loader & ml) override; - // --- per-token switch metadata (read from GGUF in load_arch_hparams) --- + // per-token switch metadata (read from GGUF in load_arch_hparams) uint32_t n_adapters = 0; // number of real adapters (excludes the zero slot) uint32_t n_slots = 0; // n_adapters + 1 (slot 0 = base/zero delta) uint32_t max_lora_rank = 0; @@ -1551,30 +1551,21 @@ struct llama_model_granite_switch : public llama_model_base { // control token id -> substitute token id (token-exchange before embedding) std::unordered_map control_token_to_substitute; - // Per-token adapter selection is computed in-graph by a single-head causal - // "router" attention whose K/V live in the model KV cache at layer index - // hparams.router_layer (== n_layer). Because that state is per-sequence (it - // lives in the cache, like vLLM's PagedAttention router), CONCURRENT requests - // are isolated for free — there is no global sticky state to leak between - // sequences (the prior POC's bug). Within one sequence the selection carries - // over once fired (flat gain, no recency) — the vLLM/HF single-switch - // contract; see the limitation note in granite_switch.cpp. Mirrors the vLLM - // backend (single.py) and HF SingleSwitch. + // the per-token adapter selection is recovered in-graph by a single-head + // causal router attention whose K/V live in the KV cache at hparams.router_layer struct graph : public llm_graph_context { graph(const llama_model & model, const llm_graph_params & params); private: - // per-token switched LoRA delta only: B_a·(A_a·x), selected per token. - // Returns {n_out, n_tokens}. Used for the fused-qkv slices where the base - // term is already computed. + // per-token switched LoRA delta only: B_a·(A_a·x) -> {n_out, n_tokens} ggml_tensor * build_switched_lora_delta( ggml_tensor * lora_a, // {n_in, max_lora_rank, N} ggml_tensor * lora_b, // {max_lora_rank, n_out, N} ggml_tensor * cur, // {n_in, n_tokens} ggml_tensor * ids); // {n_tokens} I32 adapter index per token - // per-token switched LoRA: y = W·x + B_a·(A_a·x), selected per token via mul_mat_id + // per-token switched LoRA: y = W·x + B_a·(A_a·x), selected per token ggml_tensor * build_switched_lora_mm( ggml_tensor * w, // base weight (plain 2D) ggml_tensor * lora_a, // {n_in, max_lora_rank, N} @@ -1603,16 +1594,8 @@ struct llama_model_granite_switch : public llama_model_base { }; -// Custom graph input for the per-token switch. In set_input it fills, per token: -// - sub_tokens : the token-exchange (control id -> substitute id, else the id) -// - router_ksig : +gain for a control token, -gain otherwise (router K, dim 0) -// - router_vval : float(adapter_slot) for a control token, 0 otherwise (router V, dim 0) -// - router_q : constant 1.0 (router Q, dim 0) -// The adapter index per token is NOT computed here — it is recovered in-graph by a -// single-head causal attention over (router_q, router_ksig, router_vval) whose K/V -// live in the KV cache, so the selection is per-sequence (no cross-sequence leak). -// set_input is pure per-token maps (no history / no cross-ubatch carry). -// can_reuse() returns false so these are recomputed every ubatch. +// graph input for the per-token switch: fills the substituted ids and the +// router Q/K/V signals consumed by the in-graph router attention class llm_graph_input_switch : public llm_graph_input_i { public: llm_graph_input_switch(const llama_model_granite_switch & smodel) : smodel(smodel) {} From f51c185e8d71baba805f1f07cce6119421d0266b Mon Sep 17 00:00:00 2001 From: Bar Haim Date: Wed, 24 Jun 2026 17:55:43 +0300 Subject: [PATCH 07/40] granite-switch: trim conversion comments to match native style --- conversion/granite.py | 78 +++++++++++-------------------------------- 1 file changed, 20 insertions(+), 58 deletions(-) diff --git a/conversion/granite.py b/conversion/granite.py index 736f35a9e231..cc645b3a2eef 100644 --- a/conversion/granite.py +++ b/conversion/granite.py @@ -125,57 +125,31 @@ def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iter @ModelBase.register("GraniteSwitchForCausalLM") class GraniteSwitchModel(GraniteMoeModel): - """Conversion for IBM's GraniteSwitchForCausalLM. - - Granite Switch is a dense, all-attention Granite model with N embedded LoRA - adapters selected per-token by control tokens. The on-disk checkpoint: - - - Uses a fused ``self_attn.qkv_proj`` and fused ``shared_mlp.input_linear`` - (gate+up), like GraniteMoe's shared expert, plus ``self_attn.o_proj`` and - ``shared_mlp.output_linear``. Base weights live under ``.base_layer.weight``. - - Reserves one cache slot for the (weightless) switch, so the decoder has - ``num_hidden_layers - 1`` blocks. The decoder layers are an ``nn.ModuleList`` - saved 0-based as ``model.layers.0 .. model.layers.{N-2}`` (the switch holds no - weights and is not in the list), so ``blk.{bid}`` maps straight from ``bid``. - We emit ``block_count = num_hidden_layers - 1``. - - Stacks each LoRA over the adapter dim: A ``[n_adapters, 1, max_rank, in]``, - B ``[n_adapters, 1, out_slice, max_rank]`` (lower ranks zero-padded). B is - pre-scaled by ``alpha/rank`` at compose time, so the runtime LoRA scale is - 1.0 and A is left unscaled. - - To make per-token selection branch-free with ``ggml_mul_mat_id`` (which has no - "skip/base" id), we materialize a zero adapter at slot 0 so the stacked first - dim is ``N = num_adapters + 1``: index 0 (base tokens) adds an exact-zero delta. + """Dense, all-attention Granite model with N embedded LoRA adapters selected + per-token by control tokens. Base weights live under ``.base_layer.weight``; + each LoRA is stacked over the adapter dim (B pre-scaled by ``alpha/rank``, so + the runtime scale is 1.0). A zero adapter is materialized at slot 0 so the + stacked dim is ``N = num_adapters + 1`` and base tokens add an exact-zero delta. """ model_arch = gguf.MODEL_ARCH.GRANITE_SWITCH - # ggml's granite arch uses LLAMA_ROPE_TYPE_NORM, so q/k weights must be stored - # in the interleaved layout ggml expects. HF Granite uses transformers' - # rotate_half (split-half) layout, so we apply LlamaModel's permute to the q and - # k row-blocks of the fused qkv base AND to the q/k LoRA-B output rows. We do it - # explicitly per-slice below, so disable the parent's automatic - # (separate-projection) permute. + # permute q/k per-slice below (NORM-rope layout) instead of via the parent's + # automatic separate-projection permute undo_permute = False def __init__(self, *args, **kwargs): super().__init__(*args, **kwargs) - # The (weightless) switch reserves one cache slot; the decoder has one fewer - # block than num_hidden_layers. The decoder layers are a 0-based ModuleList, - # so block ids map straight through (no index shift). Fix block_count and - # rebuild the tensor name map for the reduced count. + # the weightless switch reserves one cache slot, so the decoder has one + # fewer block than num_hidden_layers (block ids map straight through) self.block_count = self.block_count - 1 self.tensor_map = gguf.get_tensor_name_map(self.model_arch, self.block_count) self._n_adapters = int(self.hparams["num_adapters"]) self._max_lora_rank = int(self.hparams["max_lora_rank"]) - # Stacked adapter dim: +1 for the materialized zero slot at index 0. - self._n_slots = self._n_adapters + 1 + self._n_slots = self._n_adapters + 1 # +1 for the zero slot at index 0 n_head = int(self.hparams["num_attention_heads"]) n_kv_head = int(self.hparams["num_key_value_heads"]) - # The qkv projection emits projection_head_dim-wide vectors (matches the HF - # model's GraniteLoRAEmbeddedAttention.head_dim). Fall back to head_dim, then - # to hidden/heads. head_dim = ( self.hparams.get("projection_head_dim") or self.hparams.get("head_dim") @@ -188,13 +162,10 @@ def __init__(self, *args, **kwargs): self._kv_size = n_kv_head * self._head_dim def set_gguf_parameters(self): - # Emits the four Granite multipliers, shared FFN length, head counts, rope, - # vocab, etc. add_block_count uses self.block_count (already decremented). super().set_gguf_parameters() - # Granite Switch is dense (num_local_experts == 0). The llama parent may - # have emitted a nonzero expert_used_count (e.g. 2) from num_experts_per_tok; - # force the dense invariant so the loader's n_expert_used <= n_expert holds. + # the model is dense; force the dense invariant in case the llama parent + # emitted a nonzero expert_used_count from num_experts_per_tok self.gguf_writer.add_expert_count(0) self.gguf_writer.add_expert_used_count(0) @@ -207,13 +178,10 @@ def set_gguf_parameters(self): "gguf: (granite-switch) num_adapters=%s max_lora_rank=%s n_slots=%s", self._n_adapters, self._max_lora_rank, self._n_slots, ) - # NOTE: control_token_gain is intentionally not emitted — runtime adapter - # selection is exact control-token-id matching, not the attention trick. + # control_token_gain is intentionally not emitted; the runtime uses a fixed gain def _permute_qk(self, w: Tensor, n_head: int) -> Tensor: - # The same interleave permute LlamaModel applies to q/k weights to produce - # ggml's NORM-rope layout, operating on a [out, ...] slice where - # `out = n_head * head_dim`. + # the interleave permute LlamaModel applies to q/k for ggml's NORM-rope layout return LlamaModel.permute(w, n_head, n_head) def _lora_a(self, data: Tensor) -> Tensor: @@ -226,14 +194,13 @@ def _lora_b(self, data: Tensor, permute_n_head: int | None = None) -> Tensor: # on-disk B: [n_adapters, 1, out, max_rank] -> [n_adapters+1, out, max_rank] b = data.squeeze(1) if permute_n_head is not None: - # Permute the output (row) dimension of each adapter's B slice to match - # the NEOX layout used for the (permuted) q/k base weights. + # permute each adapter's B output rows to match the permuted q/k base b = torch.stack([self._permute_qk(b[i], permute_n_head) for i in range(b.shape[0])], dim=0) zero = torch.zeros_like(b[:1]) return torch.cat([zero, b], dim=0).contiguous() def tensor_force_quant(self, name, new_name, bid, n_dims): - # Keep all stacked LoRA tensors in F16 (small; selected per-token via mul_mat_id). + # keep the stacked LoRA tensors in F16 if ".lora_a" in new_name or ".lora_b" in new_name: return gguf.GGMLQuantizationType.F16 return super().tensor_force_quant(name, new_name, bid, n_dims) @@ -241,11 +208,8 @@ def tensor_force_quant(self, name, new_name, bid, n_dims): def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iterable[tuple[str, Tensor]]: T = gguf.MODEL_TENSOR - # ---- Skip the (weightless) switch + control-token buffers ---- - # The switch is a hand-constructed cumsum-attention with no learned weights; - # control-token selection/substitution is reconstructed at load time from the - # adapter id arrays we emit in set_gguf_parameters. These persistent buffers - # (model.switch.*, model.adapter_token_ids) carry no parameters to convert. + # skip the weightless switch + control-token buffers (reconstructed at + # load time from the adapter id arrays emitted in set_gguf_parameters) bare = name.split(".")[-1] if ( name.startswith("model.switch.") or name.startswith("switch.") @@ -338,8 +302,7 @@ def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iter yield (self.format_tensor_name(key, obid), data_torch) return - # ---- Everything else: embeddings, final norm (no layer index) ---- - # embed_tokens -> token_embd, model.norm -> output_norm. lm_head is tied. + # ---- Embeddings / final norm (lm_head is tied to token_embd) ---- if name in ("model.embed_tokens.weight", "embed_tokens.weight"): yield (self.format_tensor_name(T.TOKEN_EMBD), data_torch) return @@ -347,8 +310,7 @@ def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iter yield (self.format_tensor_name(T.OUTPUT_NORM), data_torch) return if name == "lm_head.weight": - # Tied to token embeddings; the runtime ties output to token_embd. - return + return # tied to token_embd raise ValueError(f"granite-switch: unhandled tensor {name!r} (bid={bid})") From 8da8df52a5d1fc34045f1ebf922f8eae43864542 Mon Sep 17 00:00:00 2001 From: Bar Haim Date: Thu, 25 Jun 2026 14:04:11 -0700 Subject: [PATCH 08/40] granite-switch: drop unused adapter_ranks metadata --- conversion/granite.py | 1 - gguf-py/gguf/constants.py | 1 - gguf-py/gguf/gguf_writer.py | 3 --- src/llama-arch.cpp | 1 - src/llama-arch.h | 1 - 5 files changed, 7 deletions(-) diff --git a/conversion/granite.py b/conversion/granite.py index cc645b3a2eef..7af46e2b8993 100644 --- a/conversion/granite.py +++ b/conversion/granite.py @@ -173,7 +173,6 @@ def set_gguf_parameters(self): self.gguf_writer.add_max_lora_rank(self._max_lora_rank) self.gguf_writer.add_adapter_token_ids(self.hparams["adapter_token_ids"]) self.gguf_writer.add_adapter_substitute_token_ids(self.hparams["adapter_substitute_token_ids"]) - self.gguf_writer.add_adapter_ranks(self.hparams["adapter_ranks"]) logger.info( "gguf: (granite-switch) num_adapters=%s max_lora_rank=%s n_slots=%s", self._n_adapters, self._max_lora_rank, self._n_slots, diff --git a/gguf-py/gguf/constants.py b/gguf-py/gguf/constants.py index d124f50796f2..d02e31a17d44 100644 --- a/gguf-py/gguf/constants.py +++ b/gguf-py/gguf/constants.py @@ -163,7 +163,6 @@ class LLM: NUM_ADAPTERS = "{arch}.num_adapters" ADAPTER_TOKEN_IDS = "{arch}.adapter_token_ids" ADAPTER_SUBSTITUTE_TOKEN_IDS = "{arch}.adapter_substitute_token_ids" - ADAPTER_RANKS = "{arch}.adapter_ranks" MAX_LORA_RANK = "{arch}.max_lora_rank" class Attention: diff --git a/gguf-py/gguf/gguf_writer.py b/gguf-py/gguf/gguf_writer.py index 2f5d6adc1dec..e76581cca541 100644 --- a/gguf-py/gguf/gguf_writer.py +++ b/gguf-py/gguf/gguf_writer.py @@ -901,9 +901,6 @@ def add_adapter_token_ids(self, ids: Sequence[int]) -> None: def add_adapter_substitute_token_ids(self, ids: Sequence[int]) -> None: self.add_array(Keys.LLM.ADAPTER_SUBSTITUTE_TOKEN_IDS.format(arch=self.arch), ids) - def add_adapter_ranks(self, ranks: Sequence[int]) -> None: - self.add_array(Keys.LLM.ADAPTER_RANKS.format(arch=self.arch), ranks) - def add_max_lora_rank(self, rank: int) -> None: self.add_uint32(Keys.LLM.MAX_LORA_RANK.format(arch=self.arch), rank) diff --git a/src/llama-arch.cpp b/src/llama-arch.cpp index 6f4218f04a5b..6ad922f363cc 100644 --- a/src/llama-arch.cpp +++ b/src/llama-arch.cpp @@ -219,7 +219,6 @@ static const std::map LLM_KV_NAMES = { { LLM_KV_NUM_ADAPTERS, "%s.num_adapters" }, { LLM_KV_ADAPTER_TOKEN_IDS, "%s.adapter_token_ids" }, { LLM_KV_ADAPTER_SUBSTITUTE_TOKEN_IDS, "%s.adapter_substitute_token_ids" }, - { LLM_KV_ADAPTER_RANKS, "%s.adapter_ranks" }, { LLM_KV_MAX_LORA_RANK, "%s.max_lora_rank" }, { LLM_KV_TOKEN_SHIFT_COUNT, "%s.token_shift_count" }, { LLM_KV_INTERLEAVE_MOE_LAYER_STEP, "%s.interleave_moe_layer_step" }, diff --git a/src/llama-arch.h b/src/llama-arch.h index fad3d33d7ed2..505ebaf07e44 100644 --- a/src/llama-arch.h +++ b/src/llama-arch.h @@ -224,7 +224,6 @@ enum llm_kv { LLM_KV_NUM_ADAPTERS, LLM_KV_ADAPTER_TOKEN_IDS, LLM_KV_ADAPTER_SUBSTITUTE_TOKEN_IDS, - LLM_KV_ADAPTER_RANKS, LLM_KV_MAX_LORA_RANK, LLM_KV_TOKEN_SHIFT_COUNT, LLM_KV_INTERLEAVE_MOE_LAYER_STEP, From 205e1fa9f8a27f74c5a7ce4b8085b6ac389f995d Mon Sep 17 00:00:00 2001 From: Bar Haim Date: Fri, 26 Jun 2026 02:41:21 +0300 Subject: [PATCH 09/40] granite-switch: rename arch to graniteswitch and drop obid alias --- conversion/granite.py | 37 ++++++++++++++++--------------------- gguf-py/gguf/constants.py | 2 +- src/llama-arch.cpp | 2 +- 3 files changed, 18 insertions(+), 23 deletions(-) diff --git a/conversion/granite.py b/conversion/granite.py index 7af46e2b8993..9f370bc047ac 100644 --- a/conversion/granite.py +++ b/conversion/granite.py @@ -174,7 +174,7 @@ def set_gguf_parameters(self): self.gguf_writer.add_adapter_token_ids(self.hparams["adapter_token_ids"]) self.gguf_writer.add_adapter_substitute_token_ids(self.hparams["adapter_substitute_token_ids"]) logger.info( - "gguf: (granite-switch) num_adapters=%s max_lora_rank=%s n_slots=%s", + "gguf: (graniteswitch) num_adapters=%s max_lora_rank=%s n_slots=%s", self._n_adapters, self._max_lora_rank, self._n_slots, ) # control_token_gain is intentionally not emitted; the runtime uses a fixed gain @@ -218,19 +218,18 @@ def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iter # ---- Attention: fused QKV ---- if "self_attn.qkv_proj" in name: - obid = bid if name.endswith("base_layer.weight"): # Fused [q|k|v] rows. Permute the q and k row-blocks (split-half rope). q, k, v = data_torch.split([self._q_size, self._kv_size, self._kv_size], dim=0) q = self._permute_qk(q, self._n_head) k = self._permute_qk(k, self._n_kv_head) fused = torch.cat([q, k, v], dim=0) - yield (self.format_tensor_name(T.ATTN_QKV, obid), fused) + yield (self.format_tensor_name(T.ATTN_QKV, bid), fused) return if "lora_A_slices." in name: slot = int(name.rsplit(".", 1)[1]) key = {0: T.ATTN_QKV_LORA_A_Q, 1: T.ATTN_QKV_LORA_A_K, 2: T.ATTN_QKV_LORA_A_V}[slot] - yield (self.format_tensor_name(key, obid, suffix=""), self._lora_a(data_torch)) + yield (self.format_tensor_name(key, bid, suffix=""), self._lora_a(data_torch)) return if "lora_B_slices." in name: slot = int(name.rsplit(".", 1)[1]) @@ -239,56 +238,53 @@ def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iter 1: (T.ATTN_QKV_LORA_B_K, self._n_kv_head), 2: (T.ATTN_QKV_LORA_B_V, None), }[slot] - yield (self.format_tensor_name(key, obid, suffix=""), self._lora_b(data_torch, ph)) + yield (self.format_tensor_name(key, bid, suffix=""), self._lora_b(data_torch, ph)) return raise ValueError(f"Unexpected qkv_proj tensor: {name}") # ---- Attention: output projection ---- if "self_attn.o_proj" in name: - obid = bid if name.endswith("base_layer.weight"): - yield (self.format_tensor_name(T.ATTN_OUT, obid), data_torch) + yield (self.format_tensor_name(T.ATTN_OUT, bid), data_torch) return if name.endswith("lora_A"): - yield (self.format_tensor_name(T.ATTN_OUT_LORA_A, obid, suffix=""), self._lora_a(data_torch)) + yield (self.format_tensor_name(T.ATTN_OUT_LORA_A, bid, suffix=""), self._lora_a(data_torch)) return if name.endswith("lora_B"): - yield (self.format_tensor_name(T.ATTN_OUT_LORA_B, obid, suffix=""), self._lora_b(data_torch)) + yield (self.format_tensor_name(T.ATTN_OUT_LORA_B, bid, suffix=""), self._lora_b(data_torch)) return raise ValueError(f"Unexpected o_proj tensor: {name}") # ---- MLP: fused gate/up (shared_mlp.input_linear) ---- if "shared_mlp.input_linear" in name: - obid = bid ffn = self.hparams["shared_intermediate_size"] if name.endswith("base_layer.weight"): gate, up = data_torch.split([ffn, ffn], dim=0) - yield (self.format_tensor_name(T.FFN_GATE, obid), gate) - yield (self.format_tensor_name(T.FFN_UP, obid), up) + yield (self.format_tensor_name(T.FFN_GATE, bid), gate) + yield (self.format_tensor_name(T.FFN_UP, bid), up) return if "lora_A_slices." in name: slot = int(name.rsplit(".", 1)[1]) key = {0: T.FFN_GATE_LORA_A, 1: T.FFN_UP_LORA_A}[slot] - yield (self.format_tensor_name(key, obid, suffix=""), self._lora_a(data_torch)) + yield (self.format_tensor_name(key, bid, suffix=""), self._lora_a(data_torch)) return if "lora_B_slices." in name: slot = int(name.rsplit(".", 1)[1]) key = {0: T.FFN_GATE_LORA_B, 1: T.FFN_UP_LORA_B}[slot] - yield (self.format_tensor_name(key, obid, suffix=""), self._lora_b(data_torch)) + yield (self.format_tensor_name(key, bid, suffix=""), self._lora_b(data_torch)) return raise ValueError(f"Unexpected shared_mlp.input_linear tensor: {name}") # ---- MLP: down (shared_mlp.output_linear) ---- if "shared_mlp.output_linear" in name: - obid = bid if name.endswith("base_layer.weight"): - yield (self.format_tensor_name(T.FFN_DOWN, obid), data_torch) + yield (self.format_tensor_name(T.FFN_DOWN, bid), data_torch) return if name.endswith("lora_A"): - yield (self.format_tensor_name(T.FFN_DOWN_LORA_A, obid, suffix=""), self._lora_a(data_torch)) + yield (self.format_tensor_name(T.FFN_DOWN_LORA_A, bid, suffix=""), self._lora_a(data_torch)) return if name.endswith("lora_B"): - yield (self.format_tensor_name(T.FFN_DOWN_LORA_B, obid, suffix=""), self._lora_b(data_torch)) + yield (self.format_tensor_name(T.FFN_DOWN_LORA_B, bid, suffix=""), self._lora_b(data_torch)) return raise ValueError(f"Unexpected shared_mlp.output_linear tensor: {name}") @@ -296,9 +292,8 @@ def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iter if bid is not None and ".layers." in name and ( "input_layernorm" in name or "post_attention_layernorm" in name ): - obid = bid key = T.ATTN_NORM if "input_layernorm" in name else T.FFN_NORM - yield (self.format_tensor_name(key, obid), data_torch) + yield (self.format_tensor_name(key, bid), data_torch) return # ---- Embeddings / final norm (lm_head is tied to token_embd) ---- @@ -311,7 +306,7 @@ def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iter if name == "lm_head.weight": return # tied to token_embd - raise ValueError(f"granite-switch: unhandled tensor {name!r} (bid={bid})") + raise ValueError(f"graniteswitch: unhandled tensor {name!r} (bid={bid})") @ModelBase.register("GraniteMoeHybridForCausalLM", "BambaForCausalLM") diff --git a/gguf-py/gguf/constants.py b/gguf-py/gguf/constants.py index d02e31a17d44..940bb67205c5 100644 --- a/gguf-py/gguf/constants.py +++ b/gguf-py/gguf/constants.py @@ -1100,7 +1100,7 @@ class MODEL_TENSOR(IntEnum): MODEL_ARCH.GRANITE: "granite", MODEL_ARCH.GRANITE_MOE: "granitemoe", MODEL_ARCH.GRANITE_HYBRID: "granitehybrid", - MODEL_ARCH.GRANITE_SWITCH: "granite-switch", + MODEL_ARCH.GRANITE_SWITCH: "graniteswitch", MODEL_ARCH.CHAMELEON: "chameleon", MODEL_ARCH.WAVTOKENIZER_DEC: "wavtokenizer-dec", MODEL_ARCH.PLM: "plm", diff --git a/src/llama-arch.cpp b/src/llama-arch.cpp index 6ad922f363cc..48148a2f3a20 100644 --- a/src/llama-arch.cpp +++ b/src/llama-arch.cpp @@ -100,7 +100,7 @@ static const std::map LLM_ARCH_NAMES = { { LLM_ARCH_GRANITE, "granite" }, { LLM_ARCH_GRANITE_MOE, "granitemoe" }, { LLM_ARCH_GRANITE_HYBRID, "granitehybrid" }, - { LLM_ARCH_GRANITE_SWITCH, "granite-switch" }, + { LLM_ARCH_GRANITE_SWITCH, "graniteswitch" }, { LLM_ARCH_CHAMELEON, "chameleon" }, { LLM_ARCH_WAVTOKENIZER_DEC, "wavtokenizer-dec" }, { LLM_ARCH_PLM, "plm" }, From 723db80bd0436d59dde98e5d51263dc64bd105f1 Mon Sep 17 00:00:00 2001 From: Bar Haim Date: Fri, 26 Jun 2026 02:50:07 +0300 Subject: [PATCH 10/40] granite-switch: fix non-ASCII comments and document router gain assumption --- src/models/granite_switch.cpp | 7 ++++--- 1 file changed, 4 insertions(+), 3 deletions(-) diff --git a/src/models/granite_switch.cpp b/src/models/granite_switch.cpp index 3d9fe7dd679d..d8e1f3e76c35 100644 --- a/src/models/granite_switch.cpp +++ b/src/models/granite_switch.cpp @@ -114,7 +114,8 @@ void llama_model_granite_switch::load_arch_tensors(llama_model_loader &) { // router K dim-0 is +GAIN for a control token and -GAIN otherwise; in the causal // softmax a single visible control token then dominates, so the readback recovers -// its adapter slot. 15.0 matches the vLLM/HF default and is F16-safe. +// its adapter slot. This assumes the HF control_token_gain default of 15.0 (the +// only value IBM ships); it is not read from the GGUF. 15.0 is also F16-safe. static constexpr float GRANITE_SWITCH_ROUTER_GAIN = 15.0f; void llm_graph_input_switch::set_input(const llama_ubatch * ubatch) { @@ -159,7 +160,7 @@ std::unique_ptr llama_model_granite_switch::build_arch_graph( return std::make_unique(*this, params); } -// per-token switched LoRA delta: B_a·(A_a·x), adapter selected per token via ids. +// per-token switched LoRA delta: B_a*(A_a*x), adapter selected per token via ids. // cur: {n_in, n_tokens}, ids: {n_tokens} -> {n_out, n_tokens} ggml_tensor * llama_model_granite_switch::graph::build_switched_lora_delta( ggml_tensor * lora_a, @@ -233,7 +234,7 @@ llama_model_granite_switch::graph::graph( // router attention: a single causal head whose K/V live in the KV cache at // layer R, recovering the per-token adapter index in-graph. Only dim 0 carries - // signal (Q[0]=1, K[0]=±gain, V[0]=slot/0), rest zero-padded; no RoPE. State is + // signal (Q[0]=1, K[0]=+/-gain, V[0]=slot/0), rest zero-padded; no RoPE. State is // isolated per-sequence because the K/V stay in the per-sequence KV cache. const int R = hparams.router_layer; GGML_ASSERT(R >= 0); From a8ee797fcbf37b422664cb1db531c52ba4bea6f9 Mon Sep 17 00:00:00 2001 From: Bar Haim Date: Sun, 28 Jun 2026 18:32:59 +0300 Subject: [PATCH 11/40] granite-switch: drop section comments from constants.py to match native style --- gguf-py/gguf/constants.py | 4 ---- 1 file changed, 4 deletions(-) diff --git a/gguf-py/gguf/constants.py b/gguf-py/gguf/constants.py index 940bb67205c5..caa55d7fe3a5 100644 --- a/gguf-py/gguf/constants.py +++ b/gguf-py/gguf/constants.py @@ -159,7 +159,6 @@ class LLM: TARGET_HIDDEN_SIZE = "{arch}.target_hidden_size" BLOCK_SIZE = "{arch}.block_size" NORM_BEFORE_RESIDUAL = "{arch}.norm_before_residual" - # Granite Switch: per-token embedded LoRA adapters NUM_ADAPTERS = "{arch}.num_adapters" ADAPTER_TOKEN_IDS = "{arch}.adapter_token_ids" ADAPTER_SUBSTITUTE_TOKEN_IDS = "{arch}.adapter_substitute_token_ids" @@ -989,7 +988,6 @@ class MODEL_TENSOR(IntEnum): A_QF_FFN_UP = auto() A_QF_FFN_DOWN = auto() A_QF_FFN_NORM = auto() - # Granite Switch: per-token embedded LoRA adapters (stacked over N = num_adapters + 1) ATTN_QKV_LORA_A_Q = auto() ATTN_QKV_LORA_B_Q = auto() ATTN_QKV_LORA_A_K = auto() @@ -1582,7 +1580,6 @@ class MODEL_TENSOR(IntEnum): MODEL_TENSOR.NEXTN_SHARED_HEAD_NORM: "blk.{bid}.nextn.shared_head_norm", MODEL_TENSOR.FC: "fc", MODEL_TENSOR.D2T: "d2t", - # Granite Switch: per-token embedded LoRA adapters MODEL_TENSOR.ATTN_QKV_LORA_A_Q: "blk.{bid}.attn_qkv.lora_a_q", MODEL_TENSOR.ATTN_QKV_LORA_B_Q: "blk.{bid}.attn_qkv.lora_b_q", MODEL_TENSOR.ATTN_QKV_LORA_A_K: "blk.{bid}.attn_qkv.lora_a_k", @@ -3717,7 +3714,6 @@ class MODEL_TENSOR(IntEnum): MODEL_TENSOR.FFN_GATE, MODEL_TENSOR.FFN_DOWN, MODEL_TENSOR.FFN_UP, - # Per-token embedded LoRA adapters MODEL_TENSOR.ATTN_QKV_LORA_A_Q, MODEL_TENSOR.ATTN_QKV_LORA_B_Q, MODEL_TENSOR.ATTN_QKV_LORA_A_K, From 4e9ac66abd3e0144ba6e45d729f86dca77c7a612 Mon Sep 17 00:00:00 2001 From: Bar Haim Date: Sun, 28 Jun 2026 18:42:15 +0300 Subject: [PATCH 12/40] granite-switch: add functional tensor block comments matching Granite4 Vision style --- gguf-py/gguf/constants.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/gguf-py/gguf/constants.py b/gguf-py/gguf/constants.py index caa55d7fe3a5..a1d45f853e04 100644 --- a/gguf-py/gguf/constants.py +++ b/gguf-py/gguf/constants.py @@ -988,6 +988,7 @@ class MODEL_TENSOR(IntEnum): A_QF_FFN_UP = auto() A_QF_FFN_DOWN = auto() A_QF_FFN_NORM = auto() + # per-token embedded LoRA adapters - Granite Switch ATTN_QKV_LORA_A_Q = auto() ATTN_QKV_LORA_B_Q = auto() ATTN_QKV_LORA_A_K = auto() @@ -1580,6 +1581,7 @@ class MODEL_TENSOR(IntEnum): MODEL_TENSOR.NEXTN_SHARED_HEAD_NORM: "blk.{bid}.nextn.shared_head_norm", MODEL_TENSOR.FC: "fc", MODEL_TENSOR.D2T: "d2t", + # per-token embedded LoRA adapters - Granite Switch MODEL_TENSOR.ATTN_QKV_LORA_A_Q: "blk.{bid}.attn_qkv.lora_a_q", MODEL_TENSOR.ATTN_QKV_LORA_B_Q: "blk.{bid}.attn_qkv.lora_b_q", MODEL_TENSOR.ATTN_QKV_LORA_A_K: "blk.{bid}.attn_qkv.lora_a_k", From 42346fd9011a7878278b74b07c350eea55d68090 Mon Sep 17 00:00:00 2001 From: Bar Haim Date: Sun, 28 Jun 2026 21:19:19 +0300 Subject: [PATCH 13/40] granite-switch: clarify n_expert_used comment State the actual constraint: mul_mat_id needs n_expert_used == 1, and since the GGUF carries expert_count = 0 the generic loader's n_expert == 0 => n_expert_used == 0 assertion has already passed by the time load_arch_hparams runs, so it is forced to 1 here. --- src/models/granite_switch.cpp | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/src/models/granite_switch.cpp b/src/models/granite_switch.cpp index d8e1f3e76c35..33c679d19e92 100644 --- a/src/models/granite_switch.cpp +++ b/src/models/granite_switch.cpp @@ -21,8 +21,10 @@ void llama_model_granite_switch::load_arch_hparams(llama_model_loader & ml) { ml.get_key(LLM_KV_EXPERT_SHARED_FEED_FORWARD_LENGTH, hparams.n_ff_shexp, /* required */ false); - // the per-token LoRA tensors use MUL_MAT_ID with n_expert_used = 1; the model - // is otherwise dense, so set it after the generic loader's n_expert checks + // per-token LoRA selection runs through MUL_MAT_ID, which needs n_expert_used == 1. + // the GGUF carries expert_count = 0 (dense model), so the generic loader's + // n_expert == 0 => n_expert_used == 0 assertion has already passed by the time + // load_arch_hparams runs; force it to 1 here for the mul_mat_id path. hparams.n_expert_used = 1; ml.get_key(LLM_KV_NUM_ADAPTERS, n_adapters); From c18749cb9886b5165abbbf3d756a415e069d7f08 Mon Sep 17 00:00:00 2001 From: Bar Haim Date: Sun, 28 Jun 2026 21:33:08 +0300 Subject: [PATCH 14/40] granite-switch: note n_layer_nextn reuse has no MTP The router carving reuses n_layer_nextn, normally the MTP/next-token count. Clarify in the comment that it is borrowed here purely as the trailing-layers lever and that there is no MTP head, to spare readers the double-take. --- src/models/granite_switch.cpp | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/src/models/granite_switch.cpp b/src/models/granite_switch.cpp index 33c679d19e92..05bcf1f2c7f6 100644 --- a/src/models/granite_switch.cpp +++ b/src/models/granite_switch.cpp @@ -48,8 +48,9 @@ void llama_model_granite_switch::load_arch_hparams(llama_model_loader & ml) { } // append one extra single-head attention layer at index R = n_real to hold - // the router K/V; bump n_layer_all so the KV cache allocates it, and set - // n_layer_nextn = 1 so n_layer() stays n_real (decoder loop unchanged) + // the router K/V; bump n_layer_all so the KV cache allocates it. reuse + // n_layer_nextn (the trailing-layers count; no MTP here) = 1 so n_layer() + // stays n_real and the decoder loop is unchanged const uint32_t n_real = hparams.n_layer(); hparams.router_layer = (int32_t) n_real; hparams.n_layer_all = n_real + 1; From 3ad6033f880c4a07e5b8aa30671ec92590fcad42 Mon Sep 17 00:00:00 2001 From: Bar Haim Date: Tue, 30 Jun 2026 11:24:58 +0300 Subject: [PATCH 15/40] granite-switch: rename source file and apply review nits --- .../{granite_switch.cpp => granite-switch.cpp} | 18 ++++++++++-------- 1 file changed, 10 insertions(+), 8 deletions(-) rename src/models/{granite_switch.cpp => granite-switch.cpp} (96%) diff --git a/src/models/granite_switch.cpp b/src/models/granite-switch.cpp similarity index 96% rename from src/models/granite_switch.cpp rename to src/models/granite-switch.cpp index 05bcf1f2c7f6..d819613c6580 100644 --- a/src/models/granite_switch.cpp +++ b/src/models/granite-switch.cpp @@ -29,10 +29,10 @@ void llama_model_granite_switch::load_arch_hparams(llama_model_loader & ml) { ml.get_key(LLM_KV_NUM_ADAPTERS, n_adapters); ml.get_key(LLM_KV_MAX_LORA_RANK, max_lora_rank); - n_slots = n_adapters + 1; + n_slots = n_adapters + 1; // n_adapters plus un-adapted base slot - std::vector token_ids; - std::vector substitute_ids; + std::vector token_ids; + std::vector substitute_ids; ml.get_arr(LLM_KV_ADAPTER_TOKEN_IDS, token_ids); ml.get_arr(LLM_KV_ADAPTER_SUBSTITUTE_TOKEN_IDS, substitute_ids); @@ -43,14 +43,16 @@ void llama_model_granite_switch::load_arch_hparams(llama_model_loader & ml) { control_token_to_substitute.clear(); for (uint32_t i = 0; i < n_adapters; ++i) { // adapter i -> stacked slot i+1 (slot 0 is the base/zero delta) - control_token_to_index[(llama_token) token_ids[i]] = (int32_t) (i + 1); - control_token_to_substitute[(llama_token) token_ids[i]] = (llama_token) substitute_ids[i]; + control_token_to_index[token_ids[i]] = (int32_t) (i + 1); + control_token_to_substitute[token_ids[i]] = substitute_ids[i]; } // append one extra single-head attention layer at index R = n_real to hold // the router K/V; bump n_layer_all so the KV cache allocates it. reuse // n_layer_nextn (the trailing-layers count; no MTP here) = 1 so n_layer() - // stays n_real and the decoder loop is unchanged + // stays n_real and the decoder loop is unchanged. note the router layer lives + // at the END (index n_real), not the start, so the regular layers keep their + // indices and the KV cache shift/defrag can skip it without renumbering const uint32_t n_real = hparams.n_layer(); hparams.router_layer = (int32_t) n_real; hparams.n_layer_all = n_real + 1; @@ -137,7 +139,7 @@ void llm_graph_input_switch::set_input(const llama_ubatch * ubatch) { const llama_token tok = ubatch->token[i]; // router K/V signals: control token -> (+gain, slot); else -> (-gain, 0) - auto it = smodel.control_token_to_index.find(tok); + const auto it = smodel.control_token_to_index.find(tok); if (it != smodel.control_token_to_index.end()) { ksig[i] = +GRANITE_SWITCH_ROUTER_GAIN; vval[i] = (float) it->second; @@ -147,7 +149,7 @@ void llm_graph_input_switch::set_input(const llama_ubatch * ubatch) { } // rewrite a control token to its substitute id before embedding - auto sit = smodel.control_token_to_substitute.find(tok); + const auto sit = smodel.control_token_to_substitute.find(tok); sub[i] = (sit != smodel.control_token_to_substitute.end()) ? (int32_t) sit->second : (int32_t) tok; From 3fdc1e4ca8e514896baa6900d5bb0d9759c8a35c Mon Sep 17 00:00:00 2001 From: Bar Haim Date: Tue, 30 Jun 2026 12:04:01 +0300 Subject: [PATCH 16/40] granite-switch: don't force LoRA tensors to F16, follow --outtype instead --- conversion/granite.py | 6 ------ 1 file changed, 6 deletions(-) diff --git a/conversion/granite.py b/conversion/granite.py index 9f370bc047ac..164fe7dedfcb 100644 --- a/conversion/granite.py +++ b/conversion/granite.py @@ -198,12 +198,6 @@ def _lora_b(self, data: Tensor, permute_n_head: int | None = None) -> Tensor: zero = torch.zeros_like(b[:1]) return torch.cat([zero, b], dim=0).contiguous() - def tensor_force_quant(self, name, new_name, bid, n_dims): - # keep the stacked LoRA tensors in F16 - if ".lora_a" in new_name or ".lora_b" in new_name: - return gguf.GGMLQuantizationType.F16 - return super().tensor_force_quant(name, new_name, bid, n_dims) - def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iterable[tuple[str, Tensor]]: T = gguf.MODEL_TENSOR From 8ded284a88d9b8c4652b601c88eac16317cff13c Mon Sep 17 00:00:00 2001 From: Bar Haim Date: Tue, 30 Jun 2026 13:07:43 +0300 Subject: [PATCH 17/40] granite-switch: drop redundant _permute_qk wrapper, call LlamaModel.permute directly --- conversion/granite.py | 14 ++++++-------- 1 file changed, 6 insertions(+), 8 deletions(-) diff --git a/conversion/granite.py b/conversion/granite.py index 164fe7dedfcb..00ec16497a24 100644 --- a/conversion/granite.py +++ b/conversion/granite.py @@ -179,10 +179,6 @@ def set_gguf_parameters(self): ) # control_token_gain is intentionally not emitted; the runtime uses a fixed gain - def _permute_qk(self, w: Tensor, n_head: int) -> Tensor: - # the interleave permute LlamaModel applies to q/k for ggml's NORM-rope layout - return LlamaModel.permute(w, n_head, n_head) - def _lora_a(self, data: Tensor) -> Tensor: # on-disk A: [n_adapters, 1, max_rank, in] -> [n_adapters+1, max_rank, in] a = data.squeeze(1) @@ -194,7 +190,8 @@ def _lora_b(self, data: Tensor, permute_n_head: int | None = None) -> Tensor: b = data.squeeze(1) if permute_n_head is not None: # permute each adapter's B output rows to match the permuted q/k base - b = torch.stack([self._permute_qk(b[i], permute_n_head) for i in range(b.shape[0])], dim=0) + # (LlamaModel.permute interleaves q/k for ggml's NORM-rope layout) + b = torch.stack([self.permute(b[i], permute_n_head, permute_n_head) for i in range(b.shape[0])], dim=0) zero = torch.zeros_like(b[:1]) return torch.cat([zero, b], dim=0).contiguous() @@ -213,10 +210,11 @@ def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iter # ---- Attention: fused QKV ---- if "self_attn.qkv_proj" in name: if name.endswith("base_layer.weight"): - # Fused [q|k|v] rows. Permute the q and k row-blocks (split-half rope). + # Fused [q|k|v] rows. Permute the q and k row-blocks for ggml's + # NORM-rope layout (LlamaModel.permute, n_head_kv == n_head per slice). q, k, v = data_torch.split([self._q_size, self._kv_size, self._kv_size], dim=0) - q = self._permute_qk(q, self._n_head) - k = self._permute_qk(k, self._n_kv_head) + q = self.permute(q, self._n_head, self._n_head) + k = self.permute(k, self._n_kv_head, self._n_kv_head) fused = torch.cat([q, k, v], dim=0) yield (self.format_tensor_name(T.ATTN_QKV, bid), fused) return From 846e6b8016110e0e66d1d67156cf5b9f58891b6c Mon Sep 17 00:00:00 2001 From: Bar Haim Date: Tue, 30 Jun 2026 13:13:14 +0300 Subject: [PATCH 18/40] granite-switch: read router gain from GGUF (control_token_gain) instead of hardcoding 15.0 --- conversion/granite.py | 7 ++++--- gguf-py/gguf/constants.py | 1 + gguf-py/gguf/gguf_writer.py | 3 +++ src/llama-arch.cpp | 1 + src/llama-arch.h | 1 + src/models/granite-switch.cpp | 11 +++++------ src/models/models.h | 1 + 7 files changed, 16 insertions(+), 9 deletions(-) diff --git a/conversion/granite.py b/conversion/granite.py index 00ec16497a24..b5ff1942cdaa 100644 --- a/conversion/granite.py +++ b/conversion/granite.py @@ -173,11 +173,12 @@ def set_gguf_parameters(self): self.gguf_writer.add_max_lora_rank(self._max_lora_rank) self.gguf_writer.add_adapter_token_ids(self.hparams["adapter_token_ids"]) self.gguf_writer.add_adapter_substitute_token_ids(self.hparams["adapter_substitute_token_ids"]) + control_token_gain = float(self.hparams.get("control_token_gain", 15.0)) + self.gguf_writer.add_control_token_gain(control_token_gain) logger.info( - "gguf: (graniteswitch) num_adapters=%s max_lora_rank=%s n_slots=%s", - self._n_adapters, self._max_lora_rank, self._n_slots, + "gguf: (graniteswitch) num_adapters=%s max_lora_rank=%s n_slots=%s control_token_gain=%s", + self._n_adapters, self._max_lora_rank, self._n_slots, control_token_gain, ) - # control_token_gain is intentionally not emitted; the runtime uses a fixed gain def _lora_a(self, data: Tensor) -> Tensor: # on-disk A: [n_adapters, 1, max_rank, in] -> [n_adapters+1, max_rank, in] diff --git a/gguf-py/gguf/constants.py b/gguf-py/gguf/constants.py index a1d45f853e04..7b6457189c25 100644 --- a/gguf-py/gguf/constants.py +++ b/gguf-py/gguf/constants.py @@ -163,6 +163,7 @@ class LLM: ADAPTER_TOKEN_IDS = "{arch}.adapter_token_ids" ADAPTER_SUBSTITUTE_TOKEN_IDS = "{arch}.adapter_substitute_token_ids" MAX_LORA_RANK = "{arch}.max_lora_rank" + CONTROL_TOKEN_GAIN = "{arch}.control_token_gain" class Attention: HEAD_COUNT = "{arch}.attention.head_count" diff --git a/gguf-py/gguf/gguf_writer.py b/gguf-py/gguf/gguf_writer.py index e76581cca541..3c7847797708 100644 --- a/gguf-py/gguf/gguf_writer.py +++ b/gguf-py/gguf/gguf_writer.py @@ -904,6 +904,9 @@ def add_adapter_substitute_token_ids(self, ids: Sequence[int]) -> None: def add_max_lora_rank(self, rank: int) -> None: self.add_uint32(Keys.LLM.MAX_LORA_RANK.format(arch=self.arch), rank) + def add_control_token_gain(self, gain: float) -> None: + self.add_float32(Keys.LLM.CONTROL_TOKEN_GAIN.format(arch=self.arch), gain) + def add_wkv_head_size(self, size: int) -> None: self.add_uint32(Keys.WKV.HEAD_SIZE.format(arch=self.arch), size) diff --git a/src/llama-arch.cpp b/src/llama-arch.cpp index 48148a2f3a20..b929f362b68c 100644 --- a/src/llama-arch.cpp +++ b/src/llama-arch.cpp @@ -220,6 +220,7 @@ static const std::map LLM_KV_NAMES = { { LLM_KV_ADAPTER_TOKEN_IDS, "%s.adapter_token_ids" }, { LLM_KV_ADAPTER_SUBSTITUTE_TOKEN_IDS, "%s.adapter_substitute_token_ids" }, { LLM_KV_MAX_LORA_RANK, "%s.max_lora_rank" }, + { LLM_KV_CONTROL_TOKEN_GAIN, "%s.control_token_gain" }, { LLM_KV_TOKEN_SHIFT_COUNT, "%s.token_shift_count" }, { LLM_KV_INTERLEAVE_MOE_LAYER_STEP, "%s.interleave_moe_layer_step" }, { LLM_KV_FULL_ATTENTION_INTERVAL, "%s.full_attention_interval" }, diff --git a/src/llama-arch.h b/src/llama-arch.h index 505ebaf07e44..4b9acd5530a7 100644 --- a/src/llama-arch.h +++ b/src/llama-arch.h @@ -225,6 +225,7 @@ enum llm_kv { LLM_KV_ADAPTER_TOKEN_IDS, LLM_KV_ADAPTER_SUBSTITUTE_TOKEN_IDS, LLM_KV_MAX_LORA_RANK, + LLM_KV_CONTROL_TOKEN_GAIN, LLM_KV_TOKEN_SHIFT_COUNT, LLM_KV_INTERLEAVE_MOE_LAYER_STEP, LLM_KV_FULL_ATTENTION_INTERVAL, diff --git a/src/models/granite-switch.cpp b/src/models/granite-switch.cpp index d819613c6580..1f6482e460a4 100644 --- a/src/models/granite-switch.cpp +++ b/src/models/granite-switch.cpp @@ -29,6 +29,7 @@ void llama_model_granite_switch::load_arch_hparams(llama_model_loader & ml) { ml.get_key(LLM_KV_NUM_ADAPTERS, n_adapters); ml.get_key(LLM_KV_MAX_LORA_RANK, max_lora_rank); + ml.get_key(LLM_KV_CONTROL_TOKEN_GAIN, router_gain, /* required */ false); n_slots = n_adapters + 1; // n_adapters plus un-adapted base slot std::vector token_ids; @@ -117,11 +118,9 @@ void llama_model_granite_switch::load_arch_tensors(llama_model_loader &) { } } -// router K dim-0 is +GAIN for a control token and -GAIN otherwise; in the causal +// router K dim-0 is +gain for a control token and -gain otherwise; in the causal // softmax a single visible control token then dominates, so the readback recovers -// its adapter slot. This assumes the HF control_token_gain default of 15.0 (the -// only value IBM ships); it is not read from the GGUF. 15.0 is also F16-safe. -static constexpr float GRANITE_SWITCH_ROUTER_GAIN = 15.0f; +// its adapter slot. The gain (smodel.router_gain) comes from the GGUF, defaulting to 15.0. void llm_graph_input_switch::set_input(const llama_ubatch * ubatch) { if (!ubatch->token) { @@ -141,10 +140,10 @@ void llm_graph_input_switch::set_input(const llama_ubatch * ubatch) { // router K/V signals: control token -> (+gain, slot); else -> (-gain, 0) const auto it = smodel.control_token_to_index.find(tok); if (it != smodel.control_token_to_index.end()) { - ksig[i] = +GRANITE_SWITCH_ROUTER_GAIN; + ksig[i] = +smodel.router_gain; vval[i] = (float) it->second; } else { - ksig[i] = -GRANITE_SWITCH_ROUTER_GAIN; + ksig[i] = -smodel.router_gain; vval[i] = 0.0f; } diff --git a/src/models/models.h b/src/models/models.h index 07627c04d4d0..ebd73b32566f 100644 --- a/src/models/models.h +++ b/src/models/models.h @@ -1545,6 +1545,7 @@ struct llama_model_granite_switch : public llama_model_base { uint32_t n_adapters = 0; // number of real adapters (excludes the zero slot) uint32_t n_slots = 0; // n_adapters + 1 (slot 0 = base/zero delta) uint32_t max_lora_rank = 0; + float router_gain = 15.0f; // control-token softmax bias magnitude (see set_input) // token id -> adapter slot (1 + adapter, so slot 0 stays "base") std::unordered_map control_token_to_index; From e5b1539b17024e2d4e49937a637300b2eb4f4fb5 Mon Sep 17 00:00:00 2001 From: Bar Haim Date: Tue, 30 Jun 2026 14:28:47 +0300 Subject: [PATCH 19/40] granite-switch: derive n_slots() --- src/models/granite-switch.cpp | 3 +-- src/models/models.h | 4 +++- 2 files changed, 4 insertions(+), 3 deletions(-) diff --git a/src/models/granite-switch.cpp b/src/models/granite-switch.cpp index 1f6482e460a4..45a60b38fd24 100644 --- a/src/models/granite-switch.cpp +++ b/src/models/granite-switch.cpp @@ -30,7 +30,6 @@ void llama_model_granite_switch::load_arch_hparams(llama_model_loader & ml) { ml.get_key(LLM_KV_NUM_ADAPTERS, n_adapters); ml.get_key(LLM_KV_MAX_LORA_RANK, max_lora_rank); ml.get_key(LLM_KV_CONTROL_TOKEN_GAIN, router_gain, /* required */ false); - n_slots = n_adapters + 1; // n_adapters plus un-adapted base slot std::vector token_ids; std::vector substitute_ids; @@ -67,7 +66,7 @@ void llama_model_granite_switch::load_arch_hparams(llama_model_loader & ml) { void llama_model_granite_switch::load_arch_tensors(llama_model_loader &) { LLAMA_LOAD_LOCALS; - const int64_t n_slots_i64 = (int64_t) n_slots; + const int64_t n_slots_i64 = (int64_t) n_slots(); const int64_t n_rank = (int64_t) max_lora_rank; const int64_t n_embd_q = n_embd_head_k * n_head; const int64_t n_embd_kv = n_embd_k_gqa; diff --git a/src/models/models.h b/src/models/models.h index ebd73b32566f..b6e05103f685 100644 --- a/src/models/models.h +++ b/src/models/models.h @@ -1543,10 +1543,12 @@ struct llama_model_granite_switch : public llama_model_base { // per-token switch metadata (read from GGUF in load_arch_hparams) uint32_t n_adapters = 0; // number of real adapters (excludes the zero slot) - uint32_t n_slots = 0; // n_adapters + 1 (slot 0 = base/zero delta) uint32_t max_lora_rank = 0; float router_gain = 15.0f; // control-token softmax bias magnitude (see set_input) + // stacked-adapter slots: n_adapters + 1 (slot 0 = base/zero delta) + uint32_t n_slots() const { return n_adapters + 1; } + // token id -> adapter slot (1 + adapter, so slot 0 stays "base") std::unordered_map control_token_to_index; // control token id -> substitute token id (token-exchange before embedding) From 37a3f873930f8760a798185007c37a5303308b92 Mon Sep 17 00:00:00 2001 From: Bar Haim Date: Tue, 30 Jun 2026 15:12:22 +0300 Subject: [PATCH 20/40] granite-switch: move llm_graph_input_switch into granite-switch.cpp --- src/models/granite-switch.cpp | 17 +++++++++++++++++ src/models/models.h | 18 ------------------ 2 files changed, 17 insertions(+), 18 deletions(-) diff --git a/src/models/granite-switch.cpp b/src/models/granite-switch.cpp index 45a60b38fd24..6d4d774bf6c9 100644 --- a/src/models/granite-switch.cpp +++ b/src/models/granite-switch.cpp @@ -117,6 +117,23 @@ void llama_model_granite_switch::load_arch_tensors(llama_model_loader &) { } } +// graph input for the per-token switch: fills the substituted ids and the +// router Q/K/V signals consumed by the in-graph router attention +class llm_graph_input_switch : public llm_graph_input_i { +public: + llm_graph_input_switch(const llama_model_granite_switch & smodel) : smodel(smodel) {} + virtual ~llm_graph_input_switch() = default; + + void set_input(const llama_ubatch * ubatch) override; + + ggml_tensor * sub_tokens = nullptr; // I32 [n_tokens] control-substituted token ids + ggml_tensor * router_ksig = nullptr; // F32 [n_tokens] router K signal (±gain) + ggml_tensor * router_vval = nullptr; // F32 [n_tokens] router V value (adapter slot / 0) + ggml_tensor * router_q = nullptr; // F32 [n_tokens] router Q value (constant 1.0) + + const llama_model_granite_switch & smodel; +}; + // router K dim-0 is +gain for a control token and -gain otherwise; in the causal // softmax a single visible control token then dominates, so the readback recovers // its adapter slot. The gain (smodel.router_gain) comes from the GGUF, defaulting to 15.0. diff --git a/src/models/models.h b/src/models/models.h index b6e05103f685..904ce5780ee3 100644 --- a/src/models/models.h +++ b/src/models/models.h @@ -1597,24 +1597,6 @@ struct llama_model_granite_switch : public llama_model_base { }; -// graph input for the per-token switch: fills the substituted ids and the -// router Q/K/V signals consumed by the in-graph router attention -class llm_graph_input_switch : public llm_graph_input_i { -public: - llm_graph_input_switch(const llama_model_granite_switch & smodel) : smodel(smodel) {} - virtual ~llm_graph_input_switch() = default; - - void set_input(const llama_ubatch * ubatch) override; - - ggml_tensor * sub_tokens = nullptr; // I32 [n_tokens] control-substituted token ids - ggml_tensor * router_ksig = nullptr; // F32 [n_tokens] router K signal (±gain) - ggml_tensor * router_vval = nullptr; // F32 [n_tokens] router V value (adapter slot / 0) - ggml_tensor * router_q = nullptr; // F32 [n_tokens] router Q value (constant 1.0) - - const llama_model_granite_switch & smodel; -}; - - struct llama_model_minicpm : public llama_model_base { llama_model_minicpm(const struct llama_model_params & params) : llama_model_base(params) {} void load_arch_hparams(llama_model_loader & ml) override; From 0173b548b18eba90065b3a6c03bd6d84408ffce0 Mon Sep 17 00:00:00 2001 From: Bar Haim Date: Tue, 30 Jun 2026 15:25:29 +0300 Subject: [PATCH 21/40] granite-switch: cut AI-style narration comments --- conversion/granite.py | 5 +--- src/llama-model.h | 1 - src/models/granite-switch.cpp | 51 +++++++++-------------------------- 3 files changed, 13 insertions(+), 44 deletions(-) diff --git a/conversion/granite.py b/conversion/granite.py index b5ff1942cdaa..be2a4a3416f1 100644 --- a/conversion/granite.py +++ b/conversion/granite.py @@ -175,10 +175,7 @@ def set_gguf_parameters(self): self.gguf_writer.add_adapter_substitute_token_ids(self.hparams["adapter_substitute_token_ids"]) control_token_gain = float(self.hparams.get("control_token_gain", 15.0)) self.gguf_writer.add_control_token_gain(control_token_gain) - logger.info( - "gguf: (graniteswitch) num_adapters=%s max_lora_rank=%s n_slots=%s control_token_gain=%s", - self._n_adapters, self._max_lora_rank, self._n_slots, control_token_gain, - ) + logger.info("gguf: (graniteswitch) num_adapters=%s max_lora_rank=%s n_slots=%s control_token_gain=%s", self._n_adapters, self._max_lora_rank, self._n_slots, control_token_gain) def _lora_a(self, data: Tensor) -> Tensor: # on-disk A: [n_adapters, 1, max_rank, in] -> [n_adapters+1, max_rank, in] diff --git a/src/llama-model.h b/src/llama-model.h index 871fa96ed50b..339d76274f57 100644 --- a/src/llama-model.h +++ b/src/llama-model.h @@ -224,7 +224,6 @@ struct llama_layer_nextn { // granite-switch: per-token embedded LoRA adapters, stacked across N = num_adapters + 1 // slots in dim 2 (slot 0 is zero-filled so base tokens add an exact-zero delta). // A: {n_embd_in, max_lora_rank, N} B: {max_lora_rank, n_embd_out, N} -// 7 injection sites (q, k, v, o, gate, up, down), each with an A and a B = 14 tensors. struct llama_layer_switch_lora { // attention: q/k/v are separate slices of the fused qkv projection (different out dims) struct ggml_tensor * a_q = nullptr; diff --git a/src/models/granite-switch.cpp b/src/models/granite-switch.cpp index 6d4d774bf6c9..003aea6521fb 100644 --- a/src/models/granite-switch.cpp +++ b/src/models/granite-switch.cpp @@ -21,10 +21,7 @@ void llama_model_granite_switch::load_arch_hparams(llama_model_loader & ml) { ml.get_key(LLM_KV_EXPERT_SHARED_FEED_FORWARD_LENGTH, hparams.n_ff_shexp, /* required */ false); - // per-token LoRA selection runs through MUL_MAT_ID, which needs n_expert_used == 1. - // the GGUF carries expert_count = 0 (dense model), so the generic loader's - // n_expert == 0 => n_expert_used == 0 assertion has already passed by the time - // load_arch_hparams runs; force it to 1 here for the mul_mat_id path. + // mul_mat_id needs n_expert_used == 1; the GGUF is dense (expert_count = 0) hparams.n_expert_used = 1; ml.get_key(LLM_KV_NUM_ADAPTERS, n_adapters); @@ -47,12 +44,9 @@ void llama_model_granite_switch::load_arch_hparams(llama_model_loader & ml) { control_token_to_substitute[token_ids[i]] = substitute_ids[i]; } - // append one extra single-head attention layer at index R = n_real to hold - // the router K/V; bump n_layer_all so the KV cache allocates it. reuse - // n_layer_nextn (the trailing-layers count; no MTP here) = 1 so n_layer() - // stays n_real and the decoder loop is unchanged. note the router layer lives - // at the END (index n_real), not the start, so the regular layers keep their - // indices and the KV cache shift/defrag can skip it without renumbering + // extra single-head attention layer at the END (index n_real) holds the router + // K/V. reusing n_layer_nextn keeps n_layer() == n_real, so the regular layers + // keep their indices and the KV cache shift/defrag skips the router layer. const uint32_t n_real = hparams.n_layer(); hparams.router_layer = (int32_t) n_real; hparams.n_layer_all = n_real + 1; @@ -73,7 +67,6 @@ void llama_model_granite_switch::load_arch_tensors(llama_model_loader &) { tok_embd = create_tensor(tn(LLM_TENSOR_TOKEN_EMBD, "weight"), {n_embd, n_vocab}, 0); - // output 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}, TENSOR_NOT_REQUIRED); if (output == NULL) { @@ -85,7 +78,6 @@ void llama_model_granite_switch::load_arch_tensors(llama_model_loader &) { layer.attn_norm = create_tensor(tn(LLM_TENSOR_ATTN_NORM, "weight", i), {n_embd}, 0); - // fused qkv base + separate o layer.wqkv = create_tensor(tn(LLM_TENSOR_ATTN_QKV, "weight", i), {n_embd, n_embd_q + 2*n_embd_kv}, 0); layer.wo = create_tensor(tn(LLM_TENSOR_ATTN_OUT, "weight", i), {n_embd_q, n_embd}, 0); @@ -95,7 +87,6 @@ void llama_model_granite_switch::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); - // stacked per-token LoRA tensors: A {n_in, max_rank, N}, B {max_rank, n_out, N} auto & sl = layer.switch_lora; sl.a_q = create_tensor(tn(LLM_TENSOR_ATTN_QKV_LORA_A_Q, i), {n_embd, n_rank, n_slots_i64}, 0); @@ -117,8 +108,6 @@ void llama_model_granite_switch::load_arch_tensors(llama_model_loader &) { } } -// graph input for the per-token switch: fills the substituted ids and the -// router Q/K/V signals consumed by the in-graph router attention class llm_graph_input_switch : public llm_graph_input_i { public: llm_graph_input_switch(const llama_model_granite_switch & smodel) : smodel(smodel) {} @@ -134,10 +123,8 @@ class llm_graph_input_switch : public llm_graph_input_i { const llama_model_granite_switch & smodel; }; -// router K dim-0 is +gain for a control token and -gain otherwise; in the causal -// softmax a single visible control token then dominates, so the readback recovers -// its adapter slot. The gain (smodel.router_gain) comes from the GGUF, defaulting to 15.0. - +// K dim-0 is +gain for a control token, -gain otherwise; the causal softmax then +// lets a single visible control token dominate so the readback recovers its slot. void llm_graph_input_switch::set_input(const llama_ubatch * ubatch) { if (!ubatch->token) { return; @@ -153,7 +140,6 @@ void llm_graph_input_switch::set_input(const llama_ubatch * ubatch) { for (int64_t i = 0; i < n_tokens; ++i) { const llama_token tok = ubatch->token[i]; - // router K/V signals: control token -> (+gain, slot); else -> (-gain, 0) const auto it = smodel.control_token_to_index.find(tok); if (it != smodel.control_token_to_index.end()) { ksig[i] = +smodel.router_gain; @@ -199,7 +185,6 @@ ggml_tensor * llama_model_granite_switch::graph::build_switched_lora_delta( return ggml_reshape_2d(ctx0, d, d->ne[0], n_tokens); } -// per-token switched matmul: base 2D weight + selected LoRA delta ggml_tensor * llama_model_granite_switch::graph::build_switched_lora_mm( ggml_tensor * w, ggml_tensor * lora_a, @@ -222,7 +207,6 @@ llama_model_granite_switch::graph::graph( GGML_ASSERT(n_embd_head == hparams.n_embd_head_k()); GGML_ASSERT(n_embd_head == n_rot); - // switch input: substituted ids + per-token router K/V/Q signals auto inp_switch = std::make_unique(smodel); inp_switch->sub_tokens = ggml_new_tensor_1d(ctx0, GGML_TYPE_I32, n_tokens); inp_switch->router_ksig = ggml_new_tensor_1d(ctx0, GGML_TYPE_F32, n_tokens); @@ -238,8 +222,7 @@ llama_model_granite_switch::graph::graph( ggml_tensor * router_q = inp_switch->router_q; res->add_input(std::move(inp_switch)); - // embed the token-exchanged ids (cannot use build_inp_embd, which embeds the - // raw ubatch tokens), then apply Granite embedding scale + // embed the substituted ids directly: build_inp_embd would embed the raw tokens ggml_tensor * inpL = ggml_get_rows(ctx0, model.tok_embd, sub_tokens); if (hparams.f_embedding_scale != 0.0f) { inpL = ggml_scale(ctx0, inpL, hparams.f_embedding_scale); @@ -252,10 +235,8 @@ llama_model_granite_switch::graph::graph( } auto * inp_attn = build_attn_inp_kv(); - // router attention: a single causal head whose K/V live in the KV cache at - // layer R, recovering the per-token adapter index in-graph. Only dim 0 carries - // signal (Q[0]=1, K[0]=+/-gain, V[0]=slot/0), rest zero-padded; no RoPE. State is - // isolated per-sequence because the K/V stay in the per-sequence KV cache. + // single causal head at layer R recovers the adapter index in-graph: only dim 0 + // carries signal (Q[0]=1, K[0]=+/-gain, V[0]=slot/0), the rest is zero-padded. const int R = hparams.router_layer; GGML_ASSERT(R >= 0); auto router_lane = [&](ggml_tensor * sig1d) { @@ -271,8 +252,7 @@ llama_model_granite_switch::graph::graph( Qr, Kr, Vr, nullptr, nullptr, nullptr, /*kq_scale=*/1.0f, /*il=*/R); cb(router_out, "router_out", R); - // readback: row 0 of router_out is the attended slot; adapter_index = - // clamp(round(row0), 0, n_adapters) as I32 for mul_mat_id (round before cast) + // row 0 of router_out is the attended slot; clamp+round to an I32 index ggml_tensor * slot_f = ggml_cont(ctx0, ggml_view_2d(ctx0, router_out, 1, n_tokens, router_out->nb[1], 0)); slot_f = ggml_reshape_1d(ctx0, slot_f, n_tokens); @@ -288,25 +268,21 @@ llama_model_granite_switch::graph::graph( for (int il = 0; il < n_layer; ++il) { ggml_tensor * inpSA = inpL; - // norm cur = build_norm(inpL, model.layers[il].attn_norm, NULL, LLM_NORM_RMS, il); cb(cur, "attn_norm", il); - // self-attention cur = build_attention_layer(cur, inp_pos, adapter_ids, inp_attn, model, n_embd_head, 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); - // reshape adapter_ids to {1, n_tokens} so get_rows selects per-token - // columns, then flatten back to 1D {n_out} + // select the same out rows from adapter_ids (2D round-trip for get_rows) const int64_t n_out = inp_out_ids->ne[0]; adapter_ids = ggml_get_rows(ctx0, ggml_reshape_2d(ctx0, adapter_ids, 1, adapter_ids->ne[0]), inp_out_ids); adapter_ids = ggml_reshape_1d(ctx0, adapter_ids, n_out); } - // ffn cur = build_layer_ffn(cur, inpSA, adapter_ids, model, il); inpL = cur; @@ -318,7 +294,6 @@ llama_model_granite_switch::graph::graph( cb(cur, "result_norm", -1); res->t_embd = cur; - // lm_head cur = build_lora_mm(model.output, cur, model.output_s); // For Granite architectures - scale logits @@ -379,8 +354,7 @@ ggml_tensor * llama_model_granite_switch::graph::build_attention_layer( const float kq_scale = hparams.f_attention_scale == 0.0f ? 1.0f/sqrtf(float(n_embd_head)) : hparams.f_attention_scale; - // wo = nullptr: build_attn returns the concatenated heads, then apply the - // switched o-proj manually + // wo = nullptr so build_attn returns concatenated heads; o-proj is switched below ggml_tensor * attn = build_attn(inp_attn, nullptr, nullptr, nullptr, Qcur, Kcur, Vcur, nullptr, nullptr, nullptr, kq_scale, il); @@ -411,7 +385,6 @@ ggml_tensor * llama_model_granite_switch::graph::build_layer_ffn( cur = build_norm(ffn_inp, layer.ffn_norm, NULL, LLM_NORM_RMS, il); cb(cur, "ffn_norm", il); - // gated SILU FFN with per-token switched LoRA on gate / up / down ggml_tensor * g = build_switched_lora_mm(layer.ffn_gate, sl.a_gate, sl.b_gate, cur, adapter_ids); ggml_tensor * u = build_switched_lora_mm(layer.ffn_up, sl.a_up, sl.b_up, cur, adapter_ids); g = ggml_silu(ctx0, g); From 9b6cb5176428a43409b9b0551eb9656e3cc387d9 Mon Sep 17 00:00:00 2001 From: Bar Haim Date: Tue, 30 Jun 2026 15:33:36 +0300 Subject: [PATCH 22/40] granite-switch: collapse multi-line comments --- conversion/granite.py | 24 +++++++----------------- 1 file changed, 7 insertions(+), 17 deletions(-) diff --git a/conversion/granite.py b/conversion/granite.py index be2a4a3416f1..340546c68a17 100644 --- a/conversion/granite.py +++ b/conversion/granite.py @@ -125,22 +125,16 @@ def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iter @ModelBase.register("GraniteSwitchForCausalLM") class GraniteSwitchModel(GraniteMoeModel): - """Dense, all-attention Granite model with N embedded LoRA adapters selected - per-token by control tokens. Base weights live under ``.base_layer.weight``; - each LoRA is stacked over the adapter dim (B pre-scaled by ``alpha/rank``, so - the runtime scale is 1.0). A zero adapter is materialized at slot 0 so the - stacked dim is ``N = num_adapters + 1`` and base tokens add an exact-zero delta. - """ + """Dense, all-attention Granite with N per-token embedded LoRA adapters, stacked + over the adapter dim with a zero adapter at slot 0 (N = num_adapters + 1).""" model_arch = gguf.MODEL_ARCH.GRANITE_SWITCH - # permute q/k per-slice below (NORM-rope layout) instead of via the parent's - # automatic separate-projection permute + # permute q/k per-slice below (NORM-rope layout), not via the parent's auto-permute undo_permute = False def __init__(self, *args, **kwargs): super().__init__(*args, **kwargs) - # the weightless switch reserves one cache slot, so the decoder has one - # fewer block than num_hidden_layers (block ids map straight through) + # the weightless switch reserves one cache slot: one fewer block than num_hidden_layers self.block_count = self.block_count - 1 self.tensor_map = gguf.get_tensor_name_map(self.model_arch, self.block_count) @@ -164,8 +158,7 @@ def __init__(self, *args, **kwargs): def set_gguf_parameters(self): super().set_gguf_parameters() - # the model is dense; force the dense invariant in case the llama parent - # emitted a nonzero expert_used_count from num_experts_per_tok + # dense model: override any expert counts the llama parent emitted self.gguf_writer.add_expert_count(0) self.gguf_writer.add_expert_used_count(0) @@ -188,7 +181,6 @@ def _lora_b(self, data: Tensor, permute_n_head: int | None = None) -> Tensor: b = data.squeeze(1) if permute_n_head is not None: # permute each adapter's B output rows to match the permuted q/k base - # (LlamaModel.permute interleaves q/k for ggml's NORM-rope layout) b = torch.stack([self.permute(b[i], permute_n_head, permute_n_head) for i in range(b.shape[0])], dim=0) zero = torch.zeros_like(b[:1]) return torch.cat([zero, b], dim=0).contiguous() @@ -196,8 +188,7 @@ def _lora_b(self, data: Tensor, permute_n_head: int | None = None) -> Tensor: def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iterable[tuple[str, Tensor]]: T = gguf.MODEL_TENSOR - # skip the weightless switch + control-token buffers (reconstructed at - # load time from the adapter id arrays emitted in set_gguf_parameters) + # skip the weightless switch + control-token buffers (rebuilt at load time) bare = name.split(".")[-1] if ( name.startswith("model.switch.") or name.startswith("switch.") @@ -208,8 +199,7 @@ def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iter # ---- Attention: fused QKV ---- if "self_attn.qkv_proj" in name: if name.endswith("base_layer.weight"): - # Fused [q|k|v] rows. Permute the q and k row-blocks for ggml's - # NORM-rope layout (LlamaModel.permute, n_head_kv == n_head per slice). + # fused [q|k|v] rows: permute q/k row-blocks for ggml's NORM-rope layout q, k, v = data_torch.split([self._q_size, self._kv_size, self._kv_size], dim=0) q = self.permute(q, self._n_head, self._n_head) k = self.permute(k, self._n_kv_head, self._n_kv_head) From 4a54c03106154abc7a8534038ca005ff539a4427 Mon Sep 17 00:00:00 2001 From: Bar Haim Date: Tue, 30 Jun 2026 15:42:43 +0300 Subject: [PATCH 23/40] granite-switch: rename control_token_* maps to adapter_token_* --- src/models/granite-switch.cpp | 24 ++++++++++++------------ src/models/models.h | 8 ++++---- 2 files changed, 16 insertions(+), 16 deletions(-) diff --git a/src/models/granite-switch.cpp b/src/models/granite-switch.cpp index 003aea6521fb..3a56ee085272 100644 --- a/src/models/granite-switch.cpp +++ b/src/models/granite-switch.cpp @@ -36,12 +36,12 @@ void llama_model_granite_switch::load_arch_hparams(llama_model_loader & ml) { GGML_ASSERT(token_ids.size() == n_adapters); GGML_ASSERT(substitute_ids.size() == n_adapters); - control_token_to_index.clear(); - control_token_to_substitute.clear(); + adapter_token_to_slot.clear(); + adapter_token_to_substitute.clear(); for (uint32_t i = 0; i < n_adapters; ++i) { // adapter i -> stacked slot i+1 (slot 0 is the base/zero delta) - control_token_to_index[token_ids[i]] = (int32_t) (i + 1); - control_token_to_substitute[token_ids[i]] = substitute_ids[i]; + adapter_token_to_slot[token_ids[i]] = (int32_t) (i + 1); + adapter_token_to_substitute[token_ids[i]] = substitute_ids[i]; } // extra single-head attention layer at the END (index n_real) holds the router @@ -115,7 +115,7 @@ class llm_graph_input_switch : public llm_graph_input_i { void set_input(const llama_ubatch * ubatch) override; - ggml_tensor * sub_tokens = nullptr; // I32 [n_tokens] control-substituted token ids + ggml_tensor * sub_tokens = nullptr; // I32 [n_tokens] adapter-substituted token ids ggml_tensor * router_ksig = nullptr; // F32 [n_tokens] router K signal (±gain) ggml_tensor * router_vval = nullptr; // F32 [n_tokens] router V value (adapter slot / 0) ggml_tensor * router_q = nullptr; // F32 [n_tokens] router Q value (constant 1.0) @@ -123,8 +123,8 @@ class llm_graph_input_switch : public llm_graph_input_i { const llama_model_granite_switch & smodel; }; -// K dim-0 is +gain for a control token, -gain otherwise; the causal softmax then -// lets a single visible control token dominate so the readback recovers its slot. +// K dim-0 is +gain for an adapter token, -gain otherwise; the causal softmax then +// lets a single visible adapter token dominate so the readback recovers its slot. void llm_graph_input_switch::set_input(const llama_ubatch * ubatch) { if (!ubatch->token) { return; @@ -140,8 +140,8 @@ void llm_graph_input_switch::set_input(const llama_ubatch * ubatch) { for (int64_t i = 0; i < n_tokens; ++i) { const llama_token tok = ubatch->token[i]; - const auto it = smodel.control_token_to_index.find(tok); - if (it != smodel.control_token_to_index.end()) { + const auto it = smodel.adapter_token_to_slot.find(tok); + if (it != smodel.adapter_token_to_slot.end()) { ksig[i] = +smodel.router_gain; vval[i] = (float) it->second; } else { @@ -149,9 +149,9 @@ void llm_graph_input_switch::set_input(const llama_ubatch * ubatch) { vval[i] = 0.0f; } - // rewrite a control token to its substitute id before embedding - const auto sit = smodel.control_token_to_substitute.find(tok); - sub[i] = (sit != smodel.control_token_to_substitute.end()) + // rewrite an adapter token to its substitute id before embedding + const auto sit = smodel.adapter_token_to_substitute.find(tok); + sub[i] = (sit != smodel.adapter_token_to_substitute.end()) ? (int32_t) sit->second : (int32_t) tok; } diff --git a/src/models/models.h b/src/models/models.h index 904ce5780ee3..90f34a85c836 100644 --- a/src/models/models.h +++ b/src/models/models.h @@ -1544,15 +1544,15 @@ struct llama_model_granite_switch : public llama_model_base { // per-token switch metadata (read from GGUF in load_arch_hparams) uint32_t n_adapters = 0; // number of real adapters (excludes the zero slot) uint32_t max_lora_rank = 0; - float router_gain = 15.0f; // control-token softmax bias magnitude (see set_input) + float router_gain = 15.0f; // routing-token softmax bias magnitude (see set_input) // stacked-adapter slots: n_adapters + 1 (slot 0 = base/zero delta) uint32_t n_slots() const { return n_adapters + 1; } // token id -> adapter slot (1 + adapter, so slot 0 stays "base") - std::unordered_map control_token_to_index; - // control token id -> substitute token id (token-exchange before embedding) - std::unordered_map control_token_to_substitute; + std::unordered_map adapter_token_to_slot; + // token id -> substitute token id (token-exchange before embedding) + std::unordered_map adapter_token_to_substitute; // the per-token adapter selection is recovered in-graph by a single-head // causal router attention whose K/V live in the KV cache at hparams.router_layer From 4675356dd1139930d6b23620195d958b58967f60 Mon Sep 17 00:00:00 2001 From: Bar Haim Date: Tue, 30 Jun 2026 18:17:31 +0300 Subject: [PATCH 24/40] granite-switch: cut noise comments --- conversion/granite.py | 6 ------ src/llama-model.h | 2 -- src/models/granite-switch.cpp | 4 ---- src/models/models.h | 1 - 4 files changed, 13 deletions(-) diff --git a/conversion/granite.py b/conversion/granite.py index 340546c68a17..885dc6a81dfc 100644 --- a/conversion/granite.py +++ b/conversion/granite.py @@ -196,7 +196,6 @@ def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iter ): return - # ---- Attention: fused QKV ---- if "self_attn.qkv_proj" in name: if name.endswith("base_layer.weight"): # fused [q|k|v] rows: permute q/k row-blocks for ggml's NORM-rope layout @@ -222,7 +221,6 @@ def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iter return raise ValueError(f"Unexpected qkv_proj tensor: {name}") - # ---- Attention: output projection ---- if "self_attn.o_proj" in name: if name.endswith("base_layer.weight"): yield (self.format_tensor_name(T.ATTN_OUT, bid), data_torch) @@ -235,7 +233,6 @@ def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iter return raise ValueError(f"Unexpected o_proj tensor: {name}") - # ---- MLP: fused gate/up (shared_mlp.input_linear) ---- if "shared_mlp.input_linear" in name: ffn = self.hparams["shared_intermediate_size"] if name.endswith("base_layer.weight"): @@ -255,7 +252,6 @@ def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iter return raise ValueError(f"Unexpected shared_mlp.input_linear tensor: {name}") - # ---- MLP: down (shared_mlp.output_linear) ---- if "shared_mlp.output_linear" in name: if name.endswith("base_layer.weight"): yield (self.format_tensor_name(T.FFN_DOWN, bid), data_torch) @@ -268,7 +264,6 @@ def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iter return raise ValueError(f"Unexpected shared_mlp.output_linear tensor: {name}") - # ---- Per-layer norms (input_layernorm / post_attention_layernorm) ---- if bid is not None and ".layers." in name and ( "input_layernorm" in name or "post_attention_layernorm" in name ): @@ -276,7 +271,6 @@ def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iter yield (self.format_tensor_name(key, bid), data_torch) return - # ---- Embeddings / final norm (lm_head is tied to token_embd) ---- if name in ("model.embed_tokens.weight", "embed_tokens.weight"): yield (self.format_tensor_name(T.TOKEN_EMBD), data_torch) return diff --git a/src/llama-model.h b/src/llama-model.h index 339d76274f57..e6f807ecad76 100644 --- a/src/llama-model.h +++ b/src/llama-model.h @@ -225,7 +225,6 @@ struct llama_layer_nextn { // slots in dim 2 (slot 0 is zero-filled so base tokens add an exact-zero delta). // A: {n_embd_in, max_lora_rank, N} B: {max_lora_rank, n_embd_out, N} struct llama_layer_switch_lora { - // attention: q/k/v are separate slices of the fused qkv projection (different out dims) struct ggml_tensor * a_q = nullptr; struct ggml_tensor * b_q = nullptr; struct ggml_tensor * a_k = nullptr; @@ -235,7 +234,6 @@ struct llama_layer_switch_lora { struct ggml_tensor * a_o = nullptr; struct ggml_tensor * b_o = nullptr; - // feed-forward struct ggml_tensor * a_gate = nullptr; struct ggml_tensor * b_gate = nullptr; struct ggml_tensor * a_up = nullptr; diff --git a/src/models/granite-switch.cpp b/src/models/granite-switch.cpp index 3a56ee085272..e25e624e2c78 100644 --- a/src/models/granite-switch.cpp +++ b/src/models/granite-switch.cpp @@ -9,7 +9,6 @@ void llama_model_granite_switch::load_arch_hparams(llama_model_loader & ml) { ml.get_key(LLM_KV_EMBEDDING_SCALE, hparams.f_embedding_scale, false); ml.get_key(LLM_KV_ATTENTION_SCALE, hparams.f_attention_scale, false); - // Granite uses rope_finetuned as a switch for rope, so default to true bool rope_finetuned = true; ml.get_key(LLM_KV_ROPE_SCALING_FINETUNED, rope_finetuned, false); hparams.rope_finetuned = rope_finetuned; @@ -149,7 +148,6 @@ void llm_graph_input_switch::set_input(const llama_ubatch * ubatch) { vval[i] = 0.0f; } - // rewrite an adapter token to its substitute id before embedding const auto sit = smodel.adapter_token_to_substitute.find(tok); sub[i] = (sit != smodel.adapter_token_to_substitute.end()) ? (int32_t) sit->second @@ -296,7 +294,6 @@ llama_model_granite_switch::graph::graph( cur = build_lora_mm(model.output, cur, model.output_s); - // For Granite architectures - scale logits cur = ggml_scale(ctx0, cur, 1.0f / hparams.f_logit_scale); cb(cur, "result_output", -1); res->t_logits = cur; @@ -375,7 +372,6 @@ ggml_tensor * llama_model_granite_switch::graph::build_layer_ffn( const auto & layer = model.layers[il]; const auto & sl = layer.switch_lora; - // For Granite architectures - scale residual if (hparams.f_residual_scale) { cur = ggml_scale(ctx0, cur, hparams.f_residual_scale); } diff --git a/src/models/models.h b/src/models/models.h index 90f34a85c836..44f62342e8aa 100644 --- a/src/models/models.h +++ b/src/models/models.h @@ -1541,7 +1541,6 @@ struct llama_model_granite_switch : public llama_model_base { void load_arch_hparams(llama_model_loader & ml) override; void load_arch_tensors(llama_model_loader & ml) override; - // per-token switch metadata (read from GGUF in load_arch_hparams) uint32_t n_adapters = 0; // number of real adapters (excludes the zero slot) uint32_t max_lora_rank = 0; float router_gain = 15.0f; // routing-token softmax bias magnitude (see set_input) From ea39782acd15c437a2a93883b35a0b8419be130f Mon Sep 17 00:00:00 2001 From: Bar Haim Date: Tue, 30 Jun 2026 19:00:04 +0300 Subject: [PATCH 25/40] granite-switch: rename embedded LoRA tensors to .lora_a/lora_b --- conversion/granite.py | 8 ++++---- gguf-py/gguf/constants.py | 36 +++++++++++++++++------------------ src/llama-arch.cpp | 24 +++++++++++------------ src/llama-arch.h | 12 ++++++------ src/models/granite-switch.cpp | 12 ++++++------ 5 files changed, 46 insertions(+), 46 deletions(-) diff --git a/conversion/granite.py b/conversion/granite.py index 885dc6a81dfc..75c147a4d99e 100644 --- a/conversion/granite.py +++ b/conversion/granite.py @@ -207,15 +207,15 @@ def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iter return if "lora_A_slices." in name: slot = int(name.rsplit(".", 1)[1]) - key = {0: T.ATTN_QKV_LORA_A_Q, 1: T.ATTN_QKV_LORA_A_K, 2: T.ATTN_QKV_LORA_A_V}[slot] + key = {0: T.ATTN_Q_LORA_A, 1: T.ATTN_K_LORA_A, 2: T.ATTN_V_LORA_A}[slot] yield (self.format_tensor_name(key, bid, suffix=""), self._lora_a(data_torch)) return if "lora_B_slices." in name: slot = int(name.rsplit(".", 1)[1]) key, ph = { - 0: (T.ATTN_QKV_LORA_B_Q, self._n_head), - 1: (T.ATTN_QKV_LORA_B_K, self._n_kv_head), - 2: (T.ATTN_QKV_LORA_B_V, None), + 0: (T.ATTN_Q_LORA_B, self._n_head), + 1: (T.ATTN_K_LORA_B, self._n_kv_head), + 2: (T.ATTN_V_LORA_B, None), }[slot] yield (self.format_tensor_name(key, bid, suffix=""), self._lora_b(data_torch, ph)) return diff --git a/gguf-py/gguf/constants.py b/gguf-py/gguf/constants.py index 7b6457189c25..03dd4f05e91b 100644 --- a/gguf-py/gguf/constants.py +++ b/gguf-py/gguf/constants.py @@ -990,12 +990,12 @@ class MODEL_TENSOR(IntEnum): A_QF_FFN_DOWN = auto() A_QF_FFN_NORM = auto() # per-token embedded LoRA adapters - Granite Switch - ATTN_QKV_LORA_A_Q = auto() - ATTN_QKV_LORA_B_Q = auto() - ATTN_QKV_LORA_A_K = auto() - ATTN_QKV_LORA_B_K = auto() - ATTN_QKV_LORA_A_V = auto() - ATTN_QKV_LORA_B_V = auto() + ATTN_Q_LORA_A = auto() + ATTN_Q_LORA_B = auto() + ATTN_K_LORA_A = auto() + ATTN_K_LORA_B = auto() + ATTN_V_LORA_A = auto() + ATTN_V_LORA_B = auto() ATTN_OUT_LORA_A = auto() ATTN_OUT_LORA_B = auto() FFN_GATE_LORA_A = auto() @@ -1583,12 +1583,12 @@ class MODEL_TENSOR(IntEnum): MODEL_TENSOR.FC: "fc", MODEL_TENSOR.D2T: "d2t", # per-token embedded LoRA adapters - Granite Switch - MODEL_TENSOR.ATTN_QKV_LORA_A_Q: "blk.{bid}.attn_qkv.lora_a_q", - MODEL_TENSOR.ATTN_QKV_LORA_B_Q: "blk.{bid}.attn_qkv.lora_b_q", - MODEL_TENSOR.ATTN_QKV_LORA_A_K: "blk.{bid}.attn_qkv.lora_a_k", - MODEL_TENSOR.ATTN_QKV_LORA_B_K: "blk.{bid}.attn_qkv.lora_b_k", - MODEL_TENSOR.ATTN_QKV_LORA_A_V: "blk.{bid}.attn_qkv.lora_a_v", - MODEL_TENSOR.ATTN_QKV_LORA_B_V: "blk.{bid}.attn_qkv.lora_b_v", + MODEL_TENSOR.ATTN_Q_LORA_A: "blk.{bid}.attn_q.lora_a", + MODEL_TENSOR.ATTN_Q_LORA_B: "blk.{bid}.attn_q.lora_b", + MODEL_TENSOR.ATTN_K_LORA_A: "blk.{bid}.attn_k.lora_a", + MODEL_TENSOR.ATTN_K_LORA_B: "blk.{bid}.attn_k.lora_b", + MODEL_TENSOR.ATTN_V_LORA_A: "blk.{bid}.attn_v.lora_a", + MODEL_TENSOR.ATTN_V_LORA_B: "blk.{bid}.attn_v.lora_b", MODEL_TENSOR.ATTN_OUT_LORA_A: "blk.{bid}.attn_output.lora_a", MODEL_TENSOR.ATTN_OUT_LORA_B: "blk.{bid}.attn_output.lora_b", MODEL_TENSOR.FFN_GATE_LORA_A: "blk.{bid}.ffn_gate.lora_a", @@ -3717,12 +3717,12 @@ class MODEL_TENSOR(IntEnum): MODEL_TENSOR.FFN_GATE, MODEL_TENSOR.FFN_DOWN, MODEL_TENSOR.FFN_UP, - MODEL_TENSOR.ATTN_QKV_LORA_A_Q, - MODEL_TENSOR.ATTN_QKV_LORA_B_Q, - MODEL_TENSOR.ATTN_QKV_LORA_A_K, - MODEL_TENSOR.ATTN_QKV_LORA_B_K, - MODEL_TENSOR.ATTN_QKV_LORA_A_V, - MODEL_TENSOR.ATTN_QKV_LORA_B_V, + MODEL_TENSOR.ATTN_Q_LORA_A, + MODEL_TENSOR.ATTN_Q_LORA_B, + MODEL_TENSOR.ATTN_K_LORA_A, + MODEL_TENSOR.ATTN_K_LORA_B, + MODEL_TENSOR.ATTN_V_LORA_A, + MODEL_TENSOR.ATTN_V_LORA_B, MODEL_TENSOR.ATTN_OUT_LORA_A, MODEL_TENSOR.ATTN_OUT_LORA_B, MODEL_TENSOR.FFN_GATE_LORA_A, diff --git a/src/llama-arch.cpp b/src/llama-arch.cpp index b929f362b68c..218136492f75 100644 --- a/src/llama-arch.cpp +++ b/src/llama-arch.cpp @@ -610,12 +610,12 @@ static const std::map LLM_TENSOR_NAMES = { { LLM_TENSOR_FC, "fc" }, { LLM_TENSOR_D2T, "d2t" }, // Granite Switch: per-token embedded LoRA adapters (no ".weight" suffix on disk) - { LLM_TENSOR_ATTN_QKV_LORA_A_Q, "blk.%d.attn_qkv.lora_a_q" }, - { LLM_TENSOR_ATTN_QKV_LORA_B_Q, "blk.%d.attn_qkv.lora_b_q" }, - { LLM_TENSOR_ATTN_QKV_LORA_A_K, "blk.%d.attn_qkv.lora_a_k" }, - { LLM_TENSOR_ATTN_QKV_LORA_B_K, "blk.%d.attn_qkv.lora_b_k" }, - { LLM_TENSOR_ATTN_QKV_LORA_A_V, "blk.%d.attn_qkv.lora_a_v" }, - { LLM_TENSOR_ATTN_QKV_LORA_B_V, "blk.%d.attn_qkv.lora_b_v" }, + { LLM_TENSOR_ATTN_Q_LORA_A, "blk.%d.attn_q.lora_a" }, + { LLM_TENSOR_ATTN_Q_LORA_B, "blk.%d.attn_q.lora_b" }, + { LLM_TENSOR_ATTN_K_LORA_A, "blk.%d.attn_k.lora_a" }, + { LLM_TENSOR_ATTN_K_LORA_B, "blk.%d.attn_k.lora_b" }, + { LLM_TENSOR_ATTN_V_LORA_A, "blk.%d.attn_v.lora_a" }, + { LLM_TENSOR_ATTN_V_LORA_B, "blk.%d.attn_v.lora_b" }, { LLM_TENSOR_ATTN_OUT_LORA_A, "blk.%d.attn_output.lora_a" }, { LLM_TENSOR_ATTN_OUT_LORA_B, "blk.%d.attn_output.lora_b" }, { LLM_TENSOR_FFN_GATE_LORA_A, "blk.%d.ffn_gate.lora_a" }, @@ -876,12 +876,12 @@ static const std::map LLM_TENSOR_INFOS = { {LLM_TENSOR_FC, {LLM_TENSOR_LAYER_OUTPUT, GGML_OP_MUL_MAT}}, {LLM_TENSOR_D2T, {LLM_TENSOR_LAYER_OUTPUT, GGML_OP_GET_ROWS}}, // granite-switch: per-token embedded LoRA adapters (selected via mul_mat_id) - {LLM_TENSOR_ATTN_QKV_LORA_A_Q, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL_MAT_ID}}, - {LLM_TENSOR_ATTN_QKV_LORA_B_Q, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL_MAT_ID}}, - {LLM_TENSOR_ATTN_QKV_LORA_A_K, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL_MAT_ID}}, - {LLM_TENSOR_ATTN_QKV_LORA_B_K, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL_MAT_ID}}, - {LLM_TENSOR_ATTN_QKV_LORA_A_V, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL_MAT_ID}}, - {LLM_TENSOR_ATTN_QKV_LORA_B_V, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL_MAT_ID}}, + {LLM_TENSOR_ATTN_Q_LORA_A, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL_MAT_ID}}, + {LLM_TENSOR_ATTN_Q_LORA_B, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL_MAT_ID}}, + {LLM_TENSOR_ATTN_K_LORA_A, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL_MAT_ID}}, + {LLM_TENSOR_ATTN_K_LORA_B, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL_MAT_ID}}, + {LLM_TENSOR_ATTN_V_LORA_A, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL_MAT_ID}}, + {LLM_TENSOR_ATTN_V_LORA_B, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL_MAT_ID}}, {LLM_TENSOR_ATTN_OUT_LORA_A, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL_MAT_ID}}, {LLM_TENSOR_ATTN_OUT_LORA_B, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL_MAT_ID}}, {LLM_TENSOR_FFN_GATE_LORA_A, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL_MAT_ID}}, diff --git a/src/llama-arch.h b/src/llama-arch.h index 4b9acd5530a7..bef4f7408424 100644 --- a/src/llama-arch.h +++ b/src/llama-arch.h @@ -618,12 +618,12 @@ enum llm_tensor { LLM_TENSOR_FC, LLM_TENSOR_D2T, // Granite Switch: per-token embedded LoRA adapters (stacked over N = num_adapters + 1) - LLM_TENSOR_ATTN_QKV_LORA_A_Q, - LLM_TENSOR_ATTN_QKV_LORA_B_Q, - LLM_TENSOR_ATTN_QKV_LORA_A_K, - LLM_TENSOR_ATTN_QKV_LORA_B_K, - LLM_TENSOR_ATTN_QKV_LORA_A_V, - LLM_TENSOR_ATTN_QKV_LORA_B_V, + LLM_TENSOR_ATTN_Q_LORA_A, + LLM_TENSOR_ATTN_Q_LORA_B, + LLM_TENSOR_ATTN_K_LORA_A, + LLM_TENSOR_ATTN_K_LORA_B, + LLM_TENSOR_ATTN_V_LORA_A, + LLM_TENSOR_ATTN_V_LORA_B, LLM_TENSOR_ATTN_OUT_LORA_A, LLM_TENSOR_ATTN_OUT_LORA_B, LLM_TENSOR_FFN_GATE_LORA_A, diff --git a/src/models/granite-switch.cpp b/src/models/granite-switch.cpp index e25e624e2c78..2950544aafac 100644 --- a/src/models/granite-switch.cpp +++ b/src/models/granite-switch.cpp @@ -88,12 +88,12 @@ void llama_model_granite_switch::load_arch_tensors(llama_model_loader &) { auto & sl = layer.switch_lora; - sl.a_q = create_tensor(tn(LLM_TENSOR_ATTN_QKV_LORA_A_Q, i), {n_embd, n_rank, n_slots_i64}, 0); - sl.b_q = create_tensor(tn(LLM_TENSOR_ATTN_QKV_LORA_B_Q, i), {n_rank, n_embd_q, n_slots_i64}, 0); - sl.a_k = create_tensor(tn(LLM_TENSOR_ATTN_QKV_LORA_A_K, i), {n_embd, n_rank, n_slots_i64}, 0); - sl.b_k = create_tensor(tn(LLM_TENSOR_ATTN_QKV_LORA_B_K, i), {n_rank, n_embd_kv, n_slots_i64}, 0); - sl.a_v = create_tensor(tn(LLM_TENSOR_ATTN_QKV_LORA_A_V, i), {n_embd, n_rank, n_slots_i64}, 0); - sl.b_v = create_tensor(tn(LLM_TENSOR_ATTN_QKV_LORA_B_V, i), {n_rank, n_embd_kv, n_slots_i64}, 0); + sl.a_q = create_tensor(tn(LLM_TENSOR_ATTN_Q_LORA_A, i), {n_embd, n_rank, n_slots_i64}, 0); + sl.b_q = create_tensor(tn(LLM_TENSOR_ATTN_Q_LORA_B, i), {n_rank, n_embd_q, n_slots_i64}, 0); + sl.a_k = create_tensor(tn(LLM_TENSOR_ATTN_K_LORA_A, i), {n_embd, n_rank, n_slots_i64}, 0); + sl.b_k = create_tensor(tn(LLM_TENSOR_ATTN_K_LORA_B, i), {n_rank, n_embd_kv, n_slots_i64}, 0); + sl.a_v = create_tensor(tn(LLM_TENSOR_ATTN_V_LORA_A, i), {n_embd, n_rank, n_slots_i64}, 0); + sl.b_v = create_tensor(tn(LLM_TENSOR_ATTN_V_LORA_B, i), {n_rank, n_embd_kv, n_slots_i64}, 0); sl.a_o = create_tensor(tn(LLM_TENSOR_ATTN_OUT_LORA_A, i), {n_embd_q, n_rank, n_slots_i64}, 0); sl.b_o = create_tensor(tn(LLM_TENSOR_ATTN_OUT_LORA_B, i), {n_rank, n_embd, n_slots_i64}, 0); From 600ef1bbbc767732f65d2a30ce8daf3de426e1bc Mon Sep 17 00:00:00 2001 From: Bar Haim Date: Wed, 1 Jul 2026 17:16:27 +0300 Subject: [PATCH 26/40] granite-switch: GGML_ASSERT token input to avoid UB on embeddings --- src/models/granite-switch.cpp | 2 ++ 1 file changed, 2 insertions(+) diff --git a/src/models/granite-switch.cpp b/src/models/granite-switch.cpp index 2950544aafac..7f89a1c98593 100644 --- a/src/models/granite-switch.cpp +++ b/src/models/granite-switch.cpp @@ -201,6 +201,8 @@ llama_model_granite_switch::graph::graph( const auto & smodel = static_cast(model); + GGML_ASSERT(ubatch.token && "granite-switch requires token input"); + const int64_t n_embd_head = hparams.n_embd_head_v(); GGML_ASSERT(n_embd_head == hparams.n_embd_head_k()); GGML_ASSERT(n_embd_head == n_rot); From 83798261acfad56332c1343dd6dc0f690240b297 Mon Sep 17 00:00:00 2001 From: Bar Haim Date: Wed, 1 Jul 2026 17:20:41 +0300 Subject: [PATCH 27/40] granite-switch: TODO for raw embedding input support --- src/models/granite-switch.cpp | 1 + 1 file changed, 1 insertion(+) diff --git a/src/models/granite-switch.cpp b/src/models/granite-switch.cpp index 7f89a1c98593..4fba3dc1e57a 100644 --- a/src/models/granite-switch.cpp +++ b/src/models/granite-switch.cpp @@ -201,6 +201,7 @@ llama_model_granite_switch::graph::graph( const auto & smodel = static_cast(model); + // TODO: support raw embedding input (multimodal / pre-embedded tokens) when needed GGML_ASSERT(ubatch.token && "granite-switch requires token input"); const int64_t n_embd_head = hparams.n_embd_head_v(); From 010f22d8a3a9a0b545af17177e956859baadc66a Mon Sep 17 00:00:00 2001 From: Bar Haim Date: Wed, 1 Jul 2026 21:06:29 +0300 Subject: [PATCH 28/40] granite-switch: collapse LoRA tensor constants to .lora_a/.lora_b suffix --- conversion/granite.py | 28 +++++++++++------------ gguf-py/gguf/constants.py | 47 +++------------------------------------ 2 files changed, 17 insertions(+), 58 deletions(-) diff --git a/conversion/granite.py b/conversion/granite.py index 75c147a4d99e..914524b47588 100644 --- a/conversion/granite.py +++ b/conversion/granite.py @@ -207,17 +207,17 @@ def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iter return if "lora_A_slices." in name: slot = int(name.rsplit(".", 1)[1]) - key = {0: T.ATTN_Q_LORA_A, 1: T.ATTN_K_LORA_A, 2: T.ATTN_V_LORA_A}[slot] - yield (self.format_tensor_name(key, bid, suffix=""), self._lora_a(data_torch)) + key = {0: T.ATTN_Q, 1: T.ATTN_K, 2: T.ATTN_V}[slot] + yield (self.format_tensor_name(key, bid, suffix=".lora_a"), self._lora_a(data_torch)) return if "lora_B_slices." in name: slot = int(name.rsplit(".", 1)[1]) key, ph = { - 0: (T.ATTN_Q_LORA_B, self._n_head), - 1: (T.ATTN_K_LORA_B, self._n_kv_head), - 2: (T.ATTN_V_LORA_B, None), + 0: (T.ATTN_Q, self._n_head), + 1: (T.ATTN_K, self._n_kv_head), + 2: (T.ATTN_V, None), }[slot] - yield (self.format_tensor_name(key, bid, suffix=""), self._lora_b(data_torch, ph)) + yield (self.format_tensor_name(key, bid, suffix=".lora_b"), self._lora_b(data_torch, ph)) return raise ValueError(f"Unexpected qkv_proj tensor: {name}") @@ -226,10 +226,10 @@ def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iter yield (self.format_tensor_name(T.ATTN_OUT, bid), data_torch) return if name.endswith("lora_A"): - yield (self.format_tensor_name(T.ATTN_OUT_LORA_A, bid, suffix=""), self._lora_a(data_torch)) + yield (self.format_tensor_name(T.ATTN_OUT, bid, suffix=".lora_a"), self._lora_a(data_torch)) return if name.endswith("lora_B"): - yield (self.format_tensor_name(T.ATTN_OUT_LORA_B, bid, suffix=""), self._lora_b(data_torch)) + yield (self.format_tensor_name(T.ATTN_OUT, bid, suffix=".lora_b"), self._lora_b(data_torch)) return raise ValueError(f"Unexpected o_proj tensor: {name}") @@ -242,13 +242,13 @@ def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iter return if "lora_A_slices." in name: slot = int(name.rsplit(".", 1)[1]) - key = {0: T.FFN_GATE_LORA_A, 1: T.FFN_UP_LORA_A}[slot] - yield (self.format_tensor_name(key, bid, suffix=""), self._lora_a(data_torch)) + key = {0: T.FFN_GATE, 1: T.FFN_UP}[slot] + yield (self.format_tensor_name(key, bid, suffix=".lora_a"), self._lora_a(data_torch)) return if "lora_B_slices." in name: slot = int(name.rsplit(".", 1)[1]) - key = {0: T.FFN_GATE_LORA_B, 1: T.FFN_UP_LORA_B}[slot] - yield (self.format_tensor_name(key, bid, suffix=""), self._lora_b(data_torch)) + key = {0: T.FFN_GATE, 1: T.FFN_UP}[slot] + yield (self.format_tensor_name(key, bid, suffix=".lora_b"), self._lora_b(data_torch)) return raise ValueError(f"Unexpected shared_mlp.input_linear tensor: {name}") @@ -257,10 +257,10 @@ def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iter yield (self.format_tensor_name(T.FFN_DOWN, bid), data_torch) return if name.endswith("lora_A"): - yield (self.format_tensor_name(T.FFN_DOWN_LORA_A, bid, suffix=""), self._lora_a(data_torch)) + yield (self.format_tensor_name(T.FFN_DOWN, bid, suffix=".lora_a"), self._lora_a(data_torch)) return if name.endswith("lora_B"): - yield (self.format_tensor_name(T.FFN_DOWN_LORA_B, bid, suffix=""), self._lora_b(data_torch)) + yield (self.format_tensor_name(T.FFN_DOWN, bid, suffix=".lora_b"), self._lora_b(data_torch)) return raise ValueError(f"Unexpected shared_mlp.output_linear tensor: {name}") diff --git a/gguf-py/gguf/constants.py b/gguf-py/gguf/constants.py index 03dd4f05e91b..dd1bdf0efe4f 100644 --- a/gguf-py/gguf/constants.py +++ b/gguf-py/gguf/constants.py @@ -989,21 +989,6 @@ class MODEL_TENSOR(IntEnum): A_QF_FFN_UP = auto() A_QF_FFN_DOWN = auto() A_QF_FFN_NORM = auto() - # per-token embedded LoRA adapters - Granite Switch - ATTN_Q_LORA_A = auto() - ATTN_Q_LORA_B = auto() - ATTN_K_LORA_A = auto() - ATTN_K_LORA_B = auto() - ATTN_V_LORA_A = auto() - ATTN_V_LORA_B = auto() - ATTN_OUT_LORA_A = auto() - ATTN_OUT_LORA_B = auto() - FFN_GATE_LORA_A = auto() - FFN_GATE_LORA_B = auto() - FFN_UP_LORA_A = auto() - FFN_UP_LORA_B = auto() - FFN_DOWN_LORA_A = auto() - FFN_DOWN_LORA_B = auto() MODEL_ARCH_NAMES: dict[MODEL_ARCH, str] = { @@ -1582,21 +1567,6 @@ class MODEL_TENSOR(IntEnum): MODEL_TENSOR.NEXTN_SHARED_HEAD_NORM: "blk.{bid}.nextn.shared_head_norm", MODEL_TENSOR.FC: "fc", MODEL_TENSOR.D2T: "d2t", - # per-token embedded LoRA adapters - Granite Switch - MODEL_TENSOR.ATTN_Q_LORA_A: "blk.{bid}.attn_q.lora_a", - MODEL_TENSOR.ATTN_Q_LORA_B: "blk.{bid}.attn_q.lora_b", - MODEL_TENSOR.ATTN_K_LORA_A: "blk.{bid}.attn_k.lora_a", - MODEL_TENSOR.ATTN_K_LORA_B: "blk.{bid}.attn_k.lora_b", - MODEL_TENSOR.ATTN_V_LORA_A: "blk.{bid}.attn_v.lora_a", - MODEL_TENSOR.ATTN_V_LORA_B: "blk.{bid}.attn_v.lora_b", - MODEL_TENSOR.ATTN_OUT_LORA_A: "blk.{bid}.attn_output.lora_a", - MODEL_TENSOR.ATTN_OUT_LORA_B: "blk.{bid}.attn_output.lora_b", - MODEL_TENSOR.FFN_GATE_LORA_A: "blk.{bid}.ffn_gate.lora_a", - MODEL_TENSOR.FFN_GATE_LORA_B: "blk.{bid}.ffn_gate.lora_b", - MODEL_TENSOR.FFN_UP_LORA_A: "blk.{bid}.ffn_up.lora_a", - MODEL_TENSOR.FFN_UP_LORA_B: "blk.{bid}.ffn_up.lora_b", - MODEL_TENSOR.FFN_DOWN_LORA_A: "blk.{bid}.ffn_down.lora_a", - MODEL_TENSOR.FFN_DOWN_LORA_B: "blk.{bid}.ffn_down.lora_b", } MODEL_TENSORS: dict[MODEL_ARCH, list[MODEL_TENSOR]] = { @@ -3712,25 +3682,14 @@ class MODEL_TENSOR(IntEnum): MODEL_TENSOR.OUTPUT, MODEL_TENSOR.ATTN_NORM, MODEL_TENSOR.ATTN_QKV, + MODEL_TENSOR.ATTN_Q, + MODEL_TENSOR.ATTN_K, + MODEL_TENSOR.ATTN_V, MODEL_TENSOR.ATTN_OUT, MODEL_TENSOR.FFN_NORM, MODEL_TENSOR.FFN_GATE, MODEL_TENSOR.FFN_DOWN, MODEL_TENSOR.FFN_UP, - MODEL_TENSOR.ATTN_Q_LORA_A, - MODEL_TENSOR.ATTN_Q_LORA_B, - MODEL_TENSOR.ATTN_K_LORA_A, - MODEL_TENSOR.ATTN_K_LORA_B, - MODEL_TENSOR.ATTN_V_LORA_A, - MODEL_TENSOR.ATTN_V_LORA_B, - MODEL_TENSOR.ATTN_OUT_LORA_A, - MODEL_TENSOR.ATTN_OUT_LORA_B, - MODEL_TENSOR.FFN_GATE_LORA_A, - MODEL_TENSOR.FFN_GATE_LORA_B, - MODEL_TENSOR.FFN_UP_LORA_A, - MODEL_TENSOR.FFN_UP_LORA_B, - MODEL_TENSOR.FFN_DOWN_LORA_A, - MODEL_TENSOR.FFN_DOWN_LORA_B, ], MODEL_ARCH.CHAMELEON: [ MODEL_TENSOR.TOKEN_EMBD, From 48fc26fa4d4b3c22288f1960cebb66bf4ea57fa1 Mon Sep 17 00:00:00 2001 From: Bar Haim Date: Thu, 2 Jul 2026 13:49:02 +0300 Subject: [PATCH 29/40] granite-switch: drop n_expert_used hack, guard mul_mat_id buft probe --- src/llama-model-loader.cpp | 3 +-- src/models/granite-switch.cpp | 3 --- 2 files changed, 1 insertion(+), 5 deletions(-) diff --git a/src/llama-model-loader.cpp b/src/llama-model-loader.cpp index 28f8bb7934bb..cea748ceb848 100644 --- a/src/llama-model-loader.cpp +++ b/src/llama-model-loader.cpp @@ -932,8 +932,7 @@ static bool weight_buft_supported(const llama_hparams & hparams, ggml_tensor * w } break; case GGML_OP_MUL_MAT_ID: { - const int n_expert_used = hparams.n_expert_used; - GGML_ASSERT(n_expert_used > 0); + const int n_expert_used = std::max(1, hparams.n_expert_used); ggml_tensor * b = ggml_new_tensor_3d(ctx, GGML_TYPE_F32, w->ne[0], n_expert_used, 512); ggml_tensor * ids = ggml_new_tensor_2d(ctx, GGML_TYPE_I32, n_expert_used, 512); op_tensor = ggml_mul_mat_id(ctx, w, b, ids); diff --git a/src/models/granite-switch.cpp b/src/models/granite-switch.cpp index 4fba3dc1e57a..76edbe31b020 100644 --- a/src/models/granite-switch.cpp +++ b/src/models/granite-switch.cpp @@ -20,9 +20,6 @@ void llama_model_granite_switch::load_arch_hparams(llama_model_loader & ml) { ml.get_key(LLM_KV_EXPERT_SHARED_FEED_FORWARD_LENGTH, hparams.n_ff_shexp, /* required */ false); - // mul_mat_id needs n_expert_used == 1; the GGUF is dense (expert_count = 0) - hparams.n_expert_used = 1; - ml.get_key(LLM_KV_NUM_ADAPTERS, n_adapters); ml.get_key(LLM_KV_MAX_LORA_RANK, max_lora_rank); ml.get_key(LLM_KV_CONTROL_TOKEN_GAIN, router_gain, /* required */ false); From 1870e92e4d9dbecab369acb9ddf3898c3f8f6d34 Mon Sep 17 00:00:00 2001 From: Bar Haim Date: Thu, 2 Jul 2026 14:37:43 +0300 Subject: [PATCH 30/40] granite-switch: stop forcing dense expert counts, read from config --- conversion/granite.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/conversion/granite.py b/conversion/granite.py index 914524b47588..743b65e9aec7 100644 --- a/conversion/granite.py +++ b/conversion/granite.py @@ -158,9 +158,9 @@ def __init__(self, *args, **kwargs): def set_gguf_parameters(self): super().set_gguf_parameters() - # dense model: override any expert counts the llama parent emitted - self.gguf_writer.add_expert_count(0) - self.gguf_writer.add_expert_used_count(0) + # dense: pin expert_used_count to 0 (config carries a leftover num_experts_per_tok) + if not self.hparams.get("num_local_experts"): + self.gguf_writer.add_expert_used_count(0) self.gguf_writer.add_num_adapters(self._n_adapters) self.gguf_writer.add_max_lora_rank(self._max_lora_rank) From c9a2850d1a3a8ac365619029a130827c364bc20a Mon Sep 17 00:00:00 2001 From: Bar Haim Date: Thu, 2 Jul 2026 15:24:42 +0300 Subject: [PATCH 31/40] granite-switch: renamed control_token_gain metadata key to router_gain --- conversion/granite.py | 6 +++--- gguf-py/gguf/constants.py | 2 +- gguf-py/gguf/gguf_writer.py | 4 ++-- src/llama-arch.cpp | 2 +- src/llama-arch.h | 2 +- src/models/granite-switch.cpp | 2 +- 6 files changed, 9 insertions(+), 9 deletions(-) diff --git a/conversion/granite.py b/conversion/granite.py index 743b65e9aec7..a35bda94478d 100644 --- a/conversion/granite.py +++ b/conversion/granite.py @@ -166,9 +166,9 @@ def set_gguf_parameters(self): self.gguf_writer.add_max_lora_rank(self._max_lora_rank) self.gguf_writer.add_adapter_token_ids(self.hparams["adapter_token_ids"]) self.gguf_writer.add_adapter_substitute_token_ids(self.hparams["adapter_substitute_token_ids"]) - control_token_gain = float(self.hparams.get("control_token_gain", 15.0)) - self.gguf_writer.add_control_token_gain(control_token_gain) - logger.info("gguf: (graniteswitch) num_adapters=%s max_lora_rank=%s n_slots=%s control_token_gain=%s", self._n_adapters, self._max_lora_rank, self._n_slots, control_token_gain) + router_gain = float(self.hparams.get("control_token_gain", 15.0)) + self.gguf_writer.add_router_gain(router_gain) + logger.info("gguf: (graniteswitch) num_adapters=%s max_lora_rank=%s n_slots=%s router_gain=%s", self._n_adapters, self._max_lora_rank, self._n_slots, router_gain) def _lora_a(self, data: Tensor) -> Tensor: # on-disk A: [n_adapters, 1, max_rank, in] -> [n_adapters+1, max_rank, in] diff --git a/gguf-py/gguf/constants.py b/gguf-py/gguf/constants.py index dd1bdf0efe4f..9ff611446cd9 100644 --- a/gguf-py/gguf/constants.py +++ b/gguf-py/gguf/constants.py @@ -163,7 +163,7 @@ class LLM: ADAPTER_TOKEN_IDS = "{arch}.adapter_token_ids" ADAPTER_SUBSTITUTE_TOKEN_IDS = "{arch}.adapter_substitute_token_ids" MAX_LORA_RANK = "{arch}.max_lora_rank" - CONTROL_TOKEN_GAIN = "{arch}.control_token_gain" + ROUTER_GAIN = "{arch}.router_gain" class Attention: HEAD_COUNT = "{arch}.attention.head_count" diff --git a/gguf-py/gguf/gguf_writer.py b/gguf-py/gguf/gguf_writer.py index 3c7847797708..a2e4e2038baa 100644 --- a/gguf-py/gguf/gguf_writer.py +++ b/gguf-py/gguf/gguf_writer.py @@ -904,8 +904,8 @@ def add_adapter_substitute_token_ids(self, ids: Sequence[int]) -> None: def add_max_lora_rank(self, rank: int) -> None: self.add_uint32(Keys.LLM.MAX_LORA_RANK.format(arch=self.arch), rank) - def add_control_token_gain(self, gain: float) -> None: - self.add_float32(Keys.LLM.CONTROL_TOKEN_GAIN.format(arch=self.arch), gain) + def add_router_gain(self, gain: float) -> None: + self.add_float32(Keys.LLM.ROUTER_GAIN.format(arch=self.arch), gain) def add_wkv_head_size(self, size: int) -> None: self.add_uint32(Keys.WKV.HEAD_SIZE.format(arch=self.arch), size) diff --git a/src/llama-arch.cpp b/src/llama-arch.cpp index 218136492f75..71ebc4bbf151 100644 --- a/src/llama-arch.cpp +++ b/src/llama-arch.cpp @@ -220,7 +220,7 @@ static const std::map LLM_KV_NAMES = { { LLM_KV_ADAPTER_TOKEN_IDS, "%s.adapter_token_ids" }, { LLM_KV_ADAPTER_SUBSTITUTE_TOKEN_IDS, "%s.adapter_substitute_token_ids" }, { LLM_KV_MAX_LORA_RANK, "%s.max_lora_rank" }, - { LLM_KV_CONTROL_TOKEN_GAIN, "%s.control_token_gain" }, + { LLM_KV_ROUTER_GAIN, "%s.router_gain" }, { LLM_KV_TOKEN_SHIFT_COUNT, "%s.token_shift_count" }, { LLM_KV_INTERLEAVE_MOE_LAYER_STEP, "%s.interleave_moe_layer_step" }, { LLM_KV_FULL_ATTENTION_INTERVAL, "%s.full_attention_interval" }, diff --git a/src/llama-arch.h b/src/llama-arch.h index bef4f7408424..80c0b9f2a836 100644 --- a/src/llama-arch.h +++ b/src/llama-arch.h @@ -225,7 +225,7 @@ enum llm_kv { LLM_KV_ADAPTER_TOKEN_IDS, LLM_KV_ADAPTER_SUBSTITUTE_TOKEN_IDS, LLM_KV_MAX_LORA_RANK, - LLM_KV_CONTROL_TOKEN_GAIN, + LLM_KV_ROUTER_GAIN, LLM_KV_TOKEN_SHIFT_COUNT, LLM_KV_INTERLEAVE_MOE_LAYER_STEP, LLM_KV_FULL_ATTENTION_INTERVAL, diff --git a/src/models/granite-switch.cpp b/src/models/granite-switch.cpp index 76edbe31b020..dff083298352 100644 --- a/src/models/granite-switch.cpp +++ b/src/models/granite-switch.cpp @@ -22,7 +22,7 @@ void llama_model_granite_switch::load_arch_hparams(llama_model_loader & ml) { ml.get_key(LLM_KV_NUM_ADAPTERS, n_adapters); ml.get_key(LLM_KV_MAX_LORA_RANK, max_lora_rank); - ml.get_key(LLM_KV_CONTROL_TOKEN_GAIN, router_gain, /* required */ false); + ml.get_key(LLM_KV_ROUTER_GAIN, router_gain, /* required */ false); std::vector token_ids; std::vector substitute_ids; From dbfd7793fb148ca25ff3768f27314aa66ea9cd30 Mon Sep 17 00:00:00 2001 From: Bar Haim Date: Thu, 2 Jul 2026 17:08:38 +0300 Subject: [PATCH 32/40] granite-switch: trim header comments to match native style --- src/llama-model.h | 5 ++--- src/models/granite-switch.cpp | 8 ++++---- src/models/models.h | 31 ++++++++++++------------------- 3 files changed, 18 insertions(+), 26 deletions(-) diff --git a/src/llama-model.h b/src/llama-model.h index e6f807ecad76..47009d69ec16 100644 --- a/src/llama-model.h +++ b/src/llama-model.h @@ -221,9 +221,8 @@ struct llama_layer_nextn { struct ggml_tensor * shared_head_norm = nullptr; }; -// granite-switch: per-token embedded LoRA adapters, stacked across N = num_adapters + 1 -// slots in dim 2 (slot 0 is zero-filled so base tokens add an exact-zero delta). -// A: {n_embd_in, max_lora_rank, N} B: {max_lora_rank, n_embd_out, N} +// granite-switch: per-token embedded LoRA adapters, stacked over slots in dim 2 +// (slot 0 is zero-filled so base tokens add an exact-zero delta) struct llama_layer_switch_lora { struct ggml_tensor * a_q = nullptr; struct ggml_tensor * b_q = nullptr; diff --git a/src/models/granite-switch.cpp b/src/models/granite-switch.cpp index dff083298352..40d87ca9a084 100644 --- a/src/models/granite-switch.cpp +++ b/src/models/granite-switch.cpp @@ -42,7 +42,7 @@ void llama_model_granite_switch::load_arch_hparams(llama_model_loader & ml) { // extra single-head attention layer at the END (index n_real) holds the router // K/V. reusing n_layer_nextn keeps n_layer() == n_real, so the regular layers - // keep their indices and the KV cache shift/defrag skips the router layer. + // keep their indices and the KV cache shift/defrag skips the router layer const uint32_t n_real = hparams.n_layer(); hparams.router_layer = (int32_t) n_real; hparams.n_layer_all = n_real + 1; @@ -220,7 +220,7 @@ llama_model_granite_switch::graph::graph( ggml_tensor * router_q = inp_switch->router_q; res->add_input(std::move(inp_switch)); - // embed the substituted ids directly: build_inp_embd would embed the raw tokens + // embed the substituted ids directly; build_inp_embd would embed the raw tokens ggml_tensor * inpL = ggml_get_rows(ctx0, model.tok_embd, sub_tokens); if (hparams.f_embedding_scale != 0.0f) { inpL = ggml_scale(ctx0, inpL, hparams.f_embedding_scale); @@ -274,7 +274,7 @@ llama_model_granite_switch::graph::graph( 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); - // select the same out rows from adapter_ids (2D round-trip for get_rows) + // keep adapter_ids aligned to the kept rows (2D round-trip for get_rows) const int64_t n_out = inp_out_ids->ne[0]; adapter_ids = ggml_get_rows(ctx0, ggml_reshape_2d(ctx0, adapter_ids, 1, adapter_ids->ne[0]), inp_out_ids); @@ -322,7 +322,7 @@ ggml_tensor * llama_model_granite_switch::graph::build_attention_layer( const int64_t n_embd_q = n_embd_head * n_head; const int64_t n_embd_kv = n_embd_head * n_head_kv; - // slice fused qkv into Q/K/V (contiguous so we can add LoRA deltas) + // slice fused qkv into Q/K/V, made contiguous so LoRA deltas can be added ggml_tensor * Qcur = ggml_cont(ctx0, ggml_view_2d(ctx0, qkv, n_embd_q, qkv->ne[1], qkv->nb[1], 0)); ggml_tensor * Kcur = ggml_cont(ctx0, ggml_view_2d(ctx0, qkv, n_embd_kv, qkv->ne[1], qkv->nb[1], n_embd_q*ggml_element_size(qkv))); ggml_tensor * Vcur = ggml_cont(ctx0, ggml_view_2d(ctx0, qkv, n_embd_kv, qkv->ne[1], qkv->nb[1], (n_embd_q + n_embd_kv)*ggml_element_size(qkv))); diff --git a/src/models/models.h b/src/models/models.h index 44f62342e8aa..28d856859eda 100644 --- a/src/models/models.h +++ b/src/models/models.h @@ -1541,39 +1541,32 @@ struct llama_model_granite_switch : public llama_model_base { void load_arch_hparams(llama_model_loader & ml) override; void load_arch_tensors(llama_model_loader & ml) override; - uint32_t n_adapters = 0; // number of real adapters (excludes the zero slot) + uint32_t n_adapters = 0; uint32_t max_lora_rank = 0; - float router_gain = 15.0f; // routing-token softmax bias magnitude (see set_input) + float router_gain = 15.0f; // stacked-adapter slots: n_adapters + 1 (slot 0 = base/zero delta) uint32_t n_slots() const { return n_adapters + 1; } - // token id -> adapter slot (1 + adapter, so slot 0 stays "base") - std::unordered_map adapter_token_to_slot; - // token id -> substitute token id (token-exchange before embedding) + std::unordered_map adapter_token_to_slot; std::unordered_map adapter_token_to_substitute; - // the per-token adapter selection is recovered in-graph by a single-head - // causal router attention whose K/V live in the KV cache at hparams.router_layer - struct graph : public llm_graph_context { graph(const llama_model & model, const llm_graph_params & params); private: - // per-token switched LoRA delta only: B_a·(A_a·x) -> {n_out, n_tokens} ggml_tensor * build_switched_lora_delta( - ggml_tensor * lora_a, // {n_in, max_lora_rank, N} - ggml_tensor * lora_b, // {max_lora_rank, n_out, N} - ggml_tensor * cur, // {n_in, n_tokens} - ggml_tensor * ids); // {n_tokens} I32 adapter index per token + ggml_tensor * lora_a, + ggml_tensor * lora_b, + ggml_tensor * cur, + ggml_tensor * ids); - // per-token switched LoRA: y = W·x + B_a·(A_a·x), selected per token ggml_tensor * build_switched_lora_mm( - ggml_tensor * w, // base weight (plain 2D) - ggml_tensor * lora_a, // {n_in, max_lora_rank, N} - ggml_tensor * lora_b, // {max_lora_rank, n_out, N} - ggml_tensor * cur, // {n_in, n_tokens} - ggml_tensor * ids); // {n_tokens} I32 adapter index per token + ggml_tensor * w, + ggml_tensor * lora_a, + ggml_tensor * lora_b, + ggml_tensor * cur, + ggml_tensor * ids); ggml_tensor * build_attention_layer( ggml_tensor * cur, From 17ad16582e7a90734e6c5b6ebe5ff86de6c4429c Mon Sep 17 00:00:00 2001 From: Bar Haim Date: Fri, 3 Jul 2026 01:06:57 +0300 Subject: [PATCH 33/40] granite-switch: collapse LoRA tensors to base name + suffix --- src/llama-arch.cpp | 30 ------------------------------ src/llama-arch.h | 15 --------------- src/llama-model-loader.cpp | 15 +++++++++++---- src/models/granite-switch.cpp | 32 ++++++++++++++++---------------- 4 files changed, 27 insertions(+), 65 deletions(-) diff --git a/src/llama-arch.cpp b/src/llama-arch.cpp index 71ebc4bbf151..59d661444419 100644 --- a/src/llama-arch.cpp +++ b/src/llama-arch.cpp @@ -609,21 +609,6 @@ static const std::map LLM_TENSOR_NAMES = { { LLM_TENSOR_MASKED_EMBD_ORDERING, "masked_embd_ordering" }, { LLM_TENSOR_FC, "fc" }, { LLM_TENSOR_D2T, "d2t" }, - // Granite Switch: per-token embedded LoRA adapters (no ".weight" suffix on disk) - { LLM_TENSOR_ATTN_Q_LORA_A, "blk.%d.attn_q.lora_a" }, - { LLM_TENSOR_ATTN_Q_LORA_B, "blk.%d.attn_q.lora_b" }, - { LLM_TENSOR_ATTN_K_LORA_A, "blk.%d.attn_k.lora_a" }, - { LLM_TENSOR_ATTN_K_LORA_B, "blk.%d.attn_k.lora_b" }, - { LLM_TENSOR_ATTN_V_LORA_A, "blk.%d.attn_v.lora_a" }, - { LLM_TENSOR_ATTN_V_LORA_B, "blk.%d.attn_v.lora_b" }, - { LLM_TENSOR_ATTN_OUT_LORA_A, "blk.%d.attn_output.lora_a" }, - { LLM_TENSOR_ATTN_OUT_LORA_B, "blk.%d.attn_output.lora_b" }, - { LLM_TENSOR_FFN_GATE_LORA_A, "blk.%d.ffn_gate.lora_a" }, - { LLM_TENSOR_FFN_GATE_LORA_B, "blk.%d.ffn_gate.lora_b" }, - { LLM_TENSOR_FFN_UP_LORA_A, "blk.%d.ffn_up.lora_a" }, - { LLM_TENSOR_FFN_UP_LORA_B, "blk.%d.ffn_up.lora_b" }, - { LLM_TENSOR_FFN_DOWN_LORA_A, "blk.%d.ffn_down.lora_a" }, - { LLM_TENSOR_FFN_DOWN_LORA_B, "blk.%d.ffn_down.lora_b" }, }; // declare information about the model weight tensors: @@ -875,21 +860,6 @@ static const std::map LLM_TENSOR_INFOS = { // eagle3 {LLM_TENSOR_FC, {LLM_TENSOR_LAYER_OUTPUT, GGML_OP_MUL_MAT}}, {LLM_TENSOR_D2T, {LLM_TENSOR_LAYER_OUTPUT, GGML_OP_GET_ROWS}}, - // granite-switch: per-token embedded LoRA adapters (selected via mul_mat_id) - {LLM_TENSOR_ATTN_Q_LORA_A, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL_MAT_ID}}, - {LLM_TENSOR_ATTN_Q_LORA_B, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL_MAT_ID}}, - {LLM_TENSOR_ATTN_K_LORA_A, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL_MAT_ID}}, - {LLM_TENSOR_ATTN_K_LORA_B, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL_MAT_ID}}, - {LLM_TENSOR_ATTN_V_LORA_A, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL_MAT_ID}}, - {LLM_TENSOR_ATTN_V_LORA_B, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL_MAT_ID}}, - {LLM_TENSOR_ATTN_OUT_LORA_A, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL_MAT_ID}}, - {LLM_TENSOR_ATTN_OUT_LORA_B, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL_MAT_ID}}, - {LLM_TENSOR_FFN_GATE_LORA_A, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL_MAT_ID}}, - {LLM_TENSOR_FFN_GATE_LORA_B, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL_MAT_ID}}, - {LLM_TENSOR_FFN_UP_LORA_A, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL_MAT_ID}}, - {LLM_TENSOR_FFN_UP_LORA_B, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL_MAT_ID}}, - {LLM_TENSOR_FFN_DOWN_LORA_A, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL_MAT_ID}}, - {LLM_TENSOR_FFN_DOWN_LORA_B, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL_MAT_ID}}, }; LLM_KV::LLM_KV(llm_arch arch, const char * suffix) : arch(arch), suffix(suffix) {} diff --git a/src/llama-arch.h b/src/llama-arch.h index 80c0b9f2a836..b3b897d7f00a 100644 --- a/src/llama-arch.h +++ b/src/llama-arch.h @@ -617,21 +617,6 @@ enum llm_tensor { LLM_TENSOR_MASKED_EMBD_ORDERING, LLM_TENSOR_FC, LLM_TENSOR_D2T, - // Granite Switch: per-token embedded LoRA adapters (stacked over N = num_adapters + 1) - LLM_TENSOR_ATTN_Q_LORA_A, - LLM_TENSOR_ATTN_Q_LORA_B, - LLM_TENSOR_ATTN_K_LORA_A, - LLM_TENSOR_ATTN_K_LORA_B, - LLM_TENSOR_ATTN_V_LORA_A, - LLM_TENSOR_ATTN_V_LORA_B, - LLM_TENSOR_ATTN_OUT_LORA_A, - LLM_TENSOR_ATTN_OUT_LORA_B, - LLM_TENSOR_FFN_GATE_LORA_A, - LLM_TENSOR_FFN_GATE_LORA_B, - LLM_TENSOR_FFN_UP_LORA_A, - LLM_TENSOR_FFN_UP_LORA_B, - LLM_TENSOR_FFN_DOWN_LORA_A, - LLM_TENSOR_FFN_DOWN_LORA_B, }; diff --git a/src/llama-model-loader.cpp b/src/llama-model-loader.cpp index cea748ceb848..4cde7ab07222 100644 --- a/src/llama-model-loader.cpp +++ b/src/llama-model-loader.cpp @@ -932,9 +932,11 @@ static bool weight_buft_supported(const llama_hparams & hparams, ggml_tensor * w } break; case GGML_OP_MUL_MAT_ID: { - const int n_expert_used = std::max(1, hparams.n_expert_used); - ggml_tensor * b = ggml_new_tensor_3d(ctx, GGML_TYPE_F32, w->ne[0], n_expert_used, 512); - ggml_tensor * ids = ggml_new_tensor_2d(ctx, GGML_TYPE_I32, n_expert_used, 512); + // Used for either MoE expert routing or embedded adapter routing + const int n_ids_used = hparams.router_layer >= 0 ? 1 : hparams.n_expert_used; + GGML_ASSERT(n_ids_used > 0); + ggml_tensor * b = ggml_new_tensor_3d(ctx, GGML_TYPE_F32, w->ne[0], n_ids_used, 512); + ggml_tensor * ids = ggml_new_tensor_2d(ctx, GGML_TYPE_I32, n_ids_used, 512); op_tensor = ggml_mul_mat_id(ctx, w, b, ids); } break; case GGML_OP_ADD: @@ -1117,15 +1119,20 @@ struct ggml_tensor * llama_model_loader::create_tensor( return nullptr; } - // tensors with "bias" suffix are always used with GGML_OP_ADD or GGML_OP_ADD_ID + // tensors with "bias" suffix are always used with GGML_OP_ADD or GGML_OP_ADD_ID; + // tensors with ".lora_a"/".lora_b" suffix are always used with GGML_OP_MUL_MAT_ID ggml_op op; bool bias = tn.suffix != nullptr && strcmp(tn.suffix, "bias") == 0; + bool lora = tn.suffix != nullptr && + (strcmp(tn.suffix, "lora_a") == 0 || strcmp(tn.suffix, "lora_b") == 0); if (bias) { if (info.op == GGML_OP_MUL_MAT_ID) { op = GGML_OP_ADD_ID; } else { op = GGML_OP_ADD; } + } else if (lora) { + op = GGML_OP_MUL_MAT_ID; } else { op = info.op; } diff --git a/src/models/granite-switch.cpp b/src/models/granite-switch.cpp index 40d87ca9a084..7eed0d279301 100644 --- a/src/models/granite-switch.cpp +++ b/src/models/granite-switch.cpp @@ -85,22 +85,22 @@ void llama_model_granite_switch::load_arch_tensors(llama_model_loader &) { auto & sl = layer.switch_lora; - sl.a_q = create_tensor(tn(LLM_TENSOR_ATTN_Q_LORA_A, i), {n_embd, n_rank, n_slots_i64}, 0); - sl.b_q = create_tensor(tn(LLM_TENSOR_ATTN_Q_LORA_B, i), {n_rank, n_embd_q, n_slots_i64}, 0); - sl.a_k = create_tensor(tn(LLM_TENSOR_ATTN_K_LORA_A, i), {n_embd, n_rank, n_slots_i64}, 0); - sl.b_k = create_tensor(tn(LLM_TENSOR_ATTN_K_LORA_B, i), {n_rank, n_embd_kv, n_slots_i64}, 0); - sl.a_v = create_tensor(tn(LLM_TENSOR_ATTN_V_LORA_A, i), {n_embd, n_rank, n_slots_i64}, 0); - sl.b_v = create_tensor(tn(LLM_TENSOR_ATTN_V_LORA_B, i), {n_rank, n_embd_kv, n_slots_i64}, 0); - - sl.a_o = create_tensor(tn(LLM_TENSOR_ATTN_OUT_LORA_A, i), {n_embd_q, n_rank, n_slots_i64}, 0); - sl.b_o = create_tensor(tn(LLM_TENSOR_ATTN_OUT_LORA_B, i), {n_rank, n_embd, n_slots_i64}, 0); - - sl.a_gate = create_tensor(tn(LLM_TENSOR_FFN_GATE_LORA_A, i), {n_embd, n_rank, n_slots_i64}, 0); - sl.b_gate = create_tensor(tn(LLM_TENSOR_FFN_GATE_LORA_B, i), {n_rank, n_ff, n_slots_i64}, 0); - sl.a_up = create_tensor(tn(LLM_TENSOR_FFN_UP_LORA_A, i), {n_embd, n_rank, n_slots_i64}, 0); - sl.b_up = create_tensor(tn(LLM_TENSOR_FFN_UP_LORA_B, i), {n_rank, n_ff, n_slots_i64}, 0); - sl.a_down = create_tensor(tn(LLM_TENSOR_FFN_DOWN_LORA_A, i), { n_ff, n_rank, n_slots_i64}, 0); - sl.b_down = create_tensor(tn(LLM_TENSOR_FFN_DOWN_LORA_B, i), {n_rank, n_embd, n_slots_i64}, 0); + sl.a_q = create_tensor(tn(LLM_TENSOR_ATTN_Q, "lora_a", i), {n_embd, n_rank, n_slots_i64}, 0); + sl.b_q = create_tensor(tn(LLM_TENSOR_ATTN_Q, "lora_b", i), {n_rank, n_embd_q, n_slots_i64}, 0); + sl.a_k = create_tensor(tn(LLM_TENSOR_ATTN_K, "lora_a", i), {n_embd, n_rank, n_slots_i64}, 0); + sl.b_k = create_tensor(tn(LLM_TENSOR_ATTN_K, "lora_b", i), {n_rank, n_embd_kv, n_slots_i64}, 0); + sl.a_v = create_tensor(tn(LLM_TENSOR_ATTN_V, "lora_a", i), {n_embd, n_rank, n_slots_i64}, 0); + sl.b_v = create_tensor(tn(LLM_TENSOR_ATTN_V, "lora_b", i), {n_rank, n_embd_kv, n_slots_i64}, 0); + + sl.a_o = create_tensor(tn(LLM_TENSOR_ATTN_OUT, "lora_a", i), {n_embd_q, n_rank, n_slots_i64}, 0); + sl.b_o = create_tensor(tn(LLM_TENSOR_ATTN_OUT, "lora_b", i), {n_rank, n_embd, n_slots_i64}, 0); + + sl.a_gate = create_tensor(tn(LLM_TENSOR_FFN_GATE, "lora_a", i), {n_embd, n_rank, n_slots_i64}, 0); + sl.b_gate = create_tensor(tn(LLM_TENSOR_FFN_GATE, "lora_b", i), {n_rank, n_ff, n_slots_i64}, 0); + sl.a_up = create_tensor(tn(LLM_TENSOR_FFN_UP, "lora_a", i), {n_embd, n_rank, n_slots_i64}, 0); + sl.b_up = create_tensor(tn(LLM_TENSOR_FFN_UP, "lora_b", i), {n_rank, n_ff, n_slots_i64}, 0); + sl.a_down = create_tensor(tn(LLM_TENSOR_FFN_DOWN, "lora_a", i), { n_ff, n_rank, n_slots_i64}, 0); + sl.b_down = create_tensor(tn(LLM_TENSOR_FFN_DOWN, "lora_b", i), {n_rank, n_embd, n_slots_i64}, 0); } } From 9fc6854ff5469552bb2f39fa4d3998068107255c Mon Sep 17 00:00:00 2001 From: Bar Haim Date: Tue, 7 Jul 2026 10:28:55 +0300 Subject: [PATCH 34/40] granite-switch: inline suffix checks in tensor op resolution --- src/llama-model-loader.cpp | 14 ++++---------- 1 file changed, 4 insertions(+), 10 deletions(-) diff --git a/src/llama-model-loader.cpp b/src/llama-model-loader.cpp index 4cde7ab07222..208117bcf420 100644 --- a/src/llama-model-loader.cpp +++ b/src/llama-model-loader.cpp @@ -1122,16 +1122,10 @@ struct ggml_tensor * llama_model_loader::create_tensor( // tensors with "bias" suffix are always used with GGML_OP_ADD or GGML_OP_ADD_ID; // tensors with ".lora_a"/".lora_b" suffix are always used with GGML_OP_MUL_MAT_ID ggml_op op; - bool bias = tn.suffix != nullptr && strcmp(tn.suffix, "bias") == 0; - bool lora = tn.suffix != nullptr && - (strcmp(tn.suffix, "lora_a") == 0 || strcmp(tn.suffix, "lora_b") == 0); - if (bias) { - if (info.op == GGML_OP_MUL_MAT_ID) { - op = GGML_OP_ADD_ID; - } else { - op = GGML_OP_ADD; - } - } else if (lora) { + if (tn.suffix != nullptr && strcmp(tn.suffix, "bias") == 0) { + op = info.op == GGML_OP_MUL_MAT_ID ? GGML_OP_ADD_ID : GGML_OP_ADD; + } else if (tn.suffix != nullptr && + (strcmp(tn.suffix, "lora_a") == 0 || strcmp(tn.suffix, "lora_b") == 0)) { op = GGML_OP_MUL_MAT_ID; } else { op = info.op; From fc5e001ea7990577898a6fbd94563bddff6333eb Mon Sep 17 00:00:00 2001 From: Bar Haim Date: Tue, 7 Jul 2026 10:36:09 +0300 Subject: [PATCH 35/40] granite-switch: drop switch-lora struct comment --- src/llama-model.h | 2 -- 1 file changed, 2 deletions(-) diff --git a/src/llama-model.h b/src/llama-model.h index 47009d69ec16..838d6dc28565 100644 --- a/src/llama-model.h +++ b/src/llama-model.h @@ -221,8 +221,6 @@ struct llama_layer_nextn { struct ggml_tensor * shared_head_norm = nullptr; }; -// granite-switch: per-token embedded LoRA adapters, stacked over slots in dim 2 -// (slot 0 is zero-filled so base tokens add an exact-zero delta) struct llama_layer_switch_lora { struct ggml_tensor * a_q = nullptr; struct ggml_tensor * b_q = nullptr; From 0cd43916209b64a7036e007ed5cd21188160ac41 Mon Sep 17 00:00:00 2001 From: Bar Haim Date: Wed, 8 Jul 2026 19:11:42 +0300 Subject: [PATCH 36/40] granite-switch: guard router layer index and inline n_slots --- src/models/granite-switch.cpp | 3 ++- src/models/models.h | 3 --- 2 files changed, 2 insertions(+), 4 deletions(-) diff --git a/src/models/granite-switch.cpp b/src/models/granite-switch.cpp index 7eed0d279301..f1f20211a38b 100644 --- a/src/models/granite-switch.cpp +++ b/src/models/granite-switch.cpp @@ -44,6 +44,7 @@ void llama_model_granite_switch::load_arch_hparams(llama_model_loader & ml) { // K/V. reusing n_layer_nextn keeps n_layer() == n_real, so the regular layers // keep their indices and the KV cache shift/defrag skips the router layer const uint32_t n_real = hparams.n_layer(); + GGML_ASSERT(n_real < LLAMA_MAX_LAYERS); hparams.router_layer = (int32_t) n_real; hparams.n_layer_all = n_real + 1; hparams.n_layer_nextn = 1; @@ -56,7 +57,7 @@ void llama_model_granite_switch::load_arch_hparams(llama_model_loader & ml) { void llama_model_granite_switch::load_arch_tensors(llama_model_loader &) { LLAMA_LOAD_LOCALS; - const int64_t n_slots_i64 = (int64_t) n_slots(); + const int64_t n_slots_i64 = (int64_t) (n_adapters + 1); // slot 0 = base/zero delta const int64_t n_rank = (int64_t) max_lora_rank; const int64_t n_embd_q = n_embd_head_k * n_head; const int64_t n_embd_kv = n_embd_k_gqa; diff --git a/src/models/models.h b/src/models/models.h index 28d856859eda..dad070c960d5 100644 --- a/src/models/models.h +++ b/src/models/models.h @@ -1545,9 +1545,6 @@ struct llama_model_granite_switch : public llama_model_base { uint32_t max_lora_rank = 0; float router_gain = 15.0f; - // stacked-adapter slots: n_adapters + 1 (slot 0 = base/zero delta) - uint32_t n_slots() const { return n_adapters + 1; } - std::unordered_map adapter_token_to_slot; std::unordered_map adapter_token_to_substitute; From f769da11f61227af655cc475b2705f38fcaac38e Mon Sep 17 00:00:00 2001 From: Bar Haim Date: Wed, 8 Jul 2026 19:43:11 +0300 Subject: [PATCH 37/40] granite-switch: group adapter metadata under {arch}.adapters.* namespace --- conversion/granite.py | 10 +++++----- gguf-py/gguf/constants.py | 10 +++++----- gguf-py/gguf/gguf_writer.py | 20 ++++++++++---------- src/llama-arch.cpp | 10 +++++----- src/llama-arch.h | 10 +++++----- src/models/granite-switch.cpp | 10 +++++----- 6 files changed, 35 insertions(+), 35 deletions(-) diff --git a/conversion/granite.py b/conversion/granite.py index a35bda94478d..956342e6d68a 100644 --- a/conversion/granite.py +++ b/conversion/granite.py @@ -162,12 +162,12 @@ def set_gguf_parameters(self): if not self.hparams.get("num_local_experts"): self.gguf_writer.add_expert_used_count(0) - self.gguf_writer.add_num_adapters(self._n_adapters) - self.gguf_writer.add_max_lora_rank(self._max_lora_rank) - self.gguf_writer.add_adapter_token_ids(self.hparams["adapter_token_ids"]) - self.gguf_writer.add_adapter_substitute_token_ids(self.hparams["adapter_substitute_token_ids"]) + self.gguf_writer.add_adapter_count(self._n_adapters) + self.gguf_writer.add_adapter_lora_rank(self._max_lora_rank) + self.gguf_writer.add_adapter_token_ids_activate(self.hparams["adapter_token_ids"]) + self.gguf_writer.add_adapter_token_ids_substitute(self.hparams["adapter_substitute_token_ids"]) router_gain = float(self.hparams.get("control_token_gain", 15.0)) - self.gguf_writer.add_router_gain(router_gain) + self.gguf_writer.add_adapter_router_gain(router_gain) logger.info("gguf: (graniteswitch) num_adapters=%s max_lora_rank=%s n_slots=%s router_gain=%s", self._n_adapters, self._max_lora_rank, self._n_slots, router_gain) def _lora_a(self, data: Tensor) -> Tensor: diff --git a/gguf-py/gguf/constants.py b/gguf-py/gguf/constants.py index 9ff611446cd9..f141cd75ec17 100644 --- a/gguf-py/gguf/constants.py +++ b/gguf-py/gguf/constants.py @@ -159,11 +159,11 @@ class LLM: TARGET_HIDDEN_SIZE = "{arch}.target_hidden_size" BLOCK_SIZE = "{arch}.block_size" NORM_BEFORE_RESIDUAL = "{arch}.norm_before_residual" - NUM_ADAPTERS = "{arch}.num_adapters" - ADAPTER_TOKEN_IDS = "{arch}.adapter_token_ids" - ADAPTER_SUBSTITUTE_TOKEN_IDS = "{arch}.adapter_substitute_token_ids" - MAX_LORA_RANK = "{arch}.max_lora_rank" - ROUTER_GAIN = "{arch}.router_gain" + ADAPTER_COUNT = "{arch}.adapters.count" + ADAPTER_TOKEN_IDS_ACTIVATE = "{arch}.adapters.token_ids_activate" + ADAPTER_TOKEN_IDS_SUBSTITUTE = "{arch}.adapters.token_ids_substitute" + ADAPTER_LORA_RANK = "{arch}.adapters.lora_rank" + ADAPTER_ROUTER_GAIN = "{arch}.adapters.router_gain" class Attention: HEAD_COUNT = "{arch}.attention.head_count" diff --git a/gguf-py/gguf/gguf_writer.py b/gguf-py/gguf/gguf_writer.py index a2e4e2038baa..2a9d743dd62a 100644 --- a/gguf-py/gguf/gguf_writer.py +++ b/gguf-py/gguf/gguf_writer.py @@ -892,20 +892,20 @@ def add_residual_scale(self, value: float) -> None: def add_embedding_scale(self, value: float) -> None: self.add_float32(Keys.LLM.EMBEDDING_SCALE.format(arch=self.arch), value) - def add_num_adapters(self, count: int) -> None: - self.add_uint32(Keys.LLM.NUM_ADAPTERS.format(arch=self.arch), count) + def add_adapter_count(self, count: int) -> None: + self.add_uint32(Keys.LLM.ADAPTER_COUNT.format(arch=self.arch), count) - def add_adapter_token_ids(self, ids: Sequence[int]) -> None: - self.add_array(Keys.LLM.ADAPTER_TOKEN_IDS.format(arch=self.arch), ids) + def add_adapter_token_ids_activate(self, ids: Sequence[int]) -> None: + self.add_array(Keys.LLM.ADAPTER_TOKEN_IDS_ACTIVATE.format(arch=self.arch), ids) - def add_adapter_substitute_token_ids(self, ids: Sequence[int]) -> None: - self.add_array(Keys.LLM.ADAPTER_SUBSTITUTE_TOKEN_IDS.format(arch=self.arch), ids) + def add_adapter_token_ids_substitute(self, ids: Sequence[int]) -> None: + self.add_array(Keys.LLM.ADAPTER_TOKEN_IDS_SUBSTITUTE.format(arch=self.arch), ids) - def add_max_lora_rank(self, rank: int) -> None: - self.add_uint32(Keys.LLM.MAX_LORA_RANK.format(arch=self.arch), rank) + def add_adapter_lora_rank(self, rank: int) -> None: + self.add_uint32(Keys.LLM.ADAPTER_LORA_RANK.format(arch=self.arch), rank) - def add_router_gain(self, gain: float) -> None: - self.add_float32(Keys.LLM.ROUTER_GAIN.format(arch=self.arch), gain) + def add_adapter_router_gain(self, gain: float) -> None: + self.add_float32(Keys.LLM.ADAPTER_ROUTER_GAIN.format(arch=self.arch), gain) def add_wkv_head_size(self, size: int) -> None: self.add_uint32(Keys.WKV.HEAD_SIZE.format(arch=self.arch), size) diff --git a/src/llama-arch.cpp b/src/llama-arch.cpp index 59d661444419..e3dd5eb14d29 100644 --- a/src/llama-arch.cpp +++ b/src/llama-arch.cpp @@ -216,11 +216,11 @@ static const std::map LLM_KV_NAMES = { { LLM_KV_TIME_DECAY_EXTRA_DIM, "%s.time_decay_extra_dim" }, { LLM_KV_RESIDUAL_SCALE, "%s.residual_scale" }, { LLM_KV_EMBEDDING_SCALE, "%s.embedding_scale" }, - { LLM_KV_NUM_ADAPTERS, "%s.num_adapters" }, - { LLM_KV_ADAPTER_TOKEN_IDS, "%s.adapter_token_ids" }, - { LLM_KV_ADAPTER_SUBSTITUTE_TOKEN_IDS, "%s.adapter_substitute_token_ids" }, - { LLM_KV_MAX_LORA_RANK, "%s.max_lora_rank" }, - { LLM_KV_ROUTER_GAIN, "%s.router_gain" }, + { LLM_KV_ADAPTER_COUNT, "%s.adapters.count" }, + { LLM_KV_ADAPTER_TOKEN_IDS_ACTIVATE, "%s.adapters.token_ids_activate" }, + { LLM_KV_ADAPTER_TOKEN_IDS_SUBSTITUTE, "%s.adapters.token_ids_substitute" }, + { LLM_KV_ADAPTER_LORA_RANK, "%s.adapters.lora_rank" }, + { LLM_KV_ADAPTER_ROUTER_GAIN, "%s.adapters.router_gain" }, { LLM_KV_TOKEN_SHIFT_COUNT, "%s.token_shift_count" }, { LLM_KV_INTERLEAVE_MOE_LAYER_STEP, "%s.interleave_moe_layer_step" }, { LLM_KV_FULL_ATTENTION_INTERVAL, "%s.full_attention_interval" }, diff --git a/src/llama-arch.h b/src/llama-arch.h index b3b897d7f00a..82040e585e36 100644 --- a/src/llama-arch.h +++ b/src/llama-arch.h @@ -221,11 +221,11 @@ enum llm_kv { LLM_KV_TIME_DECAY_EXTRA_DIM, LLM_KV_RESIDUAL_SCALE, LLM_KV_EMBEDDING_SCALE, - LLM_KV_NUM_ADAPTERS, - LLM_KV_ADAPTER_TOKEN_IDS, - LLM_KV_ADAPTER_SUBSTITUTE_TOKEN_IDS, - LLM_KV_MAX_LORA_RANK, - LLM_KV_ROUTER_GAIN, + LLM_KV_ADAPTER_COUNT, + LLM_KV_ADAPTER_TOKEN_IDS_ACTIVATE, + LLM_KV_ADAPTER_TOKEN_IDS_SUBSTITUTE, + LLM_KV_ADAPTER_LORA_RANK, + LLM_KV_ADAPTER_ROUTER_GAIN, LLM_KV_TOKEN_SHIFT_COUNT, LLM_KV_INTERLEAVE_MOE_LAYER_STEP, LLM_KV_FULL_ATTENTION_INTERVAL, diff --git a/src/models/granite-switch.cpp b/src/models/granite-switch.cpp index f1f20211a38b..fcdb85155739 100644 --- a/src/models/granite-switch.cpp +++ b/src/models/granite-switch.cpp @@ -20,14 +20,14 @@ void llama_model_granite_switch::load_arch_hparams(llama_model_loader & ml) { ml.get_key(LLM_KV_EXPERT_SHARED_FEED_FORWARD_LENGTH, hparams.n_ff_shexp, /* required */ false); - ml.get_key(LLM_KV_NUM_ADAPTERS, n_adapters); - ml.get_key(LLM_KV_MAX_LORA_RANK, max_lora_rank); - ml.get_key(LLM_KV_ROUTER_GAIN, router_gain, /* required */ false); + ml.get_key(LLM_KV_ADAPTER_COUNT, n_adapters); + ml.get_key(LLM_KV_ADAPTER_LORA_RANK, max_lora_rank); + ml.get_key(LLM_KV_ADAPTER_ROUTER_GAIN, router_gain, /* required */ false); std::vector token_ids; std::vector substitute_ids; - ml.get_arr(LLM_KV_ADAPTER_TOKEN_IDS, token_ids); - ml.get_arr(LLM_KV_ADAPTER_SUBSTITUTE_TOKEN_IDS, substitute_ids); + ml.get_arr(LLM_KV_ADAPTER_TOKEN_IDS_ACTIVATE, token_ids); + ml.get_arr(LLM_KV_ADAPTER_TOKEN_IDS_SUBSTITUTE, substitute_ids); GGML_ASSERT(token_ids.size() == n_adapters); GGML_ASSERT(substitute_ids.size() == n_adapters); From aac9a24359356507e248d41170d8ca952bf41502 Mon Sep 17 00:00:00 2001 From: Bar Haim Date: Wed, 8 Jul 2026 20:36:47 +0300 Subject: [PATCH 38/40] granite-switch: add hparams.has_rope(il) for KV-shift rope skipping --- src/llama-hparams.cpp | 10 ++++++++++ src/llama-hparams.h | 2 ++ src/llama-kv-cache.cpp | 5 +---- 3 files changed, 13 insertions(+), 4 deletions(-) diff --git a/src/llama-hparams.cpp b/src/llama-hparams.cpp index 9d0683d2fec4..67e6ce629e11 100644 --- a/src/llama-hparams.cpp +++ b/src/llama-hparams.cpp @@ -269,6 +269,16 @@ bool llama_hparams::has_kv(uint32_t il) const { return true; } +bool llama_hparams::has_rope(uint32_t il) const { + // the router layer stores adapter routing signal, not positional info, + // so it must not be RoPE-shifted + if (router_layer >= 0 && (int32_t) il == router_layer) { + return false; + } + + return true; +} + uint32_t llama_hparams::n_layer() const { return n_layer_all - n_layer_nextn; } diff --git a/src/llama-hparams.h b/src/llama-hparams.h index 7199f43f1a77..98d4f8994a05 100644 --- a/src/llama-hparams.h +++ b/src/llama-hparams.h @@ -365,6 +365,8 @@ struct llama_hparams { bool has_kv(uint32_t il) const; + bool has_rope(uint32_t il) const; + // number of effective layers (excludes nextn layers) uint32_t n_layer() const; diff --git a/src/llama-kv-cache.cpp b/src/llama-kv-cache.cpp index 3606f831f391..55f4592703bc 100644 --- a/src/llama-kv-cache.cpp +++ b/src/llama-kv-cache.cpp @@ -1925,10 +1925,7 @@ ggml_cgraph * llama_kv_cache::build_graph_shift(llm_graph_result * res, llama_co for (const auto & layer : layers) { const uint32_t il = layer.il; - // granite-switch: the router layer's K encodes the adapter slot as a - // literal dim-0 magnitude (no positional rotation), so it must NOT be - // RoPE-shifted when the KV cache is defragmented / position-shifted. - if (hparams.router_layer >= 0 && (int) il == hparams.router_layer) { + if (!hparams.has_rope(il)) { continue; } From 1a0acb54cc4a533ac6004b15fb7683023c7a2ef5 Mon Sep 17 00:00:00 2001 From: Bar Haim Date: Sun, 12 Jul 2026 15:22:36 +0300 Subject: [PATCH 39/40] granite-switch: skip arch in test-llama-archs (adapter fixture missing, TODO) --- tests/test-llama-archs.cpp | 3 +++ 1 file changed, 3 insertions(+) diff --git a/tests/test-llama-archs.cpp b/tests/test-llama-archs.cpp index f39abe773fc6..daf9fddb9f20 100644 --- a/tests/test-llama-archs.cpp +++ b/tests/test-llama-archs.cpp @@ -396,6 +396,9 @@ static bool arch_supported(const llm_arch arch) { if (arch == LLM_ARCH_GEMMA4 || arch == LLM_ARCH_GEMMA4_ASSISTANT) { return false; // FIXME @ngxson } + if (arch == LLM_ARCH_GRANITE_SWITCH) { + return false; // FIXME adapter fixture + } if (arch == LLM_ARCH_LLAMA_EMBED || arch == LLM_ARCH_GEMMA_EMBEDDING || arch == LLM_ARCH_T5ENCODER) { return false; // FIXME Embedding (?) models produce inconsistent results. } From f57f905bab2acfe3c8160c5c6a590c6c5302ecdd Mon Sep 17 00:00:00 2001 From: Bar Haim Date: Sun, 12 Jul 2026 15:44:36 +0300 Subject: [PATCH 40/40] granite-switch: Keys.Adapters namespace + simplify n_slots --- gguf-py/gguf/constants.py | 12 +++++++----- gguf-py/gguf/gguf_writer.py | 10 +++++----- src/models/granite-switch.cpp | 34 +++++++++++++++++----------------- 3 files changed, 29 insertions(+), 27 deletions(-) diff --git a/gguf-py/gguf/constants.py b/gguf-py/gguf/constants.py index f141cd75ec17..ca55d90a4d52 100644 --- a/gguf-py/gguf/constants.py +++ b/gguf-py/gguf/constants.py @@ -159,11 +159,13 @@ class LLM: TARGET_HIDDEN_SIZE = "{arch}.target_hidden_size" BLOCK_SIZE = "{arch}.block_size" NORM_BEFORE_RESIDUAL = "{arch}.norm_before_residual" - ADAPTER_COUNT = "{arch}.adapters.count" - ADAPTER_TOKEN_IDS_ACTIVATE = "{arch}.adapters.token_ids_activate" - ADAPTER_TOKEN_IDS_SUBSTITUTE = "{arch}.adapters.token_ids_substitute" - ADAPTER_LORA_RANK = "{arch}.adapters.lora_rank" - ADAPTER_ROUTER_GAIN = "{arch}.adapters.router_gain" + + class Adapters: + COUNT = "{arch}.adapters.count" + TOKEN_IDS_ACTIVATE = "{arch}.adapters.token_ids_activate" + TOKEN_IDS_SUBSTITUTE = "{arch}.adapters.token_ids_substitute" + LORA_RANK = "{arch}.adapters.lora_rank" + ROUTER_GAIN = "{arch}.adapters.router_gain" class Attention: HEAD_COUNT = "{arch}.attention.head_count" diff --git a/gguf-py/gguf/gguf_writer.py b/gguf-py/gguf/gguf_writer.py index 2a9d743dd62a..6134f1d7c33b 100644 --- a/gguf-py/gguf/gguf_writer.py +++ b/gguf-py/gguf/gguf_writer.py @@ -893,19 +893,19 @@ def add_embedding_scale(self, value: float) -> None: self.add_float32(Keys.LLM.EMBEDDING_SCALE.format(arch=self.arch), value) def add_adapter_count(self, count: int) -> None: - self.add_uint32(Keys.LLM.ADAPTER_COUNT.format(arch=self.arch), count) + self.add_uint32(Keys.Adapters.COUNT.format(arch=self.arch), count) def add_adapter_token_ids_activate(self, ids: Sequence[int]) -> None: - self.add_array(Keys.LLM.ADAPTER_TOKEN_IDS_ACTIVATE.format(arch=self.arch), ids) + self.add_array(Keys.Adapters.TOKEN_IDS_ACTIVATE.format(arch=self.arch), ids) def add_adapter_token_ids_substitute(self, ids: Sequence[int]) -> None: - self.add_array(Keys.LLM.ADAPTER_TOKEN_IDS_SUBSTITUTE.format(arch=self.arch), ids) + self.add_array(Keys.Adapters.TOKEN_IDS_SUBSTITUTE.format(arch=self.arch), ids) def add_adapter_lora_rank(self, rank: int) -> None: - self.add_uint32(Keys.LLM.ADAPTER_LORA_RANK.format(arch=self.arch), rank) + self.add_uint32(Keys.Adapters.LORA_RANK.format(arch=self.arch), rank) def add_adapter_router_gain(self, gain: float) -> None: - self.add_float32(Keys.LLM.ADAPTER_ROUTER_GAIN.format(arch=self.arch), gain) + self.add_float32(Keys.Adapters.ROUTER_GAIN.format(arch=self.arch), gain) def add_wkv_head_size(self, size: int) -> None: self.add_uint32(Keys.WKV.HEAD_SIZE.format(arch=self.arch), size) diff --git a/src/models/granite-switch.cpp b/src/models/granite-switch.cpp index fcdb85155739..3b7ff5adf32b 100644 --- a/src/models/granite-switch.cpp +++ b/src/models/granite-switch.cpp @@ -57,7 +57,7 @@ void llama_model_granite_switch::load_arch_hparams(llama_model_loader & ml) { void llama_model_granite_switch::load_arch_tensors(llama_model_loader &) { LLAMA_LOAD_LOCALS; - const int64_t n_slots_i64 = (int64_t) (n_adapters + 1); // slot 0 = base/zero delta + const int64_t n_slots = n_adapters + 1; // slot 0 = base/zero delta const int64_t n_rank = (int64_t) max_lora_rank; const int64_t n_embd_q = n_embd_head_k * n_head; const int64_t n_embd_kv = n_embd_k_gqa; @@ -86,22 +86,22 @@ void llama_model_granite_switch::load_arch_tensors(llama_model_loader &) { auto & sl = layer.switch_lora; - sl.a_q = create_tensor(tn(LLM_TENSOR_ATTN_Q, "lora_a", i), {n_embd, n_rank, n_slots_i64}, 0); - sl.b_q = create_tensor(tn(LLM_TENSOR_ATTN_Q, "lora_b", i), {n_rank, n_embd_q, n_slots_i64}, 0); - sl.a_k = create_tensor(tn(LLM_TENSOR_ATTN_K, "lora_a", i), {n_embd, n_rank, n_slots_i64}, 0); - sl.b_k = create_tensor(tn(LLM_TENSOR_ATTN_K, "lora_b", i), {n_rank, n_embd_kv, n_slots_i64}, 0); - sl.a_v = create_tensor(tn(LLM_TENSOR_ATTN_V, "lora_a", i), {n_embd, n_rank, n_slots_i64}, 0); - sl.b_v = create_tensor(tn(LLM_TENSOR_ATTN_V, "lora_b", i), {n_rank, n_embd_kv, n_slots_i64}, 0); - - sl.a_o = create_tensor(tn(LLM_TENSOR_ATTN_OUT, "lora_a", i), {n_embd_q, n_rank, n_slots_i64}, 0); - sl.b_o = create_tensor(tn(LLM_TENSOR_ATTN_OUT, "lora_b", i), {n_rank, n_embd, n_slots_i64}, 0); - - sl.a_gate = create_tensor(tn(LLM_TENSOR_FFN_GATE, "lora_a", i), {n_embd, n_rank, n_slots_i64}, 0); - sl.b_gate = create_tensor(tn(LLM_TENSOR_FFN_GATE, "lora_b", i), {n_rank, n_ff, n_slots_i64}, 0); - sl.a_up = create_tensor(tn(LLM_TENSOR_FFN_UP, "lora_a", i), {n_embd, n_rank, n_slots_i64}, 0); - sl.b_up = create_tensor(tn(LLM_TENSOR_FFN_UP, "lora_b", i), {n_rank, n_ff, n_slots_i64}, 0); - sl.a_down = create_tensor(tn(LLM_TENSOR_FFN_DOWN, "lora_a", i), { n_ff, n_rank, n_slots_i64}, 0); - sl.b_down = create_tensor(tn(LLM_TENSOR_FFN_DOWN, "lora_b", i), {n_rank, n_embd, n_slots_i64}, 0); + sl.a_q = create_tensor(tn(LLM_TENSOR_ATTN_Q, "lora_a", i), {n_embd, n_rank, n_slots}, 0); + sl.b_q = create_tensor(tn(LLM_TENSOR_ATTN_Q, "lora_b", i), {n_rank, n_embd_q, n_slots}, 0); + sl.a_k = create_tensor(tn(LLM_TENSOR_ATTN_K, "lora_a", i), {n_embd, n_rank, n_slots}, 0); + sl.b_k = create_tensor(tn(LLM_TENSOR_ATTN_K, "lora_b", i), {n_rank, n_embd_kv, n_slots}, 0); + sl.a_v = create_tensor(tn(LLM_TENSOR_ATTN_V, "lora_a", i), {n_embd, n_rank, n_slots}, 0); + sl.b_v = create_tensor(tn(LLM_TENSOR_ATTN_V, "lora_b", i), {n_rank, n_embd_kv, n_slots}, 0); + + sl.a_o = create_tensor(tn(LLM_TENSOR_ATTN_OUT, "lora_a", i), {n_embd_q, n_rank, n_slots}, 0); + sl.b_o = create_tensor(tn(LLM_TENSOR_ATTN_OUT, "lora_b", i), {n_rank, n_embd, n_slots}, 0); + + sl.a_gate = create_tensor(tn(LLM_TENSOR_FFN_GATE, "lora_a", i), {n_embd, n_rank, n_slots}, 0); + sl.b_gate = create_tensor(tn(LLM_TENSOR_FFN_GATE, "lora_b", i), {n_rank, n_ff, n_slots}, 0); + sl.a_up = create_tensor(tn(LLM_TENSOR_FFN_UP, "lora_a", i), {n_embd, n_rank, n_slots}, 0); + sl.b_up = create_tensor(tn(LLM_TENSOR_FFN_UP, "lora_b", i), {n_rank, n_ff, n_slots}, 0); + sl.a_down = create_tensor(tn(LLM_TENSOR_FFN_DOWN, "lora_a", i), { n_ff, n_rank, n_slots}, 0); + sl.b_down = create_tensor(tn(LLM_TENSOR_FFN_DOWN, "lora_b", i), {n_rank, n_embd, n_slots}, 0); } }