fix(kv-cache): gate FP8 KV cache on EP FP8-kernel capability

Address PR review: fp8_kv_cache was applied whenever GQA fusion was active
(dtype in caps.gqa_dtypes), which includes non-CUDA EPs (CPU gqa_dtypes={FLOAT},
WebGPU float/float16) that lack the FP8 GroupQueryAttention kernel. Such builds
emitted FLOAT8E4M3FN KV-cache I/O for EPs that cannot load/run it.

- Add EpCapabilities.supports_fp8_kv_cache (default False); set True only for
  CUDA (SM89+ Ada/Hopper/Blackwell ship the FP8 GQA kernel).
- Gate Fp8KvCachePass on caps.supports_fp8_kv_cache in addition to GQA activity;
  otherwise warn and ignore the request.
- Add regression test: CPU/float32 (GQA active, no FP8 kernel) keeps FLOAT KV I/O.

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
Signed-off-by: justinchuby <justinchuby@users.noreply.github.com>

Author: justinchuby

src/mobius/_execution_providers.py

@@ -93,6 +93,12 @@ class EpCapabilities:
             default) leaves ``Range`` unchanged.  Set ``False`` only when
             static cache is used on runtimes that lack a ``Range`` kernel (QNN
             HTP), where the ``Range`` node would otherwise be forced onto CPU.
+        supports_fp8_kv_cache: ``True`` allows :class:`~mobius._passes.
+            Fp8KvCachePass` to retype the ``GroupQueryAttention`` KV cache to
+            ``FLOAT8E4M3FN``.  Only CUDA (SM89+ Ada/Hopper/Blackwell) ships the
+            FP8 GQA kernel, so this defaults to ``False`` for every other EP —
+            ``--features fp8-kv-cache`` is ignored (with a warning) on EPs that
+            cannot run an FP8 KV cache, preventing invalid/unloadable models.
         supports_matmul_nbits: ``False`` converts ``com.microsoft::MatMulNBits``
             (blockwise-INT4 weight) into a standard ``DequantizeLinear`` +
             ``MatMul`` (QDQ) pair via MatMulNBitsToQDQ.  ``True`` leaves the
@@ -153,6 +159,7 @@ class EpCapabilities:
     supports_rotary_embedding: bool = True
     supports_tensor_scatter: bool = True
     supports_range: bool = True
+    supports_fp8_kv_cache: bool = False
     default_int4_accuracy_level: int = 0
     provider_options: dict[str, str] = dataclasses.field(default_factory=dict)
     enable_graph_capture: bool = False
@@ -309,6 +316,9 @@ def _register_builtins() -> None:
                 "enable_skip_layer_norm_strict_mode": "1",
             },
             supports_past_present_share_buffer=True,
+            # Only CUDA ships the FP8 (E4M3) GroupQueryAttention KV-cache kernel
+            # (SM89+ Ada/Hopper/Blackwell).
+            supports_fp8_kv_cache=True,
         ),
         EpCapabilities(
             name="dml",

src/mobius/_optimizations.py

@@ -398,9 +398,10 @@ def optimize_model(
         fp8_kv_cache: When ``True``, convert decoder
             ``GroupQueryAttention`` KV caches to ``FLOAT8E4M3FN`` (per-tensor
             E4M3) via :class:`~mobius._passes.Fp8KvCachePass`. Only applied
-            when GQA fusion is active for ``(ep, dtype)`` and
-            ``model_role == "decoder"``; otherwise a warning is emitted and the
-            request is ignored (no GQA nodes to convert).
+            when GQA fusion is active for ``(ep, dtype)``,
+            ``model_role == "decoder"``, and the EP ships the FP8 GQA kernel
+            (``caps.supports_fp8_kv_cache`` — currently CUDA only); otherwise a
+            warning is emitted and the request is ignored.
         kv_cache_scales: Optional ``layer_id -> (k_scale, v_scale)`` map of
             per-tensor FP8 scales (from offline calibration). Only used when
             ``fp8_kv_cache`` is ``True``; layers absent from the map use a unit
@@ -572,12 +573,13 @@ def _should_inline(func: ir.Function) -> bool:
     # GQA-capable EP/dtype); otherwise there is no KV-cache op to convert.
     if fp8_kv_cache:
         gqa_active = model_role == "decoder" and dtype in caps.gqa_dtypes
-        if not gqa_active:
+        if not gqa_active or not caps.supports_fp8_kv_cache:
             warnings.warn(
-                f"fp8_kv_cache=True was requested but GQA fusion is not active "
-                f"for ep={ep!r}/dtype={dtype}/role={model_role!r}. FP8 KV cache "
-                f"requires a GroupQueryAttention decoder (e.g. --execution-provider "
-                f"cuda with an fp16/bf16 dtype). Ignoring the request.",
+                f"fp8_kv_cache=True was requested but the FP8 GQA KV-cache kernel "
+                f"is not available for ep={ep!r}/dtype={dtype}/role={model_role!r}. "
+                f"FP8 KV cache requires a GroupQueryAttention decoder on an EP with "
+                f"the FP8 kernel (currently only --execution-provider cuda with an "
+                f"fp16/bf16 dtype). Ignoring the request.",
                 stacklevel=4,
             )
         else:

src/mobius/_passes/_fp8_kv_cache_test.py

@@ -145,6 +145,26 @@ def test_ignored_and_warns_on_non_gqa_ep(self):
         ins = {v.name: v for v in pkg["model"].graph.inputs}
         assert ins["past_key_values.0.key"].dtype == ir.DataType.FLOAT
 
+    def test_ignored_and_warns_on_non_fp8_ep_with_gqa(self):
+        """GQA-active but non-FP8 EP (CPU/float32) must NOT emit FP8 KV I/O.
+
+        Regression: fp8_kv_cache was previously applied whenever GQA fusion was
+        active, which includes CPU (gqa_dtypes={FLOAT}) — an EP without the FP8
+        GQA kernel — producing models that can't load/run.
+        """
+        config = _base_config(dtype=ir.DataType.FLOAT)
+        module = registry.get("qwen2")(config)
+        with pytest.warns(UserWarning, match="fp8_kv_cache=True"):
+            pkg = build_from_module(
+                module,
+                config,
+                "text-generation",
+                execution_provider="cpu",
+                fp8_kv_cache=True,
+            )
+        ins = {v.name: v for v in pkg["model"].graph.inputs}
+        assert ins["past_key_values.0.key"].dtype == ir.DataType.FLOAT
+
     def test_pass_is_idempotent(self):
         model, _config = _build_fp8_decoder()
         # Re-running the pass must not add extra inputs or re-type anything.