Support per-channel quantized export and configurable MatMulNBits acc…

src/mobius/_build_context.py

@@ -35,8 +35,10 @@
     "build_context",
     "ep_capabilities",
     "get_build_dtype",
+    "get_matmul_nbits_accuracy_level",
     "is_prefill_prefix_pruning_enabled",
     "prefill_prefix_pruning",
+    "validate_matmul_nbits_accuracy_level",
 ]
 
 _DEFAULT_CAPABILITIES = EpCapabilities(name="default")
@@ -47,6 +49,9 @@
 _current_dtype: contextvars.ContextVar[ir.DataType] = contextvars.ContextVar(
     "mobius_build_dtype", default=ir.DataType.FLOAT
 )
+_matmul_nbits_accuracy_level: contextvars.ContextVar[int | None] = contextvars.ContextVar(
+    "mobius_matmul_nbits_accuracy_level", default=None
+)
 _prune_prefill_prefix: contextvars.ContextVar[bool] = contextvars.ContextVar(
     "mobius_prune_prefill_prefix", default=False
 )
@@ -56,6 +61,8 @@
 def build_context(
     capabilities: EpCapabilities,
     dtype: ir.DataType = ir.DataType.FLOAT,
+    *,
+    matmul_nbits_accuracy_level: int | None = None,
 ) -> Iterator[None]:
     """Activate EP capabilities for the duration of graph construction.
 
@@ -66,6 +73,9 @@ def build_context(
         capabilities: EP capability descriptor to activate. Typically obtained
             via ``ep_registry.require(execution_provider)``.
         dtype: Active build dtype. Defaults to ``ir.DataType.FLOAT``.
+        matmul_nbits_accuracy_level: Override every emitted ``MatMulNBits``
+            accuracy level (0 through 4), regardless of weight bit width.
+            ``None`` preserves the EP-specific defaults.
 
     Example::
 
@@ -76,13 +86,27 @@ def build_context(
         with build_context(capabilities, ir.DataType.FLOAT16):
             pkg = task.build(module, config)
     """
+    validate_matmul_nbits_accuracy_level(matmul_nbits_accuracy_level)
     capabilities_token = _current_ep.set(capabilities)
     dtype_token = _current_dtype.set(dtype)
+    accuracy_token = _matmul_nbits_accuracy_level.set(matmul_nbits_accuracy_level)
     try:
         yield
     finally:
         _current_ep.reset(capabilities_token)
         _current_dtype.reset(dtype_token)
+        _matmul_nbits_accuracy_level.reset(accuracy_token)
+
+
+def validate_matmul_nbits_accuracy_level(level: int | None) -> None:
+    """Reject values outside the ORT ``MatMulNBits`` accuracy-level domain."""
+    if level is not None and (type(level) is not int or not 0 <= level <= 4):
+        raise ValueError("matmul_nbits_accuracy_level must be an integer from 0 to 4 or None")
+
+
+def get_matmul_nbits_accuracy_level() -> int | None:
+    """Return the explicit override, or ``None`` to preserve EP defaults."""
+    return _matmul_nbits_accuracy_level.get()
 
 
 def ep_capabilities() -> EpCapabilities:

src/mobius/_build_context_test.py

@@ -4,6 +4,7 @@
 from __future__ import annotations
 
 import threading
+from concurrent.futures import ThreadPoolExecutor
 
 import onnx_ir as ir
 import pytest
@@ -12,6 +13,7 @@
     build_context,
     ep_capabilities,
     get_build_dtype,
+    get_matmul_nbits_accuracy_level,
     is_prefill_prefix_pruning_enabled,
     prefill_prefix_pruning,
 )
@@ -27,6 +29,67 @@
 )
 
 
+@pytest.mark.parametrize("level", [None, 0, 1, 2, 3, 4])
+def test_matmul_nbits_accuracy_level_scoping(level):
+    def raise_in_inner_context():
+        with build_context(_CUDA_CAPABILITIES, matmul_nbits_accuracy_level=0):
+            assert get_matmul_nbits_accuracy_level() == 0
+            raise RuntimeError("deliberate")
+
+    assert get_matmul_nbits_accuracy_level() is None
+    with build_context(_CPU_CAPABILITIES, matmul_nbits_accuracy_level=level):
+        assert get_matmul_nbits_accuracy_level() == level
+        with build_context(_CUDA_CAPABILITIES):
+            assert get_matmul_nbits_accuracy_level() is None
+        assert get_matmul_nbits_accuracy_level() == level
+        with pytest.raises(RuntimeError, match="deliberate"):
+            raise_in_inner_context()
+        assert get_matmul_nbits_accuracy_level() == level
+    assert get_matmul_nbits_accuracy_level() is None
+
+
+@pytest.mark.parametrize("level", [-1, 5, True, False, 4.0, "4"])
+def test_matmul_nbits_accuracy_level_rejects_invalid_values_without_leaking(level):
+    with build_context(_CPU_CAPABILITIES, matmul_nbits_accuracy_level=4):
+        with (
+            pytest.raises(ValueError, match="matmul_nbits_accuracy_level"),
+            build_context(_CUDA_CAPABILITIES, matmul_nbits_accuracy_level=level),
+        ):
+            pytest.fail("Invalid context should not be entered")
+        assert get_matmul_nbits_accuracy_level() == 4
+        assert ep_capabilities() is _CPU_CAPABILITIES
+    assert get_matmul_nbits_accuracy_level() is None
+
+
+def test_matmul_nbits_accuracy_level_thread_isolation():
+    barrier = threading.Barrier(2)
+
+    def build(level):
+        with build_context(_CPU_CAPABILITIES, matmul_nbits_accuracy_level=level):
+            barrier.wait(timeout=5)
+            return get_matmul_nbits_accuracy_level()
+
+    with ThreadPoolExecutor(max_workers=2) as executor:
+        futures = [executor.submit(build, level) for level in (0, 4)]
+        assert [future.result(timeout=10) for future in futures] == [0, 4]
+    assert get_matmul_nbits_accuracy_level() is None
+
+
+def test_matmul_nbits_accuracy_level_async_isolation():
+    import asyncio
+
+    async def build(level):
+        with build_context(_CPU_CAPABILITIES, matmul_nbits_accuracy_level=level):
+            await asyncio.sleep(0)
+            return get_matmul_nbits_accuracy_level()
+
+    async def run():
+        return await asyncio.gather(build(0), build(4))
+
+    assert asyncio.run(run()) == [0, 4]
+    assert get_matmul_nbits_accuracy_level() is None
+
+
 class TestBuildContextDefaults:
     def test_default_capabilities_is_default_ep(self):
         """No context active → returns the 'default' EP descriptor."""

src/mobius/_builder.py

@@ -134,6 +134,7 @@ def build_from_module(
     task: str | ModelTask = "text-generation",
     *,
     execution_provider: str = "default",
+    matmul_nbits_accuracy_level: int | None = None,
     trace_optimization: bool = False,
     fp8_kv_cache: bool = False,
     kv_cache_scales: dict[int, tuple[float, float]] | None = None,
@@ -149,6 +150,10 @@ def build_from_module(
             precision and its ``validate`` method runs before build.
         task: Task name or :class:`ModelTask` instance.
         execution_provider: Target for EP-aware optimizations.
+        matmul_nbits_accuracy_level: Optional accuracy level (0 through 4) for
+            every emitted ``MatMulNBits``, including 2-, 4-, and 8-bit weights.
+            ``None`` preserves the EP-specific defaults. This does not change
+            weight quantization or ``GatherBlockQuantized``.
         trace_optimization: Log optimization diagnostics when true.
         fp8_kv_cache: Store supported decoder KV caches as FLOAT8E4M3FN.
         kv_cache_scales: Optional per-layer FP8 key/value scales.
@@ -172,7 +177,9 @@ def build_from_module(
     )
     _cast_module_dtype(module, dtype)
     capabilities = ep_registry.require(execution_provider)
-    with build_context(capabilities, dtype):
+    with build_context(
+        capabilities, dtype, matmul_nbits_accuracy_level=matmul_nbits_accuracy_level
+    ):
         package = resolved_task.build(module, config)
 
     for name, model in package.items():

src/mobius/_builder_test.py

@@ -22,6 +22,41 @@
 from mobius._model_package import ModelPackage
 
 
+@pytest.mark.parametrize("level", [None, 0, 4])
+@pytest.mark.parametrize("bits", [4, 8])
+def test_build_from_module_applies_accuracy_override(level, bits):
+    from mobius._builder import build_from_module
+    from mobius._configs import QuantizationConfig
+    from mobius._testing import make_config
+    from mobius.models import CausalLMModel
+
+    config = make_config(
+        quantization=QuantizationConfig(
+            quant_method="olive", bits=bits, group_size=32, sym=True
+        )
+    )
+    package = build_from_module(
+        CausalLMModel(config),
+        config,
+        execution_provider="cpu",
+        matmul_nbits_accuracy_level=level,
+    )
+    nodes = [
+        node
+        for model in package.values()
+        for node in ir.traversal.RecursiveGraphIterator(model.graph)
+        if node.op_type == "MatMulNBits"
+    ]
+    assert len(nodes) > 1
+    expected = level if level is not None else (4 if bits == 4 else None)
+    for node in nodes:
+        if expected is None:
+            assert "accuracy_level" not in node.attributes
+        else:
+            assert node.attributes["accuracy_level"].value == expected
+        assert node.attributes["bits"].value == bits
+
+
 def _make_value(name: str) -> ir.Value:
     return ir.Value(name=name)
 

src/mobius/_component_quantization.py

@@ -22,6 +22,7 @@
 
 from mobius._component_manifest import ComponentDescriptor, ComponentManifest
 from mobius._configs import BaseModelConfig, QuantizationConfig
+from mobius._quantization_utils import resolve_matmul_nbits_block_size
 from mobius._weight_utils import is_packed_quant_key
 from mobius.components import (
     ClippableLinear,
@@ -153,7 +154,8 @@ def _linear_layout_matches(
     )
     return (
         module._bits == quantization.bits
-        and module._block_size == quantization.group_size
+        and module._block_size
+        == resolve_matmul_nbits_block_size(quantization.group_size, module._k)
         and (module.zero_points is None) is quantization.sym
         and (
             module.zero_points is None or module.zero_points.dtype == expected_zero_point_dtype
@@ -626,7 +628,17 @@ def normalize_component_quantized_weights(
     component_names = tuple(component_names)
     manifest = manifest or _default_manifest(module, config, task)
     mapping = getattr(config, "component_quantization", None)
-    if mapping is None and not _uses_global_module_rules(config, manifest):
+    quantization = getattr(config, "quantization", None)
+    global_per_channel = (
+        quantization is not None
+        and quantization.quant_method in {"gptq", "awq"}
+        and quantization.group_size == -1
+    )
+    if (
+        mapping is None
+        and not _uses_global_module_rules(config, manifest)
+        and not global_per_channel
+    ):
         return state_dict
     routed = _route_component_weights(state_dict, manifest, component_names)
     result = dict(state_dict)
@@ -726,6 +738,9 @@ def normalize_component_quantized_weights(
                     kind="embedding"
                     if isinstance(local_module, QuantizedEmbedding)
                     else "linear",
+                    in_features=local_module._k
+                    if isinstance(local_module, QuantizedLinear)
+                    else None,
                 )
             )
 

src/mobius/_component_quantization_loading_test.py

@@ -210,6 +210,51 @@ def _t5_config(encoder, decoder) -> ArchitectureConfig:
     )
 
 
+@pytest.mark.parametrize("method", ["gptq", "awq"])
+@pytest.mark.parametrize("bits", [4, 8])
+def test_global_per_channel_loading_uses_each_projection_logical_width(method, bits):
+    quantization = _quantization(quant_method=method, bits=bits, group_size=-1, sym=False)
+    config = make_config(num_hidden_layers=1, intermediate_size=33, quantization=quantization)
+    module = CausalLMModel(config)
+    package = _build(module)
+    state_dict = {}
+    expected = {}
+    for projection, k, n in (("down_proj", 33, 64), ("gate_proj", 64, 33)):
+        name = f"model.layers.0.mlp.{projection}"
+        packed = torch.arange(math.ceil(k * bits / 32) * n, dtype=torch.int32).reshape(-1, n)
+        scales = _floats(1, n)
+        state_dict.update(
+            {
+                f"{name}.qweight": packed,
+                f"{name}.scales": scales,
+                f"{name}.qzeros": torch.full((1, n), 0x55555555, dtype=torch.int32),
+            }
+        )
+        blob_size = math.ceil(k * bits / 8)
+        weight = (
+            torch.stack([(packed >> shift) & 255 for shift in (0, 8, 16, 24)], dim=-1)
+            .transpose(0, 1)
+            .reshape(n, -1)[:, :blob_size]
+            .to(torch.uint8)
+            .reshape(n, 1, blob_size)
+        )
+        zero = 0x44 if method == "awq" and bits == 4 else (0x54 if method == "awq" else 0x55)
+        expected[name] = (
+            k,
+            weight,
+            scales.T,
+            torch.full((n, 1), zero, dtype=torch.uint8),
+        )
+
+    _load(module, package, state_dict)
+
+    for name, (k, weight, scales, zeros) in expected.items():
+        _assert_layout(package, "model", f"{name}.weight", bits, k)
+        _assert_bound(package, "model", f"{name}.weight", weight)
+        _assert_bound(package, "model", f"{name}.scales", scales)
+        _assert_bound(package, "model", f"{name}.zero_points", zeros)
+
+
 @pytest.mark.parametrize(("method", "bits"), [("gptq", 8), ("awq", 4)])
 def test_t5_raw_dotted_scales_survive_canonical_parameter_filter(method, bits):
     encoder = _quantization(quant_method=method, bits=bits, group_size=32)