From 1cbf8a682229e6010a17db186856a6d43a31fb64 Mon Sep 17 00:00:00 2001 From: amodhyh Date: Mon, 24 Aug 2026 18:40:27 +0530 Subject: [PATCH 1/9] feat(llama): add initial stubs for native GGUF support --- python/freetoken/models/gguf/config.py | 1 + python/freetoken/models/llama/__init__.py | 3 +- python/freetoken/models/llama/gguf.py | 42 +++++++++++++++++++++++ python/freetoken/models/llama/model.py | 5 +++ python/freetoken/models/register.py | 7 ++++ 5 files changed, 57 insertions(+), 1 deletion(-) create mode 100644 python/freetoken/models/llama/gguf.py diff --git a/python/freetoken/models/gguf/config.py b/python/freetoken/models/gguf/config.py index 63b1a18b9..bf9751bc6 100644 --- a/python/freetoken/models/gguf/config.py +++ b/python/freetoken/models/gguf/config.py @@ -18,6 +18,7 @@ # reuses the model classes but a GGUF parse_config / iter_weights). GGUF_ARCH_TO_REGISTRY: dict[str, str] = { "gemma4": "Gemma4GGUFForCausalLM", + "llama" : "LlamaGGUFForCausalLM" } diff --git a/python/freetoken/models/llama/__init__.py b/python/freetoken/models/llama/__init__.py index 9ffad0e95..e8337f148 100644 --- a/python/freetoken/models/llama/__init__.py +++ b/python/freetoken/models/llama/__init__.py @@ -1,5 +1,6 @@ from .config import parse_config from .model import LlamaForCausalLM from .weight import iter_weights +from .gguf import parse_gguf_config,iter_gguf_weights -__all__ = ["LlamaForCausalLM", "parse_config", "iter_weights"] +__all__ = ["LlamaForCausalLM", "parse_config", "iter_weights", "parse_gguf_config", "iter_gguf_weights"] diff --git a/python/freetoken/models/llama/gguf.py b/python/freetoken/models/llama/gguf.py new file mode 100644 index 000000000..dad933126 --- /dev/null +++ b/python/freetoken/models/llama/gguf.py @@ -0,0 +1,42 @@ +""" Llama GGUF adapter: build the FreeToken ``ModelConfig`` from GGUF metadata. +""" +from torch import Tensor +from typing import Iterator +from freetoken.models.config import ModelConfig + +def parse_gguf_config(shim:GGUFConfigShim) -> ModelConfig: + """Parse a GGUF config shim into a FreeToken ``ModelConfig``.""" + + # The GGUF shim is a minimal HF-config-like dict with the fields we need. + return ModelConfig( + model_type=shim.model_type, + architectures=shim.architectures, + vocab_size=shim.vocab_size, + tie_word_embeddings=shim.tie_word_embeddings, + torch_dtype="bfloat16", # GGUF weights dequantize to bf16 + ) + +def iter_gguf_weights( + model_path: str, + device, + *, + include_moe_experts: bool, + include_non_moe: bool + ) -> Iterator[tuple[str, Tensor]]: + """Iterate over GGUF weights, yielding (name, tensor) pairs for the model's parameters.""" + pass + + + +def convert_llama_to_gguf(model, config) -> None: + """Convert a FreeToken Llama model to GGUF format in-place. + + This is a no-op for non-GGUF models, and raises an error if the model is not a Llama. + """ + pass + +# ...existing code... +def is_gguf_model(config) -> bool: + """Check if the model config is for a GGUF model.""" + return getattr(config, "moe_weight_format", None) == "q4_0" +# ...existing code... \ No newline at end of file diff --git a/python/freetoken/models/llama/model.py b/python/freetoken/models/llama/model.py index a3d801712..805028e4a 100644 --- a/python/freetoken/models/llama/model.py +++ b/python/freetoken/models/llama/model.py @@ -11,6 +11,8 @@ from .attention import LlamaAttention as LlamaAttn +from freetoken.models.llama.gguf import is_gguf_model, convert_llama_to_gguf, parse_gguf_config + if TYPE_CHECKING: from freetoken.models.config import ModelConfig @@ -75,6 +77,9 @@ def __init__(self, config: ModelConfig): tied_embedding=self.model.embed_tokens if config.tie_word_embeddings else None, ) super().__init__() + if is_gguf_model(config): + convert_llama_to_gguf(self,config) + def forward(self) -> torch.Tensor: output = self.model.forward(get_global_ctx().batch.input_ids) diff --git a/python/freetoken/models/register.py b/python/freetoken/models/register.py index b94d8291b..e3ed690ae 100644 --- a/python/freetoken/models/register.py +++ b/python/freetoken/models/register.py @@ -21,6 +21,13 @@ class ModelSpec: "freetoken.models.llama", "LlamaForCausalLM", ), + # LLama architecture Entry for the GGUF support + "LlamaGGUFForCausalLM":ModelSpec( + "freetoken.models.llama", + "LlamaForCausalLM", + parse_config="parse_gguf_config", + iter_weights="iter_gguf_weights", + ), "Qwen2ForCausalLM": ModelSpec( "freetoken.models.qwen2", "Qwen2ForCausalLM", From f42698e9fbc0f66bd2f3cf8056f0479677023591 Mon Sep 17 00:00:00 2001 From: amodhyh Date: Thu, 27 Aug 2026 22:40:22 +0530 Subject: [PATCH 2/9] feat(llama): implement native GGUF support Adds parse_gguf_config to map GGUF metadata to ModelConfig, iter_gguf_weights to stream and fuse QKV/GateUp tensors, and convert_llama_to_gguf to swap dense PyTorch layers with custom GGUF GPU kernels. --- python/freetoken/models/llama/gguf.py | 157 ++++++++++++++++++++++++-- 1 file changed, 149 insertions(+), 8 deletions(-) diff --git a/python/freetoken/models/llama/gguf.py b/python/freetoken/models/llama/gguf.py index dad933126..63e2e4906 100644 --- a/python/freetoken/models/llama/gguf.py +++ b/python/freetoken/models/llama/gguf.py @@ -2,20 +2,72 @@ """ from torch import Tensor from typing import Iterator -from freetoken.models.config import ModelConfig +from freetoken.models.config import ModelConfig,RotaryConfig +from freetoken.models.gguf.config import GgufConfigShim +from freetoken.models.gguf.reader import iter_gguf_tensors -def parse_gguf_config(shim:GGUFConfigShim) -> ModelConfig: + +def parse_gguf_config(shim:GgufConfigShim) -> ModelConfig: """Parse a GGUF config shim into a FreeToken ``ModelConfig``.""" + gguf_metadata_llama = shim.metadata + + hidden_size = gguf_metadata_llama.get("llama.embedding_length") + num_of_heads = gguf_metadata_llama.get("llama.attention.head_count") + # for the rotary embedding config + head_dim = hidden_size//num_of_heads + # The GGUF shim is a minimal HF-config-like dict with the fields we need. return ModelConfig( + # directly, from shim model_type=shim.model_type, architectures=shim.architectures, vocab_size=shim.vocab_size, tie_word_embeddings=shim.tie_word_embeddings, - torch_dtype="bfloat16", # GGUF weights dequantize to bf16 + + + # from mapping GGUF keys + num_layers=shim.metadata.get("llama.block_count"), + num_qo_heads=num_of_heads, + num_kv_heads=shim.metadata.get("llama.attention.head_count_kv"), + hidden_size=hidden_size, + head_dim=head_dim, + intermediate_size=shim.metadata.get("llama.feed_forward_length"), + rms_norm_eps=shim.metadata.get("llama.attention.layer_norm_rms_epsilon"), + + # --- Rotary Config --- + rotary_config=RotaryConfig( + head_dim=head_dim, + rotary_dim=head_dim, + max_position=shim.metadata.get("llama.context_length"), + base=shim.metadata.get("llama.rope.freq_base"), + scaling=None + ), + # --- Hardcoded LLaMA Structural Defaults --- + hidden_act="silu", #llama standard activation + num_experts=1, #current llama is still a dense model(no MoE) + num_experts_per_tok=1, + moe_intermediate_size=0, + norm_topk_prob=False, + moe_weight_format="q4_0", #4bit symmetric ) +from freetoken.models.gguf.reader import iter_gguf_tensors +from freetoken.models.gguf.dequant import dequantize +import torch + +def _to_bf16(t) -> torch.Tensor: + """Dequantize a GgufTensor (F32/F16) to a dense bf16 tensor for LayerNorms.""" + flat = dequantize(t.packed().reshape(-1), t.ggml_type, torch.bfloat16) + return flat.reshape(t.shape) + +_LAYER_MAP = { + "attn_norm.weight": "input_layernorm.weight", + "ffn_norm.weight": "post_attention_layernorm.weight", + "attn_output.weight": "self_attn.o_proj.qweight", + "ffn_down.weight": "mlp.down_proj.qweight", +} + def iter_gguf_weights( model_path: str, device, @@ -24,19 +76,108 @@ def iter_gguf_weights( include_non_moe: bool ) -> Iterator[tuple[str, Tensor]]: """Iterate over GGUF weights, yielding (name, tensor) pairs for the model's parameters.""" - pass - + qkv_buf: dict[int, dict[str, torch.Tensor]] = {} + gate_up_buf: dict[int, dict[str, torch.Tensor]] = {} + + for t in iter_gguf_tensors(model_path): + name = t.name + + # Global Standalone Tensors + if name == "token_embd.weight": + yield "model.embed_tokens.qweight", t.packed() + continue + if name == "output_norm.weight": + yield "model.norm.weight", _to_bf16(t) + continue + if name == "output.weight": + yield "lm_head.qweight", t.packed() + continue + if name == "rope_freqs.weight": + continue # Recomputed dynamically in the engine + + if not name.startswith("blk."): + continue + + # Block Tensors + layer = int(name.split(".")[1]) + suffix = name.split(".", 2)[2] + base = f"model.layers.{layer}" + + # Mapped Standalones + if suffix in _LAYER_MAP: + mapped_name = _LAYER_MAP[suffix] + # Norms must be dequantized to bf16; projections stay packed as Q4_0 + if "norm" in suffix: + yield f"{base}.{mapped_name}", _to_bf16(t) + else: + yield f"{base}.{mapped_name}", t.packed() + continue + + # Fusable Tensors + if suffix == "attn_q.weight": + qkv_buf.setdefault(layer, {})["q"] = t.packed() + elif suffix == "attn_k.weight": + qkv_buf.setdefault(layer, {})["k"] = t.packed() + elif suffix == "attn_v.weight": + qkv_buf.setdefault(layer, {})["v"] = t.packed() + elif suffix == "ffn_gate.weight": + gate_up_buf.setdefault(layer, {})["gate"] = t.packed() + elif suffix == "ffn_up.weight": + gate_up_buf.setdefault(layer, {})["up"] = t.packed() + + # 3. Trigger Fusion + slots = qkv_buf.get(layer) + if slots and "q" in slots and "k" in slots and "v" in slots: + yield f"{base}.self_attn.qkv_proj.qweight", torch.cat( + [slots["q"], slots["k"], slots["v"]], dim=0 + ) + del qkv_buf[layer] + + gu = gate_up_buf.get(layer) + if gu and "gate" in gu and "up" in gu: + yield f"{base}.mlp.gate_up_proj.qweight", torch.cat( + [gu["gate"], gu["up"]], dim=0 + ) + del gate_up_buf[layer] def convert_llama_to_gguf(model, config) -> None: """Convert a FreeToken Llama model to GGUF format in-place. This is a no-op for non-GGUF models, and raises an error if the model is not a Llama. """ - pass + # to prevent the circular imports + from freetoken.layers.gguf import GGUFEmbedding,GGUFLinear + from freetoken.models.gguf.dequant import GGML_Q4_0, GGML_Q6_K + + # helper functino for swapping layers + def swap_linear(owner,attr_name, quant_type=GGML_Q4_0, has_bias=False): + old_layer=getattr(owner,attr_name) + out_features, in_features= old_layer.weight.shape + # replace the attribute with the custom GGUF kernal + setattr( + owner, + attr_name, + GGUFLinear(in_features, out_features, quant_type) + ) + inner_model=model.model + inner_model.embed_tokens=GGUFEmbedding( + num_embeddings=config.vocab_size, + embedding_dim=config.hidden_size, + quant_type=GGML_Q6_K, + embed_scale=None + ) + + for layer in inner_model.layers: + swap_linear(layer.self_attn, "qkv_proj") + swap_linear(layer.self_attn, "o_proj") + swap_linear(layer.mlp, "gate_up_proj") + swap_linear(layer.mlp, "down_proj") + + # swap the LM head + if not config.tie_word_embeddings: + swap_linear(model,"lm_head",quant_type=GGML_Q6_K) -# ...existing code... def is_gguf_model(config) -> bool: """Check if the model config is for a GGUF model.""" return getattr(config, "moe_weight_format", None) == "q4_0" -# ...existing code... \ No newline at end of file From 429ebf936a8e3a663f44b4b9a479f99b8c21b6e8 Mon Sep 17 00:00:00 2001 From: amodhyh Date: Mon, 31 Aug 2026 11:10:48 +0530 Subject: [PATCH 3/9] feat(llama): add MoE config parsing for LLaMA 4 with shared experts --- python/freetoken/models/llama/gguf.py | 72 ++++++++++++++++----------- 1 file changed, 43 insertions(+), 29 deletions(-) diff --git a/python/freetoken/models/llama/gguf.py b/python/freetoken/models/llama/gguf.py index 63e2e4906..be2ba165f 100644 --- a/python/freetoken/models/llama/gguf.py +++ b/python/freetoken/models/llama/gguf.py @@ -11,45 +11,59 @@ def parse_gguf_config(shim:GgufConfigShim) -> ModelConfig: """Parse a GGUF config shim into a FreeToken ``ModelConfig``.""" gguf_metadata_llama = shim.metadata + # model architecture + arch=shim.model_type - hidden_size = gguf_metadata_llama.get("llama.embedding_length") - num_of_heads = gguf_metadata_llama.get("llama.attention.head_count") - # for the rotary embedding config + hidden_size = gguf_metadata_llama.get(f"{arch}.embedding_length") + num_of_heads = gguf_metadata_llama.get(f"{arch}.attention.head_count") + # head_dim and rotary_dim head_dim = hidden_size//num_of_heads + + dense_ffn_size = gguf_metadata_llama.get(f"{arch}.feed_forward_length") + + expert_count = gguf_metadata_llama.get(f"{arch}.expert_count", 1) # The GGUF shim is a minimal HF-config-like dict with the fields we need. return ModelConfig( # directly, from shim - model_type=shim.model_type, - architectures=shim.architectures, - vocab_size=shim.vocab_size, - tie_word_embeddings=shim.tie_word_embeddings, + model_type = shim.model_type, + architectures = shim.architectures, + vocab_size = shim.vocab_size, + tie_word_embeddings = shim.tie_word_embeddings, + head_dim = head_dim, + hidden_size = hidden_size, - # from mapping GGUF keys - num_layers=shim.metadata.get("llama.block_count"), - num_qo_heads=num_of_heads, - num_kv_heads=shim.metadata.get("llama.attention.head_count_kv"), - hidden_size=hidden_size, - head_dim=head_dim, - intermediate_size=shim.metadata.get("llama.feed_forward_length"), - rms_norm_eps=shim.metadata.get("llama.attention.layer_norm_rms_epsilon"), +# from mapping GGUF keys + + num_layers = gguf_metadata_llama.get(f"{arch}.block_count"), + intermediate_size = gguf_metadata_llama.get(f"{arch}.feed_forward_length"), + + num_qo_heads = num_of_heads, + num_kv_heads = gguf_metadata_llama.get(f"{arch}.attention.head_count_kv"), + rms_norm_eps = gguf_metadata_llama.get(f"{arch}.attention.layer_norm_rms_epsilon"), - # --- Rotary Config --- - rotary_config=RotaryConfig( - head_dim=head_dim, - rotary_dim=head_dim, - max_position=shim.metadata.get("llama.context_length"), - base=shim.metadata.get("llama.rope.freq_base"), - scaling=None + # Rotary Config + rotary_config = RotaryConfig( + head_dim = gguf_metadata_llama.get(f"{arch}.rope.dimension_count",head_dim) , + rotary_dim = gguf_metadata_llama.get(f"{arch}.rope.dimension_count",head_dim) , #identical to head_dim + max_position = gguf_metadata_llama.get(f"{arch}.context_length"), + base = gguf_metadata_llama.get(f"{arch}.rope.freq_base"), + scaling = None ), - # --- Hardcoded LLaMA Structural Defaults --- - hidden_act="silu", #llama standard activation - num_experts=1, #current llama is still a dense model(no MoE) - num_experts_per_tok=1, - moe_intermediate_size=0, - norm_topk_prob=False, - moe_weight_format="q4_0", #4bit symmetric + + hidden_act = "silu", #llama standard activation + num_experts = expert_count, #for MoE + num_experts_per_tok = gguf_metadata_llama.get(f"{arch}.expert_used_count", 1), + norm_topk_prob = False, + + # If it's a dense model, this just falls back to the dense size. + moe_intermediate_size = gguf_metadata_llama.get(f"{arch}.expert_feed_forward_length",dense_ffn_size), + + moe_weight_format = "q4_0", #4bit symmetric + + shared_expert_intermediate_size= dense_ffn_size if expert_count> 1 else 0 + ) from freetoken.models.gguf.reader import iter_gguf_tensors From ccfe28bf2713e0bd3fec8cdf4de922f1f39fd143 Mon Sep 17 00:00:00 2001 From: amodhyh Date: Mon, 31 Aug 2026 11:55:20 +0530 Subject: [PATCH 4/9] feat(llama): add MoE and Shared Expert dual-routing to LLaMA Decoder --- python/freetoken/models/llama/model.py | 33 ++++++++++++++++++++++---- 1 file changed, 29 insertions(+), 4 deletions(-) diff --git a/python/freetoken/models/llama/model.py b/python/freetoken/models/llama/model.py index 805028e4a..42bd8da64 100644 --- a/python/freetoken/models/llama/model.py +++ b/python/freetoken/models/llama/model.py @@ -13,14 +13,31 @@ from freetoken.models.llama.gguf import is_gguf_model, convert_llama_to_gguf, parse_gguf_config +from freetoken.layers.moe import make_moe_layer + if TYPE_CHECKING: from freetoken.models.config import ModelConfig class LlamaDecoderLayer(BaseOP): def __init__(self, config: ModelConfig, layer_id: int): - self.self_attn = LlamaAttn(config, layer_id) - self.mlp = LlamaMLP(config) + self.self_attn = LlamaAttn(config, layer_id) + # build the router and the smaller experts + if config.num_experts > 1: + self.mlp = make_moe_layer( + config, + layer_id=layer_id, + weight_format= getattr(config,"moe_weight_format","bf16") + ) + # Build the massive Shared Expert + if getattr(config, "shared_expert_intermediate_size", 0) > 0: + self.shared_expert=LlamaMLP(config) + else: + self.shared_expert = None + else: + # standard dense LLaMA 1/2/3 + self.mlp = LlamaMLP(config) + self.input_layernorm = RMSNormFused( size=config.hidden_size, eps=config.rms_norm_eps, @@ -29,7 +46,6 @@ def __init__(self, config: ModelConfig, layer_id: int): size=config.hidden_size, eps=config.rms_norm_eps, ) - self._layer_id = layer_id @nvtx_annotate("Layer_{}", layer_id_field="_layer_id") @@ -41,7 +57,16 @@ def forward( x, residual = self.input_layernorm.forward(x, residual) x = self.self_attn.forward(x) x, residual = self.post_attention_layernorm.forward(x, residual) - x = self.mlp.forward(x) + + # --- NEW ROUTING LOGIC --- + if getattr(self, "shared_expert", None) is not None: + routed_out = self.mlp.forward(x) + shared_out = self.shared_expert.forward(x) + # adding the shareed expert and MoE expert together + x = routed_out + shared_out + else: + # standard dense routing + x = self.mlp.forward(x) return x, residual From 4e51005881a9b27186064581a873f959f9b84e54 Mon Sep 17 00:00:00 2001 From: amodhyh Date: Mon, 31 Aug 2026 12:17:02 +0530 Subject: [PATCH 5/9] fix(llama): instantiate and route through MoE gate --- python/freetoken/models/llama/model.py | 11 ++++++++++- 1 file changed, 10 insertions(+), 1 deletion(-) diff --git a/python/freetoken/models/llama/model.py b/python/freetoken/models/llama/model.py index 42bd8da64..d5963ec16 100644 --- a/python/freetoken/models/llama/model.py +++ b/python/freetoken/models/llama/model.py @@ -14,6 +14,7 @@ from freetoken.models.llama.gguf import is_gguf_model, convert_llama_to_gguf, parse_gguf_config from freetoken.layers.moe import make_moe_layer +from freetoken.layers.linear import LinearReplicated if TYPE_CHECKING: from freetoken.models.config import ModelConfig @@ -24,6 +25,9 @@ def __init__(self, config: ModelConfig, layer_id: int): self.self_attn = LlamaAttn(config, layer_id) # build the router and the smaller experts if config.num_experts > 1: + # router + self.router = LinearReplicated(config.hidden_size, config.num_experts, has_bias=False) + # Smaller experts self.mlp = make_moe_layer( config, layer_id=layer_id, @@ -60,7 +64,11 @@ def forward( # --- NEW ROUTING LOGIC --- if getattr(self, "shared_expert", None) is not None: - routed_out = self.mlp.forward(x) + + # routing logits for usage of the smaller experts + router_logits = self.router.forward(x) + routed_out = self.mlp.forward(x, router_logits=router_logits) + shared_out = self.shared_expert.forward(x) # adding the shareed expert and MoE expert together x = routed_out + shared_out @@ -107,6 +115,7 @@ def __init__(self, config: ModelConfig): def forward(self) -> torch.Tensor: + output = self.model.forward(get_global_ctx().batch.input_ids) logits = self.lm_head.forward(output) return logits From 707679303662d64680d60486b60a7628e0b1f404 Mon Sep 17 00:00:00 2001 From: amodhyh Date: Mon, 31 Aug 2026 22:58:51 +0530 Subject: [PATCH 6/9] feat(llama): stream MoE weights and Shared Experts from GGUF --- python/freetoken/models/llama/gguf.py | 49 +++++++++++++++++++++++---- 1 file changed, 43 insertions(+), 6 deletions(-) diff --git a/python/freetoken/models/llama/gguf.py b/python/freetoken/models/llama/gguf.py index be2ba165f..3d3e6141b 100644 --- a/python/freetoken/models/llama/gguf.py +++ b/python/freetoken/models/llama/gguf.py @@ -140,7 +140,28 @@ def iter_gguf_weights( elif suffix == "ffn_up.weight": gate_up_buf.setdefault(layer, {})["up"] = t.packed() - # 3. Trigger Fusion + # MoE and Shared Expert Tensors + # The Router Gate + elif suffix == "ffn_gate_inp.weight": + yield f"{base}.router.weight", t.packed() + + # Shared Expert (shexp) + elif suffix == "ffn_down_shexp.weight": + yield f"{base}.shared_expert.down_proj.qweight", t.packed() + elif suffix == "ffn_gate_shexp.weight": + gate_up_buf.setdefault(layer, {})["gate_shexp"] = t.packed() + elif suffix == "ffn_up_shexp.weight": + gate_up_buf.setdefault(layer, {})["up_shexp"] = t.packed() + + # Routed Experts (exps) + elif suffix == "ffn_down_exps.weight": + yield f"{base}.mlp.down_proj.qweight", t.packed() + elif suffix == "ffn_gate_exps.weight": + gate_up_buf.setdefault(layer, {})["gate_exps"] = t.packed() + elif suffix == "ffn_up_exps.weight": + gate_up_buf.setdefault(layer, {})["up_exps"] = t.packed() + + # Trigger Fusion slots = qkv_buf.get(layer) if slots and "q" in slots and "k" in slots and "v" in slots: yield f"{base}.self_attn.qkv_proj.qweight", torch.cat( @@ -149,11 +170,27 @@ def iter_gguf_weights( del qkv_buf[layer] gu = gate_up_buf.get(layer) - if gu and "gate" in gu and "up" in gu: - yield f"{base}.mlp.gate_up_proj.qweight", torch.cat( - [gu["gate"], gu["up"]], dim=0 - ) - del gate_up_buf[layer] + if gu: + if "gate" in gu and "up" in gu: + yield f"{base}.mlp.gate_up_proj.qweight", torch.cat( + [gu["gate"], gu["up"]], dim=0 + ) + del gu["gate"] + del gu["up"] + + if "gate_shexp" in gu and "up_shexp" in gu: + yield f"{base}.shared_expert.gate_up_proj.qweight", torch.cat( + [gu["gate_shexp"], gu["up_shexp"]], dim=0 + ) + del gu["gate_shexp"] + del gu["up_shexp"] + + if "gate_exps" in gu and "up_exps" in gu: + yield f"{base}.mlp.gate_up_proj.qweight", torch.cat( + [gu["gate_exps"], gu["up_exps"]], dim=0 + ) + del gu["gate_exps"] + del gu["up_exps"] def convert_llama_to_gguf(model, config) -> None: """Convert a FreeToken Llama model to GGUF format in-place. From c05412d8be542ac701fca8a93f03733a821bef3c Mon Sep 17 00:00:00 2001 From: amodhyh Date: Mon, 31 Aug 2026 23:13:45 +0530 Subject: [PATCH 7/9] fix(llama): conditionally swap MoE and Shared Expert layers for GGUF kernels --- python/freetoken/models/llama/gguf.py | 14 +++++++++++--- 1 file changed, 11 insertions(+), 3 deletions(-) diff --git a/python/freetoken/models/llama/gguf.py b/python/freetoken/models/llama/gguf.py index 3d3e6141b..af3cb7aa7 100644 --- a/python/freetoken/models/llama/gguf.py +++ b/python/freetoken/models/llama/gguf.py @@ -192,7 +192,7 @@ def iter_gguf_weights( del gu["gate_exps"] del gu["up_exps"] -def convert_llama_to_gguf(model, config) -> None: +def convert_llama_to_gguf(model, config:ModelConfig) -> None: """Convert a FreeToken Llama model to GGUF format in-place. This is a no-op for non-GGUF models, and raises an error if the model is not a Llama. @@ -220,10 +220,18 @@ def swap_linear(owner,attr_name, quant_type=GGML_Q4_0, has_bias=False): ) for layer in inner_model.layers: - swap_linear(layer.self_attn, "qkv_proj") - swap_linear(layer.self_attn, "o_proj") + swap_linear(layer.self_attn, "qkv_proj") + swap_linear(layer.self_attn, "o_proj") + + # If it's a Dense model, swap the standard MLP + if not hasattr(layer, "router"): swap_linear(layer.mlp, "gate_up_proj") swap_linear(layer.mlp, "down_proj") + + # If it's a Shared-Expert MoE model, swap the Shared Expert MLP + if getattr(layer, "shared_expert", None) is not None: + swap_linear(layer.shared_expert, "gate_up_proj") + swap_linear(layer.shared_expert, "down_proj") # swap the LM head if not config.tie_word_embeddings: From 1cabd37dc30b20ecf196c863d89bb5ccda0900f5 Mon Sep 17 00:00:00 2001 From: amodhyh Date: Mon, 31 Aug 2026 23:21:31 +0530 Subject: [PATCH 8/9] feat(llama): add tokenizer EOT tokens for LLaMA 3 and 4 --- python/freetoken/models/gguf/tokenizer.py | 10 +++++++--- 1 file changed, 7 insertions(+), 3 deletions(-) diff --git a/python/freetoken/models/gguf/tokenizer.py b/python/freetoken/models/gguf/tokenizer.py index 6d5481c17..b476db94a 100644 --- a/python/freetoken/models/gguf/tokenizer.py +++ b/python/freetoken/models/gguf/tokenizer.py @@ -13,7 +13,7 @@ from .reader import gguf_architecture, load_gguf_metadata # GGUF architecture -> transformers GGUF tokenizer-converter key. -_TOKENIZER_ARCH = {"gemma4": "gemma4_text"} +_TOKENIZER_ARCH = {"gemma4": "gemma4_text", "llama4": "llama"} def load_gguf_tokenizer(model_path: str): @@ -39,10 +39,14 @@ def tok_for(id_key: str, default: str) -> str: # gemma4 chat turns end with ; prefer it as eos so chat generation halts # (the formal is also a stop id, see gguf_eos_token_ids). turn_end = "" if "" in tokens else None + + # LLaMA 4 uses <|eot|> as the formal turn end + llama_eot = "<|eot|>" if "<|eot|>" in tokens else ("<|eot_id|>" if "<|eot_id|>" in tokens else None) + tokenizer = PreTrainedTokenizerFast( tokenizer_object=fast, bos_token=tok_for("bos_token_id", ""), - eos_token=turn_end or tok_for("eos_token_id", ""), + eos_token=llama_eot or turn_end or tok_for("eos_token_id", ""), unk_token=tok_for("unknown_token_id", ""), pad_token=tok_for("padding_token_id", ""), ) @@ -64,7 +68,7 @@ def gguf_eos_token_ids(model_path: str, tokenizer) -> set[int]: ids.add(int(eid)) # Look the stop tokens up in the vocab directly (convert_tokens_to_ids would map an # absent name to , wrongly adding it as a stop id). - for name in ("", ""): + for name in ("", "", "<|eot_id|>", "<|eot|>"): try: ids.add(tokens.index(name)) except ValueError: From cd1ff5fe0a3d00d025a5651e9e944b71485f53b0 Mon Sep 17 00:00:00 2001 From: amodhyh Date: Tue, 1 Sep 2026 13:19:57 +0530 Subject: [PATCH 9/9] feat(gguf): implement Llama-4 MoE routing and tensor filtering - Clamp RoPE max_position to 8192 to prevent massive pre-allocation OOMs - Enable MoE offload routing by setting `moe_enabled` based on expert count - Filter MoE expert tensors in `iter_gguf_weights` using include_moe_experts flags - Implement shared expert mapping and fused gate/up aggregation - Prevent 720 MiB GPU OOMs during dense pass by skipping expert allocations --- python/freetoken/models/gguf/config.py | 3 ++- python/freetoken/models/llama/gguf.py | 11 +++++++++-- 2 files changed, 11 insertions(+), 3 deletions(-) diff --git a/python/freetoken/models/gguf/config.py b/python/freetoken/models/gguf/config.py index bf9751bc6..8ec7b8b65 100644 --- a/python/freetoken/models/gguf/config.py +++ b/python/freetoken/models/gguf/config.py @@ -18,7 +18,8 @@ # reuses the model classes but a GGUF parse_config / iter_weights). GGUF_ARCH_TO_REGISTRY: dict[str, str] = { "gemma4": "Gemma4GGUFForCausalLM", - "llama" : "LlamaGGUFForCausalLM" + "llama" : "LlamaGGUFForCausalLM", + "llama4" : "LlamaGGUFForCausalLM" } diff --git a/python/freetoken/models/llama/gguf.py b/python/freetoken/models/llama/gguf.py index af3cb7aa7..c66e94217 100644 --- a/python/freetoken/models/llama/gguf.py +++ b/python/freetoken/models/llama/gguf.py @@ -47,7 +47,7 @@ def parse_gguf_config(shim:GgufConfigShim) -> ModelConfig: rotary_config = RotaryConfig( head_dim = gguf_metadata_llama.get(f"{arch}.rope.dimension_count",head_dim) , rotary_dim = gguf_metadata_llama.get(f"{arch}.rope.dimension_count",head_dim) , #identical to head_dim - max_position = gguf_metadata_llama.get(f"{arch}.context_length"), + max_position = min(gguf_metadata_llama.get(f"{arch}.context_length", 8192), 8192), base = gguf_metadata_llama.get(f"{arch}.rope.freq_base"), scaling = None ), @@ -61,6 +61,7 @@ def parse_gguf_config(shim:GgufConfigShim) -> ModelConfig: moe_intermediate_size = gguf_metadata_llama.get(f"{arch}.expert_feed_forward_length",dense_ffn_size), moe_weight_format = "q4_0", #4bit symmetric + moe_enabled = expert_count > 1, shared_expert_intermediate_size= dense_ffn_size if expert_count> 1 else 0 @@ -97,6 +98,12 @@ def iter_gguf_weights( for t in iter_gguf_tensors(model_path): name = t.name + is_moe_expert = "exps.weight" in name + if is_moe_expert and not include_moe_experts: + continue + if not is_moe_expert and not include_non_moe: + continue + # Global Standalone Tensors if name == "token_embd.weight": yield "model.embed_tokens.qweight", t.packed() @@ -219,7 +226,7 @@ def swap_linear(owner,attr_name, quant_type=GGML_Q4_0, has_bias=False): embed_scale=None ) - for layer in inner_model.layers: + for layer in inner_model.layers.op_list: swap_linear(layer.self_attn, "qkv_proj") swap_linear(layer.self_attn, "o_proj")