Fix CLI error messages for updated flags

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))