feat(kv-cache): add FP8 (E4M3) KV-cache export path - #447
Open
justinchuby wants to merge 5 commits into
Open
Conversation
Add an opt-in path that stores the com.microsoft::GroupQueryAttention past/present KV cache as FLOAT8E4M3FN instead of the model dtype, halving KV-cache memory at long context. Mirrors onnxruntime-genai's fp8_kv_cache option (microsoft/onnxruntime-genai#2351). A new post-fusion IR pass, Fp8KvCachePass, retypes each decoder GQA past_key/past_value input and present_key/present_value output to FP8, adds per-layer k_scale/v_scale FLOAT initializers at GQA input slots 12/13, and sets k_quant_type/v_quant_type="PER_TENSOR" and kv_cache_bit_width=8. The query/key/value inputs stay at the model dtype; the kernel quantizes the new K/V to FP8 on write and dequantizes on read using the scales. Scales default to a unit 1.0 (the "legacy" export shape) or come from an offline calibration file (onnxruntime-genai scales JSON format). Wired through optimize_model, build, build_from_module, and the `mobius build` CLI via --fp8-kv-cache and --kv-cache-scale-file. The option is ignored with a warning when GQA fusion is not active (e.g. non-CUDA EP or fp32), since there is no KV-cache op to convert. Verified on an H200 (SM90) with onnxruntime-gpu 1.28: the emitted FP8 GQA op loads and computes on the CUDA EP (finite fp16 output, FP8 present cache). End-to-end FP8 KV at runtime is IO-bound as device FLOAT8E4M3FN OrtValues (as onnxruntime-genai does); the Python numpy fp8 feed path is unsupported by ORT, which the runtime test works around with empty in-graph FP8 constants. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: justinchuby <justinchuby@users.noreply.github.com>
Performance Comparison
|
Address review feedback: - Fp8KvCachePass no longer early-skips a node whose past_key is already FP8. Every step is idempotent (retype no-op, scale initializers dedup by name, attributes overwrite, inputs only grow when < 14), so a re-run can never leave a node half-converted. - load_kv_cache_scale_file now rejects non-array k_scales/v_scales and any non-finite or non-positive scale, instead of silently accepting strings/dicts or 0/NaN/inf that would corrupt quantization. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: justinchuby <justinchuby@users.noreply.github.com>
Contributor
There was a problem hiding this comment.
Pull request overview
This PR adds an opt-in FP8 (E4M3FN) KV-cache export path by introducing a post-fusion IR pass that retypes com.microsoft::GroupQueryAttention cache I/O to FLOAT8E4M3FN, wires in per-layer k_scale/v_scale inputs, and sets the required quantization attributes. The feature is surfaced through optimize_model(), the Python build APIs, and the CLI mobius build command.
Changes:
- Add
Fp8KvCachePass+ JSON scale-file loader to retype KV cache I/O to FP8 and inject scale inputs/attrs. - Thread
fp8_kv_cache+kv_cache_scalesthroughoptimize_model(),build()/build_from_module(), and the CLI (--fp8-kv-cache,--kv-cache-scale-file). - Add structural + CUDA runtime tests validating graph wiring and runtime acceptance of the FP8 GQA signature.
Reviewed changes
Copilot reviewed 6 out of 6 changed files in this pull request and generated 4 comments.
Show a summary per file
| File | Description |
|---|---|
| src/mobius/_passes/_fp8_kv_cache.py | New optimization pass that converts GQA KV cache I/O to FP8 and adds scale inputs/quant attrs. |
| src/mobius/_passes/_fp8_kv_cache_test.py | New tests for pass structure, scale-file parsing, and CUDA runtime execution. |
| src/mobius/_passes/init.py | Export Fp8KvCachePass from the passes package. |
| src/mobius/_optimizations.py | Add fp8_kv_cache/kv_cache_scales options and run the new pass after fusion. |
| src/mobius/_builder.py | Plumb FP8 KV-cache options through build() / build_from_module() into optimization. |
| src/mobius/main.py | Add CLI flags and load optional scale-file for both --model and --config paths. |
Suppressed comments (2)
src/mobius/_passes/_fp8_kv_cache.py:129
- Using only
layer_idto form thekv_cache.<suffix>.k_scale/.v_scaleinitializer names risks collisions for multi-cache naming schemes (e.g.past_key_values.0.self.*andpast_key_values.0.cross.*). Since the cache value name is already available, derive the suffix from it so each distinct cache gets its own scale initializer.
present_value = node.outputs[2] if len(node.outputs) > 2 else None
_retype_fp8(present_key)
_retype_fp8(present_value)
src/mobius/_passes/_fp8_kv_cache_test.py:155
- This skip condition only checks for the CUDA EP, but the FP8 KV-cache GQA kernel is hardware-gated (SM89+). Without also checking compute capability (or skipping on kernel-not-found errors), this test can fail on older CUDA GPUs.
path = tmp_path / "bad.json"
Comment on lines
+65
to
+72
| def _layer_id_from_name(name: str | None) -> int | None: | ||
| """Return the layer index encoded in a ``past_key_values.{i}.key`` name.""" | ||
| if not name: | ||
| return None | ||
| match = _LAYER_ID_RE.search(name) | ||
| return int(match.group(1)) if match else None | ||
|
|
||
|
|
Comment on lines
+103
to
+108
| inputs = node.inputs | ||
| past_key = inputs[3] if len(inputs) > 4 else None | ||
| past_value = inputs[4] if len(inputs) > 4 else None | ||
| if past_key is None or past_value is None: | ||
| # Prompt-processing GQA without a KV cache — nothing to convert. | ||
| continue |
Comment on lines
+40
to
+43
| _FP8 = ir.DataType.FLOAT8E4M3FN | ||
| _CUDA = "CUDAExecutionProvider" | ||
| _HAS_CUDA = _CUDA in ort.get_available_providers() | ||
|
|
Contributor
There was a problem hiding this comment.
Pull request overview
Copilot reviewed 6 out of 6 changed files in this pull request and generated no new comments.
Suppressed comments (5)
src/mobius/_passes/_fp8_kv_cache.py:126
present_keyuseslen(node.outputs) > 2even though it reads index 1. This incorrectly returnsNonefor valid 2-output variants and makes the index contract inconsistent withpresent_value(index 2).
present_key = node.outputs[1] if len(node.outputs) > 2 else None
present_value = node.outputs[2] if len(node.outputs) > 2 else None
src/mobius/_passes/_fp8_kv_cache.py:15
- Docstring says the pass only applies when
past_key/past_valueare graph inputs, but the implementation converts any GQA node that has those inputs (including initializers, as in the CUDA runtime test). Updating the wording avoids misleading callers about when conversion occurs.
:func:`~mobius._optimizations.optimize_model`). For every decoder
``GroupQueryAttention`` node whose ``past_key`` / ``past_value`` inputs are
graph inputs it:
src/mobius/_passes/_fp8_kv_cache_test.py:244
- This test is gated only on the CUDA EP being available, but the FP8 KV-cache kernel is SM89+ only. On older GPUs with CUDA EP installed,
sess.run()will likely fail with a NOT_IMPLEMENTED/Fail error and turn into a flaky test failure. Wrap the execution in a try/except andpytest.skip(...)when the kernel isn't available on the current CUDA setup.
sess = ort.InferenceSession(str(path), providers=[_CUDA, "CPUExecutionProvider"])
outs = sess.run(
None,
{
"query": (np.random.randn(b, s, h * d) * 0.1).astype(np.float16),
src/mobius/_passes/_fp8_kv_cache.py:107
past_keyis guarded bylen(inputs) > 4, which is stricter than necessary for index 3 and makes the intent harder to read. Use independent bounds checks (> 3for index 3,> 4for index 4) so the logic is correct even if an upstream GQA variant ever exposes only some of the optional KV-cache inputs.
inputs = node.inputs
past_key = inputs[3] if len(inputs) > 4 else None
past_value = inputs[4] if len(inputs) > 4 else None
src/mobius/_optimizations.py:584
- The PR description/docstring says the option is ignored with a warning when GQA fusion is not active. Here,
gqa_activeonly checks EP/dtype capability, so if fusion didn't actually produce anyGroupQueryAttentionnodes (pattern miss, rule disabled, etc.), the pass becomes a silent no-op with no warning. Consider also checking that at least one GQA node exists before applying the pass, and warn otherwise.
if fp8_kv_cache:
gqa_active = model_role == "decoder" and dtype in caps.gqa_dtypes
if not gqa_active:
warnings.warn(
f"fp8_kv_cache=True was requested but GQA fusion is not active "
Copilot review feedback on the FP8 KV-cache pass:
- Enforce the documented contract: only retype KV caches that are graph
inputs or empty placeholders. A non-empty FLOAT/FLOAT16 initializer
cache is now skipped with a warning instead of being declared FP8 over
FLOAT bytes (which would corrupt the graph).
- Guard each present output on its own index (present_key at 1,
present_value at 2) so a hypothetical 2-output GQA cannot leave
present_key a different dtype from past_key.
- Derive scale-initializer names from the cache tensor names
(`{past}.fp8_scale`) rather than the numeric layer id, so distinct
caches at the same layer index (e.g. seq2seq self- vs cross-attention)
never share a scale initializer.
- Gate the CUDA runtime test on compute capability SM89+ (best-effort via
torch.cuda.get_device_capability), skipping on older GPUs / unknown
instead of running where the FP8 GQA kernel is absent.
- Add a test covering the non-empty-initializer skip contract.
Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
Signed-off-by: justinchuby <justinchuby@users.noreply.github.com>
Contributor
There was a problem hiding this comment.
Pull request overview
Copilot reviewed 6 out of 6 changed files in this pull request and generated 1 comment.
Suppressed comments (2)
src/mobius/_passes/_fp8_kv_cache_test.py:312
- The CUDA runtime test allows CPU fallback (
providers=[CUDA, CPU]), so it can still pass even if the FP8 GQA node isn’t actually supported/executed on CUDA. To make this test prove CUDA support, create the session with only the CUDA EP so session creation or execution fails if CUDA can’t run the graph.
path = tmp_path / "gqa.onnx"
ir.save(model, str(path))
sess = ort.InferenceSession(str(path), providers=[_CUDA, "CPUExecutionProvider"])
outs = sess.run(
src/mobius/_passes/_fp8_kv_cache_test.py:31
sys.path.insert(0, "tests")depends on the current working directory being the repo root; running this test from another CWD can make_test_configsfail to import. Prefer computing thetests/path relative to__file__so the import is stable.
import json
import sys
import ml_dtypes
import numpy as np
import onnx_ir as ir
import onnxruntime as ort
import pytest
sys.path.insert(0, "tests")
from _test_configs import _base_config
Comment on lines
+573
to
+584
| if fp8_kv_cache: | ||
| gqa_active = model_role == "decoder" and dtype in caps.gqa_dtypes | ||
| if not gqa_active: | ||
| 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.", | ||
| stacklevel=4, | ||
| ) | ||
| else: | ||
| Fp8KvCachePass(kv_cache_scales)(model) |
Collect the boolean build-mode toggles under a single Rust/cargo-style
`--features` option on `mobius build`. It accepts a comma-separated list
and may be repeated:
mobius build ... --features fp8-kv-cache,static-cache
mobius build ... --features fp8-kv-cache --features static-cache
Available features: static-cache, fp8-kv-cache, text-only. Each maps to
the existing boolean build attribute via the `_BUILD_FEATURES` registry;
`_resolve_build_features` folds them in at the top of `_cmd_build` before
any validation reads them. Unknown feature names are rejected with an
error that lists the valid set.
The legacy `--static-cache`, `--fp8-kv-cache`, and `--text-only` flags
remain as deprecated aliases: they still set the same behaviour but print
a deprecation notice on stderr. Companion value args (--max-seq-len,
--kv-cache-scale-file) are unchanged.
Adds CLI tests for feature parsing, comma-separated/multiple features,
unknown-feature error, static-cache equivalence, text-only/fp8 pass-
through, and the deprecation warning. Updates docs/cli_reference.md and
CHANGELOG.md.
Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
Signed-off-by: justinchuby <justinchuby@users.noreply.github.com>
Contributor
There was a problem hiding this comment.
Pull request overview
Copilot reviewed 9 out of 9 changed files in this pull request and generated no new comments.
Suppressed comments (3)
src/mobius/_passes/_fp8_kv_cache.py:80
- _retype_fp8() updates the ValueInfo dtype but leaves an empty initializer's const_value tensor dtype unchanged. This makes the graph internally type-inconsistent (initializer tensor still FLOAT/FLOAT16 while the Value is declared FP8) and forces callers/tests to manually retag the initializer before saving/running. If empty placeholder initializers are intended to be supported, the pass should also rewrite the empty const_value tensor to FP8 (safe because size==0).
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)
src/mobius/_passes/_fp8_kv_cache.py:124
- The KV-cache input detection uses
inputs[3] if len(inputs) > 4for past_key, which is an off-by-one condition (index 3 is valid when len(inputs) > 3). It currently works because the code also requires past_value, but the condition is misleading and makes the intent harder to follow. Simplify to a singlelen(inputs) <= 4guard, then index 3/4 directly.
inputs = node.inputs
past_key = inputs[3] if len(inputs) > 4 else None
past_value = inputs[4] if len(inputs) > 4 else None
if past_key is None or past_value is None:
# Prompt-processing GQA without a KV cache — nothing to convert.
src/mobius/main.py:239
- The validation error for
--kv-cache-scale-filestill points users to the deprecated--fp8-kv-cacheflag, even though--features fp8-kv-cacheis now the canonical way to enable this mode. Updating the message to reference the feature avoids directing users toward deprecated CLI usage.
fp8_kv_cache = getattr(args, "fp8_kv_cache", False)
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.")
if fp8_kv_cache and scale_file is not None:
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>
Contributor
There was a problem hiding this comment.
Pull request overview
Copilot reviewed 10 out of 10 changed files in this pull request and generated no new comments.
Suppressed comments (4)
src/mobius/_passes/_fp8_kv_cache.py:81
- _retype_fp8() updates the ValueInfo dtype but leaves an empty const_value tensor (used for placeholder initializers / constants) at its original dtype. That can create an inconsistent graph where the node input is declared FP8 but the constant/initializer is still FLOAT16/FLOAT, which can break model loading/type checking even when the tensor has zero elements.
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)
src/mobius/_passes/_fp8_kv_cache.py:137
- Fp8KvCachePass currently relies on const_value-emptiness to decide whether a KV cache is safe to retype, but it does not enforce the documented contract that past_key/past_value must be graph inputs. This allows converting caches produced by other nodes (e.g. Constant outputs), which risks retyping internal tensors that the runtime/kernel may not treat as KV-cache I/O. Add an explicit producer()==None check (graph input / initializer) before retyping.
inputs = node.inputs
past_key = inputs[3] if len(inputs) > 4 else None
past_value = inputs[4] if len(inputs) > 4 else None
if past_key is None or past_value is None:
# Prompt-processing GQA without a KV cache — nothing to convert.
continue
# Contract: only KV caches that are graph inputs (or empty
# placeholders) may be retyped. Retyping a non-empty FLOAT/FLOAT16
# initializer would declare FP8 over FLOAT bytes and corrupt the
# graph, so skip such nodes loudly rather than emit an invalid model.
if not _is_retypable_cache(past_key) or not _is_retypable_cache(past_value):
warnings.warn(
f"Fp8KvCachePass: skipping {node.name!r} — its KV cache is a "
f"non-empty initializer, not a graph input. Only graph-input "
f"KV caches can be converted to FP8.",
stacklevel=2,
)
continue
src/mobius/main.py:809
- The --kv-cache-scale-file help text says it's "Only used with --fp8-kv-cache", but that flag is deprecated and users should enable FP8 KV cache via --features fp8-kv-cache. Updating the help string will keep CLI docs consistent with the new feature mechanism.
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."
),
src/mobius/main.py:238
- The error message for using --kv-cache-scale-file without FP8 KV cache references only the deprecated flag (--fp8-kv-cache). Since --features fp8-kv-cache is now the canonical path, the message should mention it to avoid confusing users.
This issue also appears on line 804 of the same file.
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.")
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
Adds an opt-in FP8 (E4M3) KV-cache export path. When enabled, each decoder
com.microsoft::GroupQueryAttentionstores its past/present key-value cache asFLOAT8E4M3FNinstead of the model dtype (fp16/bf16), halving KV-cache memory at long context. Mirrors onnxruntime-genai'sfp8_kv_cacheoption (microsoft/onnxruntime-genai#2351).How it works
A new post-fusion IR pass
Fp8KvCachePass(src/mobius/_passes/_fp8_kv_cache.py), run after the GQA fusion rules, for every decoder GQA node whose KV cache are graph inputs:past_key/past_valueinputs andpresent_key/present_valueoutputs (all graph I/O) toFLOAT8E4M3FN, preserving shapes.k_scale/v_scalescalar FLOAT initializers at GQA input slots 12/13 (ORT ≥ 1.28 signature).k_quant_type/v_quant_type = "PER_TENSOR"andkv_cache_bit_width = 8.query/key/valuestay at the model dtype — the kernel quantizes the new K/V to FP8 on write and dequantizes on read using the scales. Scales default to a unit1.0(the "legacy" export shape), or come from an offline calibration file (onnxruntime-genai{"scales": {"k_scales": [...], "v_scales": [...]}}JSON).Usage
Also exposed programmatically on
build(...)/build_from_module(...)viafp8_kv_cache=andkv_cache_scales=. The option is ignored with a warning when GQA fusion is not active (non-GQA EP or fp32), since there is no KV-cache op to convert.Verification (H200 / SM90, onnxruntime-gpu 1.28)
test_fp8_kv_gqa_runs_on_cuda, gated on CUDA availability).FLOAT8E4M3FNOrtValues (as onnxruntime-genai does). The Python numpy fp8 feed path is unsupported by ORT; the runtime test works around it with empty in-graph FP8 constants.Tests
src/mobius/_passes/_fp8_kv_cache_test.py— 11 tests: KV I/O typed FP8, GQA scale inputs + quant attrs, default unit scales, calibrated scales, disabled path, ignored+warns on non-GQA EP, idempotency, scale-file parsing/validation, and the CUDA runtime run. Existing pass/CLI/build-graph suites pass; changed files pass ruff.