Enable support for FP4 weight quantization for non-conv layers (#34)

Enables FP4 quantization by fixing the following issues:
1. Handling subbyte tensors with safetensors, which doesn't have that support by default.
2. Update the validate_fp4_export method to enable 4D tensors as well with 32 block size.

Author: pareshs44

src/coreai_opt/_utils/torch_utils.py

@@ -439,23 +439,41 @@ def mmap_module_state_dict(module: torch.nn.Module, path: str | PathLike[str]) -
     reload it via mmap, replacing the module's parameters/buffers with mmap
     views so the in-RAM tensors can be released.
 
+    Handles ``SubbyteTensor`` subclasses (e.g. ``Float4Tensor``) that safetensors
+    cannot serialize directly by unwrapping to plain uint8 for storage and
+    re-wrapping after reload.
+
     Requires all tensors in ``module.state_dict()`` to be on CPU. Raises
     ``ValueError`` otherwise — mmap is a CPU-only mechanism
     """
+    from coreai_torch._compression._floatx import SubbyteTensor as _SubbyteTensor  # noqa: PLC0415
     from safetensors.torch import load_file, save_file  # noqa: PLC0415
 
     state_dict = module.state_dict()
+
+    # Keys whose tensors are SubbyteTensor wrappers (Float4Tensor, etc.) that
+    # safetensors cannot serialize. Track their class for re-wrapping after load.
+    subbyte_keys: dict[str, type] = {}
+    tensors_to_save: dict[str, torch.Tensor] = {}
+
     for name, tensor in state_dict.items():
         if not isinstance(tensor, torch.Tensor):
             continue
         if not is_tensor_on_cpu(tensor):
             raise ValueError(
                 f"mmap_module_state_dict requires CPU tensors; '{name}' is on {tensor.device}."
             )
+        if isinstance(tensor, _SubbyteTensor):
+            subbyte_keys[name] = type(tensor)
+            tensors_to_save[name] = tensor.elem.contiguous()
+        else:
+            tensors_to_save[name] = tensor.contiguous()
 
-    save_file(
-        {k: v.contiguous() for k, v in state_dict.items() if isinstance(v, torch.Tensor)},
-        path,
-    )
+    save_file(tensors_to_save, path)
     mmap_sd = load_file(path, device="cpu")
+
+    # Re-wrap SubbyteTensor keys from their underlying uint8 representation.
+    for name, tensor_cls in subbyte_keys.items():
+        mmap_sd[name] = tensor_cls(mmap_sd[name])
+
     module.load_state_dict(mmap_sd, assign=True)

src/coreai_opt/quantization/_export_utils.py

@@ -313,12 +313,17 @@ def validate_fp4_export(
 ) -> None:
     """Validate that FP4 quantization config is supported for MLIR export.
 
-    FP4 export only supports weight-only quantization on 2D weight tensors
-    (Linear layers) with PerBlockGranularity and block_size=32.
+    FP4 export requires weight-only quantization with ``PerBlockGranularity``
+    that resolves to blocks of size 32 along the weight's last axis and no
+    blocking along any other axis (i.e. resolved block sizes
+    ``(1, ..., 1, 32)``). This covers 2D Linear weights, which resolve to
+    ``(1, 32)``, and higher-rank weights such as MoE experts that resolve to
+    ``(1, ..., 1, 32)``.
 
     Args:
         fake_quant_mod (FakeQuantizeImplBase): The fake quantization module to validate.
-        quantized_data (torch.Tensor | None): The quantized weight tensor to validate shape.
+        quantized_data (torch.Tensor | None): The quantized weight tensor used
+            to resolve per-axis block sizes against the actual weight shape.
 
     Raises:
         ValueError: If the FP4 configuration is not supported for export.
@@ -329,20 +334,21 @@ def validate_fp4_export(
             f"Got quantization_target={fake_quant_mod.quantization_target}."
         )
 
-    if quantized_data is not None and quantized_data.ndim != 2:
+    granularity = fake_quant_mod._granularity
+    if not isinstance(granularity, PerBlockGranularity):
         raise ValueError(
-            "FP4 weight quantization export is only supported for "
-            f"2D weight tensors (Linear layers). Got {quantized_data.ndim}D tensor."
+            f"FP4 quantization requires PerBlockGranularity. Got {type(granularity).__name__}."
         )
 
-    granularity = fake_quant_mod._granularity
-    if (
-        not isinstance(granularity, PerBlockGranularity)
-        or granularity.block_size != _FP4_EXPORT_BLOCK_SIZE
-    ):
+    if quantized_data is None:
+        return
+
+    resolved_block_size = granularity.get_block_size(quantized_data.shape)
+    expected_block_size = (1,) * (quantized_data.ndim - 1) + (_FP4_EXPORT_BLOCK_SIZE,)
+    if resolved_block_size != expected_block_size:
         raise ValueError(
-            f"FP4 quantization requires PerBlockGranularity with "
-            f"block_size={_FP4_EXPORT_BLOCK_SIZE}. "
-            f"Got {type(granularity).__name__} with "
-            f"block_size={getattr(granularity, 'block_size', 'N/A')}."
+            f"FP4 export requires per-axis block sizes {expected_block_size} for a "
+            f"{quantized_data.ndim}D weight (blocks of {_FP4_EXPORT_BLOCK_SIZE} along "
+            f"the last axis, no blocking elsewhere). Got resolved block sizes "
+            f"{resolved_block_size} from granularity={granularity!r}."
         )

tests/export/test_eager_mlir_export_embedding.py

@@ -0,0 +1,98 @@
+# Copyright 2026 Apple Inc.
+#
+# Use of this source code is governed by a BSD-3-Clause license that can
+# be found in the LICENSE file or at https://opensource.org/licenses/BSD-3-Clause
+
+"""Eager Core AI export tests for FP4-quantized nn.Embedding variants.
+
+These mirror the embedding patterns seen in common LLMs:
+  - plain  : standard nn.Embedding
+  - scaled : embedding output scaled by sqrt(hidden_size)
+  - tied   : embedding weight shared with an lm_head Linear
+
+All variants use FP4 (float4_e2m1fn) weight-only quantization with symmetric
+per-block granularity (block_size=32), which is the only granularity supported
+for FP4 MLIR export.
+"""
+
+import pytest
+import torch
+from torch import nn
+
+from coreai_opt import ExportBackend
+from coreai_opt.quantization import (
+    ModuleQuantizerConfig,
+    QuantizationSpec,
+    Quantizer,
+    QuantizerConfig,
+)
+from coreai_opt.quantization.spec import PerBlockGranularity, QuantizationScheme
+from tests.models.simple import (
+    PlainEmbeddingModel,
+    ScaledEmbeddingModel,
+    TiedEmbeddingModel,
+)
+
+from . import export_utils
+
+_VOCAB_SIZE = 256
+_EMBED_DIM = 64
+_SEQ_LEN = 8
+
+
+def _fp4_config() -> QuantizerConfig:
+    """FP4 weight-only config: symmetric, per-block (block_size=32) along axis 1."""
+    return QuantizerConfig(
+        global_config=ModuleQuantizerConfig(
+            op_state_spec={
+                "weight": QuantizationSpec(
+                    dtype="float4_e2m1fn",
+                    qscheme=QuantizationScheme.SYMMETRIC,
+                    granularity=PerBlockGranularity(axis=1, block_size=32),
+                ),
+            },
+            op_input_spec=None,
+            op_output_spec=None,
+        ),
+        execution_mode="eager",
+    )
+
+
+def _run_fp4_embedding_export(model: nn.Module, expected_shift_scale_ops: int) -> None:
+    """Quantize the embedding model with FP4, finalize, and export/run on Core AI."""
+    model = model.eval().to(dtype=torch.float16)
+    input_ids = torch.randint(0, _VOCAB_SIZE, (1, _SEQ_LEN), dtype=torch.int32)
+
+    quantizer = Quantizer(model, _fp4_config())
+    prepared_model = quantizer.prepare((input_ids,))
+
+    with torch.no_grad():
+        prepared_model_output = prepared_model(input_ids)
+
+    finalized_model = quantizer.finalize(backend=ExportBackend.CoreAI)
+
+    export_utils.convert_and_verify(
+        finalized_model=finalized_model,
+        input_data=input_ids,
+        expected_ops={"constexpr_blockwise_shift_scale": expected_shift_scale_ops},
+        export_backend=ExportBackend.CoreAI,
+        prepared_model_output=prepared_model_output,
+    )
+
+
+@pytest.mark.parametrize(
+    "model_factory, expected_shift_scale_ops",
+    [
+        pytest.param(PlainEmbeddingModel, 1, id="plain"),
+        pytest.param(ScaledEmbeddingModel, 1, id="scaled"),
+        pytest.param(TiedEmbeddingModel, 2, id="tied"),
+    ],
+)
+def test_fp4_embedding_export(
+    model_factory: type[nn.Module], expected_shift_scale_ops: int
+) -> None:
+    """Eager Core AI export with FP4-quantized nn.Embedding variants."""
+    _run_fp4_embedding_export(
+        model_factory(vocab_size=_VOCAB_SIZE, embed_dim=_EMBED_DIM),
+        expected_shift_scale_ops,
+    )

tests/models/simple.py

@@ -171,3 +171,67 @@ def simple_mha_model():
 def simple_mha_model_input():
     """Fixture providing example input tensor for MHA model."""
     return torch.randn(1, 10, 64)
+
+
+class PlainEmbeddingModel(nn.Module):
+    """Simple model with standard nn.Embedding."""
+
+    def __init__(self, vocab_size: int = 256, embed_dim: int = 64) -> None:
+        super().__init__()
+        self.embed_tokens = nn.Embedding(vocab_size, embed_dim)
+
+    def forward(self, input_ids: torch.Tensor) -> torch.Tensor:
+        return self.embed_tokens(input_ids)
+
+
+class ScaledEmbeddingModel(nn.Module):
+    """Simple model with embedding output scaled by sqrt(hidden_size)."""
+
+    def __init__(self, vocab_size: int = 256, embed_dim: int = 64) -> None:
+        super().__init__()
+        self.embed_tokens = nn.Embedding(vocab_size, embed_dim)
+        self.embed_scale = embed_dim**0.5
+
+    def forward(self, input_ids: torch.Tensor) -> torch.Tensor:
+        return self.embed_tokens(input_ids) * torch.tensor(
+            self.embed_scale, dtype=self.embed_tokens.weight.dtype
+        )
+
+
+class TiedEmbeddingModel(nn.Module):
+    """Simple model with embedding weight tied with an lm_head Linear."""
+
+    def __init__(self, vocab_size: int = 256, embed_dim: int = 64) -> None:
+        super().__init__()
+        self.embed_tokens = nn.Embedding(vocab_size, embed_dim)
+        self.lm_head = nn.Linear(embed_dim, vocab_size, bias=False)
+        # Tie weights.
+        self.lm_head.weight = self.embed_tokens.weight
+
+    def forward(self, input_ids: torch.Tensor) -> torch.Tensor:
+        hidden = self.embed_tokens(input_ids)
+        return self.lm_head(hidden)
+
+
+@pytest.fixture
+def plain_embedding_model():
+    """Fixture providing a plain nn.Embedding model."""
+    return PlainEmbeddingModel()
+
+
+@pytest.fixture
+def scaled_embedding_model():
+    """Fixture providing a sqrt(hidden_size)-scaled embedding model."""
+    return ScaledEmbeddingModel()
+
+
+@pytest.fixture
+def tied_embedding_model():
+    """Fixture providing an embedding model with weight tied to an lm_head."""
+    return TiedEmbeddingModel()
+
+
+@pytest.fixture
+def embedding_model_input():
+    """Fixture providing example token-id input for the embedding models."""
+    return torch.randint(0, 256, (1, 8), dtype=torch.int32)