feat(gemma4): serve the E2B and E4B releases

docs/cli.md

@@ -71,7 +71,7 @@ ft serve --model ... --gpu GPU-9e8d7c6b  # the same card by UUID (a unique prefi
 |---|---|---|
 | `--memory-ratio` | 0.9 | Fraction of free VRAM the engine may use (weights + MoE cache + KV) |
 | `--num-pages` / `--num-tokens` | auto | KV capacity override in pages / tokens (mutually exclusive; auto sizes from VRAM left after weights and MoE cache) |
-| `--embed-weights` | gpu | Where the input token embedding table lives. `host` keeps it in pinned host RAM and each lookup reads its rows over PCIe, so its VRAM goes to the expert cache and KV. Refused when the table is tied to the LM head |
+| `--embed-weights` | gpu | Where the input token embedding tables live. `host` keeps them in pinned host RAM and each lookup reads its rows over PCIe, so their VRAM goes to the expert cache and KV. A table tied to the LM head stays on the GPU, and `host` is refused when there is no other table; Gemma-4 E2B / E4B's per-layer table is the one that moves, the way those models fit an 8 GB GPU |
 | `--page-size` | 1 | KV page size; DSV4 forces 128, the TRTLLM backend needs 16/32/64, SWA models require 1 |
 | `--cache-type` | radix | `radix` (prefix reuse; SWA/GDN-aware variants picked automatically) or `naive` |
 | `--attention-backend`, `--attn` | auto | `trtllm`/`fi`/`fa`/`triton`/`dsv4_sparse`/`dsa`; `prefill,decode` pair allowed; auto picks per model + GPU |

docs/models.md

@@ -16,7 +16,7 @@ for them; other checkpoints of the same architectures work too.
 | Qwen3-MoE | [Qwen/Qwen3-30B-A3B](https://huggingface.co/Qwen/Qwen3-30B-A3B) |
 | Qwen3-VL | [Qwen/Qwen3-VL-8B-Instruct](https://huggingface.co/Qwen/Qwen3-VL-8B-Instruct), [Qwen/Qwen3-VL-30B-A3B-Instruct](https://huggingface.co/Qwen/Qwen3-VL-30B-A3B-Instruct) |
 | gpt-oss | [openai/gpt-oss-120b](https://huggingface.co/openai/gpt-oss-120b), [openai/gpt-oss-20b](https://huggingface.co/openai/gpt-oss-20b) |
-| Gemma-4 | [google/gemma-4-26B-A4B-it](https://huggingface.co/google/gemma-4-26B-A4B-it), [nvidia/Gemma-4-26B-A4B-NVFP4](https://huggingface.co/nvidia/Gemma-4-26B-A4B-NVFP4), [google/gemma-4-12B-it](https://huggingface.co/google/gemma-4-12B-it), [nvidia/Gemma-4-31B-IT-NVFP4](https://huggingface.co/nvidia/Gemma-4-31B-IT-NVFP4) .. |
+| Gemma-4 | [google/gemma-4-E2B-it](https://huggingface.co/google/gemma-4-E2B-it), [google/gemma-4-E4B-it](https://huggingface.co/google/gemma-4-E4B-it), [google/gemma-4-26B-A4B-it](https://huggingface.co/google/gemma-4-26B-A4B-it), [nvidia/Gemma-4-26B-A4B-NVFP4](https://huggingface.co/nvidia/Gemma-4-26B-A4B-NVFP4), [google/gemma-4-12B-it](https://huggingface.co/google/gemma-4-12B-it), [nvidia/Gemma-4-31B-IT-NVFP4](https://huggingface.co/nvidia/Gemma-4-31B-IT-NVFP4) .. |
 | MiniMax-M2.5 | [nvidia/MiniMax-M2.5-NVFP4](https://huggingface.co/nvidia/MiniMax-M2.5-NVFP4) |
 | MiniMax-M3 | [nvidia/MiniMax-M3-NVFP4](https://huggingface.co/nvidia/MiniMax-M3-NVFP4) |
 | Muse-Glimmer | [meta-models/Muse-Glimmer-30B](https://huggingface.co/meta-models/Muse-Glimmer-30B), [RedHatAI/Muse-Glimmer-30B-NVFP4](https://huggingface.co/RedHatAI/Muse-Glimmer-30B-NVFP4) |
@@ -35,7 +35,7 @@ These llama.cpp architectures (`general.architecture`) are served:
 | `qwen3vlmoe`, `qwen3vl` | Qwen3-VL |
 | `qwen2` | Qwen2, Qwen2.5 |
 | `llama` | Llama, Mistral |
-| `gemma4` | Gemma-4, except E2B / E4B |
+| `gemma4` | Gemma-4 |
 | `glm4moe` | GLM-4.5, GLM-4.6, GLM-4.7 |
 | `glm5-next` | GLM-5.3-Flash |
 | `deepseek4` | DeepSeek-V4 Flash and Pro |

python/freetoken/engine/config.py

@@ -144,12 +144,17 @@ def model_spec(self) -> ModelSpec:
     @cached_property
     def active_encoders(self) -> tuple[EncoderSpec, ...]:
         """The encoder towers this process builds: the family registers them, the checkpoint config carries their section, --mm-disable did not name them."""
-        return tuple(
-            e
-            for e in self.model_spec.encoders
-            if getattr(self.hf_config, e.config_key, None) is not None
-            and e.kind not in self.mm.disabled_encoders
-        )
+        active = []
+        for e in self.model_spec.encoders:
+            section = getattr(self.hf_config, e.config_key, None)
+            if section is None or e.kind in self.mm.disabled_encoders:
+                continue
+            reason = e.unserved(section) if e.unserved is not None else None
+            if reason is not None:
+                logger.warning(f"{e.kind} input is not served on this checkpoint: {reason}")
+                continue
+            active.append(e)
+        return tuple(active)
 
     @cached_property
     def served_modalities(self) -> frozenset[str]:

python/freetoken/engine/engine.py

@@ -16,7 +16,7 @@
 from freetoken.distributed import destroy_distributed, enable_pynccl_distributed, set_tp_info
 from freetoken.gpu_select import gpu_identity
 from freetoken.checkpoint.ftw import ftw_tensor_names, is_ftw_checkpoint
-from freetoken.layers import OffloadMoELayer, VocabParallelEmbedding, iter_moe_layers, iter_ops, set_rope_device
+from freetoken.layers import OffloadMoELayer, ParallelLMHead, VocabParallelEmbedding, iter_moe_layers, iter_ops, set_rope_device
 from freetoken.layers.quantization import LayerKind, QuantBackend, QuantKind, finalize_quant, set_quant_backend
 from freetoken.moe.offload_cache import iter_offload_moe_layers
 from freetoken.mm.config import ENCODER_SECTIONS
@@ -340,13 +340,15 @@ def _host_embedding_keys(model, config: EngineConfig) -> frozenset[str]:
     """State-dict keys of the tables --embed-weights host keeps in pinned host RAM."""
     if config.embed_weights != "host":
         return frozenset()
-    if config.model_config.tie_word_embeddings:
+    # a tied LM head reads its whole table every step, so that table stays on the GPU
+    tied = {id(op.tied_embedding) for op in iter_ops(model) if isinstance(op, ParallelLMHead) and op.tied_embedding is not None}
+    tables = {id(op.weight) for op in iter_ops(model) if type(op) is VocabParallelEmbedding and id(op) not in tied}
+    keys = frozenset(key for key, t in model.state_dict().items() if id(t) in tables)
+    if not keys and tied:
         raise ValueError(
             "--embed-weights host: this model's embedding table is tied to its LM head, which reads the whole table "
             "every step; serve it with --embed-weights gpu"
         )
-    tables = {id(op.weight) for op in iter_ops(model) if type(op) is VocabParallelEmbedding}
-    keys = frozenset(key for key, t in model.state_dict().items() if id(t) in tables)
     if not keys:
         raise ValueError(
             f"--embed-weights host: {type(model).__name__} has no plain embedding table to move (a DF11 table stays "

python/freetoken/models/blocks.py

@@ -76,10 +76,18 @@ def embed_input_ids(embed_tokens, input_ids: torch.Tensor, batch: Batch) -> torc
 
 
 class GatedMLP(BaseOP):
-    def __init__(self, config: ModelConfig, *, quant_config: QuantConfig | None = None, prefix: str = ""):
+    def __init__(
+        self,
+        config: ModelConfig,
+        *,
+        intermediate_size: int | None = None,
+        quant_config: QuantConfig | None = None,
+        prefix: str = "",
+    ):
+        intermediate_size = intermediate_size or config.intermediate_size
         self.gate_up_proj = LinearColParallelMerged(
             config.hidden_size,
-            [config.intermediate_size, config.intermediate_size],
+            [intermediate_size, intermediate_size],
             has_bias=False,
             quant_config=quant_config,
             prefix=f"{prefix}.gate_up_proj",
@@ -91,7 +99,7 @@ def __init__(self, config: ModelConfig, *, quant_config: QuantConfig | None = No
             raise ValueError(f"Unsupported activation function: {config.hidden_act}")
         self.act_fn = act_fn
         self.down_proj = LinearRowParallel(
-            config.intermediate_size,
+            intermediate_size,
             config.hidden_size,
             has_bias=False,
             quant_config=quant_config,

python/freetoken/models/config.py

@@ -323,6 +323,8 @@ class ModelConfig:
     # the checkpoint's quant kind for the lm_head, detected the same way (only some NVFP4 exports quantize it)
     lm_head_quant: str = "none"
     shared_expert_intermediate_size: int = 0
+    # The dense MLP's width per layer, for a family whose layers differ (Gemma-4 E2B's double-wide KV-shared tail); () -> intermediate_size on every layer.
+    layer_intermediate_sizes: Tuple[int, ...] = ()
     use_qk_norm: bool = False
     # ----- DeepSeek/GLM-style MoE extensions (default keeps other models intact) -----
     # The first ``first_k_dense_replace`` decoder layers use a dense MLP instead of the
@@ -378,13 +380,19 @@ class ModelConfig:
     # n-gram embedding geometry and the QSA indexer scoring geometry the model module
     # needs. Opaque to model-agnostic engine code; None for every other model.
     qwen4_args: Any | None = None
+    # Gemma-4 E2B / E4B payload (Gemma4EArgs): per-layer input embeddings and KV-shared
+    # layers. Opaque to model-agnostic engine code; None for every other model.
+    gemma4_e_args: Any | None = None
     # Generic execution-path capability flags (set by a model's parse_config) so the engine and
     # factories stay model-agnostic instead of branching on dsv4_args:
     single_stream_only: bool = False  # model runs one sequence at a time -> force bs=1
     # Extra per-request tensors riding the LinearStatePool slots (see SlotStateSpec);
     # () for models without any. Requires a linear-attention group to ride on.
     slot_states: Tuple[SlotStateSpec, ...] = ()
 
+    def intermediate_size_for_layer(self, layer_id: int) -> int:
+        return self.layer_intermediate_sizes[layer_id] if self.layer_intermediate_sizes else self.intermediate_size
+
     @property
     def is_moe(self) -> bool:
         return "moe" in self.model_type or self.moe_enabled

python/freetoken/models/gemma4/attention.py

@@ -15,7 +15,13 @@
 
 
 class Gemma4Attention(BaseOP):
-    """Gemma 4 attention for one full-context or SWA layer."""
+    """Gemma 4 attention for one full-context or SWA layer.
+
+    On E2B / E4B the KV-shared tail projects q only and attends over the k, v of the last
+    earlier layer of its attention type. That layer hands its rows over in ``shared_kv``, a
+    dict living for one forward, so a captured graph records the same tensors. Each shared
+    layer still stores them in its own KV slab; reading the source's slab would save that memory.
+    """
 
     def __init__(self, config: ModelConfig, layer_id: int, *, prefix: str = ""):
         self.layer_id = layer_id
@@ -28,25 +34,35 @@ def __init__(self, config: ModelConfig, layer_id: int, *, prefix: str = ""):
         self.num_kv_heads = group.num_kv_heads
         self.num_qo_heads = config.num_qo_heads
         self.k_eq_v = isinstance(group, FullAttentionGroupConfig) and group.k_eq_v
+        e_args = config.gemma4_e_args
+        self.kv_shared = e_args is not None and e_args.is_kv_shared(layer_id)
+        self.kv_source = e_args is not None and e_args.is_kv_source(layer_id, group.layer_ids)
+        self._group_name = group.name
 
         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,
-            quant_config=config.quant,
-            prefix=f"{prefix}.qkv_proj",
-        )
+        if self.kv_shared:
+            self.q_proj = LinearReplicated(
+                config.hidden_size, self.q_dim, has_bias=False,
+                quant_config=config.quant, prefix=f"{prefix}.q_proj",
+            )
+        else:
+            self.qkv_proj = LinearQKVMerged(
+                config.hidden_size,
+                self.head_dim,
+                self.num_qo_heads,
+                self.num_kv_heads,
+                has_bias=False,
+                quant_config=config.quant,
+                prefix=f"{prefix}.qkv_proj",
+            )
+            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,
             quant_config=config.quant, prefix=f"{prefix}.o_proj",
         )
         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,
@@ -79,29 +95,38 @@ def _apply_rope(
         return q_view.view_as(q), k_view.view_as(k)
 
     @nvtx_annotate("MHA")
-    def forward(self, x: torch.Tensor) -> torch.Tensor:
+    def forward(self, x: torch.Tensor, shared_kv: dict[str, tuple[torch.Tensor, torch.Tensor]]) -> torch.Tensor:
         ctx = get_global_ctx()
         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.kv_shared:
+            q = self.q_proj.forward(x).view(T, self.num_qo_heads, self.head_dim)
+            q = self.q_norm.forward(q)
+            # the rope kernels rotate a key along with the query: give them one head of zeros
+            q, _ = self._apply_rope(positions, q, q.new_zeros(T, 1, self.head_dim))
+            k, v = shared_kv[self._group_name]
+        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)
+            k = k.reshape(T, self.num_kv_heads * self.head_dim)
+            v = v.reshape(T, self.num_kv_heads * self.head_dim)
+            if self.kv_source:
+                shared_kv[self._group_name] = (k, v)
         o = ctx.attn_backend.forward(
             q.contiguous(),
             k.contiguous(),

python/freetoken/models/gemma4/config.py

@@ -43,6 +43,34 @@ class UnifiedVisionConfig:
     text_hidden_size: int
 
 
+@dataclass(frozen=True)
+class Gemma4EArgs:
+    """What the E2B / E4B releases add to the decoder."""
+
+    # width of one layer's per-layer input; the table holds num_layers of them per token
+    per_layer_hidden_size: int
+    per_layer_vocab_size: int
+    # layers from here on project q only and reuse k, v of the last earlier layer of their attention type
+    first_kv_shared_layer: int
+
+    def is_kv_shared(self, layer_id: int) -> bool:
+        return layer_id >= self.first_kv_shared_layer
+
+    def is_kv_source(self, layer_id: int, group_layer_ids: tuple[int, ...]) -> bool:
+        return layer_id == max((i for i in group_layer_ids if i < self.first_kv_shared_layer), default=None)
+
+
+def gemma4_e_args(
+    num_layers: int, per_layer_hidden_size: int, per_layer_vocab_size: int, num_kv_shared_layers: int
+) -> Gemma4EArgs | None:
+    """None for the releases without per-layer inputs: their configs set the width to 0."""
+    if not per_layer_hidden_size:
+        if num_kv_shared_layers:
+            raise NotImplementedError("gemma4 with KV-shared layers but no per-layer inputs is not supported")
+        return None
+    return Gemma4EArgs(per_layer_hidden_size, per_layer_vocab_size, num_layers - num_kv_shared_layers)
+
+
 def _text_config(hf_config: Any) -> tuple[Any, list[str] | None, Any]:
     top_architectures = getattr(hf_config, "architectures", None)
     if hasattr(hf_config, "text_config") and hf_config.text_config is not None:
@@ -162,6 +190,15 @@ def parse_config(hf_config: Any) -> ModelConfig:
     expert_quant = detect_expert_quant(top_cfg)
     dense_quant = "nvfp4" if (expert_quant == "nvfp4" and not num_experts > 0) else "none"
 
+    # E2B doubles the MLP of its KV-shared tail. Gemma4UnifiedTextConfig omits these fields.
+    num_kv_shared_layers = getattr(cfg, "num_kv_shared_layers", None) or 0
+    first_kv_shared_layer = cfg.num_hidden_layers - num_kv_shared_layers
+    layer_intermediate_sizes = (
+        tuple(cfg.intermediate_size * (2 if i >= first_kv_shared_layer else 1) for i in range(cfg.num_hidden_layers))
+        if getattr(cfg, "use_double_wide_mlp", False) and num_kv_shared_layers
+        else ()
+    )
+
     return ModelConfig(
         num_layers=cfg.num_hidden_layers,
         num_qo_heads=cfg.num_attention_heads,
@@ -170,6 +207,7 @@ def parse_config(hf_config: Any) -> ModelConfig:
         hidden_size=cfg.hidden_size,
         vocab_size=cfg.vocab_size,
         intermediate_size=cfg.intermediate_size,
+        layer_intermediate_sizes=layer_intermediate_sizes,
         hidden_act="gelu_tanh",
         rms_norm_eps=cfg.rms_norm_eps,
         tie_word_embeddings=bool(getattr(cfg, "tie_word_embeddings", True)),
@@ -189,6 +227,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),
+        gemma4_e_args=gemma4_e_args(
+            cfg.num_hidden_layers,
+            getattr(cfg, "hidden_size_per_layer_input", None) or 0,
+            getattr(cfg, "vocab_size_per_layer_input", None) or 0,
+            num_kv_shared_layers,
+        ),
         attention_groups=(
             FullAttentionGroupConfig(
                 name="full",
@@ -211,4 +255,4 @@ def parse_config(hf_config: Any) -> ModelConfig:
     )
 
 
-__all__ = ["UnifiedVisionConfig", "VisionConfig", "parse_config"]
+__all__ = ["Gemma4EArgs", "UnifiedVisionConfig", "VisionConfig", "gemma4_e_args", "parse_config"]