Add Gemma-4 E-series (E2B/E4B) support: PLE, KV-layer sharing, double-wide MLP, and a dense GGUF parser

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

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

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:

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

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

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

python/freetoken/layers/gguf.py

@@ -22,19 +22,24 @@
     GGML_F32,
     GGML_NAME,
     GGML_Q4_0,
+    GGML_Q4_K,
+    GGML_Q5_K,
     GGML_Q6_K,
     GGML_Q8_0,
     row_bytes,
 )
 
 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

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