Fix CLI error messages for updated flags

CHANGELOG.md

@@ -38,20 +38,20 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
   Only graph-input or empty-placeholder caches are retyped — a non-empty
   initializer cache is skipped with a warning.
 
-### Text-only export for multimodal Gemma 4 (`--text-only`)
+### Text-only export for multimodal Gemma 4 (`--features text-only`)
 
 #### Added
 
-- `build(text_only=True)` and the `mobius build --text-only` CLI flag export the
+- `build(text_only=True)` and `mobius build --features text-only` export the
   **text backbone** of a unified multimodal checkpoint as a standalone
   decoder-only LLM. For `gemma4_unified` (`google/gemma-4-12B`) this remaps the
   model type to its text sibling (`gemma4_unified_text`) and strips the
   vision/audio config so the decoder fuses to `GroupQueryAttention` on
   GQA-capable execution providers (CUDA/DML) instead of the float-bias
   `Attention` path forced by the multimodal bidirectional vision-block overlay.
-  `--text-only` is rejected with `--config` / `--component` and now also bypasses
-  diffusers autodetect so `build()` validation runs (a diffusers/unsupported repo
-  raises instead of silently exporting a pipeline).
+  The `text-only` feature is rejected with `--config` / `--component` and now
+  also bypasses diffusers autodetect so `build()` validation runs (a
+  diffusers/unsupported repo raises instead of silently exporting a pipeline).
 
 #### Changed
 

docs/cli_reference.md

@@ -237,7 +237,6 @@ mobius build --model meta-llama/Llama-3.2-1B output/ \
 | `--trust-remote-code` | Trust remote code when loading the HuggingFace model config. |
 | `--component NAME` | Build only one component from a diffusers pipeline (e.g. `--component vae_decoder`). |
 | `--kv-cache-scale-file PATH` | Optional JSON file of calibrated per-layer FP8 KV-cache scales (onnxruntime-genai format). Only used with the `fp8-kv-cache` feature; without it all layers use a unit scale of 1.0. |
-| `--kv-cache-scale-file PATH` | Optional JSON file of calibrated per-layer FP8 KV-cache scales (onnxruntime-genai format). Only used with the `fp8-kv-cache` feature; without it all layers use a unit scale of 1.0. |
 
 #### Text-only example
 

src/mobius/__main__.py

@@ -177,9 +177,9 @@ def _resolve_static_cache_task(model_type: str) -> ModelTask:
     # validation reads them.
     _resolve_build_features(args)
 
-    # Validate --max-seq-len requires --static-cache
+    # Validate --max-seq-len requires the static-cache feature.
     if args.max_seq_len is not None and not args.static_cache:
-        raise SystemExit("Error: --max-seq-len can only be used with --static-cache.")
+        raise SystemExit("Error: --max-seq-len can only be used with --features static-cache.")
 
     # Validate --max-seq-len is positive
     if args.max_seq_len is not None and args.max_seq_len <= 0:
@@ -191,28 +191,28 @@ def _resolve_static_cache_task(model_type: str) -> ModelTask:
     if max_length is not None and max_length <= 0:
         raise SystemExit("Error: --max-length must be a positive integer.")
 
-    # Validate --static-cache + --task compatibility
+    # Validate static-cache + --task compatibility.
     if args.static_cache and args.task is not None:
         raise SystemExit(
-            "Error: --static-cache cannot be combined with --task. "
-            "Remove --task to use --static-cache."
+            "Error: --features static-cache cannot be combined with --task. "
+            "Remove --task to use --features static-cache."
         )
 
-    # --text-only resolution lives in build() (model_type remap + config
+    # text-only resolution lives in build() (model_type remap + config
     # stripping), which is only reached on the HuggingFace model-ID path.
     if args.text_only and args.config:
         raise SystemExit(
-            "Error: --text-only is not supported with --config (local "
-            "directory). Use --model <hf-id> --text-only instead."
+            "Error: --features text-only is not supported with --config (local "
+            "directory). Use --model <hf-id> --features text-only instead."
         )
 
-    # --component selects one component of a diffusers pipeline; --text-only
+    # --component selects one component of a diffusers pipeline; text-only
     # produces a single decoder-only model. Combining them would silently
     # filter that model away unless --component happens to be 'model'.
     if args.text_only and args.component:
         raise SystemExit(
-            "Error: --text-only is not supported with --component. "
-            "--text-only produces a single decoder-only model, while "
+            "Error: --features text-only is not supported with --component. "
+            "The text-only feature produces a single decoder-only model, while "
             "--component selects a component of a diffusers pipeline."
         )
 
@@ -225,7 +225,9 @@ def _resolve_static_cache_task(model_type: str) -> ModelTask:
     kv_cache_scales: dict[int, tuple[float, float]] | None = None
     scale_file = getattr(args, "kv_cache_scale_file", None)
     if scale_file is not None and not fp8_kv_cache:
-        raise SystemExit("Error: --kv-cache-scale-file can only be used with --fp8-kv-cache.")
+        raise SystemExit(
+            "Error: --kv-cache-scale-file can only be used with --features fp8-kv-cache."
+        )
     if fp8_kv_cache and scale_file is not None:
         from mobius._passes._fp8_kv_cache import load_kv_cache_scale_file
 
@@ -248,7 +250,7 @@ def _resolve_static_cache_task(model_type: str) -> ModelTask:
     component_filter = args.component
     execution_provider = args.execution_provider
 
-    # Auto-detect diffusers pipelines.  Skipped when --text-only is set:
+    # Auto-detect diffusers pipelines. Skipped when the text-only feature is set:
     # that flag only applies to transformers decoder exports, so we let the
     # central build() validation reject a diffusers/unsupported repo rather
     # than silently exporting a diffusion pipeline and ignoring the flag.
@@ -724,7 +726,8 @@ def main(argv: list[str] | None = None) -> None:
         default=None,
         metavar="N",
         help="Maximum sequence length for static cache buffers. "
-        "Only used with --static-cache. Defaults to max_position_embeddings from config.",
+        "Only used with --features static-cache. "
+        "Defaults to max_position_embeddings from config.",
     )
     build_parser.add_argument(
         "--ep",
@@ -761,8 +764,8 @@ def main(argv: list[str] | None = None) -> None:
         help=(
             "Optional JSON file of calibrated per-layer FP8 KV-cache scales "
             "(onnxruntime-genai format: {'scales': {'k_scales': [...], "
-            "'v_scales': [...]}}). Only used with --fp8-kv-cache; without it "
-            "all layers use a unit scale of 1.0."
+            "'v_scales': [...]}}). Only used with --features fp8-kv-cache; "
+            "without it all layers use a unit scale of 1.0."
         ),
     )
     build_parser.set_defaults(func=_cmd_build)

src/mobius/_passes/_fp8_kv_cache.py

@@ -73,11 +73,44 @@ def _layer_id_from_name(name: str | None) -> int | None:
     return int(match.group(1)) if match else None
 
 
+def _validate_scales(
+    scales: dict[int, tuple[float, float]],
+) -> dict[int, tuple[float, float]]:
+    """Return a validated copy of *scales*."""
+    validated: dict[int, tuple[float, float]] = {}
+    for layer_id, pair in scales.items():
+        try:
+            k, v = pair
+        except (TypeError, ValueError) as error:
+            raise ValueError(
+                f"kv_cache_scales[{layer_id!r}] must be a (k_scale, v_scale) "
+                f"pair, got {pair!r}."
+            ) from error
+        try:
+            k_f, v_f = float(k), float(v)
+        except (TypeError, ValueError) as error:
+            raise ValueError(
+                f"kv_cache_scales[{layer_id!r}] has a non-numeric scale "
+                f"(k={k!r}, v={v!r}); FP8 KV-cache scales must be finite and > 0."
+            ) from error
+        if not (math.isfinite(k_f) and k_f > 0.0 and math.isfinite(v_f) and v_f > 0.0):
+            raise ValueError(
+                f"kv_cache_scales[{layer_id!r}] is non-positive or non-finite "
+                f"(k={k}, v={v}); FP8 KV-cache scales must be finite and > 0."
+            )
+        validated[layer_id] = (k_f, v_f)
+    return validated
+
+
 def _retype_fp8(value: ir.Value | None) -> None:
     """Retype *value* to ``FLOAT8E4M3FN`` in place, preserving its shape."""
     if value is None:
         return
     value.type = ir.TensorType(_FP8)
+    if value.const_value is not None and value.const_value.size == 0:
+        value.const_value = ir.tensor(
+            np.zeros(tuple(value.const_value.shape), dtype=_FP8.numpy()), name=value.name
+        )
 
 
 def _is_retypable_cache(value: ir.Value) -> bool:
@@ -101,11 +134,12 @@ class Fp8KvCachePass(ir.passes.InPlacePass):
             per-tensor FP8 scales (typically produced by an offline
             calibration). Layers absent from the map — and every layer when
             *scales* is ``None`` — use a unit scale of ``1.0``.
+
     """
 
     def __init__(self, scales: dict[int, tuple[float, float]] | None = None) -> None:
         super().__init__()
-        self._scales = scales or {}
+        self._scales = _validate_scales(scales or {})
 
     def call(self, model: ir.Model) -> ir.passes.PassResult:
         graph = model.graph
@@ -217,6 +251,7 @@ def load_kv_cache_scale_file(path: str) -> dict[int, tuple[float, float]]:
     Raises:
         ValueError: If the file lacks ``scales.k_scales`` / ``scales.v_scales``
             or the two lists differ in length.
+
     """
     with open(path, encoding="utf-8") as handle:
         data = json.load(handle)
@@ -236,13 +271,5 @@ def load_kv_cache_scale_file(path: str) -> dict[int, tuple[float, float]]:
             f"{path!r}: k_scales and v_scales must have equal length "
             f"(got k={len(k_scales)}, v={len(v_scales)})."
         )
-    scales: dict[int, tuple[float, float]] = {}
-    for i, (k, v) in enumerate(zip(k_scales, v_scales)):
-        k_f, v_f = float(k), float(v)
-        if not (math.isfinite(k_f) and k_f > 0.0 and math.isfinite(v_f) and v_f > 0.0):
-            raise ValueError(
-                f"{path!r}: layer {i} has a non-positive or non-finite scale "
-                f"(k={k}, v={v}); FP8 KV-cache scales must be finite and > 0."
-            )
-        scales[i] = (k_f, v_f)
-    return scales
+    scales = {i: (k, v) for i, (k, v) in enumerate(zip(k_scales, v_scales))}
+    return _validate_scales(scales)

src/mobius/_passes/_fp8_kv_cache_test.py

@@ -20,7 +20,6 @@
 import json
 import sys
 
-import ml_dtypes
 import numpy as np
 import onnx_ir as ir
 import onnxruntime as ort
@@ -33,6 +32,7 @@
 from mobius import build_from_module
 from mobius._passes._fp8_kv_cache import (
     Fp8KvCachePass,
+    _retype_fp8,
     load_kv_cache_scale_file,
 )
 from mobius._registry import registry
@@ -78,6 +78,20 @@ def _build_fp8_decoder(fp8_kv_cache=True, kv_cache_scales=None):
 
 
 class TestFp8KvCacheGraph:
+    def test_retypes_empty_initializer_data(self):
+        value = ir.Value(
+            name="past_key",
+            type=ir.TensorType(ir.DataType.FLOAT16),
+            shape=ir.Shape([1, 1, 0, 16]),
+        )
+        value.const_value = ir.tensor(np.zeros((1, 1, 0, 16), dtype=np.float16))
+
+        _retype_fp8(value)
+
+        assert value.dtype == _FP8
+        assert value.const_value.dtype == _FP8
+        assert value.const_value.shape == ir.Shape([1, 1, 0, 16])
+
     def test_all_kv_io_typed_fp8(self):
         model, config = _build_fp8_decoder()
         ins = {v.name: v for v in model.graph.inputs}
@@ -257,6 +271,27 @@ def test_rejects_non_array(self, tmp_path):
             load_kv_cache_scale_file(str(path))
 
 
+class TestDirectScaleMapValidation:
+    @pytest.mark.parametrize(
+        "bad_scales",
+        [
+            {0: (0.0, 1.0)},
+            {0: (-1.0, 1.0)},
+            {0: (1.0, float("nan"))},
+            {0: (float("inf"), 1.0)},
+            {0: (1.0,)},
+            {0: ("invalid", 1.0)},
+        ],
+    )
+    def test_pass_rejects_invalid_scales(self, bad_scales):
+        with pytest.raises(ValueError):
+            Fp8KvCachePass(bad_scales)
+
+    def test_build_rejects_invalid_scale_map(self):
+        with pytest.raises(ValueError, match="finite and > 0"):
+            _build_fp8_decoder(kv_cache_scales={0: (0.0, 1.0)})
+
+
 @pytest.mark.skipif(not _FP8_CUDA, reason="requires CUDA FP8 GQA kernel (SM89+)")
 def test_fp8_kv_gqa_runs_on_cuda(tmp_path):
     """The emitted FP8 GQA op signature is accepted and computes on CUDA.
@@ -318,13 +353,8 @@ def const(name, arr):
     assert graph.outputs[1].dtype == _FP8
     assert node.attributes.get_int("kv_cache_bit_width") == 8
 
-    # Match the empty past constants' data to their new FP8 declared type.
-    past_key.const_value = ir.tensor(
-        np.zeros((b, kv, 0, d), dtype=ml_dtypes.float8_e4m3fn), name="past_key"
-    )
-    past_value.const_value = ir.tensor(
-        np.zeros((b, kv, 0, d), dtype=ml_dtypes.float8_e4m3fn), name="past_value"
-    )
+    assert past_key.const_value.dtype == _FP8
+    assert past_value.const_value.dtype == _FP8
 
     path = tmp_path / "gqa.onnx"
     ir.save(model, str(path))

src/mobius/integrations/modelopt/_dequant.py

@@ -54,9 +54,7 @@ def unpack_nvfp4_codes(packed_nk2: np.ndarray) -> np.ndarray:
     """
     packed = np.asarray(packed_nk2)
     if packed.ndim != 2:
-        raise ValueError(
-            f"NVFP4 packed codes must be 2D [N, K/2], got shape {packed.shape}."
-        )
+        raise ValueError(f"NVFP4 packed codes must be 2D [N, K/2], got shape {packed.shape}.")
     if packed.dtype != np.uint8:
         raise ValueError(f"NVFP4 packed codes must be uint8, got dtype {packed.dtype}.")
     packed = np.ascontiguousarray(packed)

tests/cli_test.py

@@ -76,8 +76,11 @@ def test_build_missing_model_errors(self):
             main(["build", tmpdir])  # no --model or --config
 
     def test_text_only_with_config_errors(self):
-        """--text-only is rejected on the --config (local dir) path."""
-        with tempfile.TemporaryDirectory() as tmpdir, pytest.raises(SystemExit):
+        """The text-only feature is rejected on the --config (local dir) path."""
+        with (
+            tempfile.TemporaryDirectory() as tmpdir,
+            pytest.raises(SystemExit, match=r"--features text-only.*--config"),
+        ):
             main(
                 [
                     "build",
@@ -91,8 +94,11 @@ def test_text_only_with_config_errors(self):
             )
 
     def test_text_only_with_component_errors(self):
-        """--text-only is rejected when combined with --component."""
-        with tempfile.TemporaryDirectory() as tmpdir, pytest.raises(SystemExit):
+        """The text-only feature is rejected when combined with --component."""
+        with (
+            tempfile.TemporaryDirectory() as tmpdir,
+            pytest.raises(SystemExit, match=r"--features text-only.*--component"),
+        ):
             main(
                 [
                     "build",
@@ -283,7 +289,10 @@ def test_features_unknown_errors(self):
             )
 
     def test_max_seq_len_without_static_cache_errors(self):
-        with tempfile.TemporaryDirectory() as tmpdir, pytest.raises(SystemExit):
+        with (
+            tempfile.TemporaryDirectory() as tmpdir,
+            pytest.raises(SystemExit, match=r"--features static-cache"),
+        ):
             main(
                 [
                     "build",
@@ -297,8 +306,11 @@ def test_max_seq_len_without_static_cache_errors(self):
             )
 
     def test_static_cache_with_task_errors(self):
-        """--static-cache cannot be combined with any --task."""
-        with tempfile.TemporaryDirectory() as tmpdir, pytest.raises(SystemExit):
+        """The static-cache feature cannot be combined with any --task."""
+        with (
+            tempfile.TemporaryDirectory() as tmpdir,
+            pytest.raises(SystemExit, match=r"--features static-cache.*--task"),
+        ):
             main(
                 [
                     "build",
@@ -313,6 +325,23 @@ def test_static_cache_with_task_errors(self):
                 ]
             )
 
+    def test_kv_cache_scale_file_without_fp8_feature_errors(self):
+        with (
+            tempfile.TemporaryDirectory() as tmpdir,
+            pytest.raises(SystemExit, match=r"--features fp8-kv-cache"),
+        ):
+            main(
+                [
+                    "build",
+                    "--model",
+                    "Qwen/Qwen2.5-0.5B",
+                    tmpdir,
+                    "--no-weights",
+                    "--kv-cache-scale-file",
+                    "scales.json",
+                ]
+            )
+
     def test_non_positive_max_seq_len_errors(self):
         with tempfile.TemporaryDirectory() as tmpdir, pytest.raises(SystemExit):
             main(