diff --git a/python/freetoken/attention/fi.py b/python/freetoken/attention/fi.py index c9e5853..6c8f9d1 100644 --- a/python/freetoken/attention/fi.py +++ b/python/freetoken/attention/fi.py @@ -57,7 +57,8 @@ class FIMetadata(BaseAttnMetadata): page_size: Literal[1] # currently only support page_size=1 pos_encoding_mode: str seq_lens_cpu: torch.Tensor # on cpu - dtype: torch.dtype + dtype: torch.dtype # KV-cache storage dtype (may be fp8) + q_dtype: torch.dtype # query dtype (model activation dtype, bf16) wrapper: BatchPrefillWithPagedKVCacheWrapper | BatchDecodeWithPagedKVCacheWrapper initialized: bool = False # fmt: on @@ -86,6 +87,8 @@ def __init__(self, config: ModelConfig) -> None: self.config = config self.kvcache = get_global_ctx().kv_cache + # Query dtype (bf16), distinct from self.kvcache.dtype when KV is fp8. + self.q_dtype = get_global_ctx().compute_dtype self.device = self.kvcache.device # fa2 split-KV prefill needs ``tmp_v <= qo_heads_local * padded_batch_size * # cta_tile_q * head_dim * 4`` bytes of scratch, where flashinfer's scheduler @@ -165,8 +168,7 @@ def _initialize_metadata_once(self, metadata: FIMetadata) -> None: page_size=metadata.page_size, pos_encoding_mode=metadata.pos_encoding_mode, seq_lens=metadata.seq_lens_cpu, - data_type=metadata.dtype, - q_data_type=metadata.dtype, + q_data_type=metadata.q_dtype, kv_data_type=metadata.dtype, non_blocking=True, ) @@ -182,7 +184,7 @@ def _initialize_metadata_once(self, metadata: FIMetadata) -> None: page_size=metadata.page_size, pos_encoding_mode=metadata.pos_encoding_mode, seq_lens=metadata.seq_lens_cpu, - q_data_type=metadata.dtype, + q_data_type=metadata.q_dtype, kv_data_type=metadata.dtype, non_blocking=True, causal=True, @@ -256,6 +258,7 @@ def prepare_metadata(self, batch: Batch) -> None: pos_encoding_mode="NONE", seq_lens_cpu=seq_len_cpu, dtype=self.kvcache.dtype, + q_dtype=self.q_dtype, wrapper=self.decode_wrappers if batch.is_decode else self.prefill_wrapper, ) diff --git a/python/freetoken/core.py b/python/freetoken/core.py index ef0a539..c26df06 100644 --- a/python/freetoken/core.py +++ b/python/freetoken/core.py @@ -173,6 +173,9 @@ class Context: moe_backend: BaseMoeBackend = field(init=False) moe_offload_cache: OffloadMoeCache | None = None kv_cache: BaseKVCachePool = field(init=False) + # Query/activation dtype (bf16). Distinct from kv_cache.dtype under fp8 KV; the engine + # sets it before building the attention backend, which reads it for plan()'s q_data_type. + compute_dtype: torch.dtype = field(init=False, default=torch.bfloat16) # Per-request recurrent state for GatedDeltaNet layers; set by the engine for # hybrid linear-attention models, otherwise None. linear_state_pool: LinearStatePool | None = None diff --git a/python/freetoken/engine/config.py b/python/freetoken/engine/config.py index 543012f..fba943f 100644 --- a/python/freetoken/engine/config.py +++ b/python/freetoken/engine/config.py @@ -58,6 +58,11 @@ class EngineConfig: cuda_graph_max_bs: int | None = None page_size: int = 1 memory_ratio: float = 0.9 + # KV-cache storage dtype (--kv-dtype). None -> use the model/activation ``dtype`` (bf16). + # fp8 (e4m3/e5m2) halves KV bytes/token, ~doubling the token budget that fits in memory, + # at some attention precision cost. Read via ``resolved_kv_dtype``; only the paged MHA/SWA + # KV slabs switch dtype -- queries stay in ``dtype`` and the fp8 kernels dequantize on read. + kv_dtype: torch.dtype | None = None # Hybrid GDN models default to the HybridRadixCache (cross-request GDN-state prefix reuse); # `--cache-type naive` opts out. linear_state_cache_ratio sizes the GDN snapshot cache as # ceil(ratio * max_running_req) extra slots. @@ -91,6 +96,12 @@ def model_config(self) -> ModelConfig: parse_config = _load_attr(spec.module, spec.parse_config) return parse_config(self.hf_config) + @property + def resolved_kv_dtype(self) -> torch.dtype: + """Storage dtype for the paged KV cache: the explicit ``--kv-dtype`` if set, else the + model dtype. Queries always stay in ``dtype``; this only sizes/quantizes the KV slabs.""" + return self.kv_dtype if self.kv_dtype is not None else self.dtype + @property def max_seq_len(self) -> int: if self.max_seq_len_override is not None: diff --git a/python/freetoken/engine/engine.py b/python/freetoken/engine/engine.py index ff3c985..cb05048 100644 --- a/python/freetoken/engine/engine.py +++ b/python/freetoken/engine/engine.py @@ -272,14 +272,19 @@ def _materialize_loaded_weight_state_dict( weights: Iterable[Tuple[str, torch.Tensor]], *, device: torch.device, + cpu_offload_keys: frozenset[str] = frozenset(), ) -> Dict[str, torch.Tensor]: state_dict: Dict[str, torch.Tensor] = {} for key, weight in weights: + # A model may keep specific (large, rarely-touched) weights in host RAM -- e.g. the + # Gemma-4 E-series per-layer-embedding table (~4.7 GB). Those stay on CPU so they never + # consume VRAM; the model's forward gathers from them on the host. + dst = torch.device("cpu") if key in cpu_offload_keys else device expected = model_state.get(key) if expected is None: - state_dict[key] = weight.to(device=device) + state_dict[key] = weight.to(device=dst) else: - state_dict[key] = weight.to(device=device, dtype=expected.dtype) + state_dict[key] = weight.to(device=dst, dtype=expected.dtype) return state_dict @@ -302,6 +307,8 @@ def __init__(self, config: EngineConfig): self.stream = torch.cuda.Stream() torch.cuda.set_stream(self.stream) self.dtype = config.dtype + # Paged KV storage dtype (may be fp8 via --kv-dtype); queries/activations stay self.dtype. + self.kv_dtype = config.resolved_kv_dtype self.config = config # retained for runtime cache rebuild (rebuild_runtime_cache) # KV pool family fixed at construction from the model config: its classmethods own the # page-token geometry and cost arithmetic the engine needs BEFORE the pool exists @@ -345,7 +352,7 @@ def __init__(self, config: EngineConfig): self.num_pages = self._pool_cls.solve_num_pages(config, available_memory) num_tokens = self.num_pages * config.page_size self.ctx.kv_cache = self.kv_cache = create_kv_pool( - config, self.num_pages, device=self.device, dtype=self.dtype + config, self.num_pages, device=self.device, dtype=self.kv_dtype ) # ======================= Linear (GatedDeltaNet) state initialization ======================== @@ -379,6 +386,9 @@ def __init__(self, config: EngineConfig): self.kv_cache.attach_page_table(self.page_table) # ======================= Attention & MoE backend initialization ======================== + # Query/activation dtype for attention (the fp8-KV path needs it distinct from the KV + # storage dtype); backends read it off the global context in their constructors. + self.ctx.compute_dtype = self.dtype self.ctx.attn_backend = self.attn_backend = create_attention_backend( config.attention_backend, config.model_config ) @@ -405,13 +415,25 @@ def __init__(self, config: EngineConfig): if self.linear_state_pool is not None: self.dummy_req.linear_slot_idx = self.linear_state_pool.padding_slot self.page_table[self.dummy_req.table_idx].fill_(num_tokens) # point to dummy page + # A model may forbid CUDA graphs when its forward does a host-side gather that capture + # cannot record (e.g. Gemma-4 E-series with FREETOKEN_GEMMA4_PLE_CPU=1). Force graphs + # off rather than crashing during capture on the default flags. + graph_bs = config.cuda_graph_bs + graph_max_bs = config.cuda_graph_max_bs + if not getattr(self.model, "supports_cuda_graph", True) and graph_max_bs != 0: + logger.info_rank0( + "CUDA graphs disabled: the model's forward uses a host-resident weight " + "gather that graph capture cannot record." + ) + # Mirror `--cuda-graph-max-bs 0` (bs=None so _determine_cuda_graph_bs yields []). + graph_bs, graph_max_bs = None, 0 self.graph_runner = GraphRunner( stream=self.stream, device=self.device, model=self.model, attn_backend=self.attn_backend, - cuda_graph_bs=config.cuda_graph_bs, - cuda_graph_max_bs=config.cuda_graph_max_bs, + cuda_graph_bs=graph_bs, + cuda_graph_max_bs=graph_max_bs, free_memory=init_free_memory, max_seq_len=aligned_max_seq_len, vocab_size=config.model_config.vocab_size, @@ -456,6 +478,9 @@ def _load_weight_state_dict(self, config: EngineConfig) -> Dict[str, torch.Tenso # _materialize casts each loaded tensor to its model-param dtype (model_state), so # models declaring per-tensor dtypes (e.g. DSV4's mixed fp8/fp32/bf16) are preserved; # offload models exclude experts (served from the offload cache, not dense weights). + cpu_offload_keys = frozenset( + getattr(self.model, "cpu_offloaded_weight_keys", lambda: ())() + ) return _materialize_loaded_weight_state_dict( model_state, load_weight( @@ -464,6 +489,7 @@ def _load_weight_state_dict(self, config: EngineConfig) -> Dict[str, torch.Tenso include_moe_experts=not is_offload_moe_backend(config.moe_backend), ), device=self.device, + cpu_offload_keys=cpu_offload_keys, ) def _resolve_auto_moe_cache_size(self, config: EngineConfig, banks) -> tuple[int, int, bool]: diff --git a/python/freetoken/kvcache/base.py b/python/freetoken/kvcache/base.py index ae8cf9e..5655dc3 100644 --- a/python/freetoken/kvcache/base.py +++ b/python/freetoken/kvcache/base.py @@ -22,11 +22,13 @@ def spec_kv_bytes_per_token(spec, config) -> int: per-spec arithmetic -- pool families compose it over THEIR OWN groups; no family branching here. (2 bytes/elem == the torch.bfloat16 dsa_pool.DSAKVCache._alloc hardcodes; keep the two in lockstep if the slab dtype ever changes.)""" + # KV slabs use the (possibly fp8) KV dtype; duck-typed test configs lack the property. + kv_dtype = getattr(config, "resolved_kv_dtype", None) or config.dtype per_token = ( (1 if spec.mla else 2) # MLA latent groups store one slab (V aliases K) * spec.head_dim * div_even(spec.num_kv_heads, config.tp_info.size, allow_replicate=True) - * config.dtype.itemsize + * kv_dtype.itemsize * spec.num_layers ) return per_token + spec.index_head_dim * spec.num_index_layers * 2 diff --git a/python/freetoken/kvcache/mha_pool.py b/python/freetoken/kvcache/mha_pool.py index 8ed280b..c6f9232 100644 --- a/python/freetoken/kvcache/mha_pool.py +++ b/python/freetoken/kvcache/mha_pool.py @@ -127,6 +127,13 @@ def store_kv( from freetoken.kernel import store_cache dense = self._dense(layer_id) + # fp8 KV cache: quantize the incoming bf16 k/v to the pool dtype here (scale 1.0 -- + # post-RoPE k/v magnitudes sit well inside e4m3/e5m2 range), so store_cache stays a + # same-width byte copy and the fp8 attention kernels dequantize on read. + buf_dtype = self._kv_buffer.dtype + if k.dtype != buf_dtype: + k = k.to(buf_dtype) + v = v.to(buf_dtype) store_cache( k_cache=self._k_buffer[dense].view(self._storage_shape), v_cache=self._v_buffer[dense].view(self._storage_shape), diff --git a/python/freetoken/layers/gguf.py b/python/freetoken/layers/gguf.py index ac49b1a..ffaa9a2 100644 --- a/python/freetoken/layers/gguf.py +++ b/python/freetoken/layers/gguf.py @@ -22,6 +22,8 @@ GGML_F32, GGML_NAME, GGML_Q4_0, + GGML_Q4_K, + GGML_Q5_K, GGML_Q6_K, GGML_Q8_0, row_bytes, @@ -29,12 +31,15 @@ from .base import BaseOP -# ggml type groups for kernel dispatch (subset we build kernels for). +# ggml type groups for kernel dispatch (subset we build kernels for). The borrowed ggml +# CUDA kernels dispatch on the ggml type at runtime (mmvq.cuh/mmq.cuh cover Q4_0/Q8_0 and +# the K-quants Q4_K/Q5_K/Q6_K), so the Q4_K/Q6_K mixture of a q4_K_M llama GGUF runs +# natively -- no bf16 copy of the projection weights. _UNQUANTIZED = {GGML_F32, GGML_F16, GGML_BF16} # standard + k-quants: both an MMVQ (small-batch GEMV) and MMQ (large-batch) kernel exist. -_MMVQ = {GGML_Q4_0, GGML_Q8_0, GGML_Q6_K} -_MMQ = {GGML_Q4_0, GGML_Q8_0, GGML_Q6_K} -_DEQUANT = {GGML_Q4_0, GGML_Q8_0, GGML_Q6_K} +_MMVQ = {GGML_Q4_0, GGML_Q8_0, GGML_Q4_K, GGML_Q5_K, GGML_Q6_K} +_MMQ = {GGML_Q4_0, GGML_Q8_0, GGML_Q4_K, GGML_Q5_K, GGML_Q6_K} +_DEQUANT = {GGML_Q4_0, GGML_Q8_0, GGML_Q4_K, GGML_Q5_K, GGML_Q6_K} # Below this token count, the MMVQ GEMV kernel wins (matches vLLM's heuristic). _MMVQ_SAFE = 6 diff --git a/python/freetoken/models/config.py b/python/freetoken/models/config.py index f6105e1..b545d62 100644 --- a/python/freetoken/models/config.py +++ b/python/freetoken/models/config.py @@ -269,10 +269,27 @@ class ModelConfig: output_multiplier: float | None = None vision_config: Any | None = None image_token_id: int | None = None + # ----- Gemma-4 E-series (PLE) extensions (default keeps other models intact) ----- + # Per-Layer Embeddings: an auxiliary embedding + per-layer projection that injects a + # token/context-dependent residual into every decoder layer (Gemma-4 E2B/E4B). 0 = off. + per_layer_hidden_size: int = 0 + per_layer_vocab_size: int = 0 + # KV-layer sharing: the last ``num_kv_shared_layers`` layers carry no K/V projections and + # reuse the K/V states of the last non-shared layer of the same attention type. 0 = off. + num_kv_shared_layers: int = 0 + # Double-wide MLP: KV-shared layers use 2x intermediate_size (E2B). Applies only to the + # KV-shared layers, mirroring transformers' Gemma4TextMLP. + use_double_wide_mlp: bool = False attention_groups: Tuple[AttentionGroupConfig, ...] = () has_attn_bias: bool = False has_router_bias: bool = False moe_weight_format: str | None = None + # Native-GGUF dense projections: {freetoken_module_base -> ggml_type} for every packed + # weight (attention q/k/v/o, MLP gate/up/down, token embedding). Unlike gemma4's pure + # Q4_0 GGUFs, a q4_K_M checkpoint mixes per-tensor K-quant types (Q4_K projections with + # Q6_K on some attn_v/ffn_down), so the module-swap can't hardcode one type -- it reads + # each tensor's type from here. Set by a GGUF parse_config; None for safetensors models. + gguf_quant_types: dict[str, int] | None = None swiglu_limit: float | None = None hidden_act_alpha: float = 1.702 # Full DeepseekV4Args payload for the DSV4-specific machinery (MLA sparse attention, @@ -366,6 +383,33 @@ def linear_attention_group(self) -> LinearGatedDeltaGroupConfig | None: def is_swa_layer(self, layer_id: int) -> bool: return isinstance(self.attention_group_for_layer(layer_id), SWAAttentionGroupConfig) + def first_kv_shared_layer_idx(self) -> int: + """First layer index in the KV-shared tail (Gemma-4 E-series). Equals num_layers + (i.e. no shared layers) when num_kv_shared_layers == 0.""" + if self.num_kv_shared_layers <= 0: + return self.num_layers + return self.num_layers - self.num_kv_shared_layers + + def is_kv_shared_layer(self, layer_id: int) -> bool: + """True for E-series layers that carry no K/V projections and reuse another + layer's K/V states (the last ``num_kv_shared_layers`` layers).""" + return self.num_kv_shared_layers > 0 and layer_id >= self.first_kv_shared_layer_idx() + + def is_kv_source_layer(self, layer_id: int) -> bool: + """True for the last non-shared layer of each attention type -- the layer whose + K/V the shared layers of that type reuse. No-op unless KV sharing is enabled.""" + if self.num_kv_shared_layers <= 0: + return False + first_shared = self.first_kv_shared_layer_idx() + if layer_id >= first_shared: + return False + my_name = self.attention_group_for_layer(layer_id).name + # source == the highest-index non-shared layer sharing this layer's group name + for other in range(first_shared - 1, layer_id, -1): + if self.attention_group_for_layer(other).name == my_name: + return False + return True + def is_linear_layer(self, layer_id: int) -> bool: return isinstance( self.attention_group_for_layer(layer_id), diff --git a/python/freetoken/models/gemma4/attention.py b/python/freetoken/models/gemma4/attention.py index 9103bf3..58a694e 100644 --- a/python/freetoken/models/gemma4/attention.py +++ b/python/freetoken/models/gemma4/attention.py @@ -15,7 +15,17 @@ class Gemma4Attention(BaseOP): - """Gemma 4 attention for one full-context or SWA layer.""" + """Gemma 4 attention for one full-context or SWA layer. + + Supports the E-series (E2B/E4B) KV-layer-sharing scheme: the last + ``num_kv_shared_layers`` layers carry no K/V projections and reuse the K/V + states produced by the last non-shared layer of the *same* attention type + (sliding vs full). Sharing is wired via a per-forward stash (``ctx`` buffers + keyed by attention-group name): a "source" layer writes its post-rope K and + post-norm V into the stash, and each shared layer reads them and feeds them + to its own KV cache slot. The stash is a stable, model-owned buffer so it is + captured/replayed correctly inside the decode CUDA graph. + """ def __init__(self, config: ModelConfig, layer_id: int): self.layer_id = layer_id @@ -29,19 +39,31 @@ def __init__(self, config: ModelConfig, layer_id: int): self.num_qo_heads = config.num_qo_heads self.k_eq_v = isinstance(group, FullAttentionGroupConfig) and group.k_eq_v + # --- E-series KV-layer sharing bookkeeping (no-op when num_kv_shared_layers == 0) --- + self._group_name = group.name # stash key: shared + source layers of a type must match + self.is_kv_shared = config.is_kv_shared_layer(layer_id) + self.is_kv_source = config.is_kv_source_layer(layer_id) + # Filled in by Gemma4Model after build: the model-owned {group_name: (k_buf, v_buf)} + # stash and the live token count. None on non-E-series models. + self._kv_stash: dict[str, tuple[torch.Tensor, torch.Tensor]] | None = None + self.q_dim = self.num_qo_heads * self.head_dim self.kv_dim = self.num_kv_heads * self.head_dim - self.qkv_proj = LinearQKVMerged( - config.hidden_size, - self.head_dim, - self.num_qo_heads, - self.num_kv_heads, - has_bias=False, - ) + if self.is_kv_shared: + # Shared layers project Q only (K/V come from the stash of their type). + self.q_proj = LinearReplicated(config.hidden_size, self.q_dim, has_bias=False) + else: + self.qkv_proj = LinearQKVMerged( + config.hidden_size, + self.head_dim, + self.num_qo_heads, + self.num_kv_heads, + has_bias=False, + ) + self.k_norm = GemmaRMSNorm(self.head_dim, eps=config.rms_norm_eps) + self.v_norm = GemmaRMSNorm(self.head_dim, eps=config.rms_norm_eps, with_scale=False) self.o_proj = LinearReplicated(self.q_dim, config.hidden_size, has_bias=False) self.q_norm = GemmaRMSNorm(self.head_dim, eps=config.rms_norm_eps) - self.k_norm = GemmaRMSNorm(self.head_dim, eps=config.rms_norm_eps) - self.v_norm = GemmaRMSNorm(self.head_dim, eps=config.rms_norm_eps, with_scale=False) self.attn_spec = AttentionSpec( sliding_window=group.sliding_window if self.is_swa else None, sm_scale=config.attn_sm_scale, @@ -58,6 +80,15 @@ def __init__(self, config: ModelConfig, layer_id: int): ), ) + def _apply_rope_q(self, positions: torch.Tensor, q: torch.Tensor) -> torch.Tensor: + positions = positions.reshape(-1) + if positions.device != q.device or positions.dtype != torch.long: + positions = positions.to(device=q.device, dtype=torch.long) + q_view = q.contiguous().view(q.shape[0], -1) + # rotary.forward rotates q and k together; pass q as both when we only need q. + self.rotary.forward(positions, q_view, q_view) + return q_view.view_as(q) + def _apply_rope( self, positions: torch.Tensor, @@ -78,24 +109,39 @@ def forward(self, x: torch.Tensor) -> torch.Tensor: positions = ctx.batch.positions T = x.shape[0] - qkv = self.qkv_proj.forward(x) - q_lin, k_lin, v_lin = qkv.split( - (self.q_dim, self.kv_dim, self.kv_dim), - dim=-1, - ) - del qkv - q = q_lin.view(T, self.num_qo_heads, self.head_dim) - k = k_lin.view(T, self.num_kv_heads, self.head_dim) - v = v_lin.view(T, self.num_kv_heads, self.head_dim) + if self.is_kv_shared: + q_lin = self.q_proj.forward(x) + q = q_lin.view(T, self.num_qo_heads, self.head_dim) + q = self.q_norm.forward(q) + q = self._apply_rope_q(positions, q) + k_buf, v_buf = self._kv_stash[self._group_name] + k = k_buf[:T] + v = v_buf[:T] + else: + qkv = self.qkv_proj.forward(x) + q_lin, k_lin, v_lin = qkv.split( + (self.q_dim, self.kv_dim, self.kv_dim), + dim=-1, + ) + del qkv + q = q_lin.view(T, self.num_qo_heads, self.head_dim) + k = k_lin.view(T, self.num_kv_heads, self.head_dim) + v = v_lin.view(T, self.num_kv_heads, self.head_dim) + + q = self.q_norm.forward(q) + k = self.k_norm.forward(k) + v = self.v_norm.forward(v) - q = self.q_norm.forward(q) - k = self.k_norm.forward(k) - v = self.v_norm.forward(v) + q, k = self._apply_rope(positions, q, k) - q, k = self._apply_rope(positions, q, k) + k = k.reshape(T, self.num_kv_heads * self.head_dim) + v = v.reshape(T, self.num_kv_heads * self.head_dim) + # A "source" layer publishes its full K/V for the shared layers of its type. + if self.is_kv_source and self._kv_stash is not None: + k_buf, v_buf = self._kv_stash[self._group_name] + k_buf[:T].copy_(k) + v_buf[:T].copy_(v) - k = k.reshape(T, self.num_kv_heads * self.head_dim) - v = v.reshape(T, self.num_kv_heads * self.head_dim) o = ctx.attn_backend.forward( q.contiguous(), k.contiguous(), diff --git a/python/freetoken/models/gemma4/config.py b/python/freetoken/models/gemma4/config.py index 057e0a5..6203d3a 100644 --- a/python/freetoken/models/gemma4/config.py +++ b/python/freetoken/models/gemma4/config.py @@ -180,6 +180,12 @@ def parse_config(hf_config: Any) -> ModelConfig: embedding_scale=float(cfg.hidden_size) ** 0.5, vision_config=_parse_vision_config(top_cfg, cfg.hidden_size), image_token_id=getattr(top_cfg, "image_token_id", None), + # Gemma-4 E-series (E2B/E4B): Per-Layer Embeddings + KV-layer sharing + double-wide MLP. + # These keys are absent on the dense/MoE Gemma-4 line, so getattr defaults keep it off. + per_layer_hidden_size=int(getattr(cfg, "hidden_size_per_layer_input", 0) or 0), + per_layer_vocab_size=int(getattr(cfg, "vocab_size_per_layer_input", 0) or 0), + num_kv_shared_layers=int(getattr(cfg, "num_kv_shared_layers", 0) or 0), + use_double_wide_mlp=bool(getattr(cfg, "use_double_wide_mlp", False)), attention_groups=( FullAttentionGroupConfig( name="full", diff --git a/python/freetoken/models/gemma4/gguf.py b/python/freetoken/models/gemma4/gguf.py index 437822b..50d7b16 100644 --- a/python/freetoken/models/gemma4/gguf.py +++ b/python/freetoken/models/gemma4/gguf.py @@ -57,12 +57,26 @@ def g(key: str): raise KeyError(f"missing GGUF metadata key gemma4.{key}") return val + def g_opt(key, default=None): + val = m.get(f"gemma4.{key}") + return default if val is None else val + + # Dense (E-series E2B/E4B) vs MoE (12B/26B): dense GGUFs carry no expert metadata. + num_experts = int(g_opt("expert_count", 0) or 0) + is_moe = num_experts > 0 + num_layers = int(g("block_count")) hidden = int(g("embedding_length")) num_qo_heads = int(g("attention.head_count")) - kv_per_layer = g("attention.head_count_kv") # per-layer list - # True -> sliding-window (SWA) layer, False -> full attention. - swa_pattern = [bool(x) for x in g("attention.sliding_window_pattern")] + kv_per_layer = g("attention.head_count_kv") # per-layer list OR a single scalar + # True -> sliding-window (SWA) layer, False -> full attention. Flatten single-element + # wrappers the reader may leave on each entry. + def _flat_bool(x): + while isinstance(x, (list, tuple)) and len(x) == 1: + x = x[0] + return bool(x) + + swa_pattern = [_flat_bool(x) for x in g("attention.sliding_window_pattern")] assert len(swa_pattern) == num_layers, "sliding_window_pattern length != block_count" swa_layer_ids = tuple(i for i, is_swa in enumerate(swa_pattern) if is_swa) @@ -70,8 +84,28 @@ def g(key: str): swa_head_dim = int(g("attention.key_length_swa")) full_head_dim = int(g("attention.key_length")) - swa_kv = int(kv_per_layer[swa_layer_ids[0]]) if swa_layer_ids else int(kv_per_layer[0]) - full_kv = int(kv_per_layer[full_layer_ids[0]]) if full_layer_ids else int(kv_per_layer[0]) + + def _kv_for(layer_ids): + if isinstance(kv_per_layer, (list, tuple)): + return int(kv_per_layer[layer_ids[0]] if layer_ids else kv_per_layer[0]) + return int(kv_per_layer) + + swa_kv = _kv_for(swa_layer_ids) + full_kv = _kv_for(full_layer_ids) + + # feed_forward_length is a per-layer list on the E-series (KV-shared tail is 2x wide); + # take the base width and flag double-wide when any layer is larger. Older MoE GGUFs + # store a single scalar. + _ffn = g("feed_forward_length") + _ffn = [int(x) for x in _ffn] if isinstance(_ffn, (list, tuple)) else [int(_ffn)] + intermediate = min(_ffn) + use_double_wide_mlp = max(_ffn) > intermediate + # Per-Layer Embeddings + KV-layer sharing (E-series only; absent on the MoE line). + ple_hidden = int(g_opt("embedding_length_per_layer_input", 0) or 0) + num_kv_shared_layers = int(g_opt("attention.shared_kv_layers", 0) or 0) + # Full layers reuse K as V only when the GGUF ships no separate value projection + # (older MoE gemma4); the E-series declares value_length and ships attn_v. + full_k_eq_v = g_opt("attention.value_length") is None max_pos = int(g("context_length")) full_rotary = RotaryConfig( @@ -100,24 +134,29 @@ def g(key: str): head_dim=full_head_dim, hidden_size=hidden, vocab_size=int(shim.vocab_size), - intermediate_size=int(g("feed_forward_length")), hidden_act="gelu_tanh", rms_norm_eps=float(g("attention.layer_norm_rms_epsilon")), tie_word_embeddings=bool(shim.tie_word_embeddings), rotary_config=full_rotary, - num_experts=int(g("expert_count")), - num_experts_per_tok=int(g("expert_used_count")), - moe_intermediate_size=int(g("expert_feed_forward_length")), + intermediate_size=intermediate, + num_experts=num_experts, + num_experts_per_tok=int(g_opt("expert_used_count", 0) or 0), + moe_intermediate_size=int(g_opt("expert_feed_forward_length", 0) or 0), norm_topk_prob=True, model_type="gemma4", architectures=list(shim.architectures), - moe_enabled=True, - expert_quant="q4_0", + moe_enabled=is_moe, + expert_quant="q4_0" if is_moe else "none", moe_weight_format="q4_0", use_qk_norm=True, attn_sm_scale=1.0, final_logit_softcapping=float(g("final_logit_softcapping")), embedding_scale=float(hidden) ** 0.5, + # Gemma-4 E-series (dense): Per-Layer Embeddings + KV-layer sharing + double-wide MLP. + per_layer_hidden_size=ple_hidden, + per_layer_vocab_size=int(shim.vocab_size) if ple_hidden else 0, + num_kv_shared_layers=num_kv_shared_layers, + use_double_wide_mlp=use_double_wide_mlp, attention_groups=( FullAttentionGroupConfig( name="full", @@ -125,8 +164,8 @@ def g(key: str): num_kv_heads=full_kv, head_dim=full_head_dim, rotary_config=full_rotary, - # Full layers ship no v_proj in the GGUF (k reused as v). - k_eq_v=True, + # Older MoE gemma4 GGUFs ship no v_proj (k reused as v); the E-series ships attn_v. + k_eq_v=full_k_eq_v, ), SWAAttentionGroupConfig( name="swa", @@ -157,6 +196,8 @@ def g(key: str): "post_ffw_norm_1.weight": "feed_forward.post_feedforward_layernorm_1.weight", "post_ffw_norm_2.weight": "feed_forward.post_feedforward_layernorm_2.weight", "layer_output_scale.weight": "feed_forward.layer_scalar", + # E-series Per-Layer-Embedding step (post_norm = the post-PLE RMSNorm). + "post_norm.weight": "post_per_layer_input_norm.weight", "ffn_gate_inp.weight": "feed_forward.router.proj.weight", "ffn_gate_inp.scale": "feed_forward.router.scale", "ffn_down_exps.scale": "feed_forward.router.per_expert_scale", @@ -203,22 +244,28 @@ def iter_gguf_weights( from freetoken.models.gguf.reader import iter_gguf_tensors from freetoken.utils import cached_load_hf_config - assert not include_moe_experts, ( - "gemma4 GGUF stores experts as Q4_0 and only supports the offload backend; " - "experts are loaded into the offload cache via load_q4_0_expert_sources()." + config = parse_gguf_config(cached_load_hf_config(model_path)) + # MoE gemma4 keeps its Q4_0 experts in the offload cache (include_moe_experts must be + # False). Dense E-series has no experts, so the flag is irrelevant there. + assert not (config.moe_enabled and include_moe_experts), ( + "gemma4 MoE GGUF experts are served from the offload cache " + "(include_moe_experts must be False)." ) assert include_non_moe _require_tp1("weight loading") # Full-attention layers ship no attn_v (k reused as v); SWA layers do. Knowing which # is which lets us emit the fused qkv as soon as its parts are present. - config = parse_gguf_config(cached_load_hf_config(model_path)) k_eq_v_layers = { lid for lid in range(config.num_layers) if isinstance(config.attention_group_for_layer(lid), FullAttentionGroupConfig) and config.attention_group_for_layer(lid).k_eq_v } + # E-series KV-shared tail layers carry q only (k/v reused from another layer). + kv_shared_layers = { + lid for lid in range(config.num_layers) if config.is_kv_shared_layer(lid) + } # Per-layer fusion buffers: layer -> {slot: packed[out, row_bytes]}. qkv_buf: dict[int, dict[str, torch.Tensor]] = {} @@ -237,6 +284,17 @@ def layer_of(name: str) -> int: continue if name == "rope_freqs.weight": continue # rope frequencies recomputed in-engine + # E-series Per-Layer Embeddings (top-level tensors). The big per-layer table stays a + # packed Q6_K GGUFEmbedding on GPU (~1.9 GB); the small projection/norm dequant to bf16. + if name == "per_layer_token_embd.weight": + yield "model.embed_tokens_per_layer.qweight", t.packed() + continue + if name == "per_layer_model_proj.weight": + yield "model.per_layer_model_projection.weight", _to_bf16(t) + continue + if name == "per_layer_proj_norm.weight": + yield "model.per_layer_projection_norm.weight", _to_bf16(t) + continue if not name.startswith("blk."): continue if any(name.endswith(sfx) for sfx in _EXPERT_SUFFIXES): @@ -246,6 +304,23 @@ def layer_of(name: str) -> int: suffix = name.split(".", 2)[2] # after "blk.N." base = f"model.layers.{layer}" + # E-series KV-shared layers: reuse another layer's K/V -> the model has no k/v/k_norm + # here. Drop those tensors and yield q_proj standalone (not fused into qkv). + if layer in kv_shared_layers: + if suffix in ("attn_k.weight", "attn_v.weight", "attn_k_norm.weight"): + continue + if suffix == "attn_q.weight": + yield f"{base}.self_attn.q_proj.qweight", t.packed() + continue + + # E-series Per-Layer-Embedding projections (tiny Q4_0 -> bf16 LinearReplicated). + if suffix == "inp_gate.weight": + yield f"{base}.per_layer_input_gate.weight", _to_bf16(t) + continue + if suffix == "proj.weight": + yield f"{base}.per_layer_projection.weight", _to_bf16(t) + continue + if suffix in _LAYER_SCALAR_MAP: rel = _LAYER_SCALAR_MAP[suffix] tensor = _to_bf16(t) @@ -365,8 +440,19 @@ def swap_linear(owner, attr, quant_type=GGML_Q4_0): ) inner.embed_tokens = embed - for layer in inner.layers.op_list: - swap_linear(layer.self_attn, "qkv_proj") + # E-series: the per-layer-embedding table is a packed Q6_K GGUFEmbedding on GPU (~1.9 GB) + # instead of the bf16 host-RAM CpuPerLayerEmbedding used for the safetensors path. + if getattr(inner, "_has_ple", False): + inner.embed_tokens_per_layer = GGUFEmbedding( + num_embeddings=config.per_layer_vocab_size, + embedding_dim=config.num_layers * config.per_layer_hidden_size, + quant_type=GGML_Q6_K, + embed_scale=float(config.per_layer_hidden_size) ** 0.5, + ) + + for layer_id, layer in enumerate(inner.layers.op_list): + # KV-shared layers carry a standalone q_proj (no fused qkv); others carry qkv_proj. + swap_linear(layer.self_attn, "q_proj" if config.is_kv_shared_layer(layer_id) else "qkv_proj") swap_linear(layer.self_attn, "o_proj") swap_linear(layer.feed_forward.shared_mlp, "gate_up_proj") swap_linear(layer.feed_forward.shared_mlp, "down_proj") diff --git a/python/freetoken/models/gemma4/model.py b/python/freetoken/models/gemma4/model.py index 66fcf41..df5bc36 100644 --- a/python/freetoken/models/gemma4/model.py +++ b/python/freetoken/models/gemma4/model.py @@ -1,12 +1,15 @@ from __future__ import annotations +import os from typing import TYPE_CHECKING import torch +import torch.nn.functional as F from freetoken.core import get_global_ctx from freetoken.layers import ( BaseOP, GemmaRMSNorm, + LinearReplicated, OPList, ParallelLMHead, VocabParallelEmbedding, @@ -22,18 +25,63 @@ if TYPE_CHECKING: from freetoken.models.config import ModelConfig +# Upper bound on tokens in a single forward (prefill chunk / decode batch). The E-series +# KV-sharing stash is allocated once at this size so its address is stable across CUDA-graph +# capture and replay. Matches the default --max-prefill-length. +_STASH_MAX_TOKENS = 8192 + +# The E-series per-layer-embedding table's state-dict key (kept in host RAM). +PER_LAYER_EMBED_KEY = "model.embed_tokens_per_layer.weight" + + +class CpuPerLayerEmbedding(BaseOP): + """Host-resident per-layer-embedding table (Gemma-4 E-series PLE). The table is huge + (vocab x num_layers*ple_dim ~= 4.7 GB for E2B) but each step only gathers a handful of + rows, so it stays in CPU RAM and the gather runs on the host -- freeing VRAM for the rest + of the model. Only viable because the E-series runs with CUDA graphs disabled (the CPU + gather + H2D copy cannot be captured in a graph).""" + + def __init__(self, num_embeddings: int, embedding_dim: int, embed_scale: float): + # Built on the meta device; assign-load replaces this with the real CPU tensor + # (engine keeps it on CPU via cpu_offloaded_weight_keys). + self.weight = torch.empty(num_embeddings, embedding_dim) + self._embed_scale = embed_scale + + def forward(self, input_ids: torch.Tensor) -> torch.Tensor: + idx = input_ids.to("cpu", dtype=torch.long) + rows = F.embedding(idx, self.weight) # [T, num_layers*ple_dim] on CPU + return rows * self._embed_scale # caller moves to the compute device + class Gemma4DecoderLayer(BaseOP): """Gemma 4 decoder block: attention sandwich + feed-forward sandwich, scaled by a per-layer ``layer_scalar``. The feed-forward is the dual (shared MLP || routed MoE) - branch for MoE checkpoints, or a single dense MLP branch for dense checkpoints.""" + branch for MoE checkpoints, or a single dense MLP branch for dense checkpoints. + + Gemma-4 E-series (E2B/E4B) additionally: (1) KV-shared tail layers get a 2x-wide dense + MLP (``use_double_wide_mlp``), and (2) every layer runs a Per-Layer-Embedding (PLE) step + -- gate -> gelu -> multiply by the layer's per-layer input -> project -> norm -> residual + -- inserted after the feed-forward sandwich and before the per-layer scale.""" def __init__(self, config: ModelConfig, layer_id: int): self._layer_id = layer_id self.self_attn = Gemma4Attention(config, layer_id) - self.feed_forward = ( - Gemma4MLP(config, layer_id) if config.is_moe else Gemma4DenseMLP(config) - ) + + self.has_ple = config.per_layer_hidden_size > 0 + if config.is_moe: + self.feed_forward = Gemma4MLP(config, layer_id) + else: + # E-series: KV-shared layers use a double-wide MLP; PLE runs before layer_scalar, + # so the dense MLP defers the scale to the decoder. + double_wide = ( + config.use_double_wide_mlp and config.is_kv_shared_layer(layer_id) + ) + inter = config.intermediate_size * (2 if double_wide else 1) + self.feed_forward = Gemma4DenseMLP( + config, + intermediate_size=inter, + apply_scalar=not self.has_ple, + ) eps = config.rms_norm_eps H = config.hidden_size @@ -41,6 +89,12 @@ def __init__(self, config: ModelConfig, layer_id: int): self.post_attention_layernorm = GemmaRMSNorm(H, eps=eps) self.pre_feedforward_layernorm = GemmaRMSNorm(H, eps=eps) + if self.has_ple: + P = config.per_layer_hidden_size + self.per_layer_input_gate = LinearReplicated(H, P, has_bias=False) + self.per_layer_projection = LinearReplicated(P, H, has_bias=False) + self.post_per_layer_input_norm = GemmaRMSNorm(H, eps=eps) + @nvtx_annotate("Layer_{}", layer_id_field="_layer_id") def forward(self, x: torch.Tensor) -> torch.Tensor: # --- attention sandwich --- @@ -49,7 +103,19 @@ def forward(self, x: torch.Tensor) -> torch.Tensor: h = self.self_attn.forward(h) h = self.post_attention_layernorm.forward(h) pre_ff, x = self.pre_feedforward_layernorm.forward_add_residual(h, residual) - return self.feed_forward.forward(pre_ff, x) + h = self.feed_forward.forward(pre_ff, x) + if not self.has_ple: + return h + + # --- Per-Layer-Embedding step (E-series), before the per-layer scale --- + ple_in = get_global_ctx()._gemma4_ple[:, self._layer_id, :] + g = self.per_layer_input_gate.forward(h) + g = F.gelu(g, approximate="tanh") + g = g * ple_in + g = self.per_layer_projection.forward(g) + g = self.post_per_layer_input_norm.forward(g) + h = h + g + return h * self.feed_forward.layer_scalar class Gemma4Model(BaseOP): @@ -65,14 +131,60 @@ def __init__(self, config: ModelConfig): self.norm = GemmaRMSNorm(config.hidden_size, eps=config.rms_norm_eps) self._image_token_id = config.image_token_id - def _merge_multimodal(self, input_ids: torch.Tensor, x: torch.Tensor) -> torch.Tensor: - """Scatter precomputed image soft-token embeddings at image-token positions. + # --- Gemma-4 E-series Per-Layer Embeddings (PLE) --- + self._has_ple = config.per_layer_hidden_size > 0 + # PLE per-layer table (~4.7 GB bf16 for E2B). Default: an on-GPU embedding, which is + # CUDA-graph-safe (the reason graphs are on by default). Opt in with + # FREETOKEN_GEMMA4_PLE_CPU=1 to keep it in host RAM on a VRAM-tight box -- but the + # per-forward host gather cannot be captured, so the engine then disables CUDA graphs + # for this model (see Gemma4ForCausalLM.supports_cuda_graph). The GGUF path replaces + # this with a compact Q6_K GGUFEmbedding on-GPU regardless (also graph-safe). + self._ple_on_cpu = os.getenv("FREETOKEN_GEMMA4_PLE_CPU", "0").strip().lower() in ( + "1", "true", "yes", "on" + ) + if self._has_ple: + P = config.per_layer_hidden_size + L = config.num_layers + per_layer_embed_cls = ( + CpuPerLayerEmbedding if self._ple_on_cpu else VocabParallelEmbedding + ) + self.embed_tokens_per_layer = per_layer_embed_cls( + num_embeddings=config.per_layer_vocab_size, + embedding_dim=L * P, + embed_scale=float(P) ** 0.5, + ) + self.per_layer_model_projection = LinearReplicated( + config.hidden_size, L * P, has_bias=False + ) + self.per_layer_projection_norm = GemmaRMSNorm(P, eps=config.rms_norm_eps) + self._ple_proj_scale = float(config.hidden_size) ** -0.5 + self._ple_combine_scale = 2.0**-0.5 + self._ple_hidden = P + self._num_layers = L + + # --- E-series KV-layer sharing: per-group stash buffers (lazy, fixed size) --- + self._kv_shared = config.num_kv_shared_layers > 0 + self._kv_stash: dict[str, tuple[torch.Tensor, torch.Tensor]] = {} + self._stash_kv_dim: dict[str, int] = {} + if self._kv_shared: + first_shared = config.first_kv_shared_layer_idx() + for layer_id in range(first_shared, config.num_layers): + group = config.attention_group_for_layer(layer_id) + self._stash_kv_dim[group.name] = group.num_kv_heads * group.head_dim + # share the (initially empty) stash dict with every attention module + for layer in self.layers.op_list: + layer.self_attn._kv_stash = self._kv_stash - ``mm_embeds`` (set by the scheduler from each request's vision features) is a - ``[num_image_tokens, hidden]`` tensor whose rows replace the placeholder - embeddings produced for ``image_token_id``. Only runs during prefill batches - that carry images; decode batches never do. - """ + def _ensure_stash(self, device: torch.device, dtype: torch.dtype) -> None: + if not self._kv_shared or self._kv_stash: + return + for name, kv_dim in self._stash_kv_dim.items(): + k = torch.zeros(_STASH_MAX_TOKENS, kv_dim, device=device, dtype=dtype) + v = torch.zeros(_STASH_MAX_TOKENS, kv_dim, device=device, dtype=dtype) + self._kv_stash[name] = (k, v) + + def _merge_multimodal(self, input_ids: torch.Tensor, x: torch.Tensor) -> torch.Tensor: + """Scatter precomputed image soft-token embeddings at image-token positions.""" batch = get_global_ctx().batch mm_embeds = getattr(batch, "mm_embeds", None) if mm_embeds is None or self._image_token_id is None: @@ -85,9 +197,30 @@ def _merge_multimodal(self, input_ids: torch.Tensor, x: torch.Tensor) -> torch.T ) return x.masked_scatter(mask.unsqueeze(-1), mm_embeds.to(x.dtype)) + def _compute_per_layer_inputs( + self, input_ids: torch.Tensor, inputs_embeds: torch.Tensor + ) -> torch.Tensor: + """PLE inputs [T, num_layers, ple_hidden]: token-identity embedding combined with a + context projection of ``inputs_embeds``. Mirrors transformers' get_per_layer_inputs + + project_per_layer_inputs.""" + T = input_ids.shape[0] + # embed_tokens_per_layer takes input_ids and returns [T, L*P]: an on-GPU + # VocabParallelEmbedding (default), a host-RAM CpuPerLayerEmbedding (opt-in), or a + # Q6_K GGUFEmbedding after the GGUF swap. The .to() is a no-op for the on-GPU variants. + tok = self.embed_tokens_per_layer.forward(input_ids).to(inputs_embeds.device).view( + T, self._num_layers, self._ple_hidden + ) + proj = self.per_layer_model_projection.forward(inputs_embeds) * self._ple_proj_scale + proj = proj.view(T, self._num_layers, self._ple_hidden) + proj = self.per_layer_projection_norm.forward(proj) + return (proj + tok) * self._ple_combine_scale + def forward(self, input_ids: torch.Tensor) -> torch.Tensor: x = self.embed_tokens.forward(input_ids) x = self._merge_multimodal(input_ids, x) + if self._has_ple: + self._ensure_stash(x.device, x.dtype) + get_global_ctx()._gemma4_ple = self._compute_per_layer_inputs(input_ids, x) for layer in self.layers.op_list: x = layer.forward(x) return self.norm.forward(x) @@ -115,15 +248,27 @@ def __init__(self, config: ModelConfig): if is_gguf_model(config): convert_gemma4_to_gguf(self, config) + @property + def supports_cuda_graph(self) -> bool: + """False only when the PLE table is host-resident (FREETOKEN_GEMMA4_PLE_CPU=1): the + per-forward host gather in CpuPerLayerEmbedding does a CPU<->GPU copy that CUDA graph + capture forbids (and that graph replay could not re-run anyway). The default on-GPU / + GGUF PLE embeddings are graph-safe, so graphs stay on by default.""" + return not getattr(self.model, "_ple_on_cpu", False) + + def cpu_offloaded_weight_keys(self) -> tuple[str, ...]: + """Weights the engine should keep in host RAM instead of VRAM (E-series bf16 PLE + table). Not needed for GGUF, where the table is a compact Q6_K GGUFEmbedding on GPU.""" + if self.model._has_ple and isinstance( + self.model.embed_tokens_per_layer, CpuPerLayerEmbedding + ): + return (PER_LAYER_EMBED_KEY,) + return () + @torch.inference_mode() def encode_images( self, pixel_values: torch.Tensor, image_position_ids: torch.Tensor ) -> torch.Tensor: - """Run the vision tower + projector. Returns ``[num_valid_soft_tokens, text_hidden]``. - - ``pixel_values``: ``[num_images, num_patches, 3*patch**2]``; - ``image_position_ids``: ``[num_images, num_patches, 2]`` with ``(-1, -1)`` padding. - """ features = self.vision_tower.forward(pixel_values, image_position_ids) return self.embed_vision.forward(features) diff --git a/python/freetoken/models/gemma4/moe.py b/python/freetoken/models/gemma4/moe.py index 0185561..aee9252 100644 --- a/python/freetoken/models/gemma4/moe.py +++ b/python/freetoken/models/gemma4/moe.py @@ -114,21 +114,39 @@ class Gemma4DenseMLP(BaseOP): ``post_feedforward_layernorm.``) resolves the dense checkpoint unchanged. The gated MLP is bf16, or W4A16 NVFP4 when the checkpoint quantizes the dense MLP.""" - def __init__(self, config: ModelConfig): + def __init__( + self, + config: ModelConfig, + *, + intermediate_size: int | None = None, + apply_scalar: bool = True, + ): + # E-series double-wide MLP (KV-shared layers): 2x intermediate. `apply_scalar=False` + # lets the decoder run the Per-Layer-Embedding step before the per-layer scale. + if intermediate_size is not None and intermediate_size != config.intermediate_size: + import dataclasses + + mlp_cfg = dataclasses.replace(config, intermediate_size=intermediate_size) + else: + mlp_cfg = config self.shared_mlp = ( - _Nvfp4GatedMLP(config) + _Nvfp4GatedMLP(mlp_cfg) if getattr(config, "dense_quant", "none") == "nvfp4" - else GatedMLP(config) + else GatedMLP(mlp_cfg) ) self.post_feedforward_layernorm = GemmaRMSNorm( config.hidden_size, eps=config.rms_norm_eps ) self.layer_scalar = torch.empty(1) + self._apply_scalar = apply_scalar def forward(self, pre_ff: torch.Tensor, x: torch.Tensor) -> torch.Tensor: h = self.shared_mlp.forward(pre_ff) h = self.post_feedforward_layernorm.forward(h) - return (x + h) * self.layer_scalar + out = x + h + if self._apply_scalar: + out = out * self.layer_scalar + return out __all__ = [ diff --git a/python/freetoken/models/gemma4/weight.py b/python/freetoken/models/gemma4/weight.py index 43c2355..3a5478f 100644 --- a/python/freetoken/models/gemma4/weight.py +++ b/python/freetoken/models/gemma4/weight.py @@ -164,6 +164,14 @@ def merge_info(key: str) -> tuple[str, MergeRule] | None: if isinstance(config.attention_group_for_layer(layer_id), FullAttentionGroupConfig) and config.attention_group_for_layer(layer_id).k_eq_v } + # E-series KV-shared layers carry only q_proj (no k/v). Their q_proj must NOT be merged + # into a fused qkv_proj (there is nothing to merge with) -- it loads as a standalone + # LinearReplicated. num_kv_shared_layers == 0 leaves this empty (no behavior change). + kv_shared_layers = { + layer_id + for layer_id in range(config.num_layers) + if config.is_kv_shared_layer(layer_id) + } merge_buf: dict[str, dict[str, torch.Tensor]] = {} gateup_buf: dict[str, dict[str, tuple]] = {} for file in tqdm( @@ -212,11 +220,38 @@ def merge_info(key: str) -> tuple[str, MergeRule] | None: ) continue + # E-series per-layer-embedding table (~4.7 GB) is kept in host RAM: load it + # straight to CPU so it never transiently allocates VRAM (the engine's + # cpu_offloaded_weight_keys keeps it there; the model gathers on the host). + if name.endswith("embed_tokens_per_layer.weight"): + with safetensors.safe_open(file, framework="pt", device="cpu") as fcpu: + yield name, fcpu.get_tensor(raw_name) + continue + + # E-series KV-shared layers: the checkpoint still ships k/v proj + k_norm for + # these layers, but the model reuses another layer's K/V and has no such + # weights -> drop them (and yield q_proj standalone, below). + if kv_shared_layers and ".self_attn." in name: + lm = _LAYER_INDEX_PATTERN.search(name) + if lm is not None and int(lm.group(1)) in kv_shared_layers and ( + ".self_attn.k_proj" in name + or ".self_attn.v_proj" in name + or ".self_attn.k_norm" in name + ): + continue + tensor = f.get_tensor(raw_name) if is_vision or is_expert: yield name, tensor continue + # KV-shared layers: q_proj is standalone (no k/v to fuse with). + if kv_shared_layers and ".self_attn.q_proj" in name: + lm = _LAYER_INDEX_PATTERN.search(name) + if lm is not None and int(lm.group(1)) in kv_shared_layers: + yield name, tensor + continue + info = merge_info(name) if info is None: yield name, tensor diff --git a/python/freetoken/models/gguf/config.py b/python/freetoken/models/gguf/config.py index 63b1a18..7f88f08 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/gguf/dequant.py b/python/freetoken/models/gguf/dequant.py index 77c3ea0..0ad2ee3 100644 --- a/python/freetoken/models/gguf/dequant.py +++ b/python/freetoken/models/gguf/dequant.py @@ -23,6 +23,13 @@ GGML_F16 = 1 GGML_Q4_0 = 2 GGML_Q8_0 = 8 +# K-quants: the "_M"/"_S" mixtures a q4_K_M / q6_K llama GGUF is built from. Q4_K/Q5_K +# are the packed layouts the borrowed ggml CUDA kernels dequantize on the fly (the CPU +# reference dequant below only implements the ones a load actually materializes on host: +# F32/F16 norms and Q4_0/Q6_K); the block metadata here is what row_bytes / GGUFLinear +# need to size the packed buffers. +GGML_Q4_K = 12 +GGML_Q5_K = 13 GGML_Q6_K = 14 GGML_BF16 = 30 @@ -33,6 +40,8 @@ GGML_BF16: (1, 2), GGML_Q4_0: (32, 18), GGML_Q8_0: (32, 34), + GGML_Q4_K: (256, 144), + GGML_Q5_K: (256, 176), GGML_Q6_K: (256, 210), } @@ -42,6 +51,8 @@ GGML_BF16: "BF16", GGML_Q4_0: "Q4_0", GGML_Q8_0: "Q8_0", + GGML_Q4_K: "Q4_K", + GGML_Q5_K: "Q5_K", GGML_Q6_K: "Q6_K", } @@ -143,6 +154,8 @@ def dequantize(raw: torch.Tensor, ggml_type: int, out_dtype: torch.dtype) -> tor "GGML_BF16", "GGML_Q4_0", "GGML_Q8_0", + "GGML_Q4_K", + "GGML_Q5_K", "GGML_Q6_K", "GGML_NAME", "BLOCK_SHAPE", diff --git a/python/freetoken/models/gguf/tokenizer.py b/python/freetoken/models/gguf/tokenizer.py index 6d5481c..2cd4abb 100644 --- a/python/freetoken/models/gguf/tokenizer.py +++ b/python/freetoken/models/gguf/tokenizer.py @@ -12,8 +12,9 @@ from .reader import gguf_architecture, load_gguf_metadata -# GGUF architecture -> transformers GGUF tokenizer-converter key. -_TOKENIZER_ARCH = {"gemma4": "gemma4_text"} +# GGUF architecture -> transformers GGUF tokenizer-converter key. (llama's converter key +# is just "llama"; listed explicitly even though it equals the arch and would fall through.) +_TOKENIZER_ARCH = {"gemma4": "gemma4_text", "llama": "llama"} def load_gguf_tokenizer(model_path: str): @@ -23,6 +24,12 @@ def load_gguf_tokenizer(model_path: str): meta = load_gguf_metadata(model_path) arch = gguf_architecture(model_path) conv_arch = _TOKENIZER_ARCH.get(arch, arch) + # A single architecture can ship different tokenizers: Llama-2 uses SentencePiece + # ("llama"), Llama-3 uses byte-BPE ("gpt2"/tiktoken). transformers' llama converter + # mangles the byte-BPE one (drops spaces, mis-splits tokens), so trust the GGUF's own + # declared tokenizer model and route byte-BPE to the gpt2 converter. + if conv_arch == "llama" and meta.get("tokenizer.ggml.model") == "gpt2": + conv_arch = "gpt2" tok_dict: dict[str, Any] = { k[len("tokenizer.ggml.") :]: v for k, v in meta.items() diff --git a/python/freetoken/models/llama/__init__.py b/python/freetoken/models/llama/__init__.py index 9ffad0e..abb0548 100644 --- a/python/freetoken/models/llama/__init__.py +++ b/python/freetoken/models/llama/__init__.py @@ -1,5 +1,19 @@ from .config import parse_config +from .gguf import ( + convert_llama_to_gguf, + is_gguf_model, + iter_gguf_weights, + parse_gguf_config, +) from .model import LlamaForCausalLM from .weight import iter_weights -__all__ = ["LlamaForCausalLM", "parse_config", "iter_weights"] +__all__ = [ + "LlamaForCausalLM", + "parse_config", + "iter_weights", + "parse_gguf_config", + "iter_gguf_weights", + "convert_llama_to_gguf", + "is_gguf_model", +] diff --git a/python/freetoken/models/llama/gguf.py b/python/freetoken/models/llama/gguf.py new file mode 100644 index 0000000..761172b --- /dev/null +++ b/python/freetoken/models/llama/gguf.py @@ -0,0 +1,441 @@ +"""Llama GGUF adapter: build the FreeToken ``ModelConfig`` from GGUF metadata and load +the native block-quantized weights. + +Mirrors ``gemma4/gguf.py`` but for the standard *dense* llama architecture -- there is +no MoE (routed experts / offload cache), no Per-Layer Embeddings, and no MatFormer, so +this is a much smaller adapter. ``parse_gguf_config`` maps ``llama.`` KV metadata +to the same ``ModelConfig`` ``llama.config.parse_config`` would build from a HF config; +``iter_gguf_weights`` maps GGUF tensor names to the FreeToken llama weight names keeping +the projections GGUF-block-quantized; ``convert_llama_to_gguf`` swaps the dense +``nn.Linear``/embedding for the native-quant ``GGUFLinear``/``GGUFEmbedding`` ops. + +Two things differ from the gemma4 template and are llama-GGUF specific (both flagged in +code): the checkpoint is q4_K_M, so (1) the projections carry *mixed* per-tensor K-quant +types (Q4_K, with Q6_K on some attn_v/ffn_down), read per-tensor rather than hardcoded, +and (2) because q/k/v within a layer can have different types, they are kept as separate +packed ``GGUFLinear`` sub-projections whose outputs are concatenated (weight-level +fusion, as gemma4 does for its uniform Q4_0 GGUF, would need one common packed type). +""" + +from __future__ import annotations + +from typing import TYPE_CHECKING, Iterator + +import torch + +from freetoken.layers import BaseOP +from freetoken.models.config import ModelConfig, RotaryConfig +from freetoken.models.gguf.dequant import dequantize + +if TYPE_CHECKING: + from freetoken.models.gguf.config import GgufConfigShim + + +# -------------------------------------------------------------------------------------- +# GGUF tensor name <-> FreeToken llama module base name. +# -------------------------------------------------------------------------------------- +# q/k/v and gate/up are kept as separate packed sub-projections (see module docstring), +# so each maps to its own module base under the fused ``qkv_proj`` / ``gate_up_proj`` +# container; their outputs are concatenated at forward time. +_LAYER_QUANT_MAP = { + "attn_q.weight": "self_attn.qkv_proj.q_proj", + "attn_k.weight": "self_attn.qkv_proj.k_proj", + "attn_v.weight": "self_attn.qkv_proj.v_proj", + "attn_output.weight": "self_attn.o_proj", + "ffn_gate.weight": "mlp.gate_up_proj.gate_proj", + "ffn_up.weight": "mlp.gate_up_proj.up_proj", + "ffn_down.weight": "mlp.down_proj", +} +# Per-layer 1:1 norm tensors (gguf suffix -> module-relative name). Tiny F32 -> bf16. +_LAYER_NORM_MAP = { + "attn_norm.weight": "input_layernorm.weight", + "ffn_norm.weight": "post_attention_layernorm.weight", +} + + +def _quant_base_for(name: str) -> str | None: + """FreeToken module base of a *packed* GGUF weight, or ``None`` (norms / skipped). + + Shared by ``parse_gguf_config`` (records each base's ggml type) and + ``iter_gguf_weights`` (yields ``{base}.qweight``) so the two never disagree. + """ + if name == "token_embd.weight": + return "model.embed_tokens" + if name == "output.weight": # untied LM head (absent when embeddings are tied) + return "lm_head" + if not name.startswith("blk."): + return None + parts = name.split(".", 2) + layer = parts[1] + suffix = parts[2] + rel = _LAYER_QUANT_MAP.get(suffix) + return f"model.layers.{layer}.{rel}" if rel is not None else None + + +# -------------------------------------------------------------------------------------- +# Config. +# -------------------------------------------------------------------------------------- + + +def _rope_scaling(shim: "GgufConfigShim", head_dim: int) -> dict | None: + """Recover llama3 long-context rope scaling from the baked ``rope_freqs.weight``. + + llama.cpp does not write llama3 rope params as KV metadata; like it does for gemma4's + partial rope, it bakes the per-frequency smoothing into a ``rope_freqs.weight`` divisor + tensor (ggml applies ``theta / rope_freqs[j]``). The tensor's max entry is exactly the + scaling ``factor`` (the low-frequency dims are set to ``factor``); the remaining llama3 + params are Llama-3.x's fixed constants (low_freq 1, high_freq 4, original context 8192), + read from KV only if a converter happened to emit them. FreeToken's ``"llama3"`` rope + branch reconstructs the identical smoothed frequencies from these (verified to ~1e-10). + + Returns ``None`` (plain rope) when no ``rope_freqs.weight`` is present -- a Llama-2-style + checkpoint, an un-scaled model, or a metadata-only GGUF (no tensor table). + """ + from freetoken.models.gguf.reader import iter_gguf_tensors + + m = shim.metadata + for t in iter_gguf_tensors(shim.model_path): + if t.name != "rope_freqs.weight": + continue + # F32 tensor: the packed bytes ARE the values (same trick as gemma4._full_rotary_dim). + freqs = t.packed().reshape(-1).view(torch.float32) + factor = float(freqs.max().item()) + if factor <= 1.0 + 1e-6: + return None # no long-context scaling baked in + return { + "rope_type": "llama3", + "factor": factor, + "low_freq_factor": float(m.get("llama.rope.scaling.low_freq_factor", 1.0)), + "high_freq_factor": float(m.get("llama.rope.scaling.high_freq_factor", 4.0)), + "original_max_position_embeddings": int( + m.get("llama.rope.scaling.original_context_length", 8192) + ), + } + return None + + +def parse_gguf_config(shim: "GgufConfigShim") -> ModelConfig: + m = shim.metadata + + def g(key: str): + val = m.get(f"llama.{key}") + if val is None: + raise KeyError(f"missing GGUF metadata key llama.{key}") + return val + + def g_opt(key, default=None): + val = m.get(f"llama.{key}") + return default if val is None else val + + num_layers = int(g("block_count")) + hidden = int(g("embedding_length")) + num_qo_heads = int(g("attention.head_count")) + num_kv_heads = int(g_opt("attention.head_count_kv", num_qo_heads)) + # llama is uniform GQA (unlike gemma4's per-layer kv array); a single scalar. + if isinstance(num_kv_heads, (list, tuple)): + num_kv_heads = int(num_kv_heads[0]) + head_dim = int(g_opt("attention.key_length", 0) or (hidden // num_qo_heads)) + rotary_dim = int(g_opt("rope.dimension_count", 0) or head_dim) + # feed_forward_length is a scalar for llama; guard the (unused) list form defensively. + ffn = g("feed_forward_length") + intermediate = int(ffn[0] if isinstance(ffn, (list, tuple)) else ffn) + + max_pos = int(g("context_length")) + rope_theta = float(g_opt("rope.freq_base", 10_000.0)) + rope_scaling = _rope_scaling(shim, head_dim) + rotary = RotaryConfig( + head_dim=head_dim, + rotary_dim=rotary_dim, + max_position=max_pos, + base=rope_theta, + scaling=rope_scaling, + ) + + # Per-tensor packed quant types (q4_K_M mixes Q4_K / Q6_K), keyed by module base so the + # module-swap can size each GGUFLinear/GGUFEmbedding buffer to its real type. Empty for a + # metadata-only GGUF (no tensor table) -> convert_llama_to_gguf raises with a clear note. + from freetoken.models.gguf.reader import iter_gguf_tensors + + quant_types: dict[str, int] = {} + for t in iter_gguf_tensors(shim.model_path): + base = _quant_base_for(t.name) + if base is not None: + quant_types[base] = int(t.ggml_type) + + return ModelConfig( + num_layers=num_layers, + num_qo_heads=num_qo_heads, + num_kv_heads=num_kv_heads, + head_dim=head_dim, + hidden_size=hidden, + vocab_size=int(shim.vocab_size), + intermediate_size=intermediate, + rms_norm_eps=float(g("attention.layer_norm_rms_epsilon")), + rotary_config=rotary, + hidden_act="silu", + tie_word_embeddings=bool(shim.tie_word_embeddings), + num_experts=0, + num_experts_per_tok=0, + moe_intermediate_size=0, + norm_topk_prob=False, + model_type="llama", + architectures=list(shim.architectures), + # GGUF sentinel: is_gguf_model() -> True (matches gemma4's use of this field). + moe_weight_format="q4_0", + gguf_quant_types=quant_types or None, + ) + + +# -------------------------------------------------------------------------------------- +# Weight loading: GGUF tensor names -> FreeToken llama module params. +# -------------------------------------------------------------------------------------- + + +def _to_bf16(t) -> torch.Tensor: + """Dequantize a GgufTensor (F32/F16) to a dense bf16 tensor of its torch shape.""" + flat = dequantize(t.packed().reshape(-1), t.ggml_type, torch.bfloat16) + return flat.reshape(t.shape) + + +def _unpermute_qk_rows(packed: torch.Tensor, n_head: int) -> torch.Tensor: + """Invert llama.cpp's Q/K output-row permutation on the packed (still-quantized) bytes. + + convert_hf_to_gguf permutes ``attn_q``/``attn_k`` rows ``(n_head, 2, half) -> + (n_head, half, 2)`` so ggml's adjacent-pair rope reproduces HF's split-half rope. + FreeToken applies NeoX (split-half) rope, so it needs the original HF row order -- + otherwise q/k are rotated on the wrong dim pairs and attention is scrambled (coherent + tokens leak through but generation degenerates). K-quant packs each output row as an + independent contiguous byte segment (``packed`` is ``(out_features, row_bytes)``), so we + restore the HF order by reordering whole rows -- the per-row quant blocks are untouched. + ``n_head`` is the query head count for ``attn_q`` and the KV head count for ``attn_k`` + (llama.cpp permutes K by ``n_head_kv``). + """ + out_features, row_bytes = packed.shape + half = out_features // n_head // 2 + return ( + packed.reshape(n_head, half, 2, row_bytes) + .transpose(1, 2) + .contiguous() + .reshape(out_features, row_bytes) + ) + + +def _require_tp1(what: str) -> None: + """GGUF quant layers are not tensor-parallel sharded; reject TP>1 with a clear error + (mirrors the gemma4 GGUF loader's TP=1 restriction).""" + from freetoken.distributed import get_tp_info + + if get_tp_info().size > 1: + raise NotImplementedError( + f"llama GGUF {what} currently supports TP=1 only " + "(GGUF quant layers are not tensor-parallel sharded)." + ) + + +def iter_gguf_weights( + model_path: str, + device, + *, + include_moe_experts: bool, + include_non_moe: bool, +) -> Iterator[tuple[str, torch.Tensor]]: + """Yield (param_name, tensor) for every llama param. + + Quantized projections (attention q/k/v/o, MLP gate/up/down) and the token embedding + stay in their native packed block layout and are yielded as ``.qweight`` (uint8); + the RMSNorms dequantize to bf16. Llama is dense, so there are no experts to skip and + ``include_moe_experts`` is irrelevant. Unlike gemma4, q/k/v and gate/up are NOT fused + at the weight level (their K-quant types can differ) -- each projection is emitted + standalone into its ``qkv_proj`` / ``gate_up_proj`` container sub-module. + """ + from freetoken.models.gguf.reader import iter_gguf_tensors, load_gguf_metadata + + if not include_non_moe: + return + _require_tp1("weight loading") + + # Q/K row counts for the llama.cpp rope permutation (see _unpermute_qk_rows). + meta = load_gguf_metadata(model_path) + n_qo_heads = int(meta["llama.attention.head_count"]) + n_kv_heads = int(meta.get("llama.attention.head_count_kv", n_qo_heads)) + + for t in iter_gguf_tensors(model_path): + name = t.name + 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": + # Untied LM head. convert_llama_to_gguf rejects the untied case up front, so + # this is unreachable for a tied checkpoint (Llama-3.2-1B/3B); kept for parity. + yield "lm_head.qweight", t.packed() + continue + if name == "rope_freqs.weight": + continue # rope frequencies recomputed in-engine (see _rope_scaling) + if not name.startswith("blk."): + raise ValueError(f"unmapped llama GGUF tensor: {name}") + + parts = name.split(".", 2) + layer, suffix = parts[1], parts[2] + base = f"model.layers.{layer}" + if suffix in _LAYER_NORM_MAP: + yield f"{base}.{_LAYER_NORM_MAP[suffix]}", _to_bf16(t) + continue + quant_base = _quant_base_for(name) + if quant_base is None: + raise ValueError(f"unmapped llama GGUF tensor: {name}") + packed = t.packed() + if suffix == "attn_q.weight": + packed = _unpermute_qk_rows(packed, n_qo_heads) + elif suffix == "attn_k.weight": + packed = _unpermute_qk_rows(packed, n_kv_heads) + yield f"{quant_base}.qweight", packed + + +# -------------------------------------------------------------------------------------- +# Model layer swap: dense bf16 Linear/Embedding -> native GGUF-quant ops. +# -------------------------------------------------------------------------------------- + + +def is_gguf_model(config: ModelConfig) -> bool: + """True when the model was parsed from a GGUF checkpoint (native-quant path).""" + return getattr(config, "moe_weight_format", None) == "q4_0" + + +class GGUFFusedProj(BaseOP): + """A column-parallel fused projection whose packed sub-projections may carry different + GGUF quant types. + + A q4_K_M checkpoint bumps some projections to higher precision (Q6_K attn_v), so q/k/v + within a layer need not share a quant type. Weight-level fusion (as gemma4 does for its + uniform Q4_0 GGUF, concatenating packed rows) requires one common ``row_bytes``, so + instead we hold one ``GGUFLinear`` per part and concatenate their OUTPUTS -- exactly + what a fused ``LinearQKVMerged`` / ``LinearColParallelMerged`` forward produces, at the + cost of one extra small GEMV per part. + """ + + def __init__(self, parts: list[tuple[str, BaseOP]]): + for attr, lin in parts: + setattr(self, attr, lin) + self._part_attrs = [attr for attr, _ in parts] + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return torch.cat([getattr(self, a).forward(x) for a in self._part_attrs], dim=-1) + + +class GGUFTiedLMHead: + """Tied LM head over a native GGUF embedding table (logits via ggml matmul). + + Holds only a reference to the GGUF embedding (no params of its own -> empty + state_dict), mirroring ``ParallelLMHead`` with ``tie_word_embeddings``. TP=1 only. + Same shape as gemma4's, minus the softcap (llama has none). + """ + + def __init__(self, embedding, quant_type: int): + self._embedding = embedding + self._quant_type = quant_type + + def state_dict(self, *, prefix: str = "", result=None): + return result if result is not None else {} + + def load_state_dict(self, state_dict, *, prefix: str = "", _internal: bool = False): + state_dict.pop(f"{prefix}.weight", None) + state_dict.pop(f"{prefix}.bias", None) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + from freetoken.core import get_global_ctx + from freetoken.layers.gguf import fused_mul_mat_gguf + + batch = get_global_ctx().batch + if batch.is_prefill: + indices = batch.attn_metadata.get_last_indices(batch.size) + x = x[indices].contiguous() + return fused_mul_mat_gguf(x, self._embedding.qweight, self._quant_type) + + +def convert_llama_to_gguf(model, config: ModelConfig) -> None: + """In place: replace llama's dense projections + embedding with native GGUF ops. + + Every attention (q/k/v/o) and MLP (gate/up/down) projection becomes a ``GGUFLinear`` + sized to that tensor's real K-quant type (read from ``config.gguf_quant_types``); the + token embedding becomes a ``GGUFEmbedding`` (and, tied, the LM head). RMSNorms stay + dense bf16 (F32 in the GGUF). q/k/v and gate/up live in ``GGUFFusedProj`` containers + that concatenate outputs, so mixed per-projection quant types are fine. + """ + from freetoken.layers.gguf import GGUFEmbedding, GGUFLinear + + types = config.gguf_quant_types + if not types: + raise NotImplementedError( + "llama GGUF conversion needs per-tensor quant types " + "(config.gguf_quant_types); a metadata-only / FTW-converted GGUF that strips " + "the tensor table is not supported." + ) + + def qtype(base: str) -> int: + try: + return types[base] + except KeyError as exc: + raise KeyError(f"no GGUF quant type recorded for {base}") from exc + + def lin(in_features: int, out_features: int, base: str) -> GGUFLinear: + return GGUFLinear(in_features, out_features, qtype(base), has_bias=False) + + hidden = config.hidden_size + qo_dim = config.num_qo_heads * config.head_dim + kv_dim = config.num_kv_heads * config.head_dim + inter = config.intermediate_size + + inner = model.model + embed_type = qtype("model.embed_tokens") + embed = GGUFEmbedding( + num_embeddings=config.vocab_size, + embedding_dim=hidden, + quant_type=embed_type, + embed_scale=config.embedding_scale, # None for llama + ) + inner.embed_tokens = embed + + for layer_id, layer in enumerate(inner.layers.op_list): + base = f"model.layers.{layer_id}" + attn = layer.self_attn + qkv = f"{base}.self_attn.qkv_proj" + attn.qkv_proj = GGUFFusedProj( + [ + ("q_proj", lin(hidden, qo_dim, f"{qkv}.q_proj")), + ("k_proj", lin(hidden, kv_dim, f"{qkv}.k_proj")), + ("v_proj", lin(hidden, kv_dim, f"{qkv}.v_proj")), + ] + ) + attn.o_proj = lin(qo_dim, hidden, f"{base}.self_attn.o_proj") + + mlp = layer.mlp + gu = f"{base}.mlp.gate_up_proj" + mlp.gate_up_proj = GGUFFusedProj( + [ + ("gate_proj", lin(hidden, inter, f"{gu}.gate_proj")), + ("up_proj", lin(hidden, inter, f"{gu}.up_proj")), + ] + ) + mlp.down_proj = lin(inter, hidden, f"{base}.mlp.down_proj") + + if config.tie_word_embeddings: + model.lm_head = GGUFTiedLMHead(embed, embed_type) + else: + # Untied head would need a GGUFLinear-backed lm_head that also does the prefill + # last-token gather ParallelLMHead does; Llama-3.2-1B/3B tie embeddings, so this is + # only reachable for 3.1-8B-class GGUFs. + raise NotImplementedError( + "untied llama GGUF lm_head is not yet supported (the target Llama-3.2 GGUFs " + "tie word embeddings)." + ) + + +__all__ = [ + "parse_gguf_config", + "iter_gguf_weights", + "convert_llama_to_gguf", + "is_gguf_model", +] diff --git a/python/freetoken/models/llama/model.py b/python/freetoken/models/llama/model.py index a3d8017..1da89aa 100644 --- a/python/freetoken/models/llama/model.py +++ b/python/freetoken/models/llama/model.py @@ -76,6 +76,13 @@ def __init__(self, config: ModelConfig): ) super().__init__() + # GGUF checkpoints carry native block-quantized weights: swap the dense + # projections + embedding for GGUF-quant ops. No-op for a safetensors llama. + from .gguf import convert_llama_to_gguf, is_gguf_model + + 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) diff --git a/python/freetoken/models/register.py b/python/freetoken/models/register.py index 0c033ca..87c5dc9 100644 --- a/python/freetoken/models/register.py +++ b/python/freetoken/models/register.py @@ -21,6 +21,13 @@ class ModelSpec: "freetoken.models.llama", "LlamaForCausalLM", ), + # GGUF (native Q4_K/Q6_K) dense llama: same model classes, GGUF config + weight loaders. + "LlamaGGUFForCausalLM": ModelSpec( + "freetoken.models.llama", + "LlamaForCausalLM", + parse_config="parse_gguf_config", + iter_weights="iter_gguf_weights", + ), "Qwen2ForCausalLM": ModelSpec( "freetoken.models.qwen2", "Qwen2ForCausalLM", diff --git a/python/freetoken/server/args.py b/python/freetoken/server/args.py index b2857b7..145e6c3 100644 --- a/python/freetoken/server/args.py +++ b/python/freetoken/server/args.py @@ -213,6 +213,16 @@ def _infer_reasoning_parser(model_path: str) -> str | None: help="Data type for model weights and activations. 'auto' will use FP16 for FP32/FP16 models and BF16 for BF16 models.", ) + parser.add_argument( + "--kv-dtype", + type=str, + default="auto", + choices=["auto", "bfloat16", "fp8_e4m3", "fp8_e5m2"], + help="KV-cache storage dtype. 'auto' matches the model dtype (bf16). fp8_e4m3/fp8_e5m2 " + "halve KV bytes/token (~2x the token budget that fits) with some attention-precision loss; " + "queries stay bf16 and the fp8 kernels dequantize on read.", + ) + parser.add_argument( "--tensor-parallel-size", "--tp-size", @@ -677,6 +687,16 @@ def _infer_reasoning_parser(model_path: str) -> str | None: "float32": torch.float32, } kwargs["dtype"] = DTYPE_MAP[dtype_str] if isinstance(dtype_str, str) else dtype_str + + # KV-cache dtype: "auto" -> None (falls back to the model dtype in resolved_kv_dtype). + KV_DTYPE_MAP = { + "auto": None, + "bfloat16": torch.bfloat16, + "fp8_e4m3": torch.float8_e4m3fn, + "fp8_e5m2": torch.float8_e5m2, + } + kwargs["kv_dtype"] = KV_DTYPE_MAP[kwargs["kv_dtype"]] + kwargs["tp_info"] = DistributedInfo(0, kwargs["tensor_parallel_size"]) del kwargs["tensor_parallel_size"] diff --git a/run-freetoken.sh b/run-freetoken.sh new file mode 100755 index 0000000..716f901 --- /dev/null +++ b/run-freetoken.sh @@ -0,0 +1,17 @@ +#!/usr/bin/env bash +# Convenience launcher for FreeToken in this repo's venv. +# Sets CUDA_HOME to the pip-installed CUDA 13 toolkit (nvcc/cicc/ptxas) that the +# runtime CUDA-JIT kernels need, plus PATH/LD_LIBRARY_PATH, then runs `ft`. +# +# Usage: ./run-freetoken.sh +# e.g. ./run-freetoken.sh serve --model ~/models/gpt-oss-20b +# ./run-freetoken.sh shell +# ./run-freetoken.sh bench bw +set -euo pipefail +HERE="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +VENV="$HERE/.venv" +CU="$VENV/lib/python3.13/site-packages/nvidia/cu13" +export CUDA_HOME="$CU" +export PATH="$VENV/bin:$CU/bin:$PATH" +export LD_LIBRARY_PATH="$CU/lib:${LD_LIBRARY_PATH:-}" +exec "$VENV/bin/ft" "$@"