Feat/native gguf llama

python/freetoken/models/gguf/config.py

@@ -18,6 +18,8 @@
 # reuses the model classes but a GGUF parse_config / iter_weights).
 GGUF_ARCH_TO_REGISTRY: dict[str, str] = {
     "gemma4": "Gemma4GGUFForCausalLM",
+    "llama" : "LlamaGGUFForCausalLM",
+    "llama4" : "LlamaGGUFForCausalLM"
 }
 
 

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 <turn|>; prefer it as eos so chat generation halts
     # (the formal <eos> is also a stop id, see gguf_eos_token_ids).
     turn_end = "<turn|>" if "<turn|>" 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", "<bos>"),
-        eos_token=turn_end or tok_for("eos_token_id", "<eos>"),
+        eos_token=llama_eot or turn_end or tok_for("eos_token_id", "<eos>"),
         unk_token=tok_for("unknown_token_id", "<unk>"),
         pad_token=tok_for("padding_token_id", "<pad>"),
     )
@@ -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 <unk>, wrongly adding it as a stop id).
-    for name in ("<eos>", "<turn|>"):
+    for name in ("<eos>", "<turn|>", "<|eot_id|>", "<|eot|>"):
         try:
             ids.add(tokens.index(name))
         except ValueError:

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"]

python/freetoken/models/llama/gguf.py

@@ -0,0 +1,249 @@
+""" Llama GGUF adapter: build the FreeToken ``ModelConfig`` from GGUF metadata.
+"""
+from torch import Tensor
+from typing import Iterator
+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:
+    """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(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,
+        head_dim = head_dim,
+        hidden_size = hidden_size,      
+
+
+# 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 = 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 = min(gguf_metadata_llama.get(f"{arch}.context_length", 8192), 8192),                                      
+            base = gguf_metadata_llama.get(f"{arch}.rope.freq_base"),                                           
+            scaling = None                                 
+        ),      
+
+        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    
+        moe_enabled = expert_count > 1,
+
+        shared_expert_intermediate_size= dense_ffn_size if expert_count> 1 else 0 
+
+    )
+
+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,
+        *,
+        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."""
+
+    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
+
+        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()
+            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()
+
+    # 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(
+                [slots["q"], slots["k"], slots["v"]], dim=0
+            )
+            del qkv_buf[layer]
+
+        gu = gate_up_buf.get(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: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.
+    """
+    # 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.op_list:                 
+        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:
+        swap_linear(model,"lm_head",quant_type=GGML_Q6_K)
+
+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"

python/freetoken/models/llama/model.py

@@ -11,14 +11,37 @@
 
 from .attention import LlamaAttention as LlamaAttn
 
+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
 
 
 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:
+            # 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, 
+                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,
@@ -27,7 +50,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")
@@ -39,7 +61,20 @@ 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:
+
+            # 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
+        else:
+            # standard dense routing
+            x = self.mlp.forward(x)
         return x, residual
 
 
@@ -75,8 +110,12 @@ 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)
         logits = self.lm_head.forward(output)
         return logits

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",