diff --git a/.agents/skills/mobius-onnx-export-gotchas/SKILL.md b/.agents/skills/mobius-onnx-export-gotchas/SKILL.md new file mode 100644 index 00000000..af382642 --- /dev/null +++ b/.agents/skills/mobius-onnx-export-gotchas/SKILL.md @@ -0,0 +1,287 @@ +--- +name: mobius-onnx-export-gotchas +description: Use when building/exporting ONNX models with the `mobius build` CLI (especially Phi-3 / Phi-3.5 or any model with `--execution-provider cuda` GQA fusion and/or `--static-cache`). Covers the current CLI syntax, the dtype flag values, the GQA-vs-static-cache interaction, how to verify fp16 GQA exports load in onnxruntime (the historical packed-QKV FLOAT32 load bug is fixed by the fp16 GQA fold-fix), why fp16 GQA exports need VALUE-based weight checks (corr≈1.0 / norm), not just initializer count/dtype, to catch silently-zeroed packed-QKV weights, and the graph-capture-compatibility rule for graph construction (never emit in-graph If/Loop/Scan control flow — it hard-fails session init under CUDA Graph capture). +--- + +# mobius ONNX export gotchas + +## 1. CLI syntax (editable repo differs from older docs) +`mobius build` requires `--model ` and takes the **output dir as a POSITIONAL** arg. +There is **no `-o` flag for `build`** (`-o` exists only on `build-gguf`). + +```bash +mobius build --model microsoft/Phi-3.5-mini-instruct \ + --dtype f16 --execution-provider cuda \ + --external-data onnx --trust-remote-code \ + /path/to/output_dir +``` + +- `--dtype` choices: `f16`/`float16`, `bf16`/`bfloat16`, `f32`/`float32`. **`fp16` is INVALID.** +- `--execution-provider` is an alias of `--ep`. `cuda` + fp16/bf16 triggers GQA fusion; + `default` keeps plain ONNX `Attention`. + +## 2. `--static-cache` is incompatible with GQA fusion +`--static-cache` wraps each attention with `TensorScatter` (in-place KV cache for the **ONNX Attention** +op). That breaks the pattern the GQA rewrite matches, so combining +`--execution-provider cuda --static-cache` yields **0 GroupQueryAttention + N Attention + 2N TensorScatter** +(mobius prints: "GQA fusion expected … but found 0 GroupQueryAttention and N Attention nodes"). + +- **GQA model:** `--execution-provider cuda` **alone**. GQA's shared KV buffer + (`past_present_share_buffer`) is enabled at **runtime** via IO-binding past & present to the same + OrtValue — NOT via `--static-cache`. +- **ONNX-Attention + in-place cache:** `--execution-provider default --static-cache --max-seq-len N`. + +### Static cache requires `is_causal=0` + an explicit mask (not `is_causal=1`) +The opset-24 `Attention` CUDA kernel rejects `is_causal=1` together with +`nonpad_kv_seqlen` when `S_q != total_kv` and there is no `past_key` (the +`causal_cross_no_past` guard in `attention.cc`) → ORT raises `NOT_IMPLEMENTED` at +session init. A static cache is pre-allocated to `max_seq_len`, so `S_q != total_kv` +in **both** prefill and decode and the guard always fires. + +- **Remedy:** drive the static-cache `Attention` with `is_causal=0` and supply an + explicit offset-aware causal mask (`create_static_cache_causal_mask` — keeps key + slot `j` for a query at absolute position `write_indices + t` iff `j <= write_indices + t`). + This is also **branchless**, hence graph-capture-safe — do **not** express the + prefill/decode masking difference with an `If(Greater(S_q, 1))` phase-split. See + §7 (Graph-capture compatibility) for why in-graph control flow is disallowed. + +## 3. FIXED: fp16 GQA export previously left packed-QKV weights as FLOAT32 → model wouldn't load +**Status: fixed (the fp16 GQA fold-fix).** Native fp16 Phi-3.5 GQA export now loads directly in the ORT +CUDA EP with **no manual post-cast** (32 GroupQueryAttention nodes, all-fp16 initializers). If you are on +that commit or later, you should not hit this — skip to the verification snippet below. The history is +kept here because old artifacts exported before the fix still carry fp32 packed weights. + +### Symptom (pre-fix) +For an fp16 GQA export, a folded per-layer packed QKV weight +(`..q_proj.weight__k_proj.weight__v_proj.weight__axis_0__concat`) was emitted as **FLOAT32**, while its +MatMul's other input was fp16. onnxruntime then rejected the model at load on both CPU and CUDA EPs: + +``` +Type Error: Type parameter (T) of Optype (MatMul) bound to different types +(tensor(float16) and tensor(float)) in node (node_MatMul_*) +``` + +You'd also see at save time: `The value type for shape [H, 3H] is not known. Skipping serialization`. + +### Root cause +`_cast_module_dtype` casts module params to fp16, but the resulting initializer `Value`s lose their +declared `.dtype` (it becomes `None`) while their `const_value` stays fp16. The fold passes +`FoldConcatInitializersPass` (`src/mobius/_passes/_fold_concat.py`) and `FoldTransposedInitializerPass` +(`src/mobius/_passes/_fold_transpose.py`) then defaulted the folded initializer's dtype to `FLOAT`, +serializing the packed QKV / transposed weights as fp32. + +### The fix +A shared helper `initializer_dtype()` (`src/mobius/_passes/_dtype_utils.py`) resolves the effective dtype +from the declared type, **falling back to `const_value` when the type annotation was dropped** (preferring +the data dtype and warning on stale-metadata disagreement). Both fold passes use it to stamp the correct +dtype on the new initializer's `TensorType` and `LazyTensor`, and `FoldConcatInitializersPass` now also +skips folding before weights are loaded (mirroring `FoldTransposedInitializerPass`). A regression test +loads the fp16 GQA export in the ORT CPU EP to lock this in. + +### Verify (still worth running on any fp16 build) +```python +import onnx +m = onnx.load("model.onnx", load_external_data=False) +fp32 = [i.name for i in m.graph.initializer if i.data_type == onnx.TensorProto.FLOAT] +print(len(fp32), "FLOAT32 initializers (should be 0 for fp16)") +``` + +### Convention (prevents the whole class from reappearing) +Any pass that **materializes a new initializer** must resolve its dtype via +`initializer_dtype()` (`src/mobius/_passes/_dtype_utils.py`), **never** `value.dtype or ir.DataType.FLOAT`. +The bug class originates in `_cast_module_dtype` dropping a `Value`'s declared `.dtype` (→ `None`) while its +`const_value` stays fp16; a bare `.dtype or FLOAT` fallback then silently mis-types the result as fp32. Fold +passes (`_fold_concat.py`, `_fold_transpose.py`) already follow this; mirror it in any future +initializer-producing pass. Siblings still reading `.dtype` directly remain exposed — a follow-up should +grep `_passes/` for `.dtype or ir.DataType` and consider re-stamping the type in `_cast_module_dtype` to kill +the class at source. + +### Salvaging a stale pre-fix artifact (only if re-exporting is not an option) +Prefer re-exporting on the fixed code. If you must repair an old model, cast its FLOAT32 initializers to +fp16 and re-save. **Gotcha when re-saving with external data:** if you save with `location="X.data"` and +then rename the file, the references inside `model.onnx` still point to `X.data`. Either save directly +with `location="model.onnx.data"`, or rewrite each initializer's `external_data` `location` entry. + +```python +import onnx, numpy as np +from onnx import numpy_helper, TensorProto +m = onnx.load("model.onnx", load_external_data=True) +for init in m.graph.initializer: + if init.data_type == TensorProto.FLOAT: + arr = numpy_helper.to_array(init).astype(np.float16) + init.CopyFrom(numpy_helper.from_array(arr, init.name)) +onnx.save(m, "model.onnx", save_as_external_data=True, all_tensors_to_one_file=True, + location="model.onnx.data", size_threshold=1024, convert_attribute=False) +``` + +## 4. Always validate the export in ORT before profiling +Load the model on `CUDAExecutionProvider` and run one prefill + one decode `session.run`. Confirm: +(a) the expected attention op (`com.microsoft::GroupQueryAttention` vs `ai.onnx::Attention`), +(b) finite fp16 logits, (c) no FLOAT32 initializers for an fp16 build. + +These checks are **necessary but NOT sufficient** for a fp16 GQA export — see §5. A model can pass all +three and still have silently-zeroed packed-QKV weights. + +## 5. Verifying a fp16 GQA export: use VALUE-based weight checks, NOT initializer count/dtype +**A fp16 GQA export can be all-fp16, right-count, and still all-zeros — only a corr≈1.0 / norm≈126 VALUE +check on the packed QKV proves the weights are real.** + +### Symptom +The GQA model loads cleanly (32 `GroupQueryAttention` nodes, all-fp16, finite logits) but generates +garbage (e.g. `holdou_(...artersarters`). Prefill logits come out ~3× the reference scale, with +`max|Δlogit|` ~50+ versus the reference. + +### Root cause +The packed-QKV initializer is `Transpose(Concat(q, k, v, axis=0))`. If the fold passes +(`FoldConcatInitializersPass` / `FoldTransposedInitializerPass`) leave the packed-Concat output dtype +UNKNOWN / defaulted-to-fp32 while the data is fp16, the serializer **skips** it and it loads as +**near-zero** — the weights are silently dead. (This is the §3 failure mode; the upstream fix in +The fp16 GQA fold-fix stamps the fp16 dtype at the fold-pass source. A post-hoc cast is NOT a fix — it re-corrupts.) + +### Why count/dtype checks fail (the trap) +The BROKEN export and the FIXED export can have the **same initializer count and the same fp16/fp32 dtype +ratio**, so neither is a validity signal. Worse, the fp16-init count is **not even stable across fixes** +— on Phi-3.5 it moved from ~293 down to ~197 (an unstripped intermediate carries the packed-QKV plus the +now-dead unpacked q/k/v source initializers; a safe dead-weight strip then removes the ~96 dead pre-pack +inits), with no bearing on correctness. Note the OLD broken export was *also* 197 fp16, so even a "right" +final count proves nothing. Counting initializers or checking "0 fp32 / all fp16" does **not** distinguish a +healthy model from a zeroed-weight one. §4(c) alone will pass a dead model. **Never gate on the count; +use the VALUE gate below.** + +### Canonical verification (load-bearing, not optional) +VALUE-based per-slice check on each packed-QKV initializer against its source q/k/v weights: +- per-slice correlation **≈ 1.000** (broken ≈ 0.000), AND +- packed-QKV L2 norm **≈ 126.6** at layer 0 / mean(|abs|) **≈ 0.015** (broken ≈ 0.80 / ≈ 5e-6). + +> ⚠️ **Use mean-of-ABS or norm — NEVER the signed mean.** The good model's *signed* mean is ~2.6e-6 +> (near zero, because the weights are symmetric ±), which coincidentally looks just like the broken +> model's mean(|abs|) ~5e-6. Checking signed mean would **falsely flag the good model as broken** — this +> exact confusion has already caused a false alarm in this crew. Valid discriminators: mean(|abs|) +> (good ≈ 0.015 vs broken ≈ 5e-6) or L2 norm (good ≈ 126.6 vs broken ≈ 0.80). + +Plus an end-to-end next-token greedy-argmax parity check vs the `attn_dynamic` reference (expect +**~19–20 / 20**). Isolated single-token divergences are fp16 dead-ties (reference top1−top2 gap = 0.0000), +not bugs. Optional hardening: assert **0 unused initializers** and that all N packed-QKV initializers are +present, to catch dead-weight OVER-stripping. + +QA's `gqa_weight_integrity_gate.py` (`--self-check --strip-audit --scan-all`, per-layer corr/norm) +implements exactly this gate. + +## 6. FIXED: GQA `present.*` KV-cache outputs declared the wrong `head_dim` +**Status: fixed (the GQA present-KV shape fix).** A native fp16 GQA export now declares +`present.{i}.{key,value}` with the correct `head_dim`, symmetric to its `past_key_values.{i}.*` inputs. + +### Symptom (pre-fix) +The graph **output** `present.{i}.key/value` declared the wrong `head_dim` (e.g. `32` instead of the real +`96` on Phi-3.5) while the matching `past_key_values.{i}.*` **input** was correct (`96`). At load ORT logged +(once per key+value per layer — 64 on Phi-3.5): + +``` +[W ...MergeShapeInfo] Error merging shape info for output. 'present.0.key' +source:{-1,32,-1,96} target:{-1,32,-1,32}. Falling back to lenient merge. +``` + +Runtime still produced correct (96-wide) arrays via lenient merge, but any consumer that **trusts declared +shapes** (e.g. `onnxruntime-genai`) would see inconsistent past-vs-present KV cache types. + +### Root cause +`GroupQueryAttention`'s contrib-op shape inference mis-derives the present `head_dim` (it does **not** +reproduce on the plain `Attention` op, which infers correctly). `_register_kv_cache_outputs` +(`src/mobius/tasks/_cache_utils.py`) added the present outputs with **no explicit shape**, so the buggy +inference won. + +### The fix +`_register_kv_cache_outputs` now opt-in **stamps** `present.{i}.{key,value}` shape+dtype symmetric to the +past inputs when the caller passes `batch`/`num_kv_heads`/`key_head_dim`/`value_head_dim`/`total_seq_len`/ +`dtype` (wired from `_causal_lm.py`). Omitting them preserves inference-only behavior, so the other ~10 +callers are unaffected. The stamp survives `SymbolicShapeInferencePass` (policy `refine` only tightens +unknown dims; it won't replace a concrete `96` with a conflicting `32`). + +### Verify +```python +import onnx +m = onnx.load("model.onnx", load_external_data=False) +d = lambda vi: [(x.dim_param or x.dim_value) for x in vi.type.tensor_type.shape.dim] +o = {v.name: v for v in m.graph.output} +print("present.0.key:", d(o["present.0.key"])) # head_dim must equal the past input's (e.g. 96, NOT 32) +``` + +### Known remaining (separate, pre-existing, harmless) +ORT still logs ~32 `Error merging shape info ... source:{-1,-1,3072} target:{-1,-1,1024}` warnings on the +GQA op's **internal hidden-state output** value_info (`v_*.GroupQueryAttention_*_0`, `1024`=32×32 vs the +correct `3072`=32×96). That value is **not** a declared graph I/O — runtime is correct and `onnxruntime-genai` +does not trust it — so it does not bite shape-trusting consumers the way the present-output bug did. Tracked +as a follow-up in the GQA rewrite emission path (not the KV-cache output path). + +## 7. Graph-capture compatibility: never emit in-graph control flow + +This is general guidance for **anyone constructing model graphs in mobius**, broader +than the static-cache case below — it constrains how *every* exported graph must be +built. + +### Principle: exported graphs must run under graph capture, not just eager + +mobius ships models for **onnxruntime / onnxruntime-genai**. On the production and edge +execution paths these run under **CUDA Graph capture**: + +- CUDA EP with `enable_cuda_graph=1`, +- DML EP **always** captures, +- NvTensorRtRtx EP captures (with shared buffers), +- onnxruntime-genai's default decode design uses `past_present_share_buffer` + static + shapes, purpose-built for capture. + +Capture records the kernel-launch sequence once and replays it, eliminating the +per-launch overhead that dominates the decode loop. So an exported graph **must** be +capture-compatible. **Do not ship a model that only runs in eager.** + +### Hard constraint: control-flow nodes make session init FAIL under capture + +When a graph-capture-enabled EP loads a model that contains **control-flow nodes**, ORT +**hard-fails session init** — it does not silently fall back to eager. + +- Source: onnxruntime `core/session/inference_session.cc`. `HasControlflowNodes(graph)` + returns true for any node that owns a subgraph (**`If` / `Loop` / `Scan`**), and ORT + then returns `ORT_MAKE_STATUS(..., FAIL, "This session cannot use the graph capture + feature as the model has control flow nodes which can't be supported by ")`. +- It is **branch-agnostic**: an `If` that always takes the same branch at runtime + **still** triggers the failure — the node is physically present in the graph, which is + all the check inspects. +- Secondary break: capture requires a single partition / no `Memcpy` nodes; control-flow + plumbing that lands partly on the CPU EP also defeats capture even where the primary + check would not. + +### Empirically confirmed + +A per-layer-`If` static-cache export loaded on **CUDA EP with `enable_cuda_graph=1`** +fails with `FAIL: ... model has control flow nodes which can't be supported by +CUDAExecutionProvider`. The **same** model loads fine in **eager** (`enable_cuda_graph=0`). +DML and NvTensorRtRtx follow the same EP-agnostic gate (source-reasoned; DML not +runtime-tested on Linux). + +### Measurement caveat + +A control-flow node's **own** GPU cost measured in **eager** (e.g. a CUPTI per-node +profile) can read as ~0 and is **not** representative: in capture mode that same node +prevents capture *entirely*, so its real cost is "the whole model can't be captured." +Always evaluate export choices against the **capture path**, not eager. + +### Rule for graph construction + +Avoid **in-graph, data-dependent control flow** (`If` / `Loop` / `Scan`). Express +phase- or shape-dependent behavior via either: + +1. **Host-side dispatch** — separate prefill/decode `Run` calls (which onnxruntime-genai + already does), or +2. **Branchless formulations** — compute the same result without a subgraph. + +Concrete: static-cache causal masking must use **`is_causal=0` + an explicit +offset-aware mask** (branchless), **not** an `If(Greater(S_q, 1))` phase-split. + +### Cautionary example (why this PR ships the branchless path) + +A per-layer `If(Greater(S_q, 1))` phase-split (prefill-masked / decode-maskless) was +numerically correct and even kept decode on the Flash kernel — but it **fails to load +under graph capture** and was abandoned for exactly that reason. The always-masked +**`is_causal=0` + explicit-mask** path (decode runs on Memory-Efficient Attention) is +**branchless and capture-safe**, and that is what this PR ships. See §2 / the +`create_static_cache_causal_mask` helper for the masking mechanics. diff --git a/CHANGELOG.md b/CHANGELOG.md index 83edfb24..d6bca84b 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -7,6 +7,96 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ## [Unreleased] +### KV-cache present-shape: fail-closed on partial parameter sets + +#### Fixed + +- `_register_kv_cache_outputs` now **raises `ValueError`** when given a partial + set of present-shape parameters (1–5 of the six `batch`, `num_kv_heads`, + `key_head_dim`, `value_head_dim`, `total_seq_len`, `dtype`) instead of logging + a warning and proceeding. A partial set is always a wiring slip with no + legitimate use; the previous fail-open shipped a structurally-wrong model + (mis-derived `GroupQueryAttention` present `head_dim`) with only a log line. + Passing all six (stamp) or none (infer) is unaffected. (closes #341) + +### Static-cache mask: build-once hoist + +#### Performance + +- The static-cache causal mask is now built **once** and shared (same + `ir.Value`) across all decoder layers instead of being rebuilt per layer. The + mask is layer-invariant (it depends only on `S_q`, `max_seq` and + `write_indices`, all identical across layers). A new + `StaticCacheState.causal_mask` field threads the shared value from + `_make_static_cache_inputs`; `_apply_attention` consumes it and keeps a + fallback that builds the mask on demand for direct callers. This removes + `~(mask_nodes)·(num_layers − 1)` duplicated nodes from the graph, scaling with + layer count (e.g. ~16 nodes × 31 saved on a 32-layer model). + +### Static-cache Attention Causal-Mask Fix + +#### Fixed + +- Static-cache exports (`--static-cache`, the ONNX `Attention` + `TensorScatter` + in-place KV-cache path) now drive the `Attention` op with `is_causal=0` plus an + explicit offset-aware causal mask instead of `is_causal=1`. The opset-24 + `Attention` CUDA kernel rejects `is_causal=1` together with `nonpad_kv_seqlen` + when the query length differs from the total KV length and there is no + `past_key` (the `causal_cross_no_past` guard in `attention.cc`). Because the + static cache is pre-allocated to `max_seq_len`, `S_q != total_kv` in **both** + prefill and decode, so that guard fired and ORT raised `NOT_IMPLEMENTED` at + session init. A new `create_static_cache_causal_mask` helper builds a mask that + keeps key slot `j` for a query at absolute position `write_indices[b] + t` iff + `j <= write_indices[b] + t`, serving prefill (triangular) and decode (prefix) + with a single rule. The mask is **causal-only**; the `nonpad_kv_seqlen` + key-bound is enforced by the ORT `Attention` kernel itself (external-cache + input #6, on both CUDA and CPU EPs), so the mask does not re-encode it. + Non-compact / interior padding holes are out of contract (a scalar + `nonpad_kv_seqlen` cannot express them). The formulation is + branchless (no `If` phase-split), so the exported graph stays compatible with + CUDA Graph capture (see the export-gotchas skill, "Graph-capture + compatibility"); decode and prefill both run on the Memory-Efficient Attention + kernel. + +--- + +### fp16 GQA Export Fix + +#### Fixed + +- Native fp16 GroupQueryAttention exports (e.g. `microsoft/Phi-3.5-mini-instruct` + with `--dtype f16 --execution-provider cuda`) no longer emit fp32 packed-QKV / + transposed weights. Previously the fold passes (`FoldConcatInitializersPass`, + `FoldTransposedInitializerPass`) defaulted a folded initializer's dtype to + `FLOAT` when the source `Value`'s declared type had been dropped during fp16 + casting, producing a model onnxruntime rejected at load with a + `MatMul` type-parameter error (`tensor(float16)` vs `tensor(float)`) on both + CPU and CUDA EPs. A new `mobius._passes._dtype_utils.initializer_dtype()` + helper now resolves the effective dtype from `const_value` when the type + annotation is missing, so fp16 GQA models load directly with no manual + post-cast. + +--- + +### GQA Present KV-Cache Shape Fix + +#### Fixed + +- GroupQueryAttention exports now declare correct `present.{i}.key` / + `present.{i}.value` graph-output shapes and dtype. The GQA contrib op's shape + inference mis-derived the present KV `head_dim` (e.g. 32 instead of 96 on + `microsoft/Phi-3.5-mini-instruct`), so the present KV-cache outputs declared a + `head_dim` inconsistent with the (correct) `past_key_values` inputs. ORT logged + `Error merging shape info ... lenient merge` (64 warnings on Phi-3.5) and any + consumer that chains `present` → `past` and trusts declared shapes (e.g. + `onnxruntime-genai`) saw mismatched past-vs-present KV cache types. This is a + metadata / declared-shape correction only — runtime numerics are unchanged + (weights byte-identical, next-token parity 20/20). `_register_kv_cache_outputs` + now stamps the present KV outputs symmetric to the past inputs. Affects + GQA-fusion packed-QKV exports (Phi-3.5, Llama-3.2, Qwen2, Mistral, Phi-3-GQA). + +--- + ### WebGPU Shape Op Support #### Changed diff --git a/src/mobius/_passes/_dtype_utils.py b/src/mobius/_passes/_dtype_utils.py new file mode 100644 index 00000000..d5368dd2 --- /dev/null +++ b/src/mobius/_passes/_dtype_utils.py @@ -0,0 +1,55 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT License. + +"""Shared dtype helpers for graph passes that materialize new initializers. + +Passes such as :class:`~mobius._passes.FoldConcatInitializersPass` and +:class:`~mobius._passes.FoldTransposedInitializerPass` pre-compute new +initializers from existing ones. They must stamp the *correct* dtype on the +result, otherwise an fp16 model can silently end up with fp32 weights that +onnxruntime rejects at load time (a MatMul binding fp16 and fp32 to the same +type parameter ``T``). +""" + +from __future__ import annotations + +import logging + +import onnx_ir as ir + +logger = logging.getLogger(__name__) + + +def initializer_dtype(value: ir.Value) -> ir.DataType | None: + """Return the effective dtype of an initializer ``value``. + + Uses the value's declared ``type`` dtype, but falls back to the dtype of its + ``const_value`` when the type annotation is missing. When both are present + but disagree, the ``const_value`` dtype wins (it is the data actually + serialized) and a warning is logged, since a healthy initializer should + never have a declared type that contradicts its data. + + Graph building can drop the declared ``type`` on an initializer while its + actual tensor data (``const_value``) still carries the correct dtype. In + that situation, defaulting to ``ir.DataType.FLOAT`` would emit fp32 weights + into an otherwise fp16 model. Reading the dtype from ``const_value`` keeps + folded initializers consistent with the weights they are derived from. + + Returns ``None`` only when neither the declared type nor ``const_value`` is + available; callers decide on a final fallback. + """ + declared = value.dtype + const_dtype = value.const_value.dtype if value.const_value is not None else None + + if declared is not None and const_dtype is not None and declared != const_dtype: + logger.warning( + "Initializer %r declares dtype %s but its data is %s; using the data " + "dtype. This indicates stale type metadata.", + value.name, + declared, + const_dtype, + ) + return const_dtype + if declared is not None: + return declared + return const_dtype diff --git a/src/mobius/_passes/_dtype_utils_test.py b/src/mobius/_passes/_dtype_utils_test.py new file mode 100644 index 00000000..4267eb7e --- /dev/null +++ b/src/mobius/_passes/_dtype_utils_test.py @@ -0,0 +1,49 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT License. + +"""Tests for the shared initializer dtype helper.""" + +from __future__ import annotations + +import logging + +import numpy as np +import onnx_ir as ir + +from mobius._passes._dtype_utils import initializer_dtype + + +def _value(dtype: ir.DataType | None, const: np.ndarray | None) -> ir.Value: + v = ir.Value(name="w") + if dtype is not None: + v.dtype = dtype + if const is not None: + v.const_value = ir.tensor(const) + return v + + +class TestInitializerDtype: + def test_uses_declared_dtype_when_present(self): + v = _value(ir.DataType.FLOAT16, np.ones((2,), np.float16)) + assert initializer_dtype(v) == ir.DataType.FLOAT16 + + def test_falls_back_to_const_value_when_declared_missing(self): + """The core fix: a dropped declared type must not hide the real dtype.""" + v = _value(None, np.ones((2,), np.float16)) + assert v.dtype is None + assert initializer_dtype(v) == ir.DataType.FLOAT16 + + def test_const_value_wins_on_disagreement(self, caplog): + """Stale declared metadata must not override the serialized data dtype.""" + v = _value(ir.DataType.FLOAT, np.ones((2,), np.float16)) + with caplog.at_level(logging.WARNING, logger="mobius._passes._dtype_utils"): + assert initializer_dtype(v) == ir.DataType.FLOAT16 + assert any( + record.levelno == logging.WARNING + and "stale type metadata" in record.getMessage().lower() + for record in caplog.records + ), "expected a warning when declared dtype disagrees with const_value" + + def test_returns_none_when_nothing_available(self): + v = _value(None, None) + assert initializer_dtype(v) is None diff --git a/src/mobius/_passes/_fold_concat.py b/src/mobius/_passes/_fold_concat.py index b8f2eed9..7035f1e9 100644 --- a/src/mobius/_passes/_fold_concat.py +++ b/src/mobius/_passes/_fold_concat.py @@ -27,6 +27,8 @@ import numpy as np import onnx_ir as ir +from mobius._passes._dtype_utils import initializer_dtype + logger = logging.getLogger(__name__) @@ -63,6 +65,19 @@ def call(self, model: ir.Model) -> ir.passes.PassResult: if not all(v is not None and v.is_initializer() for v in inputs): continue + # All inputs must have loaded weights. Folding before weights are + # loaded would force the dtype to default to FLOAT and bake a wrong + # type into the packed initializer. Mirror FoldTransposedInitializer's + # guard and skip with a warning so the Concat survives for a later run. + if any(v.const_value is None for v in inputs): # type: ignore[union-attr] + logger.warning( + "FoldConcatInitializersPass: skipping Concat %r — an input " + "initializer has no const_value (pass ran before weights " + "were loaded, or a weight is missing).", + node.name, + ) + continue + axis_attr = node.attributes.get("axis") axis: int = axis_attr.value if axis_attr is not None else 0 @@ -79,8 +94,10 @@ def call(self, model: ir.Model) -> ir.passes.PassResult: # Require uniform dtype — mixed-dtype concat is unusual and likely # a modelling error; skip and warn rather than silently produce - # a wrong result. - dtypes = [v.dtype for v in inputs] # type: ignore[union-attr] + # a wrong result. Resolve each dtype from the declared type or, when + # that is missing, from the loaded ``const_value`` so an fp16 weight + # whose type annotation was dropped is not mistaken for fp32. + dtypes = [initializer_dtype(v) for v in inputs] # type: ignore[union-attr] if len(set(dtypes)) > 1: logger.warning( "FoldConcatInitializersPass: skipping Concat with mixed dtypes %s" @@ -89,6 +106,10 @@ def call(self, model: ir.Model) -> ir.passes.PassResult: ) continue + # Authoritative dtype for the packed initializer. Falls back to FLOAT + # only when neither the declared type nor const_value is available. + packed_dtype = dtypes[0] or ir.DataType.FLOAT + # Build a name for the packed initializer from the input names and axis. # The name encodes both the ordered input names and the axis, so two Concat # nodes with the same inputs in the same order along the same axis will @@ -104,7 +125,12 @@ def call(self, model: ir.Model) -> ir.passes.PassResult: if packed_name in model.graph.initializers: existing_val = model.graph.initializers[packed_name] out_val.replace_all_uses_with(existing_val, replace_graph_outputs=True) - model.graph.remove(node) + # safe=True detaches the node from its inputs before removal so + # the folded q/k/v initializers' use lists are cleared. Otherwise + # `inp.uses()` still points at the removed Concat and the dead + # pre-pack weights survive RemoveUnusedNodesPass (~1.8 GB of + # orphaned weights serialized into the fp16 GQA model). + model.graph.remove(node, safe=True) folded_count += 1 modified = True @@ -117,7 +143,12 @@ def call(self, model: ir.Model) -> ir.passes.PassResult: ) continue - new_val = ir.Value(name=packed_name, shape=out_shape, type=out_val.type) + # Stamp the resolved dtype on the new initializer's type so the + # declared type, the LazyTensor, and the materialized data all agree + # (out_val.type may be missing after stage-2 rewrites). + new_val = ir.Value( + name=packed_name, shape=out_shape, type=ir.TensorType(packed_dtype) + ) captured_inputs = list(inputs) # capture for closure @@ -133,13 +164,17 @@ def _make_packed( arrays.append(v.const_value.numpy()) return ir.tensor(np.concatenate(arrays, axis=ax)) - dtype = inputs[0].dtype or ir.DataType.FLOAT # type: ignore[union-attr] - new_val.const_value = ir.LazyTensor(_make_packed, dtype=dtype, shape=out_shape) + new_val.const_value = ir.LazyTensor( + _make_packed, dtype=packed_dtype, shape=out_shape + ) model.graph.initializers[new_val.name] = new_val out_val.replace_all_uses_with(new_val, replace_graph_outputs=True) - model.graph.remove(node) + # safe=True detaches the node from its inputs before removal so the + # folded q/k/v initializers' use lists are cleared; otherwise the + # dead pre-pack weights survive RemoveUnusedNodesPass. + model.graph.remove(node, safe=True) folded_count += 1 modified = True diff --git a/src/mobius/_passes/_fold_concat_test.py b/src/mobius/_passes/_fold_concat_test.py index 492538d6..3cb237fa 100644 --- a/src/mobius/_passes/_fold_concat_test.py +++ b/src/mobius/_passes/_fold_concat_test.py @@ -124,6 +124,61 @@ def test_dynamic_input_not_folded(self): result = FoldConcatInitializersPass()(model) assert not result.modified + def test_folded_inputs_detached_so_dce_strips_dead_weights(self): + """After folding, the source initializers must be detached and DCE-able. + + FoldConcat removes the Concat node; if it does not detach the node from + its inputs (``safe=True``), the source q/k/v initializers keep a stale + use pointing at the removed node, so RemoveUnusedNodesPass cannot strip + them. For fp16 GQA that leaves ~1.8 GB of dead pre-pack weights in the + exported model. + """ + from onnx_ir.passes import common as common_passes + + a = np.ones((4, 8), dtype=np.float16) + b = np.ones((4, 8), dtype=np.float16) * 2 + model, init_vals = _make_concat_model([a, b], axis=0) + # The Concat output is the only graph output; route it through a MatMul + # so the packed initializer stays live while the sources become dead. + packed_out = model.graph.outputs[0] + hidden = ir.Value( + name="hidden", + shape=ir.Shape(["N", 8]), + type=ir.TensorType(ir.DataType.FLOAT16), + ) + matmul = ir.Node("", "MatMul", inputs=[hidden, packed_out], num_outputs=1) + mm_out = matmul.outputs[0] + mm_out.shape = ir.Shape(["N", 8]) + mm_out.dtype = ir.DataType.FLOAT16 + model.graph.append(matmul) + model.graph.inputs.append(hidden) + model.graph.outputs[0] = mm_out + + FoldConcatInitializersPass()(model) + + # Sources must be detached (zero uses) so DCE can remove them. + for v in init_vals: + assert len(list(v.uses())) == 0, ( + f"{v.name} still has uses after fold — node not detached" + ) + + common_passes.RemoveUnusedNodesPass()(model) + remaining = set(model.graph.initializers) + assert "init_0" not in remaining and "init_1" not in remaining, ( + f"Dead pre-pack weights survived DCE: {remaining}" + ) + assert "init_0__init_1__axis_0__concat" in remaining, ( + f"the live packed result must NOT be stripped by DCE: {remaining}" + ) + # The survived packed weight must retain its exact values, not just its + # name — guards against a future DCE that mutates retained tensors. + packed = model.graph.initializers["init_0__init_1__axis_0__concat"] + np.testing.assert_array_equal( + packed.const_value.numpy(), + np.concatenate([a, b], axis=0), + err_msg="packed-QKV values corrupted by DCE", + ) + def test_uses_lazy_tensor(self): """The packed initializer wraps sources in a LazyTensor. @@ -467,6 +522,56 @@ def test_concat_with_no_inputs_skipped(self): result = FoldConcatInitializersPass()(model) assert not result.modified + def test_packed_dtype_follows_const_value_when_declared_dtype_missing(self): + """Regression: fp16 weights with an unset declared dtype must fold to fp16. + + After ``_cast_module_dtype`` casts parameters to fp16, the resulting + initializer Values keep an fp16 ``const_value`` but lose their declared + ``.dtype`` (it becomes ``None``). The pass must resolve the packed dtype + from ``const_value`` rather than defaulting to FLOAT, otherwise GQA's + packed-QKV weight is serialized as fp32 and ORT rejects the model with a + fp16/fp32 MatMul type-mismatch error. + """ + q = ir.Value(name="q_weight") + q.const_value = ir.tensor(np.ones((4, 3), dtype=np.float16)) + k = ir.Value(name="k_weight") + k.const_value = ir.tensor(np.ones((4, 3), dtype=np.float16)) + v = ir.Value(name="v_weight") + v.const_value = ir.tensor(np.ones((4, 3), dtype=np.float16)) + # Critically: do NOT set .dtype — declared dtype stays None. + assert q.dtype is None and k.dtype is None and v.dtype is None + + concat_node = ir.Node( + "", + "Concat", + inputs=[q, k, v], + attributes=[ir.Attr("axis", ir.AttributeType.INT, 0)], + num_outputs=1, + ) + graph = ir.Graph( + inputs=[ir.Value(name="x")], + outputs=[concat_node.outputs[0]], + nodes=[concat_node], + name="test_graph", + opset_imports={"": 20}, + ) + graph.register_initializer(q) + graph.register_initializer(k) + graph.register_initializer(v) + model = ir.Model(graph, ir_version=10) + + result = FoldConcatInitializersPass()(model) + assert result.modified + + packed = model.graph.initializers["q_weight__k_weight__v_weight__axis_0__concat"] + assert packed.dtype == ir.DataType.FLOAT16, ( + f"Packed initializer declared dtype should be FLOAT16, got {packed.dtype}" + ) + assert packed.const_value.dtype == ir.DataType.FLOAT16, ( + f"Packed LazyTensor dtype should be FLOAT16, got {packed.const_value.dtype}" + ) + assert packed.const_value.numpy().dtype == np.float16 + def test_shape_computed_from_inputs_when_output_shape_missing(self): """When the Concat output has no shape, shape is inferred from input shapes.""" a = np.ones((4, 3), dtype=np.float32) @@ -512,3 +617,275 @@ def test_shape_computed_from_inputs_when_output_shape_missing(self): # Shape should be computed from inputs: (4,3) + (4,3) along axis 0 → (8,3). packed = model.graph.initializers[packed_name] assert list(packed.shape) == [8, 3] + + +class TestFoldConcatOrtLoad: + """End-to-end regression: a folded fp16 GQA-style graph must load in ORT. + + Reproduces the original failure where the packed-QKV initializer was + serialized as fp32 while the rest of the graph was fp16, causing ORT to + reject the model with:: + + Type parameter (T) of Optype (MatMul) bound to different types + (tensor(float16) and tensor(float)) + + The failure occurs on the CPU EP too, so no GPU is needed. + """ + + def test_folded_fp16_concat_matmul_loads_in_ort(self, tmp_path): + import onnxruntime as ort + + # Three fp16 weights with UNSET declared dtype (as after _cast_module_dtype), + # but with shape retained so they serialize before they are pruned. + q = ir.Value(name="q_weight", shape=ir.Shape([8, 4])) + q.const_value = ir.tensor(np.random.randn(8, 4).astype(np.float16)) + k = ir.Value(name="k_weight", shape=ir.Shape([8, 4])) + k.const_value = ir.tensor(np.random.randn(8, 4).astype(np.float16)) + v = ir.Value(name="v_weight", shape=ir.Shape([8, 4])) + v.const_value = ir.tensor(np.random.randn(8, 4).astype(np.float16)) + + concat_node = ir.Node( + "", + "Concat", + inputs=[q, k, v], + attributes=[ir.Attr("axis", ir.AttributeType.INT, 1)], + num_outputs=1, + ) + packed_out = concat_node.outputs[0] # shape (8, 12), packed QKV weight + packed_out.shape = ir.Shape([8, 12]) + + # hidden @ packed_qkv : (N, 8) @ (8, 12) -> (N, 12), all fp16. + hidden = ir.Value( + name="hidden", shape=ir.Shape(["N", 8]), type=ir.TensorType(ir.DataType.FLOAT16) + ) + matmul_node = ir.Node("", "MatMul", inputs=[hidden, packed_out], num_outputs=1) + out = matmul_node.outputs[0] + out.shape = ir.Shape(["N", 12]) + out.dtype = ir.DataType.FLOAT16 + + graph = ir.Graph( + inputs=[hidden], + outputs=[out], + nodes=[concat_node, matmul_node], + name="qkv_matmul", + opset_imports={"": 20}, + ) + graph.register_initializer(q) + graph.register_initializer(k) + graph.register_initializer(v) + model = ir.Model(graph, ir_version=10) + + FoldConcatInitializersPass()(model) + + # Drop the now-unused source initializers (a DCE step runs after folding + # in the real export pipeline). Without this they linger as dead weights. + for dead in ("q_weight", "k_weight", "v_weight"): + del model.graph.initializers[dead] + + model_path = tmp_path / "fp16_qkv.onnx" + ir.save(model, model_path) + + # Before the fix this raises a fp16/fp32 MatMul type-mismatch at load. + sess = ort.InferenceSession(str(model_path), providers=["CPUExecutionProvider"]) + feed = {"hidden": np.random.randn(2, 8).astype(np.float16)} + (result,) = sess.run(None, feed) + assert result.shape == (2, 12) + assert result.dtype == np.float16 + + def test_folded_fp16_packed_qkv_values_survive_serialization(self, tmp_path): + """Value gate: packed-QKV must keep its fp16 *values* through serialize→reload. + + The other fold tests assert the packed value in memory and that ORT can + load+run the model, but none compare the *serialized* packed-QKV weight + to its source q/k/v projections. The original garbage-export bug + corrupted bytes at serialization (fp16 data written under a defaulted + FLOAT32 dtype → near-zero/garbage) — a failure an in-memory + ``const_value`` check cannot see, and which a load+run check misses + because the model still loads and produces a right-shaped fp16 output. + + This round-trips through the production save path (``ir.save`` with + external data, exactly like the real fp16 export's model.onnx + + model.onnx.data) and asserts the reloaded packed weight matches its + sources per-slice, catching a numerically-corrupt pack that still has + the right count and dtype (the corr≈0 / norm≈0.8 failure mode). + """ + import onnxruntime as ort + + rng = np.random.default_rng(0) + # fp16 q/k/v projection weights with UNSET declared dtype (the state left + # by _cast_module_dtype). Sized > 256 bytes so ir.save externalizes them, + # exercising the same external-data path as the real export. Column-concat + # (axis=1) mirrors the real packed-QKV layout [Q | K | V]. + q_arr = rng.standard_normal((16, 16)).astype(np.float16) + k_arr = rng.standard_normal((16, 16)).astype(np.float16) + v_arr = rng.standard_normal((16, 16)).astype(np.float16) + q = ir.Value(name="q_weight", shape=ir.Shape([16, 16])) + q.const_value = ir.tensor(q_arr) + k = ir.Value(name="k_weight", shape=ir.Shape([16, 16])) + k.const_value = ir.tensor(k_arr) + v = ir.Value(name="v_weight", shape=ir.Shape([16, 16])) + v.const_value = ir.tensor(v_arr) + + concat_node = ir.Node( + "", + "Concat", + inputs=[q, k, v], + attributes=[ir.Attr("axis", ir.AttributeType.INT, 1)], + num_outputs=1, + ) + packed_out = concat_node.outputs[0] # (16, 48) packed QKV weight + packed_out.shape = ir.Shape([16, 48]) + + hidden = ir.Value( + name="hidden", shape=ir.Shape(["N", 16]), type=ir.TensorType(ir.DataType.FLOAT16) + ) + matmul_node = ir.Node("", "MatMul", inputs=[hidden, packed_out], num_outputs=1) + out = matmul_node.outputs[0] + out.shape = ir.Shape(["N", 48]) + out.dtype = ir.DataType.FLOAT16 + + graph = ir.Graph( + inputs=[hidden], + outputs=[out], + nodes=[concat_node, matmul_node], + name="qkv_matmul", + opset_imports={"": 20}, + ) + graph.register_initializer(q) + graph.register_initializer(k) + graph.register_initializer(v) + model = ir.Model(graph, ir_version=10) + + FoldConcatInitializersPass()(model) + # DCE step the real export runs after folding. + for dead in ("q_weight", "k_weight", "v_weight"): + del model.graph.initializers[dead] + + # Serialize through the production path: external data, exactly as the + # real fp16 export writes model.onnx + model.onnx.data. This is where the + # original dtype bug corrupted the packed bytes. + model_path = tmp_path / "model.onnx" + ir.save(model, model_path, external_data="model.onnx.data") + assert (tmp_path / "model.onnx.data").exists(), ( + "weights should be externalized, exercising the real export's save path" + ) + + reloaded = ir.load(model_path) + packed_name = "q_weight__k_weight__v_weight__axis_1__concat" + packed = reloaded.graph.initializers[packed_name].const_value.numpy() + assert packed.dtype == np.float16, ( + f"reloaded packed-QKV must serialize as fp16, got {packed.dtype}" + ) + expected = np.concatenate([q_arr, k_arr, v_arr], axis=1) + assert packed.shape == expected.shape + + # Per-slice discriminator (mirrors QA's weight-integrity gate): each q/k/v + # column block of the reloaded pack must match its source projection. A + # serialize-time dtype corruption surfaces here as corr≈0 / norm collapse + # even when the pack's shape and dtype look correct. + blocks = { + "q": (packed[:, 0:16], q_arr), + "k": (packed[:, 16:32], k_arr), + "v": (packed[:, 32:48], v_arr), + } + for name, (got, src) in blocks.items(): + got32 = got.astype(np.float32).ravel() + src32 = src.astype(np.float32).ravel() + corr = np.corrcoef(got32, src32)[0, 1] + assert corr >= 0.99, ( + f"{name}-slice packed-QKV corr={corr:.4f} (<0.99) — corrupted pack" + ) + src_norm = np.linalg.norm(src32) + rel_err = abs(np.linalg.norm(got32) - src_norm) / src_norm + assert rel_err <= 0.02, ( + f"{name}-slice packed-QKV norm rel_err={rel_err:.4f} (>2%) — corrupted pack" + ) + # Non-degenerate magnitude. Use mean of ABSOLUTE values, not signed + # mean: symmetric fp16 weights have a signed mean ~1e-6 that is + # indistinguishable from a broken near-zero tensor, whereas mean|abs| + # separates cleanly (healthy ~0.0x vs unserialized ~1e-6). This also + # catches the near-zero failure mode robustly where corr is undefined + # (a constant/zero slice has zero variance → corrcoef is nan). + mean_abs = float(np.abs(got32).mean()) + assert mean_abs >= 1e-3, ( + f"{name}-slice packed-QKV mean|abs|={mean_abs:.2e} (~0) — near-zero/unserialized" + ) + + # fp16 concat is lossless, so the strongest assert also holds end-to-end. + np.testing.assert_array_equal( + packed, + expected, + err_msg="packed-QKV bytes corrupted by the serialization round-trip", + ) + + # Functional check: ORT inference on the reloaded model must match a numpy + # reference, proving the packed weight is correct in use, not just on disk. + sess = ort.InferenceSession(str(model_path), providers=["CPUExecutionProvider"]) + feed = {"hidden": rng.standard_normal((2, 16)).astype(np.float16)} + (result,) = sess.run(None, feed) + reference = feed["hidden"].astype(np.float32) @ expected.astype(np.float32) + np.testing.assert_allclose(result.astype(np.float32), reference, rtol=1e-2, atol=1e-2) + + def test_value_gate_catches_corrupted_packed_slice(self, tmp_path): + """Negative control: the per-slice value gate must FAIL on a bad pack. + + Proves the discriminators in + ``test_folded_fp16_packed_qkv_values_survive_serialization`` actually + catch the original failure signature — a packed-QKV with a near-zero + slice (corr≈0 / norm collapse) — and that the corruption survives the + serialize→reload round-trip rather than being silently "healed" by + save/load. Without this negative control a future change could neuter + the value asserts and still go green. + """ + rng = np.random.default_rng(1) + q_arr = rng.standard_normal((16, 16)).astype(np.float16) + k_arr = rng.standard_normal((16, 16)).astype(np.float16) + v_arr = rng.standard_normal((16, 16)).astype(np.float16) + + # Poison the K slice to zero — the exact "unserialized / near-zero packed + # weight" signature of the original garbage export. + poisoned = np.concatenate([q_arr, k_arr, v_arr], axis=1).copy() + poisoned[:, 16:32] = 0 + + packed = ir.Value( + name="packed_qkv", + shape=ir.Shape([16, 48]), + type=ir.TensorType(ir.DataType.FLOAT16), + ) + packed.const_value = ir.tensor(poisoned) + graph = ir.Graph( + inputs=[], + outputs=[packed], + nodes=[], + name="poisoned_pack", + opset_imports={"": 20}, + ) + graph.register_initializer(packed) + model = ir.Model(graph, ir_version=10) + + model_path = tmp_path / "model.onnx" + ir.save(model, model_path, external_data="model.onnx.data") + reloaded = ir.load(model_path) + got = reloaded.graph.initializers["packed_qkv"].const_value.numpy() + + # The corruption must survive the round-trip (save/load must not "fix" it). + np.testing.assert_array_equal(got[:, 16:32], np.zeros((16, 16), dtype=np.float16)) + + # The gate's discriminators must flag the zeroed K slice. mean|abs| is the + # robust primary signal: corr is undefined (nan) for a zero-variance slice, + # which is exactly why a corr-only gate would be unsafe here. + k_block = got[:, 16:32].astype(np.float32).ravel() + k_ref = k_arr.astype(np.float32).ravel() + assert np.abs(k_block).mean() < 1e-3, ( + "mean|abs| discriminator must flag a zeroed packed slice" + ) + rel_err = abs(np.linalg.norm(k_block) - np.linalg.norm(k_ref)) / np.linalg.norm(k_ref) + assert rel_err > 0.02, "norm rel_err discriminator must flag a zeroed packed slice" + + # The untouched Q and V slices must still read as healthy — the gate is + # specific to the corrupted slice, not a blanket failure. + for sl, ref in ((got[:, 0:16], q_arr), (got[:, 32:48], v_arr)): + sl32 = sl.astype(np.float32).ravel() + assert np.abs(sl32).mean() >= 1e-3 + corr = np.corrcoef(sl32, ref.astype(np.float32).ravel())[0, 1] + assert corr >= 0.99 diff --git a/src/mobius/_passes/_fold_dtype_e2e_test.py b/src/mobius/_passes/_fold_dtype_e2e_test.py new file mode 100644 index 00000000..0b98f59c --- /dev/null +++ b/src/mobius/_passes/_fold_dtype_e2e_test.py @@ -0,0 +1,358 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT License. + +"""End-to-end regression tests for fp16 dtype preservation through the fold passes. + +These complement the unit-level coverage in ``_dtype_utils_test.py`` and the +pass-level coverage in ``_fold_concat_test.py`` / ``_fold_transpose_test.py`` by +driving the *real* export pipeline (``build_from_module`` + ``apply_weights``) +rather than hand-built single-pass graphs. + +The df203cc fix has two parts, and this module guards both at the export level: + +1. **dtype stamping (type-stamping path).** The fold passes stamp the *resolved* + fp16 dtype on the new packed / transposed initializer (its ``TensorType`` and + its ``LazyTensor``) instead of inheriting the unset type of the rewrite-produced + ``Concat`` intermediate or defaulting to ``FLOAT``. Guarded by + ``test_packed_qkv_weights_stay_fp16_after_export`` and + ``test_all_folded_matmul_weights_stay_fp16_after_export``: in the real fp16 GQA + export the PackQKV rewrite emits a ``MatMul(hidden, Transpose(Concat(W_q, W_k, + W_v)))`` whose ``Concat`` output carries no declared dtype, so without the + stamping the chained FoldConcat -> FoldTranspose widens the packed-QKV weight + to fp32. + +2. **const_value fallback.** ``_dtype_utils.initializer_dtype`` resolves the dtype + from ``const_value`` when an initializer's declared ``type`` was dropped during + graph building. Guarded by + ``test_dropped_declared_dtype_falls_back_to_fp16_const``, which reproduces that + dropped-dtype condition on the packed-QKV weights (as the original regression + exhibited) so the const_value fallback is the only thing keeping the folded + weights fp16. + +Either failure silently widens an fp16 model to fp32 packed weights, which +onnxruntime rejects at load time with a fp16/fp32 ``MatMul`` type-mismatch +(``Type parameter (T) of Optype (MatMul) bound to different types``). + +The in-memory dtype assertions read the IR ``const_value`` metadata. The +regression's *ground-truth* symptom, however, is a serialize-time corruption: +the fold passes write the fp16 *bytes* under a fp32 declaration, so the persisted +initializer is a fp32-declared tensor backed by half-width data. An in-memory +``const_value.numpy()`` can stay fp16 and miss this (only the declared dtype +widens), so ``test_packed_qkv_survives_serialization_roundtrip`` and the +serialize check in ``test_dropped_declared_dtype_falls_back_to_fp16_const`` drive +the real export artifact (``ir.save`` with external data — the same +``model.onnx`` + ``model.onnx.data`` layout the fp16 export produces), reload it, +and assert the folded weights reload as fp16 with their bytes intact. + +The models are tiny and fully synthetic (no HuggingFace download, no GPU and no +onnxruntime execution — only the ONNX graph is constructed, serialized and +inspected), so the tests fit the per-PR CI tier. +""" + +from __future__ import annotations + +from pathlib import Path + +import numpy as np +import onnx_ir as ir +import pytest +import torch + +from mobius._builder import build_from_module +from mobius._configs import ArchitectureConfig +from mobius._optimizations import fold_initializers_after_weights +from mobius._registry import registry +from mobius._weight_loading import apply_weights + + +def _make_fp16_llama_config() -> ArchitectureConfig: + """A tiny llama config whose Q/K/V weights are packable (no QK norm). + + ``dtype=FLOAT16`` routes the build through the fp16 packed-QKV path that the + fold passes must preserve. + """ + return ArchitectureConfig( + # Structural invariants that make the fp16 PackQKV fold fire: + # num_attention_heads * head_dim == hidden_size (4 * 16 == 64), and + # num_key_value_heads < num_attention_heads (2 < 4, i.e. GQA). + # The remaining fields are arbitrary small values kept lightweight. + hidden_size=64, + intermediate_size=128, + num_attention_heads=4, + num_key_value_heads=2, + head_dim=16, + num_hidden_layers=2, + vocab_size=256, + max_position_embeddings=128, + hidden_act="silu", + rms_norm_eps=1e-6, + rope_type="default", + rope_theta=10000.0, + pad_token_id=0, + dtype=ir.DataType.FLOAT16, + ) + + +def _build_fp16_decoder(config: ArchitectureConfig) -> ir.Model: + """Build an fp16 llama decoder graph with packed QKV, weights not yet loaded. + + Uses the ``cuda`` execution provider so the GQA + PackQKV rewrites pack + Q/K/V into a single ``MatMul(hidden, Transpose(Concat(...)))``. Only the ONNX + graph is constructed — no GPU or onnxruntime execution is required. The fold + passes have not run yet (they run when weights are loaded). + """ + return build_from_module( + registry.get("llama")(config), + config, + execution_provider="cuda", + )["model"] + + +def _fp16_weight_tensors(model: ir.Model) -> dict[str, torch.Tensor]: + """fp16 tensors for every uninitialised parameter, mirroring an fp16 checkpoint. + + Weight initializers always carry fully-static integer shapes (they are + concrete parameter tensors, never symbolic), so ``int(d)`` is safe here. + """ + return { + name: torch.randn(*[int(d) for d in init.shape]).to(torch.float16) + for name, init in model.graph.initializers.items() + if init.const_value is None + } + + +def _matmul_weight_initializers(model: ir.Model) -> list[ir.Value]: + """Return every ``MatMul`` second input that is a graph initializer.""" + initializers = model.graph.initializers + weights: list[ir.Value] = [] + for node in model.graph: + if node.op_type != "MatMul": + continue + weight = node.inputs[1] + if weight is not None and weight.name in initializers: + weights.append(weight) + return weights + + +def _assert_matmul_weights_roundtrip_as_fp16(model: ir.Model, tmp_path: Path) -> None: + """Assert every folded MatMul weight reloads as fp16 with its bytes intact. + + Serializes ``model`` through the production external-data path. This is the + ground-truth df203cc guard. The regression writes the packed / transposed fp16 *bytes* under a fp32 declaration, a corruption an in-memory + ``const_value`` check can mask (the backing array stays fp16 while only the + declared dtype widens) but that surfaces on serialize -> reload as a + fp32-declared initializer whose bytes no longer match the fp16 source. + ``ir.save`` with ``external_data`` mirrors the real fp16 export's + ``model.onnx`` + ``model.onnx.data`` layout, exactly where the dtype bug + corrupts the weight bytes. + """ + before = {w.name: w.const_value.numpy() for w in _matmul_weight_initializers(model)} + assert before, "Expected the export to contain MatMul weight initializers to round-trip" + + model_path = tmp_path / "model.onnx" + ir.save(model, model_path, external_data="model.onnx.data") + assert (tmp_path / "model.onnx.data").exists(), ( + "fp16 weights should be externalized, exercising the real export's " + "model.onnx + model.onnx.data save path where the dtype bug corrupts bytes" + ) + + reloaded = ir.load(model_path) + for weight in _matmul_weight_initializers(reloaded): + assert weight.dtype == ir.DataType.FLOAT16, ( + f"Reloaded MatMul weight {weight.name!r} serialized as {weight.dtype} " + f"(expected FLOAT16): fp16 bytes were written under a fp32 declaration " + f"— the df203cc serialize-time corruption." + ) + np.testing.assert_array_equal( + weight.const_value.numpy(), + before[weight.name], + err_msg=( + f"Reloaded MatMul weight {weight.name!r} bytes changed across " + f"serialize -> reload; the fp16 packed weight was silently corrupted " + f"(fp16 data written under a fp32 dtype)." + ), + ) + + +@pytest.fixture +def fp16_export() -> tuple[ArchitectureConfig, ir.Model]: + """A real fp16 packed-QKV export: build + ``apply_weights`` (folds inside). + + Function-scoped (a fresh build per test) on purpose: the serialize-roundtrip + test calls ``ir.save`` on the model, and on some ``onnx_ir`` versions that + offloads the initializers' ``const_value`` to external tensors in place. A + shared (module-scoped) model could then leak that mutated/externalized state + into another test, and under ``pytest-xdist`` the tests' execution order is + not guaranteed — making such cross-test contamination order-dependent and + flaky. A fresh model per test keeps each test hermetic and deterministic. The + build is tiny, so the rebuild cost is negligible. + """ + config = _make_fp16_llama_config() + model = _build_fp16_decoder(config) + # apply_weights assigns the fp16 weights and then runs + # fold_initializers_after_weights, folding the Transpose/Concat weight nodes. + apply_weights(model, _fp16_weight_tensors(model)) + return config, model + + +class TestFp16FoldDtypeE2E: + def test_packed_qkv_weights_stay_fp16_after_export( + self, fp16_export: tuple[ArchitectureConfig, ir.Model] + ) -> None: + """An fp16 packed-QKV export keeps fp16 packed/transposed weights. + + Guards the df203cc dtype-stamping path: the chained + ``FoldConcatInitializersPass`` (the Q/K/V pack) and + ``FoldTransposedInitializerPass`` (the weight transpose) must carry the + fp16 dtype through the rewrite-produced ``Concat`` intermediate instead + of defaulting it to ``FLOAT``. + """ + config, model = fp16_export + + # The packed path must actually have run, otherwise the dtype assertion + # below would pass vacuously (e.g. if the build stopped packing QKV). + gqa_nodes = [n for n in model.graph if n.op_type == "GroupQueryAttention"] + assert len(gqa_nodes) == config.num_hidden_layers, ( + "Expected one packed GroupQueryAttention per layer; the fp16 GQA " + "PackQKV path was not exercised, so this regression guard is vacuous." + ) + + for gqa in gqa_nodes: + matmul = gqa.inputs[0].producer() + assert matmul is not None and matmul.op_type == "MatMul", ( + "Packed GQA projection should be produced by a MatMul" + ) + weight = matmul.inputs[1] + assert weight is not None, "Packed projection MatMul is missing its weight input" + # FoldConcat + FoldTranspose ran inside apply_weights, so the packed + # weight is now a plain folded initializer (no producer node left). + assert weight.producer() is None, ( + "Packed QKV weight should be a folded initializer after export, " + "not the output of a residual Transpose/Concat node" + ) + assert weight.const_value is not None + assert weight.const_value.dtype == ir.DataType.FLOAT16, ( + f"Packed QKV weight {weight.name!r} widened to " + f"{weight.const_value.dtype} (expected FLOAT16); the fp16 fold " + f"dtype fix (df203cc) has regressed." + ) + assert weight.dtype == ir.DataType.FLOAT16 + + def test_all_folded_matmul_weights_stay_fp16_after_export( + self, fp16_export: tuple[ArchitectureConfig, ir.Model] + ) -> None: + """Every folded MatMul weight in an fp16 export stays fp16. + + Broader guard than the packed-QKV case: it also covers the standalone + ``o_proj`` and MLP weight transposes folded by + ``FoldTransposedInitializerPass``. A single fp32-widened weight here is + exactly what makes onnxruntime reject the fp16 model at load time. + """ + _, model = fp16_export + + weights = _matmul_weight_initializers(model) + assert weights, "Expected the export to contain MatMul weight initializers" + # Anti-vacuity: at least one weight must be a folded transposed initializer + # (FoldTransposedInitializerPass names them ``..._t``). Without this, the + # test could pass simply because no folding ran. + assert any(weight.name.endswith("_t") for weight in weights), ( + "Expected at least one folded transposed weight (``..._t``); the fold " + "passes did not run, so this regression guard would be vacuous." + ) + + # fp32 widening (FLOAT) is the specific documented failure mode of the + # df203cc regression — the fold passes defaulting a dropped declared dtype + # to FLOAT — so we flag exactly that rather than any non-fp16 dtype. + widened = [ + weight.name + for weight in weights + if weight.const_value is not None and weight.const_value.dtype == ir.DataType.FLOAT + ] + assert not widened, ( + f"fp16 export produced fp32 MatMul weights after folding: {widened}. " + f"The fold passes must preserve fp16 (df203cc)." + ) + + def test_packed_qkv_survives_serialization_roundtrip( + self, fp16_export: tuple[ArchitectureConfig, ir.Model], tmp_path: Path + ) -> None: + """The fp16 export's folded weights survive serialize -> reload as fp16. + + Ground-truth guard for the dtype-stamping path. The in-memory assertions + above read the IR ``const_value`` metadata; this drives the real export + artifact (``ir.save`` with external data, like the production + ``model.onnx`` + ``model.onnx.data``) and reloads it, the only check that + catches the df203cc serialize-time fp16-under-fp32 byte corruption — an + in-memory ``const_value.numpy()`` can stay fp16 and miss it. + """ + _, model = fp16_export + _assert_matmul_weights_roundtrip_as_fp16(model, tmp_path) + + def test_dropped_declared_dtype_falls_back_to_fp16_const(self, tmp_path: Path) -> None: + """Folding fp16 weights whose declared dtype was dropped stays fp16. + + Guards the df203cc ``initializer_dtype`` const_value fallback. The + original regression arose because graph building left the packed-QKV + weights with fp16 ``const_value`` but no declared ``type``; the fold + passes then defaulted to ``FLOAT``. Here we reproduce that exact + condition on a real fp16 export — load fp16 weights, clear the declared + type on the packed-QKV ``Concat`` inputs, then run the real fold + orchestration — so the const_value fallback is the *only* thing that can + keep the folded weights fp16. + + ``fold_initializers_after_weights`` is invoked directly (rather than via + ``apply_weights``) so the dropped-dtype condition can be injected between + loading the weights and folding; it is the same orchestration + ``apply_weights`` runs internally. + """ + config = _make_fp16_llama_config() + model = _build_fp16_decoder(config) + + # Load fp16 weights onto the initializers (mirrors apply_weights' assign + # step) without yet folding. + for name, tensor in _fp16_weight_tensors(model).items(): + model.graph.initializers[name].const_value = ir.tensor(tensor.numpy()) + + # Reproduce the dropped-declared-dtype condition: clear the declared type + # on the packed-QKV weight initializers (the Concat inputs) while keeping + # their fp16 const_value. + dropped: list[str] = [] + for node in model.graph: + if node.op_type != "Concat": + continue + for inp in node.inputs: + if ( + inp is not None + and inp.name in model.graph.initializers + and inp.const_value is not None + ): + inp.type = None + dropped.append(inp.name) + assert dropped, ( + "Expected packed-QKV Concat inputs to drop declared dtype on; the " + "fp16 PackQKV path was not exercised, so this guard would be vacuous." + ) + + fold_initializers_after_weights(model) + + gqa_nodes = [n for n in model.graph if n.op_type == "GroupQueryAttention"] + assert len(gqa_nodes) == config.num_hidden_layers + for gqa in gqa_nodes: + matmul = gqa.inputs[0].producer() + assert matmul is not None and matmul.op_type == "MatMul", ( + "Packed GQA projection should be produced by a MatMul" + ) + weight = matmul.inputs[1] + assert weight is not None, "Packed projection MatMul is missing its weight input" + assert weight.const_value is not None + assert weight.const_value.dtype == ir.DataType.FLOAT16, ( + f"Packed QKV weight {weight.name!r} widened to " + f"{weight.const_value.dtype} when its declared dtype was dropped; " + f"the initializer_dtype const_value fallback (df203cc) regressed." + ) + + # Ground-truth check: the const_value fallback must also hold through the + # real serialize -> reload path. Reverting only the ``initializer_dtype`` + # call-sites widens these dropped-dtype weights and corrupts their bytes + # here even though the type-stamp keeps the realistic-export tests green. + _assert_matmul_weights_roundtrip_as_fp16(model, tmp_path) diff --git a/src/mobius/_passes/_fold_transpose.py b/src/mobius/_passes/_fold_transpose.py index 402cad7c..e7f6cae3 100644 --- a/src/mobius/_passes/_fold_transpose.py +++ b/src/mobius/_passes/_fold_transpose.py @@ -26,6 +26,8 @@ import onnx_ir as ir +from mobius._passes._dtype_utils import initializer_dtype + logger = logging.getLogger(__name__) @@ -107,6 +109,11 @@ def call(self, model: ir.Model) -> ir.passes.PassResult: new_name = f"{inp.name}_t" + # Resolve the dtype from the declared type or, when it was + # dropped during graph building, from the loaded const_value, so + # a transposed fp16 weight is not silently emitted as fp32. + transposed_dtype = initializer_dtype(inp) or ir.DataType.FLOAT + # Derive shape of the transposed tensor from the Transpose output. # Fall back to computing from the input shape if shape inference # did not propagate to this new node (e.g. after stage-2 rewrites). @@ -115,7 +122,11 @@ def call(self, model: ir.Model) -> ir.passes.PassResult: perm = list(perm_attr.value) t_shape = ir.Shape([inp.shape[p] for p in perm]) - new_val = ir.Value(name=new_name, shape=t_shape, type=inp.type) + new_val = ir.Value( + name=new_name, + shape=t_shape, + type=ir.TensorType(transposed_dtype), + ) # Create a LazyTensor that transposes the original data on demand. # The actual numpy transposition is deferred until serialization, @@ -123,7 +134,7 @@ def call(self, model: ir.Model) -> ir.passes.PassResult: src = inp # captured for the closure below new_val.const_value = ir.LazyTensor( lambda s=src: ir.tensor(s.const_value.numpy().T), # type: ignore[union-attr] - dtype=inp.dtype or ir.DataType.FLOAT, + dtype=transposed_dtype, shape=t_shape, name=new_val.name, ) diff --git a/src/mobius/_passes/_fold_transpose_test.py b/src/mobius/_passes/_fold_transpose_test.py index 93c74b37..0e3def2a 100644 --- a/src/mobius/_passes/_fold_transpose_test.py +++ b/src/mobius/_passes/_fold_transpose_test.py @@ -508,3 +508,53 @@ def test_original_weight_pruned_after_full_pipeline(): assert "weight_t" not in model.graph.initializers transpose_nodes = [n for n in model.graph.all_nodes() if n.op_type == "Transpose"] assert len(transpose_nodes) == 1 + + +class TestFoldTransposeDtype: + """Regression: a transposed fp16 weight with unset declared dtype stays fp16.""" + + def test_transposed_dtype_follows_const_value_when_declared_dtype_missing(self): + """fp16 weights lose their declared dtype after _cast_module_dtype. + + The pass must resolve the transposed initializer's dtype from + ``const_value`` rather than defaulting to FLOAT, otherwise a fp16 weight + is emitted as fp32 and ORT rejects the downstream MatMul. + """ + weight = ir.Value(name="weight") + weight.const_value = ir.tensor(np.arange(12, dtype=np.float16).reshape(4, 3)) + # Do NOT set .dtype — declared dtype stays None. + assert weight.dtype is None + + x = ir.Value(name="x") + x.shape = ir.Shape([2, 4]) + x.dtype = ir.DataType.FLOAT16 + + transpose_node = ir.Node( + "", + "Transpose", + inputs=[weight], + attributes=[ir.Attr("perm", ir.AttributeType.INTS, [1, 0])], + num_outputs=1, + ) + w_t = transpose_node.outputs[0] + matmul_node = ir.Node("", "MatMul", inputs=[x, w_t], num_outputs=1) + + graph = ir.Graph( + inputs=[x], + outputs=[matmul_node.outputs[0]], + nodes=[transpose_node, matmul_node], + name="test_graph", + opset_imports={"": 20}, + ) + graph.register_initializer(weight) + model = ir.Model(graph, ir_version=10) + + result = FoldTransposedInitializerPass()(model) + assert result.modified + + packed = model.graph.initializers["weight_t"] + assert packed.dtype == ir.DataType.FLOAT16, ( + f"Transposed initializer declared dtype should be FLOAT16, got {packed.dtype}" + ) + assert packed.const_value.dtype == ir.DataType.FLOAT16 + assert packed.const_value.numpy().dtype == np.float16 diff --git a/src/mobius/components/_attention.py b/src/mobius/components/_attention.py index 9c076d2e..7b5fce70 100644 --- a/src/mobius/components/_attention.py +++ b/src/mobius/components/_attention.py @@ -10,7 +10,7 @@ from onnxscript import OpBuilder, nn from mobius._configs import ArchitectureConfig -from mobius.components._common import Linear +from mobius.components._common import Linear, create_static_cache_causal_mask from mobius.components._rms_norm import OffsetRMSNorm, RMSNorm from mobius.components._rotary_embedding import apply_rotary_pos_emb @@ -66,12 +66,20 @@ class StaticCacheState(NamedTuple): value_cache: Pre-allocated value cache [B, max_seq, kv_hidden] 3D. write_indices: Position to write new tokens [B] int64. nonpad_kv_seqlen: Valid KV length per batch entry [B] int64. + causal_mask: Prebuilt causal mask [B, 1, S_q, max_seq] bool, shared + (same ``ir.Value``) across all layers. The mask is layer-invariant + (depends only on S_q, max_seq and write_indices), so the task layer + builds it once and threads the same Value to every layer instead of + rebuilding it per layer. When ``None`` (e.g. a direct + ``_apply_attention`` caller that did not hoist the mask) the + attention path builds it on demand. """ key_cache: ir.Value value_cache: ir.Value write_indices: ir.Value nonpad_kv_seqlen: ir.Value + causal_mask: ir.Value | None = None def _apply_attention( @@ -99,14 +107,19 @@ def _apply_attention( Static cache mode (``static_cache is not None``): Scatters new key/value into the static cache via TensorScatter, - then attends over the full cache using ``nonpad_kv_seqlen``. - Also uses ``is_causal=1``. + then attends over the full cache using ``nonpad_kv_seqlen`` with + ``is_causal=0`` plus an explicit causal mask derived from + ``write_indices`` (see :func:`create_static_cache_causal_mask`). + ``is_causal=1`` cannot be used here: the opset-24 Attention kernel + rejects it together with ``nonpad_kv_seqlen`` when ``S_q`` differs + from the (pre-allocated) cache length, which is always the case. Returns ``(attn_output, updated_key_cache, updated_value_cache)``. Note: - Both paths set ``is_causal=1`` on the Attention op, which enables - built-in causal masking. This means ``attn_mask`` should encode - only padding information (as a bool mask), not causality. + The dynamic path sets ``is_causal=1`` so callers only provide a + bool padding mask. The static path instead uses ``is_causal=0`` + with a full causal+padding mask because the external-cache kernel + does not accept ``is_causal=1`` alongside ``nonpad_kv_seqlen``. Note: ``nonpad_kv_seqlen`` (input #6) is only valid in static cache mode @@ -138,28 +151,56 @@ def _apply_attention( axis=1, ) # [B, max_seq, kv_hidden] - # Attend over the full cache. We pass None for attn_mask and use - # is_causal=1 instead — the Attention op handles causal + padding - # masking internally via is_causal + nonpad_kv_seqlen. Using - # create_attention_bias() here would produce incorrect causality - # during prefill because it cannot represent the relationship - # between query positions and the full cache length. + # Attend over the full cache with is_causal=0 plus an explicit + # causal mask. The opset-24 Attention CUDA kernel rejects + # is_causal=1 together with nonpad_kv_seqlen when S_q != total_kv + # and there is no past_key (the causal_cross_no_past guard in + # attention.cc). With a pre-allocated [B, max_seq, ...] cache the + # query length never equals the total cache length, so that guard + # fires in BOTH prefill and decode. Per ORT's own guidance we set + # is_causal=0 and pass an explicit causal mask built from + # write_indices: a query token at offset t attends to cache slots + # j <= write_indices[b] + t. This single rule serves prefill + # (write_indices=0 -> triangular) and decode (write_indices=N, + # S_q=1 -> keep slots 0..N), and subsumes padding so it stays + # consistent with nonpad_kv_seqlen. + # + # nonpad_kv_seqlen is still passed: it selects the external-cache + # (TensorScatter) kernel path. When both nonpad_kv_seqlen and an + # attn_mask are present, ORT skips Flash and routes to the + # memory-efficient / unfused path, which converts the bool mask to + # an additive bias (attention.cc ConvertAttnMaskToBias). + # + # This single always-masked (branchless) formulation deliberately + # trades decode-on-Flash for graph-capture compatibility: an + # If(Greater(S_q, 1)) phase-split could keep decode maskless on Flash, + # but emits in-graph control flow that fails CUDA Graph capture at + # session init (see the graph-capture compatibility gotcha skill). # - # NOTE: The ONNX Attention spec supports attn_mask alongside - # nonpad_kv_seqlen for custom masking (e.g., user-defined masks - # beyond causal + padding). Currently we rely on is_causal=1 + - # nonpad_kv_seqlen for standard LLM causal + padding masking. # TODO(titaiwang): Support user-provided attn_mask in external # cache mode for advanced use cases (e.g., prefix masking, # document boundaries in batched inference). # TODO(titaiwang): Support sliding window (circular cache mode) # with static cache for long-context models that use local # attention windows. - attn_output, _, _ = op.Attention( + # The causal mask is layer-invariant, so the task layer normally builds + # it ONCE and shares the same ir.Value across every layer via + # static_cache.causal_mask (graph-size / perf win that scales with + # layer count). Fall back to building it here for direct callers + # (e.g. unit tests) that did not hoist the mask. + static_causal_mask = static_cache.causal_mask + if static_causal_mask is None: + static_causal_mask = create_static_cache_causal_mask( + op, + query, + updated_k, + static_cache.write_indices, + ) + attn_output = op.Attention( query, updated_k, updated_v, - None, # no attn_mask — is_causal handles masking + static_causal_mask, # explicit causal mask (is_causal=0) None, # no past_key (full cache is already provided) None, # no past_value static_cache.nonpad_kv_seqlen, @@ -167,8 +208,11 @@ def _apply_attention( kv_num_heads=num_key_value_heads, scale=scale, softcap=softcap, - is_causal=1, - _outputs=3, + is_causal=0, + # Only attn_output is consumed; the updated cache comes from the + # TensorScatter writes (updated_k/updated_v), not the op's present + # outputs, so request a single output instead of discarding two. + _outputs=1, ) return attn_output, updated_k, updated_v diff --git a/src/mobius/components/_attention_test.py b/src/mobius/components/_attention_test.py index a60cd7c2..ab6bd247 100644 --- a/src/mobius/components/_attention_test.py +++ b/src/mobius/components/_attention_test.py @@ -460,3 +460,128 @@ def test_direct_gqa_and_rewrite_rule_produce_same_structure(self): # Neither path should leave any standard Attention nodes assert ops_direct.get("Attention", 0) == 0 assert ops_rewrite.get("Attention", 0) == 0 + + +class TestApplyAttentionStaticCacheFallback: + """`_apply_attention` static-cache mask fallback when ``causal_mask`` is None. + + Production hoists a single shared mask via ``StaticCacheState.causal_mask``. + Direct callers (e.g. unit tests) may leave ``causal_mask=None``; the + fallback must then build the mask on demand via the 4-arg + ``create_static_cache_causal_mask`` path, yielding a graph equivalent to the + hoisted path. This covers the otherwise-untested ``None`` branch. + """ + + NUM_HEADS = 2 + KV_HEADS = 2 + HEAD_DIM = 8 + BATCH = 1 + S_Q = 2 + MAX_SEQ = 4 + + def _build_static_cache_attention(self, *, hoist: bool): + """Build a graph that calls ``_apply_attention`` in static-cache mode. + + When ``hoist`` is True the mask is built once up front and threaded + through ``StaticCacheState.causal_mask`` (production path); when False + the field is ``None`` so ``_apply_attention`` must build it on demand + (fallback path). Returns ``(graph, hoisted_mask_or_None)``. + """ + from mobius.components._attention import ( + StaticCacheState, + _apply_attention, + ) + from mobius.components._common import create_static_cache_causal_mask + + hidden = self.NUM_HEADS * self.HEAD_DIM + builder, op, graph = create_test_builder() + query = create_test_input(builder, "query", [self.BATCH, self.S_Q, hidden]) + key = create_test_input(builder, "key", [self.BATCH, self.S_Q, hidden]) + value = create_test_input(builder, "value", [self.BATCH, self.S_Q, hidden]) + key_cache = create_test_input(builder, "key_cache", [self.BATCH, self.MAX_SEQ, hidden]) + value_cache = create_test_input( + builder, "value_cache", [self.BATCH, self.MAX_SEQ, hidden] + ) + write_indices = create_test_input( + builder, "write_indices", [self.BATCH], ir.DataType.INT64 + ) + nonpad_kv_seqlen = create_test_input( + builder, "nonpad_kv_seqlen", [self.BATCH], ir.DataType.INT64 + ) + + hoisted_mask = None + if hoist: + hoisted_mask = create_static_cache_causal_mask(op, query, key_cache, write_indices) + + static_cache = StaticCacheState( + key_cache=key_cache, + value_cache=value_cache, + write_indices=write_indices, + nonpad_kv_seqlen=nonpad_kv_seqlen, + causal_mask=hoisted_mask, + ) + _apply_attention( + op, + query, + key, + value, + None, + None, + None, + num_attention_heads=self.NUM_HEADS, + num_key_value_heads=self.KV_HEADS, + scale=self.HEAD_DIM**-0.5, + static_cache=static_cache, + ) + return graph, hoisted_mask + + def test_fallback_builds_causal_mask_when_none(self): + """causal_mask=None → Attention attn_mask is built via the 4-arg path.""" + graph, _ = self._build_static_cache_attention(hoist=False) + + attention_nodes = [n for n in graph if n.op_type == "Attention"] + assert len(attention_nodes) == 1 + + attn_mask_input = attention_nodes[0].inputs[3] + assert attn_mask_input is not None and attn_mask_input.name != "", ( + f"fallback should connect an explicit causal attn_mask, got {attn_mask_input}" + ) + producer = attn_mask_input.producer() + assert producer is not None and producer.op_type == "GreaterOrEqual", ( + "fallback should build the mask via create_static_cache_causal_mask " + f"(GreaterOrEqual root), got {None if producer is None else producer.op_type}" + ) + # The mask construction subgraph is present exactly once (built on demand): + # one GreaterOrEqual root over two Range index vectors (q_offsets, key_positions). + assert count_op_type(graph, "GreaterOrEqual") == 1 + assert count_op_type(graph, "Range") == 2 + + def test_hoisted_mask_is_consumed_not_rebuilt(self): + """causal_mask present → the exact shared Value is consumed, no rebuild.""" + graph, hoisted_mask = self._build_static_cache_attention(hoist=True) + + attention_nodes = [n for n in graph if n.op_type == "Attention"] + assert len(attention_nodes) == 1 + # Identity: the hoisted mask Value is consumed directly. + assert attention_nodes[0].inputs[3] is hoisted_mask + # Still exactly one mask root — the hoisted one, not a second rebuild. + assert count_op_type(graph, "GreaterOrEqual") == 1 + + def test_fallback_graph_equivalent_to_hoisted(self): + """The fallback-built mask is equivalent to the hoisted mask. + + Both paths call ``create_static_cache_causal_mask`` exactly once with the + same args and run the same ``_apply_attention`` body, so the resulting + graphs must have identical op-type multisets. + """ + from collections import Counter + + graph_fallback, _ = self._build_static_cache_attention(hoist=False) + graph_hoist, _ = self._build_static_cache_attention(hoist=True) + + counts_fallback = Counter(n.op_type for n in graph_fallback) + counts_hoist = Counter(n.op_type for n in graph_hoist) + assert counts_fallback == counts_hoist, ( + "fallback-built mask graph differs from the hoisted-mask graph: " + f"{dict(counts_fallback)} vs {dict(counts_hoist)}" + ) diff --git a/src/mobius/components/_common.py b/src/mobius/components/_common.py index 2d4836f5..da5f4a0b 100644 --- a/src/mobius/components/_common.py +++ b/src/mobius/components/_common.py @@ -321,3 +321,103 @@ def create_sliding_window_mask( # Combine with padding mask padding_mask = op.Cast(op.Unsqueeze(attention_mask, [1]), to=ir.DataType.BOOL) return op.And(within_window, padding_mask) + + +def create_static_cache_causal_mask( + op: OpBuilder, + query: ir.Value, + key_cache: ir.Value, + write_indices: ir.Value, +): + """Build a causal attention mask for the static (TensorScatter) KV cache. + + The opset-24 ONNX ``Attention`` CUDA kernel rejects ``is_causal=1`` + together with ``nonpad_kv_seqlen`` whenever the query length differs + from the total KV length and there is no ``past_key`` (see + ``onnxruntime/core/providers/cuda/llm/attention.cc`` — the + ``causal_cross_no_past`` guard). In the static cache the KV buffers + are pre-allocated to ``max_seq_len``, so ``S_q != total_kv`` in **both** + prefill and decode. The fix, per ORT's own guidance, is to drive the + Attention op with ``is_causal=0`` and supply an explicit causal mask. + + A query token at position ``t`` within this step writes into cache slot + ``write_indices[b] + t`` (``write_indices[b]`` is the number of valid + cache tokens *before* this step). Causality means it may attend to + every cache slot ``j`` with ``j <= write_indices[b] + t``. This single + rule serves both phases: + + * **Prefill** (``write_indices=0``, ``S_q=N``): triangular causal mask + ``j <= t``. + * **Decode** (``write_indices=N``, ``S_q=1``): keep slots ``j <= N``, + i.e. all previously written tokens plus the just-written one. + + **Cache contract — compact / right-trimmed.** The mask is *purely + positional*: it keeps slot ``j`` iff ``j <= write_indices[b] + t`` and + never reads ``nonpad_kv_seqlen`` in the mask math. Correctness therefore + relies on the **compact-cache invariant** — valid tokens occupy cache + slots ``[0, write_indices[b] + t]`` contiguously, with padding only on the + right, i.e. ``nonpad_kv_seqlen[b] == write_indices[b] + S_q`` per batch. + Under that invariant the padding slots ``j >= nonpad_kv_seqlen[b]`` are + always greater than ``write_indices[b] + t``, so they are masked out as a + side effect. The padding bound itself is **enforced by the ORT Attention + kernel** via ``nonpad_kv_seqlen`` (external-cache input #6): the kernel + clamps attention to key slots ``j < nonpad_kv_seqlen[b]`` independently of + this attn_mask, on both the CUDA and CPU EPs (verified: poisoning slots + ``[nonpad, max_seq)`` yields bit-identical logits). This mask is therefore + **causal-only** — it must not re-encode the nonpad bound, since a + ``j < nonpad_kv_seqlen[b]`` term here would merely duplicate input #6 and + add dead graph nodes. mobius's ``TensorScatter`` writes with sequential + ``write_indices`` satisfy the compact invariant for standard batched + generation (per-row prefill scatters ``[0, prompt_len)``). + + Ragged / left / interior padding is **out of contract**: a scalar + ``nonpad_kv_seqlen`` can only express a compact valid prefix ``[0, nonpad)`` + and cannot describe interior holes, so neither this mask nor the kernel's + input #6 can represent them. This export does not support ragged caches. + + The mask is 4D ``[batch, 1, S_q, max_seq]`` rather than 3D on purpose: + ``ConvertAttnMaskToBias`` (attention.cc) treats a 3D mask as + ``[heads, q, kv]`` and broadcasts over the batch dimension, which would + be incorrect because ``write_indices`` is per-batch. A 4D mask with a + leading batch dim is honored per-batch (and ``dim1 == 1`` broadcasts + over heads). + + Args: + op: The OpBuilder. + query: Query sequence source ``[batch, S_q, hidden]`` (3D) **or** + ``[batch, S_q]`` (2D, e.g. ``input_ids``); only dim 1 is read, + to derive the query sequence length ``S_q``. + key_cache: Pre-allocated key cache ``[batch, max_seq, kv_hidden]``; + dim 1 supplies the total KV length ``max_seq``. + write_indices: Per-batch write start position ``[batch]`` INT64 — + the number of valid cache tokens before this step. + + Returns: + Bool mask ``[batch, 1, S_q, max_seq]``. ``True`` = attend, + ``False`` = mask out. + """ + zero = op.Constant(value_int=0) + one = op.Constant(value_int=1) + + # Scalar S_q (query length) and total KV length (max_seq) for Range. + q_len = op.Squeeze(op.Shape(query, start=1, end=2), op.Constant(value_ints=[0])) + total_kv = op.Squeeze(op.Shape(key_cache, start=1, end=2), op.Constant(value_ints=[0])) + + # Per-step query offsets 0..S_q-1 and key slot indices 0..max_seq-1. + q_offsets = op.Range(zero, q_len, one) # [S_q] int64 + key_positions = op.Range(zero, total_kv, one) # [max_seq] int64 + + # Absolute query positions: write_indices[b] + t → [batch, S_q]. + query_positions = op.Add( + op.Unsqueeze(write_indices, [1]), # [batch, 1] + op.Unsqueeze(q_offsets, [0]), # [1, S_q] + ) + + # Reshape for broadcasting to [batch, 1, S_q, max_seq]. + query_positions = op.Unsqueeze(query_positions, [1, 3]) # [batch, 1, S_q, 1] + key_positions = op.Unsqueeze(key_positions, [0, 1, 2]) # [1, 1, 1, max_seq] + + # Keep key slot j for query at position p iff j <= p (positional causal + # bound only; the nonpad/padding key-bound is enforced by the ORT Attention + # kernel via nonpad_kv_seqlen, input #6 — this mask must not re-encode it). + return op.GreaterOrEqual(query_positions, key_positions) diff --git a/src/mobius/components/_common_test.py b/src/mobius/components/_common_test.py index 3ebc1bf9..05d1dad1 100644 --- a/src/mobius/components/_common_test.py +++ b/src/mobius/components/_common_test.py @@ -5,7 +5,9 @@ from __future__ import annotations +import numpy as np import onnx_ir as ir +import onnxruntime as ort from mobius._testing import count_op_type, create_test_builder, create_test_input from mobius.components._common import ( @@ -13,9 +15,37 @@ Linear, create_attention_bias, create_padding_mask, + create_static_cache_causal_mask, ) +def _run_static_cache_mask(max_seq: int, query_len: int, write_index: int) -> np.ndarray: + """Build the static-cache causal mask subgraph and evaluate it on CPU. + + Returns the bool mask array of shape ``[1, 1, query_len, max_seq]`` where + ``True`` means the query token attends to that cache slot. + """ + builder, op, graph = create_test_builder() + query = create_test_input(builder, "query", [1, "S_q", 32], dtype=ir.DataType.FLOAT) + key_cache = create_test_input( + builder, "key_cache", [1, max_seq, 16], dtype=ir.DataType.FLOAT + ) + write_indices = create_test_input(builder, "write_indices", [1], dtype=ir.DataType.INT64) + mask = create_static_cache_causal_mask(op, query, key_cache, write_indices) + mask.name = "mask" + graph.outputs.append(mask) + proto = ir.to_proto(ir.Model(graph, ir_version=10)) + session = ort.InferenceSession( + proto.SerializeToString(), providers=["CPUExecutionProvider"] + ) + feeds = { + "query": np.zeros((1, query_len, 32), np.float32), + "key_cache": np.zeros((1, max_seq, 16), np.float32), + "write_indices": np.array([write_index], np.int64), + } + return session.run(None, feeds)[0] + + class TestLinear: def test_linear_with_bias(self): linear = Linear(64, 128, bias=True) @@ -178,3 +208,65 @@ def test_uses_simpler_ops_than_attention_bias(self): assert count_op_type(graph_bias, "CumSum") >= 1 assert count_op_type(graph_bias, "GreaterOrEqual") >= 1 assert count_op_type(graph_bias, "Where") >= 1 + + +class TestCreateStaticCacheCausalMask: + """Value-level checks for the static-cache causal mask. + + Query token ``t`` writes to absolute cache slot ``write_index + t`` and must + attend to every slot ``j <= write_index + t`` (causal), while padded/future + slots stay masked. This is verified by running the subgraph on CPU. + """ + + def test_structure_uses_range_and_greater_or_equal(self): + builder, op, graph = create_test_builder() + query = create_test_input(builder, "query", [1, "S_q", 32], dtype=ir.DataType.FLOAT) + key_cache = create_test_input( + builder, "key_cache", [1, 8, 16], dtype=ir.DataType.FLOAT + ) + write_indices = create_test_input( + builder, "write_indices", [1], dtype=ir.DataType.INT64 + ) + mask = create_static_cache_causal_mask(op, query, key_cache, write_indices) + assert mask is not None + assert count_op_type(graph, "Range") >= 1 + assert count_op_type(graph, "GreaterOrEqual") >= 1 + + def test_mask_shape_is_4d_batch_one_head(self): + mask = _run_static_cache_mask(max_seq=8, query_len=4, write_index=0) + # [B, 1, S_q, max_seq] — dim1=1 broadcasts over heads, dim0=batch. + assert mask.shape == (1, 1, 4, 8) + assert mask.dtype == bool + + def test_prefill_is_triangular(self): + """write_index=0: each query token t attends to slots 0..t only.""" + mask = _run_static_cache_mask(max_seq=8, query_len=4, write_index=0) + expected = np.array( + [ + [1, 0, 0, 0, 0, 0, 0, 0], + [1, 1, 0, 0, 0, 0, 0, 0], + [1, 1, 1, 0, 0, 0, 0, 0], + [1, 1, 1, 1, 0, 0, 0, 0], + ], + dtype=bool, + ) + np.testing.assert_array_equal(mask[0, 0], expected) + + def test_decode_keeps_prefix_through_write_index(self): + """Single decode token at slot N attends to slots 0..N, masks the rest.""" + mask = _run_static_cache_mask(max_seq=8, query_len=1, write_index=3) + expected = np.array([[1, 1, 1, 1, 0, 0, 0, 0]], dtype=bool) + np.testing.assert_array_equal(mask[0, 0], expected) + + def test_chunked_prefill_offset_write_index(self): + """Second prefill chunk (write_index=2, S_q=3) stays causal at absolute pos.""" + mask = _run_static_cache_mask(max_seq=8, query_len=3, write_index=2) + expected = np.array( + [ + [1, 1, 1, 0, 0, 0, 0, 0], + [1, 1, 1, 1, 0, 0, 0, 0], + [1, 1, 1, 1, 1, 0, 0, 0], + ], + dtype=bool, + ) + np.testing.assert_array_equal(mask[0, 0], expected) diff --git a/src/mobius/tasks/_cache_utils.py b/src/mobius/tasks/_cache_utils.py index 41cf5071..d8c156b5 100644 --- a/src/mobius/tasks/_cache_utils.py +++ b/src/mobius/tasks/_cache_utils.py @@ -114,13 +114,69 @@ def _register_kv_cache_outputs( present_key_values: list[tuple[ir.Value, ir.Value]], *, prefix: str = "present", + batch: ir.SymbolicDim | str | None = None, + num_kv_heads: int | None = None, + key_head_dim: int | None = None, + value_head_dim: int | None = None, + total_seq_len: ir.SymbolicDim | str | int | None = None, + dtype: ir.DataType | None = None, ) -> None: """Name and register KV cache outputs on the graph. - Output shapes and dtypes are inferred by the shape inference pass - that runs during model optimization. + When every present-shape parameter (``batch``, ``num_kv_heads``, + ``key_head_dim``, ``value_head_dim``, ``total_seq_len``, ``dtype``) is + supplied, each ``present.{i}.{key,value}`` output is stamped with an + explicit ``[batch, num_kv_heads, total_seq_len, head_dim]`` type, symmetric + to the ``past_key_values`` inputs created by :func:`_make_kv_cache_inputs`. + + This explicit stamp is required for ``com.microsoft::GroupQueryAttention`` + exports: that contrib op's shape inference mis-derives the present + ``head_dim`` (it divides the *packed* QKV query hidden by + ``num_heads + 2 * kv_num_heads`` and lands on the wrong value), so the + ``present.*`` outputs would otherwise declare the wrong ``head_dim`` even + though the kernel produces correct data at runtime. The mismatch makes ORT + log a shape-merge warning and breaks present->past chaining in consumers + such as onnxruntime-genai that trust the declared shapes. The plain ONNX + ``Attention`` op infers the present shape correctly, so the stamp is a + no-op there. + + When the parameters are omitted, output shapes/dtypes are left to the + shape inference pass that runs during model optimization. + + The present-shape parameters are all-or-nothing by design: pass every one + to stamp the explicit type, or none to opt out and infer. A *partial* set + is always a wiring slip (a caller wired some dims but dropped others) and is + never legitimate, so it is rejected fail-closed: a partial set raises + :class:`ValueError` naming the provided and missing parameters, rather than + silently falling back to the known-wrong inference path. """ + params = { + "batch": batch, + "num_kv_heads": num_kv_heads, + "key_head_dim": key_head_dim, + "value_head_dim": value_head_dim, + "total_seq_len": total_seq_len, + "dtype": dtype, + } + provided = [name for name, value in params.items() if value is not None] + stamp = len(provided) == len(params) + if provided and not stamp: + missing = [name for name in params if params[name] is None] + raise ValueError( + f"_register_kv_cache_outputs received a partial set of present-shape " + f"parameters (provided {provided}, missing {missing}); these are " + f"all-or-nothing. Pass all six to stamp explicit present.* types " + f"(required for correct GroupQueryAttention head_dim), or none to opt " + f"out and infer." + ) for i, (present_key, present_value) in enumerate(present_key_values): + if stamp: + present_key.shape = ir.Shape([batch, num_kv_heads, total_seq_len, key_head_dim]) + present_key.type = ir.TensorType(dtype) + present_value.shape = ir.Shape( + [batch, num_kv_heads, total_seq_len, value_head_dim] + ) + present_value.type = ir.TensorType(dtype) builder.add_output(present_key, f"{prefix}.{i}.key") builder.add_output(present_value, f"{prefix}.{i}.value") diff --git a/src/mobius/tasks/_cache_utils_test.py b/src/mobius/tasks/_cache_utils_test.py new file mode 100644 index 00000000..780e1cfa --- /dev/null +++ b/src/mobius/tasks/_cache_utils_test.py @@ -0,0 +1,124 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT License. + +"""Tests for KV-cache I/O helpers in :mod:`mobius.tasks._cache_utils`.""" + +from __future__ import annotations + +import onnx_ir as ir +import pytest + +from mobius.tasks._base import _make_graph +from mobius.tasks._cache_utils import _register_kv_cache_outputs + + +def _present_pair(name: str, wrong_head_dim: int) -> tuple[ir.Value, ir.Value]: + """A present key/value pair carrying a deliberately wrong inferred shape. + + Mimics ``com.microsoft::GroupQueryAttention`` shape inference, which + mis-derives ``head_dim`` (e.g. 32 instead of 96) for the present outputs. + """ + key = ir.Value(name=f"{name}_key") + value = ir.Value(name=f"{name}_value") + for v in (key, value): + v.shape = ir.Shape(["batch", 32, "seq", wrong_head_dim]) + v.type = ir.TensorType(ir.DataType.FLOAT16) + return key, value + + +def _dims(value: ir.Value) -> list[object]: + return [d if isinstance(d, int) else str(d) for d in value.shape] + + +class TestRegisterKVCacheOutputs: + def test_stamps_explicit_shapes_over_wrong_inference(self): + """Explicit params must override mis-inferred present shapes (GQA bug).""" + _, builder = _make_graph() + pairs = [_present_pair("present.0", wrong_head_dim=32)] + + _register_kv_cache_outputs( + builder, + pairs, + batch=ir.SymbolicDim("batch"), + num_kv_heads=32, + key_head_dim=96, + value_head_dim=96, + total_seq_len="past_sequence_len + sequence_len", + dtype=ir.DataType.FLOAT16, + ) + + key, value = pairs[0] + assert _dims(key) == ["batch", 32, "past_sequence_len + sequence_len", 96] + assert _dims(value) == ["batch", 32, "past_sequence_len + sequence_len", 96] + assert key.dtype == ir.DataType.FLOAT16 + assert value.dtype == ir.DataType.FLOAT16 + + def test_distinct_key_value_head_dims(self): + """MLA-style caches may use different key/value head dims.""" + _, builder = _make_graph() + pairs = [_present_pair("present.0", wrong_head_dim=32)] + + _register_kv_cache_outputs( + builder, + pairs, + batch=ir.SymbolicDim("batch"), + num_kv_heads=16, + key_head_dim=192, + value_head_dim=128, + total_seq_len="past_sequence_len + sequence_len", + dtype=ir.DataType.FLOAT16, + ) + + key, value = pairs[0] + assert _dims(key)[1] == 16 and _dims(key)[3] == 192 + assert _dims(value)[1] == 16 and _dims(value)[3] == 128 + + def test_no_params_leaves_shapes_untouched(self): + """Without shape params the helper must not stamp (inference path).""" + _, builder = _make_graph() + pairs = [_present_pair("present.0", wrong_head_dim=32)] + + _register_kv_cache_outputs(builder, pairs) + + key, _ = pairs[0] + # Unchanged: still the (wrong) pre-existing inferred shape. + assert _dims(key) == ["batch", 32, "seq", 32] + + def test_partial_params_raise(self): + """A partial present-shape set is a wiring slip -> must raise. + + This is the fail-closed regression proof: the exact input that + previously warned-and-proceeded (shipping a structurally-wrong model + with mis-derived GroupQueryAttention ``head_dim``) now raises before + any output is registered. + """ + _, builder = _make_graph() + pairs = [_present_pair("present.0", wrong_head_dim=32)] + + with pytest.raises(ValueError, match="partial set of present-shape parameters") as exc: + _register_kv_cache_outputs( + builder, + pairs, + batch=ir.SymbolicDim("batch"), + num_kv_heads=32, + # key_head_dim / value_head_dim / total_seq_len / dtype omitted + ) + + # The message must name every omitted parameter so the slip is diagnosable. + for missing in ("key_head_dim", "value_head_dim", "total_seq_len", "dtype"): + assert missing in str(exc.value) + + def test_registers_named_outputs(self): + """Outputs are registered with the conventional present.{i}.* names.""" + _, builder = _make_graph() + pairs = [_present_pair("a", 32), _present_pair("b", 32)] + + _register_kv_cache_outputs(builder, pairs) + + names = [v.name for v in builder.graph.outputs] + assert names == [ + "present.0.key", + "present.0.value", + "present.1.key", + "present.1.value", + ] diff --git a/src/mobius/tasks/_causal_lm.py b/src/mobius/tasks/_causal_lm.py index 471d350b..4a2ca585 100644 --- a/src/mobius/tasks/_causal_lm.py +++ b/src/mobius/tasks/_causal_lm.py @@ -11,6 +11,7 @@ from mobius._configs import ArchitectureConfig from mobius._model_package import ModelPackage from mobius.components._attention import StaticCacheState +from mobius.components._common import create_static_cache_causal_mask from mobius.tasks._base import ( ModelTask, _make_graph, @@ -61,7 +62,10 @@ class CausalLMTask(ModelTask): - logits: FLOAT - updated_key_cache.{i} / updated_value_cache.{i}: FLOAT - No ``attention_mask`` input — causal masking uses ``is_causal=1``. + No ``attention_mask`` input — the static cache drives the + ``Attention`` op with ``is_causal=0`` and an explicit offset-aware + causal mask built from ``write_indices``/``nonpad_kv_seqlen`` (see + :func:`mobius.components._common.create_static_cache_causal_mask`). The module's ``forward()`` must accept ``(op, input_ids, attention_mask, position_ids, past_key_values)`` @@ -131,6 +135,7 @@ def build( config.dtype, batch, max_seq_len, + input_ids, ) else: past_seq_len = ir.SymbolicDim("past_sequence_len") @@ -153,6 +158,10 @@ def build( num_kv_cache_heads = ( config.num_attention_heads if use_mla else config.num_key_value_heads ) + kv_key_head_dim = ( + (config.qk_nope_head_dim or 0) + (config.qk_rope_head_dim or 0) + ) or config.head_dim + kv_value_head_dim = config.v_head_dim or config.head_dim past_key_values = _make_kv_cache_inputs( builder, @@ -162,9 +171,8 @@ def build( config.dtype, batch, past_seq_len, - key_head_dim=((config.qk_nope_head_dim or 0) + (config.qk_rope_head_dim or 0)) - or None, - value_head_dim=config.v_head_dim or None, + key_head_dim=kv_key_head_dim, + value_head_dim=kv_value_head_dim, ) logits, present_key_values = module( @@ -184,9 +192,19 @@ def build( present_key_values, ) else: + # Stamp explicit present shapes symmetric to the past inputs so the + # com.microsoft::GroupQueryAttention export declares the correct + # head_dim (its contrib-op shape inference otherwise mis-derives it + # from the packed QKV hidden). total_seq = past + current sequence. _register_kv_cache_outputs( builder, present_key_values, + batch=batch, + num_kv_heads=num_kv_cache_heads, + key_head_dim=kv_key_head_dim, + value_head_dim=kv_value_head_dim, + total_seq_len="past_sequence_len + sequence_len", + dtype=config.dtype, ) return ModelPackage({"model": _make_model(graph)}, config=config) @@ -271,15 +289,24 @@ def _make_static_cache_inputs( dtype: ir.DataType, batch: ir.SymbolicDim, max_seq_len: int, + query_seq_source: ir.Value, ) -> list[StaticCacheState]: """Create static KV cache inputs for ``num_layers`` layers. Uses ``builder.input()`` to create and register graph inputs directly. + Args: + query_seq_source: Any ``[batch, S_q, ...]`` or ``[batch, S_q]`` Value + (``input_ids`` works) — only dim 1 is read, to derive the query + sequence length ``S_q`` for the shared causal mask. + Returns: A list of :class:`StaticCacheState` tuples for passing to the - module via ``past_key_values``. + module via ``past_key_values``. Every entry carries the SAME + ``causal_mask`` Value (the mask is layer-invariant, so it is built + once here and shared across layers). """ + op = builder.op kv_hidden = num_key_value_heads * head_dim cache_pairs: list[tuple[ir.Value, ir.Value]] = [] @@ -296,6 +323,12 @@ def _make_static_cache_inputs( ) cache_pairs.append((key_cache, value_cache)) + # A zero-layer model has no cache and no mask to build; return early + # before indexing cache_pairs[0][0] (which would raise IndexError) and + # before registering shared inputs that would have no consumer. + if not cache_pairs: + return [] + # Shared inputs across all layers write_indices = builder.input( "write_indices", @@ -308,7 +341,20 @@ def _make_static_cache_inputs( shape=[batch], ) - # Build StaticCacheState for each layer (shared indices) + # Build the static-cache causal mask ONCE — it is layer-invariant + # (depends only on S_q, max_seq and write_indices, all identical across + # layers). The first layer's key_cache supplies max_seq (identical for + # every layer). The same ir.Value is then threaded to every layer's + # StaticCacheState below, so the graph holds a single shared mask subgraph + # instead of one rebuilt copy per layer. + causal_mask = create_static_cache_causal_mask( + op, + query_seq_source, + cache_pairs[0][0], + write_indices, + ) + + # Build StaticCacheState for each layer (shared indices + shared mask Value) static_caches: list[StaticCacheState] = [] for key_cache, value_cache in cache_pairs: static_caches.append( @@ -317,6 +363,7 @@ def _make_static_cache_inputs( value_cache=value_cache, write_indices=write_indices, nonpad_kv_seqlen=nonpad_kv_seqlen, + causal_mask=causal_mask, ) ) @@ -358,10 +405,13 @@ def _validate_static_cache_support(module: nn.Module) -> None: ``create_attention_bias()``, so ``attention_mask=None`` would fail. Needs position embedding adaptation. - - **Falcon (ALiBi)**: The ALiBi variant uses ``is_causal=0`` with a - position-dependent bias that encodes both causal masking and - distance-based attention decay. This is fundamentally - incompatible with the ``is_causal=1`` static cache pattern. + - **Falcon (ALiBi)**: The ALiBi variant replaces a plain causal mask + with a position-dependent *additive* bias that encodes both causal + masking and distance-based attention decay. The static cache path + already occupies the ``attn_mask`` input with its own explicit + offset-aware causal mask (``is_causal=0``), so ALiBi support would + require folding the distance-decay bias into that static-mask + construction rather than passing a separate bias. Raises: TypeError: If any decoder layer is not a supported type. diff --git a/src/mobius/tasks/_task_test.py b/src/mobius/tasks/_task_test.py index d50f4337..cf2fdd0a 100644 --- a/src/mobius/tasks/_task_test.py +++ b/src/mobius/tasks/_task_test.py @@ -91,6 +91,96 @@ def test_build_outputs(self): assert "present.0.key" in output_names assert "present.0.value" in output_names + def test_present_outputs_match_past_inputs(self): + """present.{i}.* must match the corresponding past_key_values.{i}.* inputs. + + They must declare the same kv_heads/head_dim/dtype as the corresponding + past_key_values.{i}.* inputs, with total (past+current) sequence length. + Guards the GQA present-head_dim export bug, where the contrib-op shape + inference would otherwise mis-declare head_dim. + """ + config = make_config() + module = CausalLMModel(config) + pkg = CausalLMTask().build(module, config) + model = pkg["model"] + inputs = {v.name: v for v in model.graph.inputs} + outputs = {v.name: v for v in model.graph.outputs} + + def dims(value): + return [d if isinstance(d, int) else str(d) for d in value.shape] + + for i in range(config.num_hidden_layers): + for kind in ("key", "value"): + past = inputs[f"past_key_values.{i}.{kind}"] + present = outputs[f"present.{i}.{kind}"] + past_dims = dims(past) + present_dims = dims(present) + # batch, kv_heads, head_dim must match the past input exactly. + assert present_dims[0] == past_dims[0] + assert present_dims[1] == past_dims[1] + assert present_dims[3] == past_dims[3] + # present covers past + current tokens. + assert present_dims[2] == "past_sequence_len + sequence_len" + assert present.dtype == past.dtype + + def test_present_outputs_match_past_inputs_mla_distinct_head_dims(self): + """present/past KV symmetry holds when key_head_dim != value_head_dim. + + The default-config sibling test only exercises + ``key_head_dim == value_head_dim == head_dim``, so it cannot catch a + call-site wiring slip that swaps or drops one of the two distinct dims. + An MLA (DeepSeek-style) config gives the keys a head dim of + ``qk_nope_head_dim + qk_rope_head_dim`` and the values a *different* + ``v_head_dim``; this guards the present-KV head_dim stamping fix for the + very case it was written for (MLA distinct key/value head dims). + """ + config = make_config( + q_lora_rank=16, + kv_lora_rank=32, + qk_nope_head_dim=12, + qk_rope_head_dim=4, + v_head_dim=8, + ) + module = CausalLMModel(config) + pkg = CausalLMTask().build(module, config) + model = pkg["model"] + inputs = {v.name: v for v in model.graph.inputs} + outputs = {v.name: v for v in model.graph.outputs} + + # Precondition: the config must actually produce distinct key/value + # head dims, otherwise this test degenerates into the default case. + key_head_dim = inputs["past_key_values.0.key"].shape[3] + value_head_dim = inputs["past_key_values.0.value"].shape[3] + assert key_head_dim != value_head_dim, ( + "MLA config did not yield distinct key/value head dims " + f"(key={key_head_dim}, value={value_head_dim}); the test would not " + "exercise the distinct-head-dim path it is meant to guard" + ) + + def dims(value): + return [d if isinstance(d, int) else str(d) for d in value.shape] + + for i in range(config.num_hidden_layers): + for kind in ("key", "value"): + past = inputs[f"past_key_values.{i}.{kind}"] + present = outputs[f"present.{i}.{kind}"] + # Fail clearly (not with a cryptic TypeError inside dims()) if the + # present-KV head_dim stamp left the shape undeclared. + assert present.shape is not None, ( + f"present.{i}.{kind} has no declared shape — the present-KV " + "head_dim stamp did not run for the MLA distinct-head-dim path" + ) + past_dims = dims(past) + present_dims = dims(present) + # batch, kv_heads, head_dim must match the past input exactly -- + # including the per-kind (key vs value) distinct head_dim. + assert present_dims[0] == past_dims[0] + assert present_dims[1] == past_dims[1] + assert present_dims[3] == past_dims[3] + # present covers past + current tokens. + assert present_dims[2] == "past_sequence_len + sequence_len" + assert present.dtype == past.dtype + def test_build_producer_info(self): config = make_config() module = CausalLMModel(config) diff --git a/tests/build_graph_test.py b/tests/build_graph_test.py index d6d2f9c8..c4b421f4 100644 --- a/tests/build_graph_test.py +++ b/tests/build_graph_test.py @@ -4485,8 +4485,9 @@ def test_static_cache_graph_inputs(self): assert "input_ids" in input_names assert "position_ids" in input_names - # No attention_mask in static cache mode — causal masking is - # handled by is_causal=1 on the Attention op. + # No attention_mask in static cache mode — causal masking is handled + # by is_causal=0 plus an explicit offset-aware causal mask built from + # write_indices (create_static_cache_causal_mask). assert "attention_mask" not in input_names # Per-layer static cache inputs @@ -4554,8 +4555,14 @@ def test_static_cache_graph_validates(self): proto = ir.serde.serialize_model(model) assert len(proto.SerializeToString()) > 0 - def test_static_cache_attention_is_causal(self): - """Verify Attention ops use is_causal=1 in static cache mode.""" + def test_static_cache_attention_causal_via_explicit_mask(self): + """Verify Attention ops use is_causal=0 in static cache mode. + + The opset-24 Attention CUDA kernel rejects is_causal=1 together + with nonpad_kv_seqlen when S_q != total_kv with no past_key (always + true for a pre-allocated cache). The static path therefore uses + is_causal=0 plus an explicit causal mask. + """ model, config = self._build_static_cache_model() attention_nodes = [n for n in model.graph if n.op_type == "Attention"] @@ -4566,24 +4573,95 @@ def test_static_cache_attention_is_causal(self): assert is_causal is not None, ( f"Attention node {node.name} missing is_causal attribute" ) - assert is_causal.as_int() == 1, ( - f"Attention node {node.name} should have is_causal=1" + assert is_causal.as_int() == 0, ( + f"Attention node {node.name} should have is_causal=0 " + f"(causality is supplied via an explicit attn_mask)" ) - def test_static_cache_attention_no_attn_mask_input(self): - """Verify Attention ops do NOT receive attn_mask in static cache mode.""" + def test_static_cache_attention_has_causal_mask_input(self): + """Verify Attention ops receive an explicit causal attn_mask. + + With is_causal=0, causality must come from input 3 (attn_mask). + The mask is produced by create_static_cache_causal_mask, whose + final op is a GreaterOrEqual yielding a bool mask. Because the mask is + layer-invariant it is hoisted and built once, so every Attention node + must reference the *same* mask ``ir.Value`` (identity), not a per-layer + rebuild. + """ model, config = self._build_static_cache_model() attention_nodes = [n for n in model.graph if n.op_type == "Attention"] assert len(attention_nodes) == config.num_hidden_layers + shared_mask = None for node in attention_nodes: - # Input 3 (0-indexed) is attn_mask — should be empty/None + # Input 3 (0-indexed) is attn_mask — should be a real value now. attn_mask_input = node.inputs[3] - assert attn_mask_input is None or attn_mask_input.name == "", ( - f"Attention node {node.name} should not have attn_mask " - f"connected, but got input: {attn_mask_input}" + assert attn_mask_input is not None and attn_mask_input.name != "", ( + f"Attention node {node.name} should have an explicit causal " + f"attn_mask connected, but got: {attn_mask_input}" + ) + producer = attn_mask_input.producer() + assert producer is not None and producer.op_type == "GreaterOrEqual", ( + f"Attention node {node.name} attn_mask should be produced by " + f"the causal-mask GreaterOrEqual, got " + f"{None if producer is None else producer.op_type}" ) + # Hoist invariant: all layers share the identical mask Value. + if shared_mask is None: + shared_mask = attn_mask_input + else: + assert attn_mask_input is shared_mask, ( + "Attention nodes reference different mask Values — the " + "layer-invariant static-cache mask was rebuilt per layer " + "instead of being hoisted/shared" + ) + + def test_static_cache_mask_built_once(self): + """The layer-invariant mask is built once and shared, not per layer. + + The mask depends only on S_q, max_seq and write_indices — all identical + across layers — so the construction ops must appear exactly once + regardless of layer count, and every Attention node must consume that + single shared mask Value. + """ + model, config = self._build_static_cache_model() + num_layers = config.num_hidden_layers + assert num_layers >= 2, "Need >=2 layers to prove the mask is not per-layer" + + op_type_counts: dict[str, int] = {} + for node in model.graph: + op_type_counts[node.op_type] = op_type_counts.get(node.op_type, 0) + 1 + + # Mask root op appears ONCE (not once per layer). + assert op_type_counts.get("GreaterOrEqual", 0) == 1, ( + f"Expected exactly 1 GreaterOrEqual (single shared mask root), got " + f"{op_type_counts.get('GreaterOrEqual', 0)} — mask appears to be " + f"rebuilt per layer" + ) + + # Every Attention node references the one shared mask Value. + attention_nodes = [n for n in model.graph if n.op_type == "Attention"] + assert len(attention_nodes) == num_layers + mask_values = {id(n.inputs[3]) for n in attention_nodes} + assert len(mask_values) == 1, ( + f"Attention nodes reference {len(mask_values)} distinct mask Values; " + f"expected 1 shared hoisted mask across all {num_layers} layers" + ) + + def test_static_cache_graph_contains_causal_mask_ops(self): + """Static cache graph must contain the causal-mask construction ops. + + Guards against regressing back to the is_causal=1 / no-mask form + that ORT rejects: the graph must build per-step query positions + (Range + Add) feeding a GreaterOrEqual mask. + """ + model, _ = self._build_static_cache_model() + op_types = [n.op_type for n in model.graph] + assert "Range" in op_types, "Causal mask should use Range for positions" + assert "GreaterOrEqual" in op_types, ( + "Causal mask should compare query/key positions with GreaterOrEqual" + ) def test_static_cache_moe_graph_builds(self): """Build a MoE model (qwen2_moe) with static cache.""" diff --git a/tests/static_cache_decode_test.py b/tests/static_cache_decode_test.py new file mode 100644 index 00000000..a48202f7 --- /dev/null +++ b/tests/static_cache_decode_test.py @@ -0,0 +1,465 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT License. + +"""End-to-end regression test for the static (TensorScatter) KV cache. + +mobius used to emit the static-cache ``Attention`` op with ``is_causal=1`` +alongside ``nonpad_kv_seqlen``. The opset-24 ONNX ``Attention`` CUDA kernel +rejects that combination whenever the query length differs from the total KV +length and there is no ``past_key`` (the ``causal_cross_no_past`` guard in +``onnxruntime/core/providers/cuda/llm/attention.cc``). Because the static +cache is pre-allocated to ``max_seq_len``, that guard fires in **both** +prefill (``S_q = N``) and decode (``S_q = 1``), raising ``NOT_IMPLEMENTED``. + +The fix sets ``is_causal=0`` and supplies an explicit causal mask +(:func:`mobius.components._common.create_static_cache_causal_mask`). This test +exercises the actual ONNX Runtime kernel for both phases so the regression +cannot silently come back. It requires the CUDA Execution Provider because +``TensorScatter`` and the external-cache ``Attention`` path are CUDA-only. + +The model is built fp32: the ``is_causal`` guard fires in ORT *before* kernel +dtype dispatch, so fp32 exercises the same external-cache code path as a +production fp16 export while avoiding the cos/sin-cache dtype casting that +only the full CLI build pipeline applies. A raw ``InferenceSession`` is used +(instead of the ``OnnxModelSession`` test helper) to keep this regression +guard free of the optional ``onnxruntime-easy`` dependency. +""" + +from __future__ import annotations + +import tempfile +from pathlib import Path + +import numpy as np +import onnx_ir as ir +import onnxruntime as ort +import pytest +from _test_configs import _base_config + +from mobius._registry import registry +from mobius.tasks import CausalLMTask + +pytestmark = pytest.mark.skipif( + "CUDAExecutionProvider" not in ort.get_available_providers(), + reason="static-cache TensorScatter / external-cache Attention are CUDA-only", +) + +_MAX_SEQ_LEN = 16 +_MODEL_TYPE = "qwen2" +_CACHE_DTYPE = np.float32 + +# CUDA attention logits are not bit-reproducible across kernels/drivers, so the +# "unchanged" (equality) poison controls compare with a tolerance rather than +# exact equality. The "changed" (negative) controls instead use exact +# inequality (``not np.array_equal``): any bit of change proves the slot was +# attended, and the poison magnitude (50.0) guarantees a large, unambiguous +# change there. That same 50.0 poison also far exceeds the 1e-5 tolerance +# below, so a genuine leak into a masked slot would move the logits well +# outside the band — the tolerant "unchanged" check still cannot mask a real +# leak. +_LOGIT_RTOL = 1e-5 +_LOGIT_ATOL = 1e-5 + + +def _fill_random_weights(model: ir.Model, rng: np.random.Generator) -> None: + """Fill empty initializers with small random values of their dtype. + + The graph is built without real weights; ORT still needs concrete + initializers to run. Small values keep logits finite and well-scaled. + """ + for initializer in model.graph.initializers.values(): + if initializer.const_value is not None: + continue + shape = initializer.shape + dims = [d if isinstance(d, int) else 1 for d in shape] if shape else [1] + dtype = initializer.dtype or ir.DataType.FLOAT + np_dtype = dtype.numpy() + if np.issubdtype(np_dtype, np.floating): + data = (rng.standard_normal(dims) * 0.02).astype(np_dtype) + else: + data = np.zeros(dims, dtype=np_dtype) + initializer.const_value = ir.Tensor(data) + + +def _build_static_cache_session( + tmp_dir: str, +) -> tuple[ort.InferenceSession, object]: + """Build a tiny static-cache qwen2 graph and load it on CUDA.""" + config = _base_config() + module = registry.get(_MODEL_TYPE)(config) + task = CausalLMTask(static_cache=True, max_seq_len=_MAX_SEQ_LEN) + model = task.build(module, config)["model"] + _fill_random_weights(model, np.random.default_rng(0)) + + model_path = str(Path(tmp_dir) / "model.onnx") + ir.save(model, model_path, external_data="model.onnx.data") + session = ort.InferenceSession( + model_path, + providers=["CUDAExecutionProvider", "CPUExecutionProvider"], + ) + assert "CUDAExecutionProvider" in session.get_providers(), ( + "static-cache regression test must run on CUDA" + ) + return session, config + + +def _empty_caches(num_layers: int, kv_hidden: int) -> dict[str, np.ndarray]: + """Zeroed ``[1, max_seq, kv_hidden]`` cache buffers for every layer.""" + feeds: dict[str, np.ndarray] = {} + for layer in range(num_layers): + zeros = np.zeros((1, _MAX_SEQ_LEN, kv_hidden), dtype=_CACHE_DTYPE) + feeds[f"key_cache.{layer}"] = zeros.copy() + feeds[f"value_cache.{layer}"] = zeros.copy() + return feeds + + +def _carry_caches(outputs: dict[str, np.ndarray], num_layers: int) -> dict[str, np.ndarray]: + """Feed the prefill ``updated_*`` caches back in as decode inputs.""" + feeds: dict[str, np.ndarray] = {} + for layer in range(num_layers): + feeds[f"key_cache.{layer}"] = outputs[f"updated_key_cache.{layer}"] + feeds[f"value_cache.{layer}"] = outputs[f"updated_value_cache.{layer}"] + return feeds + + +def test_static_cache_prefill_and_decode_run_on_cuda(): + """Prefill (S_q>1) and decode (S_q=1) both run without NOT_IMPLEMENTED. + + This is the regression guard the codebase previously lacked: it loads + a real static-cache graph and runs it on ``CUDAExecutionProvider`` for + both phases. Reverting to ``is_causal=1`` (no mask) makes ORT raise + ``NOT_IMPLEMENTED`` here, failing the test. + """ + with tempfile.TemporaryDirectory() as tmp_dir: + session, config = _build_static_cache_session(tmp_dir) + num_layers = config.num_hidden_layers + kv_hidden = config.num_key_value_heads * config.head_dim + vocab = config.vocab_size + output_names = [out.name for out in session.get_outputs()] + rng = np.random.default_rng(1) + + # --- Prefill: write N tokens from slot 0 (S_q = N != max_seq). --- + prefill_len = 4 + prefill_feeds: dict[str, np.ndarray] = { + "input_ids": rng.integers(0, vocab, size=(1, prefill_len), dtype=np.int64), + "position_ids": np.arange(prefill_len, dtype=np.int64)[None, :], + "write_indices": np.array([0], dtype=np.int64), + "nonpad_kv_seqlen": np.array([prefill_len], dtype=np.int64), + } + prefill_feeds.update(_empty_caches(num_layers, kv_hidden)) + + prefill_out = dict(zip(output_names, session.run(output_names, prefill_feeds))) + prefill_logits = prefill_out["logits"] + assert prefill_logits.shape == (1, prefill_len, vocab) + assert np.isfinite(prefill_logits).all(), "prefill logits must be finite" + + # --- Decode: one token at slot N (S_q = 1 != max_seq). --- + decode_feeds: dict[str, np.ndarray] = { + "input_ids": rng.integers(0, vocab, size=(1, 1), dtype=np.int64), + "position_ids": np.array([[prefill_len]], dtype=np.int64), + "write_indices": np.array([prefill_len], dtype=np.int64), + "nonpad_kv_seqlen": np.array([prefill_len + 1], dtype=np.int64), + } + decode_feeds.update(_carry_caches(prefill_out, num_layers)) + + decode_out = dict(zip(output_names, session.run(output_names, decode_feeds))) + decode_logits = decode_out["logits"] + assert decode_logits.shape == (1, 1, vocab) + assert np.isfinite(decode_logits).all(), "decode logits must be finite" + + # The decode step must have scattered its key into slot N (the + # previously-empty tail), confirming the in-place cache advanced. + advanced_key = decode_out["updated_key_cache.0"] + assert advanced_key.shape == (1, _MAX_SEQ_LEN, kv_hidden) + assert np.any(advanced_key[0, prefill_len] != 0), ( + "decode should scatter the new key into cache slot N" + ) + + +def test_static_cache_decode_mask_bounds_attention_to_frontier_on_cuda(): + """Always-masked decode does not attend to padding slots beyond ``nonpad``. + + The static cache exports the attention with ``is_causal=0`` plus an explicit + :func:`mobius.components._common.create_static_cache_causal_mask`, which keeps + key slot ``j`` for a query at absolute position ``p = write_indices[b] + t`` + iff ``j <= p``. For a single-token decode (``S_q=1``, ``write_indices=N``) + the causal frontier ``p = N`` coincides with the padding boundary, so this + test exercises the **padding** side of the mask: cache slots + ``j >= nonpad_kv_seqlen`` are unwritten padding and must never reach the + softmax. + + The intra-sequence **causal** side — a *written*, within-``nonpad`` key that + is in the future of an earlier query row — is a distinct bound that a + single-token decode cannot isolate; it is covered separately by + :func:`test_static_cache_prefill_causal_mask_blocks_future_keys_within_nonpad_on_cuda`. + + Checked two ways: + + * **Padding (negative) control:** poisoning every padding slot + ``j >= nonpad`` with large garbage must NOT change the decode logits — + those slots are masked out, so the result is bit-identical to the + clean-cache decode. + * **In-range (positive) control:** poisoning a slot strictly inside the + frontier MUST change the decode logits — proving the decode genuinely + attends to in-range keys, so the negative control is meaningful rather than + passing vacuously because decode ignores the cache. + """ + with tempfile.TemporaryDirectory() as tmp_dir: + session, config = _build_static_cache_session(tmp_dir) + num_layers = config.num_hidden_layers + kv_hidden = config.num_key_value_heads * config.head_dim + vocab = config.vocab_size + output_names = [out.name for out in session.get_outputs()] + rng = np.random.default_rng(7) + + # Prefill four real tokens into slots 0..3 to populate the cache. + prefill_len = 4 + prefill_feeds: dict[str, np.ndarray] = { + "input_ids": rng.integers(0, vocab, size=(1, prefill_len), dtype=np.int64), + "position_ids": np.arange(prefill_len, dtype=np.int64)[None, :], + "write_indices": np.array([0], dtype=np.int64), + "nonpad_kv_seqlen": np.array([prefill_len], dtype=np.int64), + } + prefill_feeds.update(_empty_caches(num_layers, kv_hidden)) + prefill_out = dict(zip(output_names, session.run(output_names, prefill_feeds))) + + # Decode one token into slot 4 at offset>0: write_indices=4, so the + # causal frontier is j <= 4 (slots 0..4 valid; nonpad=5 marks the tail). + nonpad = prefill_len + 1 + decode_inputs: dict[str, np.ndarray] = { + "input_ids": rng.integers(0, vocab, size=(1, 1), dtype=np.int64), + "position_ids": np.array([[prefill_len]], dtype=np.int64), + "write_indices": np.array([prefill_len], dtype=np.int64), + "nonpad_kv_seqlen": np.array([nonpad], dtype=np.int64), + } + + clean_feeds = {**decode_inputs, **_carry_caches(prefill_out, num_layers)} + baseline = dict(zip(output_names, session.run(output_names, clean_feeds))) + + # Padding (negative) control: poison every padding slot (j >= nonpad). + # These are unwritten padding; the mask excludes them, so the decode + # logits must be unchanged. + out_of_range = _carry_caches(prefill_out, num_layers) + for layer in range(num_layers): + for name in (f"key_cache.{layer}", f"value_cache.{layer}"): + buf = out_of_range[name].copy() + buf[:, nonpad:, :] = _CACHE_DTYPE(50.0) + out_of_range[name] = buf + out_of_range_feeds = {**decode_inputs, **out_of_range} + masked_out = dict(zip(output_names, session.run(output_names, out_of_range_feeds))) + + np.testing.assert_allclose( + masked_out["logits"], + baseline["logits"], + rtol=_LOGIT_RTOL, + atol=_LOGIT_ATOL, + err_msg=( + "decode logits changed when padding slots (j >= nonpad_kv_seqlen) " + "were poisoned — the static-cache mask is not excluding unwritten " + "padding from attention" + ), + ) + + # Positive control: poison an in-range slot (0, strictly inside the + # frontier). It is attended, so the decode logits MUST change — proving + # the negative control above is a live guard, not a no-op. + in_range = _carry_caches(prefill_out, num_layers) + for layer in range(num_layers): + for name in (f"key_cache.{layer}", f"value_cache.{layer}"): + buf = in_range[name].copy() + buf[:, 0, :] = _CACHE_DTYPE(50.0) + in_range[name] = buf + in_range_feeds = {**decode_inputs, **in_range} + attended = dict(zip(output_names, session.run(output_names, in_range_feeds))) + + # Exact inequality is correct for this negative control: any change at + # all proves decode attended the in-range slot (no fp-tolerance band). + assert not np.array_equal(baseline["logits"], attended["logits"]), ( + "decode logits were unchanged when an in-frontier cache slot (0) was " + "poisoned — decode is not attending to valid in-range keys, so the " + "out-of-range guard would pass vacuously" + ) + + +def test_static_cache_prefill_causal_mask_blocks_future_keys_within_nonpad_on_cuda(): + """The causal mask blocks *future* keys that are valid (within ``nonpad``). + + The decode guard above only exercises the *padding* side of the mask (slots + ``j >= nonpad``). This test isolates the orthogonal **causal** side: a key + slot that is genuinely written and within ``nonpad`` — so the padding bound + alone would admit it — but lies in the *future* of an earlier query row, and + so must be excluded by the explicit causal mask ``j <= write_indices + t``. + A single-token decode cannot probe this (it has no within-``nonpad`` future + slot); a multi-row step at a non-terminal offset can. + + Setup: prefill positions 0..3 to populate slots 0..3, then re-run a 2-token + block at positions 1,2 (``write_indices=1``, ``nonpad=4`` so every slot 0..3 + is valid, not padding). This block scatters into slots 1,2 only, leaving the + carried slots 0 and 3 untouched and poisonable: + + * **Causal (negative) control:** slot 3 is valid (within ``nonpad``) but in + the future of both query rows (positions 1 and 2 < 3). Poisoning it must + NOT change either row's logits — only the causal mask, not the padding + bound, can exclude a within-``nonpad`` slot. + * **In-range (positive) control:** slot 0 is in the causal past of both rows + and is carried (not rewritten). Poisoning it MUST change both rows' + logits, proving the rows genuinely attend their causal history (so the + negative control is not vacuous). + """ + with tempfile.TemporaryDirectory() as tmp_dir: + session, config = _build_static_cache_session(tmp_dir) + num_layers = config.num_hidden_layers + kv_hidden = config.num_key_value_heads * config.head_dim + vocab = config.vocab_size + output_names = [out.name for out in session.get_outputs()] + rng = np.random.default_rng(11) + + # Populate slots 0..3 with real keys (positions 0..3). + seed_len = 4 + seed_feeds: dict[str, np.ndarray] = { + "input_ids": rng.integers(0, vocab, size=(1, seed_len), dtype=np.int64), + "position_ids": np.arange(seed_len, dtype=np.int64)[None, :], + "write_indices": np.array([0], dtype=np.int64), + "nonpad_kv_seqlen": np.array([seed_len], dtype=np.int64), + } + seed_feeds.update(_empty_caches(num_layers, kv_hidden)) + seed_out = dict(zip(output_names, session.run(output_names, seed_feeds))) + + # Re-run a 2-token block at positions 1,2. write_indices=1 scatters into + # slots 1,2 only; nonpad=4 keeps slots 0..3 all valid (none are padding). + block_inputs: dict[str, np.ndarray] = { + "input_ids": rng.integers(0, vocab, size=(1, 2), dtype=np.int64), + "position_ids": np.array([[1, 2]], dtype=np.int64), + "write_indices": np.array([1], dtype=np.int64), + "nonpad_kv_seqlen": np.array([seed_len], dtype=np.int64), + } + + clean_feeds = {**block_inputs, **_carry_caches(seed_out, num_layers)} + baseline = dict(zip(output_names, session.run(output_names, clean_feeds))) + + # Causal (negative) control: poison slot 3 — valid (within nonpad) but in + # the future of both query rows (positions 1, 2). Not rewritten by this + # block (write region is {1, 2}), so the poison survives the scatter. + future = _carry_caches(seed_out, num_layers) + for layer in range(num_layers): + for name in (f"key_cache.{layer}", f"value_cache.{layer}"): + buf = future[name].copy() + buf[:, 3, :] = _CACHE_DTYPE(50.0) + future[name] = buf + future_feeds = {**block_inputs, **future} + future_poisoned = dict(zip(output_names, session.run(output_names, future_feeds))) + + np.testing.assert_allclose( + future_poisoned["logits"], + baseline["logits"], + rtol=_LOGIT_RTOL, + atol=_LOGIT_ATOL, + err_msg=( + "block logits changed when a valid within-nonpad but causally-future " + "key (slot 3, future of query positions 1 and 2) was poisoned — the " + "explicit causal mask is not enforcing j <= write_indices + t" + ), + ) + + # In-range (positive) control: poison slot 0 — causal past of both rows + # and carried (not rewritten) — so it must change the logits. + past = _carry_caches(seed_out, num_layers) + for layer in range(num_layers): + for name in (f"key_cache.{layer}", f"value_cache.{layer}"): + buf = past[name].copy() + buf[:, 0, :] = _CACHE_DTYPE(50.0) + past[name] = buf + past_feeds = {**block_inputs, **past} + past_poisoned = dict(zip(output_names, session.run(output_names, past_feeds))) + + # Exact inequality is correct for this negative control: any change at + # all proves the rows attended their causal history (no tolerance band). + assert not np.array_equal(baseline["logits"], past_poisoned["logits"]), ( + "block logits were unchanged when a causal-past key (slot 0) was " + "poisoned — the query rows are not attending their causal history, so " + "the future-key guard would pass vacuously" + ) + + +def test_static_cache_decode_mask_is_per_batch_on_cuda(): + """Each batch row's causal frontier follows *its own* ``write_indices``. + + The 4D ``[batch, 1, S_q, max_seq]`` mask exists so that per-row + ``write_indices`` are honored independently (a 3D mask would broadcast one + frontier across the whole batch). This is the value-level guard for that: + two rows decode with *different* ``write_indices`` over an identical cache, + then a single shared slot that is in one row's frontier but the other's + future is poisoned. Only the row that legitimately attends the slot may + change; the other row must be untouched — which can only hold if the mask + is built per batch. + """ + with tempfile.TemporaryDirectory() as tmp_dir: + session, config = _build_static_cache_session(tmp_dir) + num_layers = config.num_hidden_layers + kv_hidden = config.num_key_value_heads * config.head_dim + vocab = config.vocab_size + output_names = [out.name for out in session.get_outputs()] + rng = np.random.default_rng(23) + batch = 2 + + # Prefill five real tokens into slots 0..4 for both rows. + prefill_len = 5 + prefill_feeds: dict[str, np.ndarray] = { + "input_ids": rng.integers(0, vocab, size=(batch, prefill_len), dtype=np.int64), + "position_ids": np.tile(np.arange(prefill_len, dtype=np.int64), (batch, 1)), + "write_indices": np.zeros((batch,), dtype=np.int64), + "nonpad_kv_seqlen": np.full((batch,), prefill_len, dtype=np.int64), + } + for layer in range(num_layers): + zeros = np.zeros((batch, _MAX_SEQ_LEN, kv_hidden), dtype=_CACHE_DTYPE) + prefill_feeds[f"key_cache.{layer}"] = zeros.copy() + prefill_feeds[f"value_cache.{layer}"] = zeros.copy() + prefill_out = dict(zip(output_names, session.run(output_names, prefill_feeds))) + + # Decode one token per row with DIFFERENT write_indices: row 0's frontier + # is j <= 2, row 1's is j <= 4. Each writes its new token into its own + # write slot (2 / 4), so the shared probe slot 3 is left untouched by the + # scatter in both rows. Per the compact-cache invariant nonpad == write + # + S_q, so nonpad = [3, 5]. + decode_inputs: dict[str, np.ndarray] = { + "input_ids": rng.integers(0, vocab, size=(batch, 1), dtype=np.int64), + "position_ids": np.array([[2], [4]], dtype=np.int64), + "write_indices": np.array([2, 4], dtype=np.int64), + "nonpad_kv_seqlen": np.array([3, 5], dtype=np.int64), + } + + clean_feeds = {**decode_inputs, **_carry_caches(prefill_out, num_layers)} + baseline = dict(zip(output_names, session.run(output_names, clean_feeds))) + + # Poison the shared probe slot 3 in both rows. Slot 3 is in row 1's + # frontier (3 <= 4) but row 0's future (3 > 2), and is padding for row 0 + # (3 >= nonpad 3) yet valid for row 1 (3 < nonpad 5). + poisoned = _carry_caches(prefill_out, num_layers) + probe_slot = 3 + for layer in range(num_layers): + for name in (f"key_cache.{layer}", f"value_cache.{layer}"): + buf = poisoned[name].copy() + buf[:, probe_slot, :] = _CACHE_DTYPE(50.0) + poisoned[name] = buf + poisoned_feeds = {**decode_inputs, **poisoned} + poisoned_out = dict(zip(output_names, session.run(output_names, poisoned_feeds))) + + # Row 0 must be untouched: slot 3 is beyond its per-row frontier. + np.testing.assert_allclose( + poisoned_out["logits"][0], + baseline["logits"][0], + rtol=_LOGIT_RTOL, + atol=_LOGIT_ATOL, + err_msg=( + "row 0 logits changed when slot 3 was poisoned, but slot 3 is beyond " + "row 0's frontier (write_indices=2) — the mask is not applying " + "per-batch write_indices (it leaked row 1's wider frontier onto row 0)" + ), + ) + # Row 1 must change: slot 3 is within its per-row frontier and attended. + # Exact inequality is correct for this negative control (no tolerance band). + assert not np.array_equal(baseline["logits"][1], poisoned_out["logits"][1]), ( + "row 1 logits were unchanged when slot 3 (within its frontier, " + "write_indices=4) was poisoned — row 1 is not attending an in-frontier " + "key, so the per-batch row 0 control would pass vacuously" + )