Quantize Gemma4 per-layer embedding tables for WebGPU (INT4/INT8 configurable)

src/mobius/_builder.py

@@ -354,6 +354,7 @@ def build(
     execution_provider: str = "default",
     trace_optimization: bool = False,
     text_only: bool = False,
+    embedding_bits: int | None = None,
 ) -> ModelPackage:
     """Build an ONNX :class:`ModelPackage` from a HuggingFace model ID.
 
@@ -576,6 +577,23 @@ def build(
         # states.  See ``ArchitectureConfig.output_layer_indices``.
         config = dataclasses.replace(config, output_layer_indices=list(output_layer_indices))
 
+    if embedding_bits is not None:
+        from mobius._configs import QuantizationConfig
+
+        if config.quantization is None:
+            config = dataclasses.replace(
+                config,
+                quantization=QuantizationConfig(
+                    bits=embedding_bits, group_size=32, quant_method="mobius",
+                    sym=False, quantize_embeddings=True,
+                ),
+            )
+        else:
+            qc = dataclasses.replace(
+                config.quantization, bits=embedding_bits, quantize_embeddings=True,
+            )
+            config = dataclasses.replace(config, quantization=qc)
+
     if task is None:
         task = _default_task_for_model(model_type)
 

src/mobius/_weight_utils.py

@@ -658,6 +658,108 @@ def preprocess_gptq_weights(
     return result
 
 
+def quantize_embedding_rtn(
+    weight: torch.Tensor,
+    bits: int = 4,
+    block_size: int = 32,
+    symmetric: bool = False,
+) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor | None]:
+    """Block-wise RTN quantization of a 2D embedding table.
+
+    Produces the packed uint8 format expected by GatherBlockQuantized /
+    QuantizedEmbedding: quantize along axis=1 (embedding dim), gather along
+    axis=0 (vocabulary).
+
+    Args:
+        weight: [num_embeddings, embedding_dim] float tensor.
+        bits: Quantization bit-width (4 or 8).
+        block_size: Number of elements per quantization group.
+        symmetric: If True, use symmetric quantization (no zero_points).
+
+    Returns:
+        (qweight, scales, zero_points) where:
+        - qweight: [num_embeddings, embedding_dim * bits // 8] uint8
+        - scales: [num_embeddings, n_blocks] same dtype as input
+        - zero_points: [num_embeddings, ceil(n_blocks * bits / 8)] uint8,
+          or None if symmetric.
+    """
+    assert bits in (4, 8), f"bits must be 4 or 8, got {bits}"
+    assert weight.ndim == 2
+    num_embeddings, embedding_dim = weight.shape
+    assert embedding_dim % block_size == 0, (
+        f"embedding_dim ({embedding_dim}) must be divisible by block_size ({block_size})"
+    )
+
+    n_blocks = embedding_dim // block_size
+    w = weight.float().numpy()
+
+    import numpy as np
+
+    # Reshape to [num_embeddings, n_blocks, block_size] for per-block quantization
+    blocks = w.reshape(num_embeddings, n_blocks, block_size)
+
+    if symmetric:
+        qmax = (1 << (bits - 1)) - 1  # 7 for 4-bit
+        qmin = -(1 << (bits - 1))     # -8 for 4-bit
+        abs_max = np.maximum(np.abs(blocks.max(axis=2, keepdims=True)),
+                             np.abs(blocks.min(axis=2, keepdims=True)))
+        scales_np = abs_max / float(qmax)
+        scales_np = np.where(scales_np == 0, 1.0, scales_np)
+        quantized = np.clip(np.round(blocks / scales_np), qmin, qmax).astype(np.int8)
+        # Convert signed to unsigned for packing: [-8,7] -> [0,15]
+        quantized_unsigned = (quantized.astype(np.int16) + (1 << (bits - 1))).astype(np.uint8)
+        scales_np = scales_np.squeeze(2)  # [num_embeddings, n_blocks]
+        zero_points_np = None
+    else:
+        qmax = (1 << bits) - 1  # 15 for 4-bit
+        block_min = blocks.min(axis=2, keepdims=True)
+        block_max = blocks.max(axis=2, keepdims=True)
+        # Ensure range includes zero
+        block_min = np.minimum(block_min, 0.0)
+        block_max = np.maximum(block_max, 0.0)
+        scales_np = (block_max - block_min) / float(qmax)
+        scales_np = np.where(scales_np == 0, 1.0, scales_np)
+        zp = np.clip(np.round(-block_min / scales_np), 0, qmax).astype(np.uint8)
+        quantized_unsigned = np.clip(
+            np.round(blocks / scales_np + zp), 0, qmax
+        ).astype(np.uint8)
+        scales_np = scales_np.squeeze(2)  # [num_embeddings, n_blocks]
+        zero_points_np = zp.squeeze(2)     # [num_embeddings, n_blocks]
+
+    # Pack into uint8
+    if bits == 4:
+        # Flatten quantized to [num_embeddings, embedding_dim]
+        flat = quantized_unsigned.reshape(num_embeddings, embedding_dim)
+        # Pack pairs of 4-bit values into uint8 (low nibble first)
+        packed = (flat[:, 0::2] & 0x0F) | ((flat[:, 1::2] & 0x0F) << 4)
+        qweight = packed.astype(np.uint8)  # [num_embeddings, embedding_dim // 2]
+
+        if zero_points_np is not None:
+            # Pack zero_points the same way
+            if n_blocks % 2 == 0:
+                zp_packed = (zero_points_np[:, 0::2] & 0x0F) | ((zero_points_np[:, 1::2] & 0x0F) << 4)
+            else:
+                # Pad to even count
+                padded = np.pad(zero_points_np, ((0, 0), (0, 1)), constant_values=0)
+                zp_packed = (padded[:, 0::2] & 0x0F) | ((padded[:, 1::2] & 0x0F) << 4)
+            zero_points_out = torch.from_numpy(zp_packed.astype(np.uint8))
+        else:
+            zero_points_out = None
+    else:
+        # 8-bit: no packing
+        qweight = quantized_unsigned.reshape(num_embeddings, embedding_dim)
+        if zero_points_np is not None:
+            zero_points_out = torch.from_numpy(zero_points_np.astype(np.uint8))
+        else:
+            zero_points_out = None
+
+    return (
+        torch.from_numpy(qweight),
+        torch.from_numpy(scales_np).to(weight.dtype),
+        zero_points_out,
+    )
+
+
 def preprocess_olive_weights(
     state_dict: dict[str, torch.Tensor],
     bits: int = 4,

src/mobius/models/gemma4.py

@@ -41,6 +41,7 @@
     Embedding,
     LayerNorm,
     Linear,
+    QuantizedEmbedding,
     RMSNorm,
     create_attention_bias,
     initialize_rope,
@@ -55,6 +56,27 @@
     from mobius.components._attention import GQAContext
 
 
+class QuantizedScaledWordEmbedding(QuantizedEmbedding):
+    """Quantized embedding with scaling — INT4 GatherBlockQuantized + scale multiply."""
+
+    def __init__(
+        self,
+        num_embeddings: int,
+        embedding_dim: int,
+        padding_idx: int,
+        embed_scale: float = 1.0,
+        bits: int = 4,
+        block_size: int = 32,
+        has_zero_point: bool = True,
+    ):
+        super().__init__(num_embeddings, embedding_dim, bits, block_size, has_zero_point, padding_idx)
+        self.embed_scale = embed_scale
+
+    def forward(self, op: OpBuilder, input_ids: ir.Value):
+        embeddings = super().forward(op, input_ids)
+        return op.Mul(embeddings, self.embed_scale)
+
+
 def _dtype_safe_compress(
     op: OpBuilder, data: ir.Value, condition: ir.Value, *, axis: int
 ) -> ir.Value:
@@ -1573,17 +1595,34 @@ def __init__(self, config: Gemma4Config):
             # 256 MiB limit; ~128 MiB each vs ~4.7 GB fused).
             # Only the table actually called in forward() is realized as an
             # ONNX initializer, so the unused one adds no graph weight.
-            self.embed_tokens_per_layer_split = nn.ModuleList(
-                [
-                    Gemma3TextScaledWordEmbedding(
-                        vocab_per_layer,
-                        self._per_layer_dim,
-                        config.pad_token_id,
-                        embed_scale=float(self._per_layer_dim**0.5),
-                    )
-                    for _ in range(self._num_layers)
-                ]
-            )
+            qc = getattr(config, "quantization", None)
+            if qc is not None and getattr(qc, "quantize_embeddings", False):
+                self.embed_tokens_per_layer_split = nn.ModuleList(
+                    [
+                        QuantizedScaledWordEmbedding(
+                            vocab_per_layer,
+                            self._per_layer_dim,
+                            config.pad_token_id,
+                            embed_scale=float(self._per_layer_dim**0.5),
+                            bits=qc.bits,
+                            block_size=qc.group_size,
+                            has_zero_point=not qc.sym,
+                        )
+                        for _ in range(self._num_layers)
+                    ]
+                )
+            else:
+                self.embed_tokens_per_layer_split = nn.ModuleList(
+                    [
+                        Gemma3TextScaledWordEmbedding(
+                            vocab_per_layer,
+                            self._per_layer_dim,
+                            config.pad_token_id,
+                            embed_scale=float(self._per_layer_dim**0.5),
+                        )
+                        for _ in range(self._num_layers)
+                    ]
+                )
             self.per_layer_model_projection = Linear(
                 config.hidden_size,
                 config.num_hidden_layers * self._per_layer_dim,
@@ -2015,8 +2054,25 @@ def preprocess_weights(
                     f"got {fused.shape[1]}"
                 )
                 chunks = fused.chunk(num_layers, dim=1)
-                for i, chunk in enumerate(chunks):
-                    state_dict[f"model.embed_tokens_per_layer_split.{i}.weight"] = chunk
+                qc = getattr(self.config, "quantization", None)
+                quantize_per_layer = qc is not None and getattr(qc, "quantize_embeddings", False)
+                if quantize_per_layer:
+                    from mobius._weight_utils import quantize_embedding_rtn
+
+                    for i, chunk in enumerate(chunks):
+                        qweight, scales, zero_points = quantize_embedding_rtn(
+                            chunk.contiguous(),
+                            bits=qc.bits,
+                            block_size=qc.group_size,
+                            symmetric=qc.sym,
+                        )
+                        state_dict[f"model.embed_tokens_per_layer_split.{i}.qweight"] = qweight
+                        state_dict[f"model.embed_tokens_per_layer_split.{i}.scales"] = scales
+                        if zero_points is not None:
+                            state_dict[f"model.embed_tokens_per_layer_split.{i}.zero_points"] = zero_points
+                else:
+                    for i, chunk in enumerate(chunks):
+                        state_dict[f"model.embed_tokens_per_layer_split.{i}.weight"] = chunk
         return state_dict
 
 
@@ -2809,8 +2865,25 @@ def preprocess_weights(
                     f"got {fused.shape[1]}"
                 )
                 chunks = fused.chunk(num_layers, dim=1)
-                for i, chunk in enumerate(chunks):
-                    renamed[f"decoder.model.embed_tokens_per_layer_split.{i}.weight"] = chunk
+                qc = getattr(self.config, "quantization", None)
+                quantize_per_layer = qc is not None and getattr(qc, "quantize_embeddings", False)
+                if quantize_per_layer:
+                    from mobius._weight_utils import quantize_embedding_rtn
+
+                    for i, chunk in enumerate(chunks):
+                        qweight, scales, zero_points = quantize_embedding_rtn(
+                            chunk.contiguous(),
+                            bits=qc.bits,
+                            block_size=qc.group_size,
+                            symmetric=qc.sym,
+                        )
+                        renamed[f"decoder.model.embed_tokens_per_layer_split.{i}.qweight"] = qweight
+                        renamed[f"decoder.model.embed_tokens_per_layer_split.{i}.scales"] = scales
+                        if zero_points is not None:
+                            renamed[f"decoder.model.embed_tokens_per_layer_split.{i}.zero_points"] = zero_points
+                else:
+                    for i, chunk in enumerate(chunks):
+                        renamed[f"decoder.model.embed_tokens_per_layer_split.{i}.weight"] = chunk
             # Re-route the projection weights from embedding.* → decoder.model.*
             for k in list(renamed.keys()):
                 if k.startswith(

src/mobius/tasks/_gemma4.py

@@ -25,7 +25,7 @@
 from onnxscript import GraphBuilder, nn
 
 from mobius._build_context import ep_capabilities
-from mobius._configs import Gemma4Config
+from mobius._configs import Gemma4Config, QuantizationConfig
 from mobius._model_package import ModelPackage
 from mobius.tasks._base import (
     ModelTask,
@@ -222,6 +222,52 @@ def build(
             config.split_per_layer_embedding = fused_bytes > caps.max_buffer_size
         else:
             config.split_per_layer_embedding = False
+        # When splitting per-layer embeddings, quantize them to reduce size.
+        # Bits default to 4 (INT4) but can be overridden via embedding_bits
+        # in mobius.build() / the MobiusBuilder Olive pass config.
+        if config.split_per_layer_embedding:
+            qc = config.quantization
+            if qc is not None and qc.quantize_embeddings:
+                # embedding_bits was set by the caller — use its bits value
+                pass
+            elif qc is not None:
+                # Existing quantization config without embedding quantization —
+                # enable it, defaulting to INT4
+                config.quantization.quantize_embeddings = True
+            else:
+                config.quantization = QuantizationConfig(
+                    bits=4, group_size=32, quant_method="mobius", sym=False,
+                    quantize_embeddings=True,
+                )
+            # The module was already constructed before build() runs, so the
+            # per-layer embeddings are plain Embedding (Gather).  Replace them
+            # with QuantizedScaledWordEmbedding (GatherBlockQuantized).
+            from mobius.models.gemma4 import QuantizedScaledWordEmbedding
+
+            qc = config.quantization
+            per_layer_dim = getattr(config, "hidden_size_per_layer_input", 0)
+            vocab_per_layer = getattr(config, "vocab_size_per_layer_input", 0)
+            import numpy as np
+
+            embed_scale = float(np.float16(per_layer_dim**0.5))
+            module.decoder.model.embed_tokens_per_layer_split = nn.ModuleList(
+                [
+                    QuantizedScaledWordEmbedding(
+                        vocab_per_layer,
+                        per_layer_dim,
+                        config.pad_token_id,
+                        embed_scale=embed_scale,
+                        bits=qc.bits,
+                        block_size=qc.group_size,
+                        has_zero_point=not qc.sym,
+                    )
+                    for _ in range(config.num_hidden_layers)
+                ]
+            )
+            # Cast scales to model dtype (modules created after _cast_module_dtype)
+            from mobius._builder import _cast_module_dtype
+
+            _cast_module_dtype(module.decoder.model.embed_tokens_per_layer_split, config.dtype)
         models: dict[str, ir.Model] = {}
         models["decoder"] = self._build_decoder(module.decoder, config)
         models["vision_encoder"] = self._build_vision(module.vision_encoder, config)