From e9f89962122abdd3e30e827c176bc48f850f4855 Mon Sep 17 00:00:00 2001
From: Fei Chen <feich@microsoft.com>
Date: Fri, 10 Jul 2026 12:02:26 +0800
Subject: [PATCH 1/5] Quantize per-layer embedding tables to INT4 for WebGPU
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

The 35 per-layer embedding tables ([262144, 256] FP16 each) take 4.48 GB
unquantized — 77% of the decoder.  Quantize them to INT4 using
GatherBlockQuantized, reducing the decoder from 5.8 GB to 2.7 GB and
total model from 6.9 GB to 3.8 GB.

Co-Authored-By: Claude <noreply@anthropic.com>
---
 src/mobius/_weight_utils.py | 102 +++++++++++++++++++++++++++++++++++
 src/mobius/models/gemma4.py | 103 ++++++++++++++++++++++++++++++------
 src/mobius/tasks/_gemma4.py |  41 +++++++++++++-
 3 files changed, 230 insertions(+), 16 deletions(-)

diff --git a/src/mobius/_weight_utils.py b/src/mobius/_weight_utils.py
index 57f6c4c3..643417dd 100644
--- a/src/mobius/_weight_utils.py
+++ b/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,
diff --git a/src/mobius/models/gemma4.py b/src/mobius/models/gemma4.py
index a4557b2a..cdeb7049 100644
--- a/src/mobius/models/gemma4.py
+++ b/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(
diff --git a/src/mobius/tasks/_gemma4.py b/src/mobius/tasks/_gemma4.py
index 1e048b73..5bd7d9c5 100644
--- a/src/mobius/tasks/_gemma4.py
+++ b/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,45 @@ 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, also quantize them to INT4 to
+        # reduce the 4.5 GB FP16 tables to ~560 MB.
+        if config.split_per_layer_embedding:
+            if config.quantization is None:
+                config.quantization = QuantizationConfig(
+                    bits=4, group_size=32, quant_method="mobius", sym=False,
+                    quantize_embeddings=True,
+                )
+            else:
+                config.quantization.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)

From ded37e9ed1c9f89d240228766efa71884a62e2fc Mon Sep 17 00:00:00 2001
From: Fei Chen <feich@microsoft.com>
Date: Fri, 10 Jul 2026 12:43:09 +0800
Subject: [PATCH 2/5] Make embedding quantization bits configurable via
 embedding_bits param

Add embedding_bits parameter to mobius.build() that threads through to
the Gemma4 task's per-layer embedding quantization. When set, the
specified bit-width (4 or 8) is used for GatherBlockQuantized embedding
tables instead of the previous hardcoded INT4. When omitted, the
default INT4 behavior is preserved.

Co-Authored-By: Claude <noreply@anthropic.com>
---
 src/mobius/_builder.py      | 18 ++++++++++++++++++
 src/mobius/tasks/_gemma4.py | 17 ++++++++++++-----
 2 files changed, 30 insertions(+), 5 deletions(-)

diff --git a/src/mobius/_builder.py b/src/mobius/_builder.py
index a080f36a..3e484f78 100644
--- a/src/mobius/_builder.py
+++ b/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)
 
diff --git a/src/mobius/tasks/_gemma4.py b/src/mobius/tasks/_gemma4.py
index 5bd7d9c5..b4e623a8 100644
--- a/src/mobius/tasks/_gemma4.py
+++ b/src/mobius/tasks/_gemma4.py
@@ -222,16 +222,23 @@ 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, also quantize them to INT4 to
-        # reduce the 4.5 GB FP16 tables to ~560 MB.
+        # 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:
-            if config.quantization is None:
+            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,
                 )
-            else:
-                config.quantization.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 c54866bf350a82c15ac14d9661cdbc1a64308104 Mon Sep 17 00:00:00 2001
From: Fei Chen <feich@microsoft.com>
Date: Tue, 14 Jul 2026 16:31:37 +0800
Subject: [PATCH 3/5] Address PR review: validation, error handling, docstring,
 and tests

- Add embedding_bits to build() docstring with description and valid values
- Validate embedding_bits is 4 or 8 in build(), raise ValueError otherwise
- Replace assert-based validation in quantize_embedding_rtn with ValueError
- Add .detach().cpu() before .numpy() to handle non-CPU tensors
- Add bits validation (4/8) in both gemma4.py preprocess_weights paths
- Fix _gemma4.py to explicitly set bits/group_size/sym when enabling
  quantize_embeddings on an existing QuantizationConfig
- Add 9 unit tests for quantize_embedding_rtn covering shapes, dtypes,
  error cases, and roundtrip accuracy

Co-Authored-By: Claude <noreply@anthropic.com>
---
 src/mobius/_builder.py           |  9 ++++
 src/mobius/_weight_utils.py      | 15 +++---
 src/mobius/_weight_utils_test.py | 89 ++++++++++++++++++++++++++++++++
 src/mobius/models/gemma4.py      |  8 +++
 src/mobius/tasks/_gemma4.py      | 11 +++-
 5 files changed, 124 insertions(+), 8 deletions(-)

diff --git a/src/mobius/_builder.py b/src/mobius/_builder.py
index 3e484f78..32b5631c 100644
--- a/src/mobius/_builder.py
+++ b/src/mobius/_builder.py
@@ -418,6 +418,13 @@ def build(
             Raises :class:`ValueError` if the resolved ``model_type`` has no
             text-only sibling. Currently supported for ``gemma4_unified``
             (``google/gemma-4-12B``).
+        embedding_bits: Quantization bit-width for per-layer embedding tables
+            (4 or 8). When set, large embedding tables that exceed the EP's
+            buffer limit are block-quantized with ``GatherBlockQuantized``
+            instead of stored at full precision. Only affects models that
+            split per-layer embeddings (e.g. Gemma4 on WebGPU). When ``None``
+            (default), the model task decides — currently defaults to INT4
+            when splitting is required.
 
     Returns:
         A :class:`ModelPackage` containing the built model(s).
@@ -578,6 +585,8 @@ def build(
         config = dataclasses.replace(config, output_layer_indices=list(output_layer_indices))
 
     if embedding_bits is not None:
+        if embedding_bits not in (4, 8):
+            raise ValueError(f"embedding_bits must be 4 or 8, got {embedding_bits}")
         from mobius._configs import QuantizationConfig
 
         if config.quantization is None:
diff --git a/src/mobius/_weight_utils.py b/src/mobius/_weight_utils.py
index 643417dd..d940ab79 100644
--- a/src/mobius/_weight_utils.py
+++ b/src/mobius/_weight_utils.py
@@ -683,15 +683,18 @@ def quantize_embedding_rtn(
         - 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
+    if bits not in (4, 8):
+        raise ValueError(f"bits must be 4 or 8, got {bits}")
+    if weight.ndim != 2:
+        raise ValueError(f"weight must be 2D, got {weight.ndim}D")
     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})"
-    )
+    if embedding_dim % block_size != 0:
+        raise ValueError(
+            f"embedding_dim ({embedding_dim}) must be divisible by block_size ({block_size})"
+        )
 
     n_blocks = embedding_dim // block_size
-    w = weight.float().numpy()
+    w = weight.detach().cpu().float().numpy()
 
     import numpy as np
 
diff --git a/src/mobius/_weight_utils_test.py b/src/mobius/_weight_utils_test.py
index 14b05e57..8207bb6c 100644
--- a/src/mobius/_weight_utils_test.py
+++ b/src/mobius/_weight_utils_test.py
@@ -1049,3 +1049,92 @@ def test_indivisible_mp_num_raises(self):
         """mp_num that doesn't divide hidden raises ValueError."""
         with pytest.raises(ValueError, match="divisible"):
             split_codegen_qkv(torch.zeros(96, 32), num_heads=4, head_dim=8, mp_num=3)
+
+
+class TestQuantizeEmbeddingRTN:
+    """Tests for quantize_embedding_rtn."""
+
+    def setup_method(self):
+        from mobius._weight_utils import quantize_embedding_rtn
+
+        self.quantize = quantize_embedding_rtn
+
+    def test_int4_asymmetric_shapes(self):
+        """INT4 asymmetric: qweight is packed to half, scales/zp have correct shapes."""
+        weight = torch.randn(128, 64, dtype=torch.float16)
+        qweight, scales, zp = self.quantize(weight, bits=4, block_size=32, symmetric=False)
+        assert qweight.shape == (128, 32)  # 64 * 4 / 8 = 32
+        assert qweight.dtype == torch.uint8
+        assert scales.shape == (128, 2)  # 64 / 32 = 2 blocks
+        assert scales.dtype == torch.float16
+        assert zp is not None
+        assert zp.shape == (128, 1)  # 2 blocks packed: ceil(2*4/8) = 1
+        assert zp.dtype == torch.uint8
+
+    def test_int8_asymmetric_shapes(self):
+        """INT8 asymmetric: qweight is same shape (no packing), zp unpacked."""
+        weight = torch.randn(128, 64, dtype=torch.float16)
+        qweight, scales, zp = self.quantize(weight, bits=8, block_size=32, symmetric=False)
+        assert qweight.shape == (128, 64)
+        assert qweight.dtype == torch.uint8
+        assert scales.shape == (128, 2)
+        assert zp is not None
+        assert zp.shape == (128, 2)
+        assert zp.dtype == torch.uint8
+
+    def test_int4_symmetric_no_zero_points(self):
+        """Symmetric mode returns None for zero_points."""
+        weight = torch.randn(64, 32, dtype=torch.float16)
+        qweight, scales, zp = self.quantize(weight, bits=4, block_size=32, symmetric=True)
+        assert zp is None
+        assert qweight.shape == (64, 16)
+        assert scales.shape == (64, 1)
+
+    def test_int8_symmetric_no_zero_points(self):
+        """INT8 symmetric returns None for zero_points."""
+        weight = torch.randn(64, 32, dtype=torch.float16)
+        qweight, scales, zp = self.quantize(weight, bits=8, block_size=32, symmetric=True)
+        assert zp is None
+        assert qweight.shape == (64, 32)
+        assert scales.shape == (64, 1)
+
+    def test_invalid_bits_raises(self):
+        """bits not in (4, 8) raises ValueError."""
+        weight = torch.randn(32, 32, dtype=torch.float16)
+        with pytest.raises(ValueError, match="bits must be 4 or 8"):
+            self.quantize(weight, bits=3)
+
+    def test_non_2d_raises(self):
+        """Non-2D input raises ValueError."""
+        weight = torch.randn(4, 8, 16, dtype=torch.float16)
+        with pytest.raises(ValueError, match="2D"):
+            self.quantize(weight, bits=4)
+
+    def test_indivisible_dim_raises(self):
+        """embedding_dim not divisible by block_size raises ValueError."""
+        weight = torch.randn(32, 30, dtype=torch.float16)
+        with pytest.raises(ValueError, match="divisible by block_size"):
+            self.quantize(weight, bits=4, block_size=32)
+
+    def test_roundtrip_accuracy_int8(self):
+        """INT8 quantize-dequantize error is small."""
+        torch.manual_seed(42)
+        weight = torch.randn(64, 64, dtype=torch.float16)
+        qweight, scales, zp = self.quantize(weight, bits=8, block_size=32, symmetric=False)
+        # Dequantize manually
+        import numpy as np
+
+        q = qweight.numpy().astype(np.float32).reshape(64, 2, 32)
+        s = scales.numpy().astype(np.float32).reshape(64, 2, 1)
+        z = zp.numpy().astype(np.float32).reshape(64, 2, 1)
+        deq = (q - z) * s
+        deq = deq.reshape(64, 64)
+        error = np.abs(deq - weight.float().numpy()).mean()
+        assert error < 0.02, f"Mean abs error too large: {error}"
+
+    def test_cpu_tensor_from_cuda_mock(self):
+        """Handles non-CPU tensors gracefully (detach().cpu() path)."""
+        weight = torch.randn(32, 32, dtype=torch.float16)
+        # Just verify it doesn't crash — the .detach().cpu() handles it
+        qweight, scales, zp = self.quantize(weight, bits=4, block_size=32)
+        assert qweight.shape == (32, 16)
diff --git a/src/mobius/models/gemma4.py b/src/mobius/models/gemma4.py
index cdeb7049..6a168102 100644
--- a/src/mobius/models/gemma4.py
+++ b/src/mobius/models/gemma4.py
@@ -2057,6 +2057,10 @@ def preprocess_weights(
                 qc = getattr(self.config, "quantization", None)
                 quantize_per_layer = qc is not None and getattr(qc, "quantize_embeddings", False)
                 if quantize_per_layer:
+                    if qc.bits not in (4, 8):
+                        raise ValueError(
+                            f"quantize_embeddings requires bits=4 or bits=8, got {qc.bits}"
+                        )
                     from mobius._weight_utils import quantize_embedding_rtn
 
                     for i, chunk in enumerate(chunks):
@@ -2868,6 +2872,10 @@ def preprocess_weights(
                 qc = getattr(self.config, "quantization", None)
                 quantize_per_layer = qc is not None and getattr(qc, "quantize_embeddings", False)
                 if quantize_per_layer:
+                    if qc.bits not in (4, 8):
+                        raise ValueError(
+                            f"quantize_embeddings requires bits=4 or bits=8, got {qc.bits}"
+                        )
                     from mobius._weight_utils import quantize_embedding_rtn
 
                     for i, chunk in enumerate(chunks):
diff --git a/src/mobius/tasks/_gemma4.py b/src/mobius/tasks/_gemma4.py
index b4e623a8..07462c44 100644
--- a/src/mobius/tasks/_gemma4.py
+++ b/src/mobius/tasks/_gemma4.py
@@ -229,11 +229,18 @@ def build(
             qc = config.quantization
             if qc is not None and qc.quantize_embeddings:
                 # embedding_bits was set by the caller — use its bits value
-                pass
+                if qc.bits not in (4, 8):
+                    raise ValueError(
+                        f"quantize_embeddings requires bits=4 or bits=8, got {qc.bits}"
+                    )
             elif qc is not None:
                 # Existing quantization config without embedding quantization —
-                # enable it, defaulting to INT4
+                # enable it with INT4 (the existing bits may be for MatMul
+                # quantization which is unrelated, so override for embeddings)
                 config.quantization.quantize_embeddings = True
+                config.quantization.bits = 4
+                config.quantization.group_size = 32
+                config.quantization.sym = False
             else:
                 config.quantization = QuantizationConfig(
                     bits=4, group_size=32, quant_method="mobius", sym=False,

From a1d6f94b1f07375ad536ab94976b2625025a8e26 Mon Sep 17 00:00:00 2001
From: justinchuby <justinchuby@users.noreply.github.com>
Date: Wed, 29 Jul 2026 17:25:40 +0000
Subject: [PATCH 4/5] Fix embedding_bits quantization scope

Restrict embedding_bits to Gemma4 per-layer embedding tables so ordinary text models are not opted into quantized embeddings or linears. Defer per-layer quantized table construction until splitting is active and use a valid embedding block size for smaller configs. Add regression coverage for embedding_bits validation and no-op behavior on regular text models.\n\nCo-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

Signed-off-by: justinchuby <justinchuby@users.noreply.github.com>
---
 src/mobius/_builder.py       | 44 +++++++++++++++++-----------
 src/mobius/_exporter_test.py | 29 +++++++++++++++++++
 src/mobius/models/gemma4.py  | 56 ++++++++++++++++++++++++++++--------
 src/mobius/tasks/_gemma4.py  | 13 +++++++--
 4 files changed, 111 insertions(+), 31 deletions(-)

diff --git a/src/mobius/_builder.py b/src/mobius/_builder.py
index 0875e5b5..7198651b 100644
--- a/src/mobius/_builder.py
+++ b/src/mobius/_builder.py
@@ -425,7 +425,8 @@ def build(
             instead of stored at full precision. Only affects models that
             split per-layer embeddings (e.g. Gemma4 on WebGPU). When ``None``
             (default), the model task decides — currently defaults to INT4
-            when splitting is required.
+            when splitting is required. Ignored by models without per-layer
+            embedding tables.
 
     Returns:
         A :class:`ModelPackage` containing the built model(s).
@@ -588,21 +589,32 @@ def build(
     if embedding_bits is not None:
         if embedding_bits not in (4, 8):
             raise ValueError(f"embedding_bits must be 4 or 8, got {embedding_bits}")
-        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)
+        # ``embedding_bits`` is only for Gemma4's per-layer embedding table.  Do not
+        # attach a QuantizationConfig to ordinary text models, because that changes
+        # their regular token embedding/Linear modules.
+        if getattr(config, "hidden_size_per_layer_input", 0) and getattr(
+            config, "vocab_size_per_layer_input", 0
+        ):
+            from mobius._configs import QuantizationConfig
+
+            if config.quantization is None:
+                config = dataclasses.replace(
+                    config,
+                    quantization=QuantizationConfig(
+                        bits=embedding_bits,
+                        group_size=32,
+                        quant_method="none",
+                        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)
diff --git a/src/mobius/_exporter_test.py b/src/mobius/_exporter_test.py
index 557560b8..b662602c 100644
--- a/src/mobius/_exporter_test.py
+++ b/src/mobius/_exporter_test.py
@@ -375,3 +375,32 @@ def test_build_sets_graph_name(self, monkeypatch):
         )
         pkg = build("fake/model-id", load_weights=False)
         assert pkg["model"].graph.name == "fake/model-id/model"
+
+    def test_embedding_bits_ignored_for_regular_text_models(self, monkeypatch):
+        """embedding_bits must not opt ordinary text models into quantized modules."""
+        import transformers
+
+        fake_config = self._mock_hf_config()
+        monkeypatch.setattr(
+            transformers.AutoConfig,
+            "from_pretrained",
+            lambda *a, **kw: fake_config,
+        )
+        pkg = build("fake/model-id", embedding_bits=4, load_weights=False)
+        op_types = {node.op_type for node in pkg["model"].graph}
+        assert "MatMulNBits" not in op_types
+        assert "GatherBlockQuantized" not in op_types
+        assert pkg.config.quantization is None
+
+    def test_embedding_bits_validates_bit_width(self, monkeypatch):
+        """embedding_bits only accepts the bit-widths supported by GatherBlockQuantized."""
+        import transformers
+
+        fake_config = self._mock_hf_config()
+        monkeypatch.setattr(
+            transformers.AutoConfig,
+            "from_pretrained",
+            lambda *a, **kw: fake_config,
+        )
+        with pytest.raises(ValueError, match="embedding_bits must be 4 or 8"):
+            build("fake/model-id", embedding_bits=3, load_weights=False)
diff --git a/src/mobius/models/gemma4.py b/src/mobius/models/gemma4.py
index f3c07d2d..91b121e6 100644
--- a/src/mobius/models/gemma4.py
+++ b/src/mobius/models/gemma4.py
@@ -213,7 +213,9 @@ def __init__(
         block_size: int = 32,
         has_zero_point: bool = True,
     ):
-        super().__init__(num_embeddings, embedding_dim, bits, block_size, has_zero_point, padding_idx)
+        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):
@@ -221,6 +223,13 @@ def forward(self, op: OpBuilder, input_ids: ir.Value):
         return op.Mul(embeddings, self.embed_scale)
 
 
+def _per_layer_embedding_block_size(embedding_dim: int, group_size: int) -> int:
+    """Use one valid block for small test embeddings whose dim is below the real group."""
+    if group_size > embedding_dim >= 16 and embedding_dim & (embedding_dim - 1) == 0:
+        return embedding_dim
+    return group_size
+
+
 def _dtype_safe_compress(
     op: OpBuilder, data: ir.Value, condition: ir.Value, *, axis: int
 ) -> ir.Value:
@@ -1822,7 +1831,14 @@ def __init__(self, config: Gemma4Config):
             # Only the table actually called in forward() is realized as an
             # ONNX initializer, so the unused one adds no graph weight.
             qc = getattr(config, "quantization", None)
-            if qc is not None and getattr(qc, "quantize_embeddings", False):
+            if (
+                qc is not None
+                and getattr(qc, "quantize_embeddings", False)
+                and config.split_per_layer_embedding
+            ):
+                block_size = _per_layer_embedding_block_size(
+                    self._per_layer_dim, qc.group_size
+                )
                 self.embed_tokens_per_layer_split = nn.ModuleList(
                     [
                         QuantizedScaledWordEmbedding(
@@ -1831,7 +1847,7 @@ def __init__(self, config: Gemma4Config):
                             config.pad_token_id,
                             embed_scale=float(self._per_layer_dim**0.5),
                             bits=qc.bits,
-                            block_size=qc.group_size,
+                            block_size=block_size,
                             has_zero_point=not qc.sym,
                         )
                         for _ in range(self._num_layers)
@@ -2387,7 +2403,9 @@ def preprocess_weights(
                 )
                 chunks = fused.chunk(num_layers, dim=1)
                 qc = getattr(self.config, "quantization", None)
-                quantize_per_layer = qc is not None and getattr(qc, "quantize_embeddings", False)
+                quantize_per_layer = qc is not None and getattr(
+                    qc, "quantize_embeddings", False
+                )
                 if quantize_per_layer:
                     if qc.bits not in (4, 8):
                         raise ValueError(
@@ -2395,17 +2413,20 @@ def preprocess_weights(
                         )
                     from mobius._weight_utils import quantize_embedding_rtn
 
+                    block_size = _per_layer_embedding_block_size(per_layer_dim, qc.group_size)
                     for i, chunk in enumerate(chunks):
                         qweight, scales, zero_points = quantize_embedding_rtn(
                             chunk.contiguous(),
                             bits=qc.bits,
-                            block_size=qc.group_size,
+                            block_size=block_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
+                            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
@@ -3209,7 +3230,9 @@ def preprocess_weights(
                 )
                 chunks = fused.chunk(num_layers, dim=1)
                 qc = getattr(self.config, "quantization", None)
-                quantize_per_layer = qc is not None and getattr(qc, "quantize_embeddings", False)
+                quantize_per_layer = qc is not None and getattr(
+                    qc, "quantize_embeddings", False
+                )
                 if quantize_per_layer:
                     if qc.bits not in (4, 8):
                         raise ValueError(
@@ -3217,20 +3240,29 @@ def preprocess_weights(
                         )
                     from mobius._weight_utils import quantize_embedding_rtn
 
+                    block_size = _per_layer_embedding_block_size(per_layer_dim, qc.group_size)
                     for i, chunk in enumerate(chunks):
                         qweight, scales, zero_points = quantize_embedding_rtn(
                             chunk.contiguous(),
                             bits=qc.bits,
-                            block_size=qc.group_size,
+                            block_size=block_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
+                        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
+                            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
+                        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(
diff --git a/src/mobius/tasks/_gemma4.py b/src/mobius/tasks/_gemma4.py
index 216d64e8..476ee1db 100644
--- a/src/mobius/tasks/_gemma4.py
+++ b/src/mobius/tasks/_gemma4.py
@@ -460,13 +460,19 @@ def build(
                 config.quantization.sym = False
             else:
                 config.quantization = QuantizationConfig(
-                    bits=4, group_size=32, quant_method="mobius", sym=False,
+                    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
+            from mobius.models.gemma4 import (
+                QuantizedScaledWordEmbedding,
+                _per_layer_embedding_block_size,
+            )
 
             qc = config.quantization
             per_layer_dim = getattr(config, "hidden_size_per_layer_input", 0)
@@ -474,6 +480,7 @@ def build(
             import numpy as np
 
             embed_scale = float(np.float16(per_layer_dim**0.5))
+            block_size = _per_layer_embedding_block_size(per_layer_dim, qc.group_size)
             module.decoder.model.embed_tokens_per_layer_split = nn.ModuleList(
                 [
                     QuantizedScaledWordEmbedding(
@@ -482,7 +489,7 @@ def build(
                         config.pad_token_id,
                         embed_scale=embed_scale,
                         bits=qc.bits,
-                        block_size=qc.group_size,
+                        block_size=block_size,
                         has_zero_point=not qc.sym,
                     )
                     for _ in range(config.num_hidden_layers)

From 839ccf3b8e2f406366444711b0d37aa8690e9a04 Mon Sep 17 00:00:00 2001
From: justinchuby <justinchuby@users.noreply.github.com>
Date: Wed, 29 Jul 2026 17:38:24 +0000
Subject: [PATCH 5/5] Isolate Gemma4 per-layer embedding quantization

Keep embedding_bits and split per-layer embedding quantization separate from the shared text QuantizationConfig so decoder Linear, LM head, and token embedding quantization are not changed incidentally. Add dedicated Gemma4 config fields for per-layer embedding quantization and a regression test that split defaults do not mutate an existing quantization config.\n\nCo-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

Signed-off-by: justinchuby <justinchuby@users.noreply.github.com>
---
 src/mobius/_builder.py           | 32 +++++---------
 src/mobius/_configs/_base.py     | 14 +++++-
 src/mobius/models/gemma4.py      | 73 +++++++++++++++++---------------
 src/mobius/models/gemma4_test.py | 45 +++++++++++++++++++-
 src/mobius/tasks/_gemma4.py      | 42 +++++++++---------
 5 files changed, 124 insertions(+), 82 deletions(-)

diff --git a/src/mobius/_builder.py b/src/mobius/_builder.py
index 7198651b..7e8dc044 100644
--- a/src/mobius/_builder.py
+++ b/src/mobius/_builder.py
@@ -592,29 +592,17 @@ def build(
         # ``embedding_bits`` is only for Gemma4's per-layer embedding table.  Do not
         # attach a QuantizationConfig to ordinary text models, because that changes
         # their regular token embedding/Linear modules.
-        if getattr(config, "hidden_size_per_layer_input", 0) and getattr(
-            config, "vocab_size_per_layer_input", 0
+        if (
+            hasattr(config, "per_layer_embedding_bits")
+            and getattr(config, "hidden_size_per_layer_input", 0)
+            and getattr(config, "vocab_size_per_layer_input", 0)
         ):
-            from mobius._configs import QuantizationConfig
-
-            if config.quantization is None:
-                config = dataclasses.replace(
-                    config,
-                    quantization=QuantizationConfig(
-                        bits=embedding_bits,
-                        group_size=32,
-                        quant_method="none",
-                        sym=False,
-                        quantize_embeddings=True,
-                    ),
-                )
-            else:
-                qc = dataclasses.replace(
-                    config.quantization,
-                    bits=embedding_bits,
-                    quantize_embeddings=True,
-                )
-                config = dataclasses.replace(config, quantization=qc)
+            config = dataclasses.replace(
+                config,
+                per_layer_embedding_bits=embedding_bits,
+                per_layer_embedding_group_size=32,
+                per_layer_embedding_sym=False,
+            )
 
     if task is None:
         task = _default_task_for_model(model_type)
diff --git a/src/mobius/_configs/_base.py b/src/mobius/_configs/_base.py
index 7ebbf390..d7f225d7 100644
--- a/src/mobius/_configs/_base.py
+++ b/src/mobius/_configs/_base.py
@@ -324,7 +324,9 @@ def _extract_vision_config(config, parent_config, model_type: str) -> dict:
     hooks live under :mod:`mobius._configs.per_model` and are
     registered with :mod:`mobius._configs._extractors` at import time.
     """
-    from mobius._configs import per_model  # noqa: F401 - imported for registration side effect
+    from mobius._configs import (
+        per_model,  # ruff: ignore[unused-import] - imported for registration side effect
+    )
     from mobius._configs._extractors import extract_vision_config as _dispatch
 
     return _dispatch(config, parent_config, model_type)
@@ -337,7 +339,9 @@ def _extract_audio_config(config, parent_config, model_type: str) -> dict:
     hooks live under :mod:`mobius._configs.per_model` and are
     registered with :mod:`mobius._configs._extractors` at import time.
     """
-    from mobius._configs import per_model  # noqa: F401 - imported for registration side effect
+    from mobius._configs import (
+        per_model,  # ruff: ignore[unused-import] - imported for registration side effect
+    )
     from mobius._configs._extractors import extract_audio_config as _dispatch
 
     return _dispatch(config, parent_config, model_type)
@@ -1593,6 +1597,12 @@ class Gemma4Config(VisionLanguageConfig):
     # is too small for the fused [V, L*D] per-layer embedding table, to split
     # it into L separate [V, D] tables that each fit within the EP's buffer limit.
     split_per_layer_embedding: bool = False
+    # Dedicated quantization settings for the split per-layer embedding table.
+    # Kept separate from ``quantization`` so embedding_bits does not affect text
+    # token embeddings, LM head, or decoder Linear projection quantization.
+    per_layer_embedding_bits: int | None = None
+    per_layer_embedding_group_size: int = 32
+    per_layer_embedding_sym: bool = False
 
     @classmethod
     def from_transformers(cls, config, parent_config=None) -> Gemma4Config:
diff --git a/src/mobius/models/gemma4.py b/src/mobius/models/gemma4.py
index 91b121e6..66899b36 100644
--- a/src/mobius/models/gemma4.py
+++ b/src/mobius/models/gemma4.py
@@ -230,6 +230,26 @@ def _per_layer_embedding_block_size(embedding_dim: int, group_size: int) -> int:
     return group_size
 
 
+def _per_layer_embedding_quantization(config: Gemma4Config) -> tuple[int, int, bool] | None:
+    """Return split per-layer embedding quantization as (bits, group_size, symmetric)."""
+    bits = getattr(config, "per_layer_embedding_bits", None)
+    if bits is None:
+        quantization_config = getattr(config, "quantization", None)
+        if quantization_config is None or not getattr(
+            quantization_config, "quantize_embeddings", False
+        ):
+            return None
+        bits = quantization_config.bits
+        group_size = quantization_config.group_size
+        symmetric = quantization_config.sym
+    else:
+        group_size = getattr(config, "per_layer_embedding_group_size", 32)
+        symmetric = getattr(config, "per_layer_embedding_sym", False)
+    if bits not in (4, 8):
+        raise ValueError(f"quantize_embeddings requires bits=4 or bits=8, got {bits}")
+    return bits, group_size, symmetric
+
+
 def _dtype_safe_compress(
     op: OpBuilder, data: ir.Value, condition: ir.Value, *, axis: int
 ) -> ir.Value:
@@ -1830,15 +1850,10 @@ 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.
-            qc = getattr(config, "quantization", None)
-            if (
-                qc is not None
-                and getattr(qc, "quantize_embeddings", False)
-                and config.split_per_layer_embedding
-            ):
-                block_size = _per_layer_embedding_block_size(
-                    self._per_layer_dim, qc.group_size
-                )
+            per_layer_quant = _per_layer_embedding_quantization(config)
+            if per_layer_quant is not None and config.split_per_layer_embedding:
+                bits, group_size, symmetric = per_layer_quant
+                block_size = _per_layer_embedding_block_size(self._per_layer_dim, group_size)
                 self.embed_tokens_per_layer_split = nn.ModuleList(
                     [
                         QuantizedScaledWordEmbedding(
@@ -1846,9 +1861,9 @@ def __init__(self, config: Gemma4Config):
                             self._per_layer_dim,
                             config.pad_token_id,
                             embed_scale=float(self._per_layer_dim**0.5),
-                            bits=qc.bits,
+                            bits=bits,
                             block_size=block_size,
-                            has_zero_point=not qc.sym,
+                            has_zero_point=not symmetric,
                         )
                         for _ in range(self._num_layers)
                     ]
@@ -2402,24 +2417,18 @@ def preprocess_weights(
                     f"got {fused.shape[1]}"
                 )
                 chunks = fused.chunk(num_layers, dim=1)
-                qc = getattr(self.config, "quantization", None)
-                quantize_per_layer = qc is not None and getattr(
-                    qc, "quantize_embeddings", False
-                )
-                if quantize_per_layer:
-                    if qc.bits not in (4, 8):
-                        raise ValueError(
-                            f"quantize_embeddings requires bits=4 or bits=8, got {qc.bits}"
-                        )
+                per_layer_quant = _per_layer_embedding_quantization(self.config)
+                if per_layer_quant is not None:
+                    bits, group_size, symmetric = per_layer_quant
                     from mobius._weight_utils import quantize_embedding_rtn
 
-                    block_size = _per_layer_embedding_block_size(per_layer_dim, qc.group_size)
+                    block_size = _per_layer_embedding_block_size(per_layer_dim, group_size)
                     for i, chunk in enumerate(chunks):
                         qweight, scales, zero_points = quantize_embedding_rtn(
                             chunk.contiguous(),
-                            bits=qc.bits,
+                            bits=bits,
                             block_size=block_size,
-                            symmetric=qc.sym,
+                            symmetric=symmetric,
                         )
                         state_dict[f"model.embed_tokens_per_layer_split.{i}.qweight"] = qweight
                         state_dict[f"model.embed_tokens_per_layer_split.{i}.scales"] = scales
@@ -3229,24 +3238,18 @@ def preprocess_weights(
                     f"got {fused.shape[1]}"
                 )
                 chunks = fused.chunk(num_layers, dim=1)
-                qc = getattr(self.config, "quantization", None)
-                quantize_per_layer = qc is not None and getattr(
-                    qc, "quantize_embeddings", False
-                )
-                if quantize_per_layer:
-                    if qc.bits not in (4, 8):
-                        raise ValueError(
-                            f"quantize_embeddings requires bits=4 or bits=8, got {qc.bits}"
-                        )
+                per_layer_quant = _per_layer_embedding_quantization(self.config)
+                if per_layer_quant is not None:
+                    bits, group_size, symmetric = per_layer_quant
                     from mobius._weight_utils import quantize_embedding_rtn
 
-                    block_size = _per_layer_embedding_block_size(per_layer_dim, qc.group_size)
+                    block_size = _per_layer_embedding_block_size(per_layer_dim, group_size)
                     for i, chunk in enumerate(chunks):
                         qweight, scales, zero_points = quantize_embedding_rtn(
                             chunk.contiguous(),
-                            bits=qc.bits,
+                            bits=bits,
                             block_size=block_size,
-                            symmetric=qc.sym,
+                            symmetric=symmetric,
                         )
                         renamed[f"decoder.model.embed_tokens_per_layer_split.{i}.qweight"] = (
                             qweight
diff --git a/src/mobius/models/gemma4_test.py b/src/mobius/models/gemma4_test.py
index 66e24e19..165fd19a 100644
--- a/src/mobius/models/gemma4_test.py
+++ b/src/mobius/models/gemma4_test.py
@@ -8,7 +8,7 @@
 import onnx_ir as ir
 import torch
 
-from mobius._configs import AudioConfig, Gemma4Config
+from mobius._configs import AudioConfig, Gemma4Config, QuantizationConfig
 from mobius.models.gemma4 import Gemma4CausalLMModel, Gemma4EmbeddingModel, Gemma4Model
 
 
@@ -120,6 +120,49 @@ def test_per_expert_scale_not_folded(self):
         assert torch.allclose(result[key], torch.ones(4))
 
 
+class TestGemma4PerLayerEmbeddingQuantization:
+    """Per-layer embedding quantization stays independent from text quantization."""
+
+    def test_split_defaults_do_not_mutate_existing_quantization_config(self, monkeypatch):
+        import dataclasses
+
+        import mobius.tasks._gemma4 as gemma4_task_module
+        from mobius.tasks._gemma4 import Gemma4Task
+
+        class _TinyCaps:
+            max_buffer_size = 1
+
+        monkeypatch.setattr(gemma4_task_module, "ep_capabilities", lambda: _TinyCaps())
+
+        config = _tiny_gemma4_config(
+            enable_moe_block=False,
+            hidden_size_per_layer_input=32,
+            vocab_size_per_layer_input=64,
+        )
+        config = dataclasses.replace(
+            config,
+            quantization=QuantizationConfig(
+                bits=8,
+                group_size=64,
+                quant_method="none",
+                sym=True,
+                quantize_embeddings=False,
+            ),
+        )
+
+        pkg = Gemma4Task().build(Gemma4Model(config), config)
+
+        assert config.quantization is not None
+        assert config.quantization.bits == 8
+        assert config.quantization.group_size == 64
+        assert config.quantization.sym is True
+        assert config.quantization.quantize_embeddings is False
+        assert config.per_layer_embedding_bits == 4
+        assert config.per_layer_embedding_group_size == 32
+        assert config.per_layer_embedding_sym is False
+        assert any(node.op_type == "GatherBlockQuantized" for node in pkg["decoder"].graph)
+
+
 class TestGemma4EmbeddingModel:
     def test_reuses_token_id_fields_for_masking(self):
         config = _tiny_gemma4_config(
diff --git a/src/mobius/tasks/_gemma4.py b/src/mobius/tasks/_gemma4.py
index 476ee1db..a54c8502 100644
--- a/src/mobius/tasks/_gemma4.py
+++ b/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, QuantizationConfig
+from mobius._configs import Gemma4Config
 from mobius._model_package import ModelPackage
 from mobius._pipeline_contract import (
     declare_component_presence,
@@ -444,28 +444,24 @@ def build(
         # 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
+            if getattr(config, "per_layer_embedding_bits", None) is not None:
+                if config.per_layer_embedding_bits not in (4, 8):
+                    raise ValueError(
+                        "quantize_embeddings requires bits=4 or bits=8, "
+                        f"got {config.per_layer_embedding_bits}"
+                    )
+            elif qc is not None and qc.quantize_embeddings:
                 if qc.bits not in (4, 8):
                     raise ValueError(
                         f"quantize_embeddings requires bits=4 or bits=8, got {qc.bits}"
                     )
-            elif qc is not None:
-                # Existing quantization config without embedding quantization —
-                # enable it with INT4 (the existing bits may be for MatMul
-                # quantization which is unrelated, so override for embeddings)
-                config.quantization.quantize_embeddings = True
-                config.quantization.bits = 4
-                config.quantization.group_size = 32
-                config.quantization.sym = False
+                config.per_layer_embedding_bits = qc.bits
+                config.per_layer_embedding_group_size = qc.group_size
+                config.per_layer_embedding_sym = qc.sym
             else:
-                config.quantization = QuantizationConfig(
-                    bits=4,
-                    group_size=32,
-                    quant_method="mobius",
-                    sym=False,
-                    quantize_embeddings=True,
-                )
+                config.per_layer_embedding_bits = 4
+                config.per_layer_embedding_group_size = 32
+                config.per_layer_embedding_sym = False
             # The module was already constructed before build() runs, so the
             # per-layer embeddings are plain Embedding (Gather).  Replace them
             # with QuantizedScaledWordEmbedding (GatherBlockQuantized).
@@ -474,13 +470,15 @@ def build(
                 _per_layer_embedding_block_size,
             )
 
-            qc = config.quantization
+            bits = config.per_layer_embedding_bits
+            group_size = config.per_layer_embedding_group_size
+            symmetric = config.per_layer_embedding_sym
             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))
-            block_size = _per_layer_embedding_block_size(per_layer_dim, qc.group_size)
+            block_size = _per_layer_embedding_block_size(per_layer_dim, group_size)
             module.decoder.model.embed_tokens_per_layer_split = nn.ModuleList(
                 [
                     QuantizedScaledWordEmbedding(
@@ -488,9 +486,9 @@ def build(
                         per_layer_dim,
                         config.pad_token_id,
                         embed_scale=embed_scale,
-                        bits=qc.bits,
+                        bits=bits,
                         block_size=block_size,
-                        has_zero_point=not qc.sym,
+                        has_zero_point=not symmetric,
                     )
                     for _ in range(config.num_hidden_layers)
                 ]