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..0ff4c8b5 --- /dev/null +++ b/.agents/skills/mobius-onnx-export-gotchas/SKILL.md @@ -0,0 +1,200 @@ +--- +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 as of df203cc), and 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. +--- + +# 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`. + +## 3. FIXED: fp16 GQA export previously left packed-QKV weights as FLOAT32 → model wouldn't load +**Status: fixed as of commit `df203cc`.** 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 +`df203cc` 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 as of commit `cf6c5c4`.** 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). diff --git a/.github/workflows/main.yml b/.github/workflows/main.yml index e1b7e259..70e35c7a 100644 --- a/.github/workflows/main.yml +++ b/.github/workflows/main.yml @@ -292,6 +292,10 @@ jobs: -k "smollm-135m or albert-base or t5-small or dinov2-small or wav2vec2-base or whisper-tiny or test_gemma3_multimodal or test_qwen35 or test_qwen3_next or test_deepseek or test_sam_vit or test_ocr2" \ --cov=src --cov-report=xml --cov-branch --junitxml junit.xml timeout-minutes: 15 + - name: Static-cache decode runtime tests + run: | + pytest tests/static_cache_decode_test.py -v --tb=short + timeout-minutes: 10 - name: Upload coverage to Codecov if: always() uses: codecov/codecov-action@v6 diff --git a/CHANGELOG.md b/CHANGELOG.md index 83edfb24..7fefb3d2 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -7,6 +7,79 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ## [Unreleased] +### Static-cache Attention Phase-split Export + +#### Added + +- Causal-LM static-cache exports now emit ONNX `Attention` with `is_causal=0` + wrapped in a per-layer phase-split `If(Greater(seq_len, 1))` subgraph: prefill + takes an explicit-causal-mask branch while single-token decode takes a maskless + branch (which lets onnxruntime route decode to the Flash kernel). This is a + user-visible exported-graph change — it adds a fixed +10 top-level nodes per + static-cache export (`llama` 58→68, `qwen2` 58→68, `phi3` 56→66), which the + benchmark regression comparator recognizes as an intended structural change + (see `tests/benchmark_compare.py`). Lands alongside the GQA present-KV shape + fix below as part of the same export-correctness PR. + +--- + +### Developer Tooling + +#### Internal + +- The Architecture-Diff CI tool (`src/mobius/_graph_diff.py`) now recurses into + `If` / `Loop` / `Scan` `GRAPH`-typed attributes instead of collapsing them to a + bare type marker, so changes *inside* control-flow subgraphs are visible in the + diff. A dedicated `subgraph_structure_change` (MODERATE) severity is emitted for + structural in-branch deltas (node added/removed/rewired, branch added/removed), + while pure inner-attribute tweaks stay `changed_attrs` (MINOR). This PR introduces + the first control-flow/subgraph *structure* into the exported graphs — the + static-cache per-layer phase-split `If(Greater(seq_len, 1))`. The previous + top-level diff already reports those newly-added `If` nodes; what it could not do + is see *inside* a branch. Once these `If` nodes exist, any future change within + their branches would be invisible to the old diff (which collapsed each subgraph + to a type marker), so the recursion is added now to keep such in-branch changes + visible going forward. Developer-tooling only; no exported-graph or runtime impact. + +--- + +### 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). + +--- + +### 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. + +--- + ### WebGPU Shape Op Support #### Changed diff --git a/pyproject.toml b/pyproject.toml index fdcf48d2..0ccb5f9b 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -15,7 +15,7 @@ license = "MIT" dependencies = [ "huggingface_hub", "numpy>=1.24.0", - "onnx_ir>=0.1.0", + "onnx_ir>=0.1.2", "onnx-shape-inference>=0.1.9", "onnxscript>=0.7.0", "safetensors", diff --git a/src/mobius/_graph_diff.py b/src/mobius/_graph_diff.py index f85ea1e3..3cd78ae5 100644 --- a/src/mobius/_graph_diff.py +++ b/src/mobius/_graph_diff.py @@ -54,11 +54,23 @@ def _dtype_str(value: ir.Value) -> str: return "UNKNOWN" +# Sentinel keys marking a recursively-canonicalised subgraph payload inside an +# attribute's comparable value. diff_graphs uses these to render subgraph +# deltas readably instead of dumping a raw nested canonical dict. +_SUBGRAPH_KEY = "__subgraph__" +_SUBGRAPHS_KEY = "__subgraphs__" + + def _attr_to_comparable(attr: ir.Attr) -> Any: """Convert an attribute to a JSON-serialisable comparable value. - Graph and tensor attributes are reduced to their type string so - that canonicalisation stays lightweight. + Most attributes reduce to their scalar / list value. GRAPH-typed + attributes (``If``'s ``then_branch`` / ``else_branch``, ``Loop`` / + ``Scan`` bodies) are *recursively canonicalised* so subgraph structure + participates in the diff — without this, per-layer phase-split ``If`` + subgraphs are invisible to the architecture diff. Remaining opaque + types (TENSOR, SPARSE_TENSOR, TYPE_PROTO, …) are recorded as their type + string to keep canonicalisation lightweight. """ simple_types = { ir.AttributeType.FLOAT, @@ -74,7 +86,14 @@ def _attr_to_comparable(attr: ir.Attr) -> Any: if isinstance(v, tuple): return list(v) return v - # For TENSOR, GRAPH etc. just record the type + # Recurse into subgraphs so their node structure is compared, not collapsed. + # Subgraphs reuse canonicalize_graph, so inner node/value names are ignored + # the same way top-level ones are (see canonicalize_graph's name-independence). + if attr.type == ir.AttributeType.GRAPH: + return {_SUBGRAPH_KEY: canonicalize_graph(attr.value)} + if attr.type == ir.AttributeType.GRAPHS: + return {_SUBGRAPHS_KEY: [canonicalize_graph(g) for g in attr.value]} + # For TENSOR, SPARSE_TENSOR, TYPE_PROTO, … just record the type. return f"<{attr.type.name}>" @@ -203,13 +222,92 @@ def _describe_port_diff(base_port: dict, head_port: dict) -> str: return "; ".join(parts) or "changed" +def _is_subgraph_payload(value: Any) -> bool: + """True if *value* is a recursively-canonicalised subgraph payload.""" + return isinstance(value, dict) and (_SUBGRAPH_KEY in value or _SUBGRAPHS_KEY in value) + + +def _subgraph_list(value: Any) -> list[dict]: + """Extract the list of subgraph canonical forms from an attr payload.""" + if isinstance(value, dict): + if _SUBGRAPH_KEY in value: + return [value[_SUBGRAPH_KEY]] + if _SUBGRAPHS_KEY in value: + return list(value[_SUBGRAPHS_KEY]) + return [] + + +# Nested sub-change types that make a subgraph delta *structurally* significant +# (a node/branch was added/removed, rewired, or the subgraph interface moved), +# as opposed to a mere inner-attribute tweak. Any of these promotes the +# containing attribute to a subgraph_structure_change → MODERATE. Note this is +# uniformly MODERATE: unlike a *top-level* interface_change (MAJOR, an external +# model-contract break), a subgraph's interface is internal control-flow plumbing, +# so it stays MODERATE here — a deliberate asymmetry. +# ``subgraph_structure_change`` is included so structural significance +# *propagates* upward through nested subgraphs (e.g. an If inside an If). +_STRUCTURAL_SUB_TYPES = frozenset( + { + "added_node", + "removed_node", + "changed_connectivity", + "interface_change", + "subgraph_structure_change", + } +) + + +def _describe_subgraph_attr_change(key: str, base_val: Any, head_val: Any) -> tuple[str, bool]: + """Describe a GRAPH-typed attribute change, recursing into the subgraph(s). + + Returns ``(detail, structural)`` where *detail* is a readable summary + that surfaces the nested diff (e.g. ``"then_branch: node[0] Concat: + axis: 0 → 1"`` or ``"then_branch: + Mul; - Add"``) and *structural* is + True when the nested delta adds/removes a node or branch, rewires + connectivity, or changes the subgraph interface — i.e. a change that + should outrank a pure inner-attribute tweak. + """ + base_subs = _subgraph_list(base_val) + head_subs = _subgraph_list(head_val) + count = max(len(base_subs), len(head_subs)) + structural = False + parts: list[str] = [] + for idx in range(count): + bs = base_subs[idx] if idx < len(base_subs) else None + hs = head_subs[idx] if idx < len(head_subs) else None + # The attribute *key* already names a single subgraph (then_branch, + # body, …); only disambiguate by index when there are several (GRAPHS). + label = "" if count == 1 else f"subgraph[{idx}]" + if bs is None or hs is None: + # A whole branch/body was added or removed. + structural = True + verb = "added" if bs is None else "removed" + parts.append(f"{label} {verb}".strip()) + continue + sub_changes = diff_graphs(bs, hs) + if not sub_changes: + continue + if {c["type"] for c in sub_changes} & _STRUCTURAL_SUB_TYPES: + structural = True + inner = "; ".join(c["details"] for c in sub_changes) + parts.append(f"{label}: {inner}" if label else inner) + detail = f"{key}: " + ("; ".join(parts) if parts else "subgraph changed") + return detail, structural + + def diff_graphs(base: dict, head: dict) -> list[dict[str, Any]]: """Compare two canonical graph representations. Returns a list of change dicts. Each dict has a ``"type"`` key with one of: ``"added_node"``, ``"removed_node"``, ``"changed_attrs"``, + ``"subgraph_structure_change"``, ``"changed_connectivity"``, ``"interface_change"``, ``"initializer_change"``. A ``"details"`` key carries human-readable information about the change. + + ``subgraph_structure_change`` is emitted for a GRAPH-typed attribute + (``If`` branches, ``Loop`` / ``Scan`` bodies) whose nested graph gains + or loses a node/branch, is rewired, or changes interface; a subgraph + delta that only tweaks an inner attribute stays ``changed_attrs``. """ changes: list[dict[str, Any]] = [] @@ -304,19 +402,32 @@ def diff_graphs(base: dict, head: dict) -> list[dict[str, Any]]: if bn["attributes"] != hn["attributes"]: ba = bn["attributes"] ha = hn["attributes"] - attr_details: list[str] = [] + plain_details: list[str] = [] all_keys = sorted(set(ba) | set(ha)) for k in all_keys: bv = ba.get(k) hv = ha.get(k) - if bv != hv: - attr_details.append(f"{k}: {bv!r} → {hv!r}") - changes.append( - { - "type": "changed_attrs", - "details": (f"node[{i}] {bn['op_type']}: " + ", ".join(attr_details)), - } - ) + if bv == hv: + continue + if _is_subgraph_payload(bv) or _is_subgraph_payload(hv): + detail, structural = _describe_subgraph_attr_change(k, bv, hv) + changes.append( + { + "type": ( + "subgraph_structure_change" if structural else "changed_attrs" + ), + "details": f"node[{i}] {bn['op_type']}: {detail}", + } + ) + else: + plain_details.append(f"{k}: {bv!r} → {hv!r}") + if plain_details: + changes.append( + { + "type": "changed_attrs", + "details": (f"node[{i}] {bn['op_type']}: " + ", ".join(plain_details)), + } + ) if bn["input_ids"] != hn["input_ids"]: changes.append( { @@ -343,7 +454,12 @@ def _change_status(change_list: list[dict[str, Any]]) -> str: types = {c["type"] for c in change_list} if types & {"interface_change"}: return _STATUS_MAJOR - if types & {"added_node", "removed_node", "changed_connectivity"}: + if types & { + "added_node", + "removed_node", + "changed_connectivity", + "subgraph_structure_change", + }: return _STATUS_MODERATE if types & {"changed_attrs", "initializer_change"}: return _STATUS_MINOR @@ -459,6 +575,7 @@ def _sha_link(sha: str) -> str: removed = [c for c in change_list if c["type"] == "removed_node"] attrs = [c for c in change_list if c["type"] == "changed_attrs"] connectivity = [c for c in change_list if c["type"] == "changed_connectivity"] + subgraph = [c for c in change_list if c["type"] == "subgraph_structure_change"] iface = [c for c in change_list if c["type"] == "interface_change"] inits = [c for c in change_list if c["type"] == "initializer_change"] @@ -474,6 +591,12 @@ def _sha_link(sha: str) -> str: lines.append(f"- `{c['details']}`") lines.append("") + if subgraph: + lines.append("**Subgraph structure changes:**") + for c in subgraph: + lines.append(f"- `{c['details']}`") + lines.append("") + if attrs: lines.append("**Modified attributes:**") for c in attrs: diff --git a/src/mobius/_graph_diff_test.py b/src/mobius/_graph_diff_test.py index 2f87842e..bcb94ffd 100644 --- a/src/mobius/_graph_diff_test.py +++ b/src/mobius/_graph_diff_test.py @@ -12,6 +12,8 @@ import onnx_ir as ir from mobius._graph_diff import ( + _SUBGRAPH_KEY, + _SUBGRAPHS_KEY, canonicalize_graph, diff_graphs, render_markdown, @@ -611,3 +613,376 @@ def test_structural_change_detected(self) -> None: head = _add_relu_graph() changes = diff_graphs(canonicalize_graph(base), canonicalize_graph(head)) assert len(changes) > 0 + + +# ------------------------------------------------------------------ +# Subgraph recursion — helpers (graphs containing an If with subgraphs) +# ------------------------------------------------------------------ + + +def _if_graph( + *, + then_op: str = "Add", + else_op: str = "Sub", + then_axis: int | None = None, + then_domain: str = "", + then_node_name: str = "then_op", + else_node_name: str = "else_op", +) -> ir.Graph: + """Return a graph: y = If(cond) {then: then_op(x,x)} {else: else_op(x,x)}. + + ``then_axis`` (when given) attaches an ``axis`` attribute to the + then-branch node so inner-attribute deltas can be exercised. + ``then_domain`` sets the then-branch node's op domain (diff_graphs does + not compare domain, so this exercises the 'subgraph changed' fallback). + ``*_node_name`` allow renaming inner nodes to prove names are ignored. + """ + x = ir.val("x", type=ir.TensorType(ir.DataType.FLOAT), shape=ir.Shape([1, 4])) + cond = ir.val("cond", type=ir.TensorType(ir.DataType.BOOL), shape=ir.Shape([])) + + then_attrs = [ir.AttrInt64("axis", then_axis)] if then_axis is not None else [] + then_node = ir.Node(then_domain, then_op, [x, x], then_attrs, name=then_node_name) + then_out = then_node.outputs[0] + then_out.name = "then_out" + then_g = ir.Graph([], [then_out], nodes=[then_node], name="then_branch") + + else_node = ir.Node("", else_op, [x, x], name=else_node_name) + else_out = else_node.outputs[0] + else_out.name = "else_out" + else_g = ir.Graph([], [else_out], nodes=[else_node], name="else_branch") + + if_node = ir.Node( + "", + "If", + [cond], + [ir.AttrGraph("then_branch", then_g), ir.AttrGraph("else_branch", else_g)], + name="if_node", + ) + if_out = if_node.outputs[0] + if_out.name = "y" + return ir.Graph([x, cond], [if_out], nodes=[if_node]) + + +# ------------------------------------------------------------------ +# Subgraph recursion — tests (architect D12: descend into subgraphs) +# ------------------------------------------------------------------ + + +def _nested_if_graph( + *, inner_then_op: str = "Add", inner_then_axis: int | None = None +) -> ir.Graph: + """Return a graph whose If then-branch itself contains an If (depth 2). + + ``inner_then_axis`` attaches an ``axis`` attribute to the innermost + then-branch node, to exercise a *pure inner-attribute* delta nested + two levels deep (which must stay MINOR, not promote to structural). + """ + x = ir.val("x", type=ir.TensorType(ir.DataType.FLOAT), shape=ir.Shape([1, 4])) + cond = ir.val("cond", type=ir.TensorType(ir.DataType.BOOL), shape=ir.Shape([])) + + # Innermost two branches. + inner_then_attrs = ( + [ir.AttrInt64("axis", inner_then_axis)] if inner_then_axis is not None else [] + ) + inner_then = ir.Node("", inner_then_op, [x, x], inner_then_attrs, name="inner_then") + it_out = inner_then.outputs[0] + it_out.name = "it_out" + inner_then_g = ir.Graph([], [it_out], nodes=[inner_then], name="inner_then_branch") + + inner_else = ir.Node("", "Sub", [x, x], name="inner_else") + ie_out = inner_else.outputs[0] + ie_out.name = "ie_out" + inner_else_g = ir.Graph([], [ie_out], nodes=[inner_else], name="inner_else_branch") + + inner_if = ir.Node( + "", + "If", + [cond], + [ + ir.AttrGraph("then_branch", inner_then_g), + ir.AttrGraph("else_branch", inner_else_g), + ], + name="inner_if", + ) + inner_if_out = inner_if.outputs[0] + inner_if_out.name = "inner_if_out" + # Outer then-branch wraps the inner If; outer else-branch is a plain op. + outer_then_g = ir.Graph([], [inner_if_out], nodes=[inner_if], name="outer_then") + + outer_else = ir.Node("", "Mul", [x, x], name="outer_else") + oe_out = outer_else.outputs[0] + oe_out.name = "oe_out" + outer_else_g = ir.Graph([], [oe_out], nodes=[outer_else], name="outer_else") + + outer_if = ir.Node( + "", + "If", + [cond], + [ + ir.AttrGraph("then_branch", outer_then_g), + ir.AttrGraph("else_branch", outer_else_g), + ], + name="outer_if", + ) + out = outer_if.outputs[0] + out.name = "y" + return ir.Graph([x, cond], [out], nodes=[outer_if]) + + +def _graphs_attr_graph(*, body_ops: list[str]) -> ir.Graph: + """Return a graph with a node carrying a GRAPHS-plural attribute. + + GRAPHS is rare for standard ops (``If`` / ``Loop`` / ``Scan`` bodies are + GRAPH singular), so this uses a synthetic op purely to exercise the + plural recursion branch in ``_attr_to_comparable``. + """ + x = ir.val("x", type=ir.TensorType(ir.DataType.FLOAT), shape=ir.Shape([1, 4])) + + bodies = [] + for idx, op in enumerate(body_ops): + n = ir.Node("", op, [x, x], name=f"body_{idx}") + o = n.outputs[0] + o.name = f"body_out_{idx}" + bodies.append(ir.Graph([], [o], nodes=[n], name=f"body_g_{idx}")) + + multi = ir.Node("", "CustomMulti", [x], [ir.AttrGraphs("bodies", bodies)], name="cm") + out = multi.outputs[0] + out.name = "y" + return ir.Graph([x], [out], nodes=[multi]) + + +def _if_graph_branch_output_dtype(*, then_dtype: ir.DataType) -> ir.Graph: + """If-graph whose then-branch *output dtype* varies (subgraph interface).""" + x = ir.val("x", type=ir.TensorType(ir.DataType.FLOAT), shape=ir.Shape([1, 4])) + cond = ir.val("cond", type=ir.TensorType(ir.DataType.BOOL), shape=ir.Shape([])) + + cast = ir.Node("", "Cast", [x], [ir.AttrInt64("to", int(then_dtype))], name="cast") + then_out = cast.outputs[0] + then_out.name = "then_out" + then_out.type = ir.TensorType(then_dtype) + then_out.shape = ir.Shape([1, 4]) + then_g = ir.Graph([], [then_out], nodes=[cast], name="then_branch") + + else_node = ir.Node("", "Identity", [x], name="else_node") + else_out = else_node.outputs[0] + else_out.name = "else_out" + else_out.type = ir.TensorType(ir.DataType.FLOAT) + else_out.shape = ir.Shape([1, 4]) + else_g = ir.Graph([], [else_out], nodes=[else_node], name="else_branch") + + if_node = ir.Node( + "", + "If", + [cond], + [ir.AttrGraph("then_branch", then_g), ir.AttrGraph("else_branch", else_g)], + name="if_node", + ) + if_out = if_node.outputs[0] + if_out.name = "y" + return ir.Graph([x, cond], [if_out], nodes=[if_node]) + + +class TestSubgraphRecursion: + """canonicalize_graph / diff_graphs descend into GRAPH-typed attrs.""" + + def test_canonicalize_descends_into_subgraphs(self) -> None: + """GRAPH attrs are recursively canonicalised, not collapsed to a type.""" + canon = canonicalize_graph(_if_graph(then_op="Add", else_op="Sub")) + if_attrs = canon["nodes"][0]["attributes"] + # Not collapsed to the old "" placeholder. + assert if_attrs["then_branch"] != "" + assert if_attrs["else_branch"] != "" + # Nested canonical form carries the subgraph's op sequence. + then_sub = if_attrs["then_branch"][_SUBGRAPH_KEY] + else_sub = if_attrs["else_branch"][_SUBGRAPH_KEY] + assert then_sub["op_sequence"] == ["Add"] + assert else_sub["op_sequence"] == ["Sub"] + + def test_identical_subgraphs_no_diff(self) -> None: + """Two structurally identical If-graphs produce no diff.""" + base = canonicalize_graph(_if_graph(then_op="Add", else_op="Sub")) + head = canonicalize_graph(_if_graph(then_op="Add", else_op="Sub")) + assert diff_graphs(base, head) == [] + + def test_subgraph_node_names_ignored(self) -> None: + """Renaming inner subgraph nodes does not produce a spurious diff.""" + base = canonicalize_graph(_if_graph(then_node_name="a", else_node_name="b")) + head = canonicalize_graph(_if_graph(then_node_name="x", else_node_name="y")) + assert diff_graphs(base, head) == [] + + def test_then_branch_op_delta_is_structural(self) -> None: + """A differing then-branch op is a MODERATE subgraph_structure_change.""" + base = canonicalize_graph(_if_graph(then_op="Add", else_op="Sub")) + head = canonicalize_graph(_if_graph(then_op="Mul", else_op="Sub")) + changes = diff_graphs(base, head) + structural = [c for c in changes if c["type"] == "subgraph_structure_change"] + assert structural, "op delta inside a branch must be structural" + assert any("then_branch" in c["details"] for c in structural) + # The op-level subgraph delta is surfaced (added/removed inner node). + assert any("Add" in c["details"] and "Mul" in c["details"] for c in structural) + + def test_else_branch_op_delta_is_structural(self) -> None: + """A differing else-branch op is detected independently of then-branch.""" + base = canonicalize_graph(_if_graph(then_op="Add", else_op="Sub")) + head = canonicalize_graph(_if_graph(then_op="Add", else_op="Div")) + changes = diff_graphs(base, head) + structural = [c for c in changes if c["type"] == "subgraph_structure_change"] + assert any("else_branch" in c["details"] for c in structural) + + def test_structural_subgraph_change_is_moderate(self) -> None: + """render_markdown rates a subgraph structure change as MODERATE (🟡).""" + base = canonicalize_graph(_if_graph(then_op="Add", else_op="Sub")) + head = canonicalize_graph(_if_graph(then_op="Mul", else_op="Sub")) + changes = diff_graphs(base, head) + diffs = { + "model": { + "sub": { + "changes": changes, + "_base_ops": base["op_sequence"], + "_head_ops": head["op_sequence"], + "_base_node_count": len(base["nodes"]), + "_head_node_count": len(head["nodes"]), + } + } + } + md = render_markdown(diffs) + assert "🟡" in md + assert "Subgraph structure changes" in md + + def test_inner_attribute_delta_is_minor_and_surfaced(self) -> None: + """A pure inner-attribute tweak (Concat axis) stays changed_attrs/MINOR. + + The actual nested detail (``axis: 0 → 1``) must be surfaced, not an + opaque ``"N sub-change(s)"`` summary. + """ + base = canonicalize_graph(_if_graph(then_op="Concat", else_op="Sub", then_axis=0)) + head = canonicalize_graph(_if_graph(then_op="Concat", else_op="Sub", then_axis=1)) + changes = diff_graphs(base, head) + # Not promoted to structural — it's an inner-attr-only change. + assert not any(c["type"] == "subgraph_structure_change" for c in changes) + attr_changes = [c for c in changes if c["type"] == "changed_attrs"] + assert attr_changes + detail = " ".join(c["details"] for c in attr_changes) + assert "then_branch" in detail + # The real inner delta is surfaced. + assert "axis" in detail and "0" in detail and "1" in detail + + def test_nested_subgraph_delta_detected(self) -> None: + """A delta in an If-inside-an-If (depth 2) propagates to the top diff.""" + base = canonicalize_graph(_nested_if_graph(inner_then_op="Add")) + head = canonicalize_graph(_nested_if_graph(inner_then_op="Mul")) + changes = diff_graphs(base, head) + structural = [c for c in changes if c["type"] == "subgraph_structure_change"] + assert structural, "innermost branch op delta must surface at the top" + assert any("then_branch" in c["details"] for c in structural) + + def test_nested_pure_attr_delta_stays_minor(self) -> None: + """A pure inner-attr tweak nested two levels deep stays changed_attrs/MINOR. + + (Code-review NIT: lock the severity boundary — nesting must not + spuriously promote a non-structural change to structural.) + """ + base = canonicalize_graph(_nested_if_graph(inner_then_op="Concat", inner_then_axis=0)) + head = canonicalize_graph(_nested_if_graph(inner_then_op="Concat", inner_then_axis=1)) + changes = diff_graphs(base, head) + assert not any(c["type"] == "subgraph_structure_change" for c in changes) + attr_changes = [c for c in changes if c["type"] == "changed_attrs"] + assert attr_changes + # The innermost axis delta is still surfaced through both nesting levels. + detail = " ".join(c["details"] for c in attr_changes) + assert "then_branch" in detail and "axis" in detail + + def test_subgraph_changed_fallback_path(self) -> None: + """A subgraph delta invisible to diff_graphs hits the readable fallback. + + (Code-review NIT: lock the ``"subgraph changed"`` fallback.) The + then-branch node's *domain* differs; canonicalize records domain so + the attribute payloads differ, but diff_graphs does not compare + domain, so the nested diff is empty → the summary falls back to + ``"subgraph changed"`` and is classified MINOR (changed_attrs). + """ + base = canonicalize_graph(_if_graph(then_op="Add", then_domain="")) + head = canonicalize_graph(_if_graph(then_op="Add", then_domain="custom.domain")) + changes = diff_graphs(base, head) + assert not any(c["type"] == "subgraph_structure_change" for c in changes) + attr_changes = [c for c in changes if c["type"] == "changed_attrs"] + assert any("subgraph changed" in c["details"] for c in attr_changes) + + def test_graphs_plural_attr_recursed_and_diffed(self) -> None: + """GRAPHS-plural attrs are canonicalised and per-body deltas detected.""" + base = canonicalize_graph(_graphs_attr_graph(body_ops=["Add", "Add"])) + head = canonicalize_graph(_graphs_attr_graph(body_ops=["Add", "Mul"])) + # Plural payload is recursed, not collapsed. + bodies_attr = base["nodes"][0]["attributes"]["bodies"] + assert bodies_attr != "" + changes = diff_graphs(base, head) + structural = [c for c in changes if c["type"] == "subgraph_structure_change"] + assert structural + # Plural payload genuinely drove the AttributeType.GRAPHS branch. + assert _SUBGRAPHS_KEY in bodies_attr + assert len(bodies_attr[_SUBGRAPHS_KEY]) == 2 + # The differing body is identified by index. + assert any("subgraph[1]" in c["details"] for c in structural) + + def test_graphs_plural_branch_count_change_is_structural(self) -> None: + """Adding/removing a body in a GRAPHS-plural attr is structural.""" + base = canonicalize_graph(_graphs_attr_graph(body_ops=["Add"])) + head = canonicalize_graph(_graphs_attr_graph(body_ops=["Add", "Mul"])) + changes = diff_graphs(base, head) + structural = [c for c in changes if c["type"] == "subgraph_structure_change"] + assert structural + assert any("added" in c["details"] or "removed" in c["details"] for c in structural) + + def test_subgraph_interface_change_is_structural(self) -> None: + """A subgraph *interface* (output dtype) change is structural.""" + base = canonicalize_graph(_if_graph_branch_output_dtype(then_dtype=ir.DataType.FLOAT)) + head = canonicalize_graph( + _if_graph_branch_output_dtype(then_dtype=ir.DataType.FLOAT16) + ) + changes = diff_graphs(base, head) + structural = [c for c in changes if c["type"] == "subgraph_structure_change"] + assert structural, "subgraph interface change must be structural" + assert any("then_branch" in c["details"] for c in structural) + + def test_subgraph_delta_does_not_perturb_top_level_ops(self) -> None: + """The top-level op sequence is unchanged when only a subgraph differs.""" + base = canonicalize_graph(_if_graph(then_op="Add", else_op="Sub")) + head = canonicalize_graph(_if_graph(then_op="Mul", else_op="Sub")) + # Both still have exactly one top-level If node. + assert base["op_sequence"] == ["If"] == head["op_sequence"] + changes = diff_graphs(base, head) + # No spurious added/removed at the top level. + assert not any(c["type"] in {"added_node", "removed_node"} for c in changes) + + def test_top_level_op_swap_not_double_counted(self) -> None: + """An op_type swap at the same position is structural-only, never MINOR. + + (Reviewer 6def2895 (a): pick one classification — structural wins.) + """ + + def _single(op: str) -> ir.Graph: + x = ir.val("x", type=ir.TensorType(ir.DataType.FLOAT), shape=ir.Shape([1, 4])) + n = ir.Node("", op, [x, x], name="n") + o = n.outputs[0] + o.name = "y" + return ir.Graph([x], [o], nodes=[n]) + + changes = diff_graphs( + canonicalize_graph(_single("Add")), canonicalize_graph(_single("Mul")) + ) + types = {c["type"] for c in changes} + assert types == {"added_node", "removed_node"} + assert "changed_attrs" not in types + + def test_subgraph_op_swap_is_structural_only(self) -> None: + """MEA↔Flash-style op swap *inside* a branch is structural, not MINOR. + + A same-position op_type swap within a subgraph must classify as + subgraph_structure_change only — not additionally as changed_attrs. + """ + base = canonicalize_graph(_if_graph(then_op="Add", else_op="Sub")) + head = canonicalize_graph(_if_graph(then_op="Mul", else_op="Sub")) + changes = diff_graphs(base, head) + # Exactly one change for the swapped branch, classified structural. + assert [c["type"] for c in changes] == ["subgraph_structure_change"] + assert not any(c["type"] == "changed_attrs" for c in changes) 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..4073991b 100644 --- a/src/mobius/components/_attention.py +++ b/src/mobius/components/_attention.py @@ -7,12 +7,123 @@ from typing import NamedTuple import onnx_ir as ir -from onnxscript import OpBuilder, nn +from onnxscript import GraphBuilder, OpBuilder, nn from mobius._configs import ArchitectureConfig -from mobius.components._common import Linear +from mobius._constants import OPSET_VERSION +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 +from mobius.components._scan_utils import rename_subgraph_values + + +def _attend_over_static_cache( + op: OpBuilder, + query: ir.Value, + key_cache: ir.Value, + value_cache: ir.Value, + write_indices: ir.Value, + nonpad_kv_seqlen: ir.Value, + num_attention_heads: int, + num_key_value_heads: int, + scale: float | None, + softcap: float | None, +) -> ir.Value: + """Attend over the static KV cache, phase-split for decode kernel fidelity. + + The opset-24 ONNX ``Attention`` CUDA kernel rejects ``is_causal=1`` + together with ``nonpad_kv_seqlen`` when ``S_q != total_kv`` with no + ``past_key`` (the ``causal_cross_no_past`` guard in ``attention.cc``). + With a pre-allocated ``[B, max_seq, ...]`` cache that condition holds in + both prefill and decode, so ``is_causal=0`` must be used. + + Both phases run with ``is_causal=0`` and keep ``nonpad_kv_seqlen`` (which + selects the external-cache / TensorScatter kernel path). The phases + differ in whether an explicit ``attn_mask`` is supplied, because the + *presence* of ``attn_mask`` — regardless of its contents — disables Flash + Attention in ORT (``attn_mask != nullptr`` routes to the slower memory- + efficient or unfused path; see the kernel-selection cascade in + ``attention.cc``). The phase split exists precisely to pay that + Flash→MEA latency cost only where it is unavoidable (multi-token prefill) + and never on the per-token decode hot path: + + * **Multi-token step** (``S_q > 1``: prefill or speculative/chunked decode): + needs intra-query causality, so it passes an explicit causal mask built + from ``write_indices`` (:func:`create_static_cache_causal_mask`) and + therefore runs on the memory-efficient path. This is unavoidable — the + static buffer makes ``K_seq == total``, so Flash prefill is blocked by + the same guard regardless — and is the cheap, amortized path anyway. + * **Single-token decode** (``S_q == 1``): a lone query needs no + intra-query causal mask; ``nonpad_kv_seqlen`` alone bounds attention to + the valid prefix ``0..write_indices[b]``. Omitting ``attn_mask`` keeps + this hot path on Flash / XQA — the kernel the GQA variant also uses, + so the comparison stays apples-to-apples. + + The two phases are emitted as the branches of an ``If`` keyed on + ``Shape(query)[1] > 1`` so a single exported graph serves both, while + decode structurally omits the mask input. + + Returns: + The attention output for the active phase, shape ``[B, S_q, hidden]``. + """ + seq_len = op.Squeeze(op.Shape(query, start=1, end=2), op.Constant(value_ints=[0])) + is_multi_token_step = op.Greater(seq_len, op.Constant(value_int=1)) + + def _build_attention_branch(name: str, use_causal_mask: bool) -> ir.Graph: + branch = ir.Graph([], [], nodes=[], name=name, opset_imports={"": OPSET_VERSION}) + branch_op = GraphBuilder(branch).op + attn_mask = ( + create_static_cache_causal_mask(branch_op, query, key_cache, write_indices) + if use_causal_mask + else None + ) + attn_output, _, _ = branch_op.Attention( + query, + key_cache, + value_cache, + attn_mask, + None, # no past_key (full cache is already provided) + None, # no past_value + nonpad_kv_seqlen, + q_num_heads=num_attention_heads, + kv_num_heads=num_key_value_heads, + scale=scale, + softcap=softcap, + is_causal=0, + _outputs=3, + ) + # Prefix internal node/value names so the two branches stay in SSA + # form when merged under the parent graph, then pin the branch + # output name (the If wires branches by output position). + # + # Ordering is load-bearing — do NOT reorder these three lines: + # 1. rename first: rename_subgraph_values renames node OUTPUT value + # names but skips graph.inputs/outputs. Running it before the + # output is appended (branch.outputs is still empty) means + # attn_output is renamed here as an internal node output, not + # protected as a graph output. + # 2. pin the name next: this deterministic name must be set AFTER + # the rename so the rename does not clobber it, and BEFORE the + # append so the If sees a stable output name. + # 3. append last: registers the now-stable value as the branch + # output the parent If wires by position. + # Outer-scope captures (query/key_cache/value_cache/nonpad_kv_seqlen/ + # write_indices) are only ever node INPUTS inside the branch, never + # outputs, so the rename never touches them and implicit-input capture + # stays intact. + rename_subgraph_values(branch, f"{name}_") + attn_output.name = f"{name}_attn_output" + branch.outputs.append(attn_output) + return branch + + prefill_branch = _build_attention_branch("static_cache_prefill", use_causal_mask=True) + decode_branch = _build_attention_branch("static_cache_decode", use_causal_mask=False) + return op.If( + is_multi_token_step, + then_branch=prefill_branch, + else_branch=decode_branch, + _outputs=1, + ) class GQAContext(NamedTuple): @@ -99,14 +210,23 @@ 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``. Causality is phase-split (see + :func:`_attend_over_static_cache`): multi-token steps (``S_q > 1``) + use an explicit causal mask derived from ``write_indices`` (memory- + efficient path), while single-token decode (``S_q == 1``) omits the + mask to stay on Flash/XQA. ``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`` + because the external-cache kernel does not accept ``is_causal=1`` + alongside ``nonpad_kv_seqlen``; causality is supplied per phase via + an explicit mask (prefill) or ``nonpad_kv_seqlen`` alone (decode). Note: ``nonpad_kv_seqlen`` (input #6) is only valid in static cache mode @@ -138,37 +258,31 @@ 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. Both phases use is_causal=0 (the + # opset-24 Attention kernel rejects is_causal=1 + nonpad_kv_seqlen + # for a pre-allocated cache) and keep nonpad_kv_seqlen to select the + # external-cache kernel path. Causality is enforced per-phase: the + # multi-token branch supplies an explicit causal mask (MEA path), + # while single-token decode omits the mask to stay on Flash/XQA. + # See _attend_over_static_cache for the full rationale. # - # 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( + attn_output = _attend_over_static_cache( + op, query, updated_k, updated_v, - None, # no attn_mask — is_causal handles masking - None, # no past_key (full cache is already provided) - None, # no past_value + static_cache.write_indices, static_cache.nonpad_kv_seqlen, - q_num_heads=num_attention_heads, - kv_num_heads=num_key_value_heads, + num_attention_heads=num_attention_heads, + num_key_value_heads=num_key_value_heads, scale=scale, softcap=softcap, - is_causal=1, - _outputs=3, ) return attn_output, updated_k, updated_v diff --git a/src/mobius/components/_common.py b/src/mobius/components/_common.py index 2d4836f5..ef16c5c5 100644 --- a/src/mobius/components/_common.py +++ b/src/mobius/components/_common.py @@ -321,3 +321,81 @@ 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. + + Because padding slots ``j >= nonpad_kv_seqlen[b]`` are always greater + than ``write_indices[b] + t``, they are masked out too, so the mask is + consistent with the ``nonpad_kv_seqlen`` bounds (which are still passed + to the Attention op to select the external-cache kernel path). + + 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 tensor ``[batch, S_q, hidden]``; only dims 0/1 are + 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 (causal + padding). + 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..3e001e7d 100644 --- a/src/mobius/tasks/_cache_utils.py +++ b/src/mobius/tasks/_cache_utils.py @@ -11,6 +11,7 @@ from __future__ import annotations +import logging from typing import NamedTuple import onnx_ir as ir @@ -18,6 +19,8 @@ from mobius._configs import BaseModelConfig +logger = logging.getLogger(__name__) + _FUNCTIONS_DOMAIN = "com.microsoft" # Cache state pair: (key, value) or (conv_state, ssm_state) for stateful @@ -114,13 +117,71 @@ 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 almost always a wiring slip (a caller wired some dims but dropped + others), so it is treated conservatively -- the stamp is skipped, the + shapes fall back to the known-wrong inference path, and a warning is + logged naming the missing parameters so the slip is loud rather than silent. """ + 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] + logger.warning( + "_register_kv_cache_outputs received a partial set of present-shape " + "parameters (provided %s, missing %s); these are all-or-nothing, so " + "the explicit present.* stamp is SKIPPED and shapes fall back to " + "inference (which mis-derives head_dim for GroupQueryAttention). " + "Pass all six parameters to stamp, or none to opt out.", + provided, + missing, + ) 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..64924ead --- /dev/null +++ b/src/mobius/tasks/_cache_utils_test.py @@ -0,0 +1,125 @@ +# 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 logging + +import onnx_ir as ir + +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, caplog): + """Without shape params the helper must not stamp (inference path).""" + _, builder = _make_graph() + pairs = [_present_pair("present.0", wrong_head_dim=32)] + + with caplog.at_level(logging.WARNING, logger="mobius.tasks._cache_utils"): + _register_kv_cache_outputs(builder, pairs) + + key, _ = pairs[0] + # Unchanged: still the (wrong) pre-existing inferred shape. + assert _dims(key) == ["batch", 32, "seq", 32] + # Opting out (zero params) is intentional and must stay silent. + assert caplog.text == "" + + def test_partial_params_do_not_stamp(self, caplog): + """A partial set falls back to inference AND warns about the slip.""" + _, builder = _make_graph() + pairs = [_present_pair("present.0", wrong_head_dim=32)] + + with caplog.at_level(logging.WARNING, logger="mobius.tasks._cache_utils"): + _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 + ) + + key, _ = pairs[0] + assert _dims(key) == ["batch", 32, "seq", 32] + assert "partial set of present-shape parameters" in caplog.text + # The warning must name the omitted parameters so the slip is diagnosable. + for missing in ("key_head_dim", "value_head_dim", "total_seq_len", "dtype"): + assert missing in caplog.text + + 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..972437c7 100644 --- a/src/mobius/tasks/_causal_lm.py +++ b/src/mobius/tasks/_causal_lm.py @@ -153,6 +153,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 +166,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 +187,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) diff --git a/src/mobius/tasks/_task_test.py b/src/mobius/tasks/_task_test.py index d50f4337..ef81e1c6 100644 --- a/src/mobius/tasks/_task_test.py +++ b/src/mobius/tasks/_task_test.py @@ -91,6 +91,38 @@ 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 outputs must match their past-input KV metadata. + + present.{i}.* 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_build_producer_info(self): config = make_config() module = CausalLMModel(config) diff --git a/tests/benchmark_compare.py b/tests/benchmark_compare.py index c9b9fcf9..5d8a2e1b 100644 --- a/tests/benchmark_compare.py +++ b/tests/benchmark_compare.py @@ -33,6 +33,30 @@ "num_nodes": (0.05, 0.10), } +# Intended structural changes that would otherwise trip a deterministic-metric +# blocker, keyed by model display key -> metric -> the exact (baseline, current) +# values to waive. ABSOLUTE counts (not a relative delta) are pinned on purpose: +# the waiver fires only for this exact base->current transition, so once the +# change merges into the base branch (baseline becomes the new value) the entry +# can NEVER match again and is truly self-cleaning — a later regression that +# happens to add the same number of nodes still blocks. Only the exact pinned +# transition is waived; any other base/current pair falls through to the normal +# threshold logic and can block, so this never masks an accidental regression. +# +# PR #328 (static-cache export): the is_causal=0 static-cache path wraps each +# decoder layer's attention in a phase-split ``If(Greater(seq_len, 1))`` subgraph +# (prefill takes the explicit-causal-mask branch, decode the maskless branch) +# plus the prefill mask-build nodes. This intentionally raises num_nodes by a +# fixed +10 top-level nodes per static-cache export. These are correct, not a +# regression. Safe to delete this whole table after PR #328 merges. (If an +# unrelated opset/ORT/exporter change shifts the baseline counts before merge, +# update the pinned values here — a mismatch fails closed with a RED blocker.) +EXPECTED_CHANGES: dict[str, dict[str, tuple[int, int]]] = { + "llama (static-cache)": {"num_nodes": (58, 68)}, + "qwen2 (static-cache)": {"num_nodes": (58, 68)}, + "phi3 (static-cache)": {"num_nodes": (56, 66)}, +} + _GITHUB_REPO_URL = "https://github.com/onnxruntime/mobius" @@ -58,7 +82,12 @@ def compare(current_path: str, baseline_path: str) -> tuple[str, bool]: continue delta_pct = (curr_val - base_val) / base_val warn_t, block_t = THRESHOLDS[metric] - if delta_pct > block_t: + expected = EXPECTED_CHANGES.get(model, {}).get(metric) + if expected is not None and (base_val, curr_val) == expected: + # Intended structural change (e.g. the static-cache phase-split + # If+mask). Waive only this exact pinned base->current transition. + status = "\U0001f7e6" # blue square: accepted intended change + elif delta_pct > block_t: status = "\U0001f534" # red circle has_blocker = True elif delta_pct > warn_t: @@ -105,7 +134,14 @@ def _sha_link(sha: str) -> str: elif any(r[5] == "\u26a0\ufe0f" for r in rows): md += "\n> Warning: minor regressions detected. Review flagged metrics.\n" else: - md += "\n> No performance regressions.\n" + md += "\n> No blocking regressions.\n" + + if any(r[5] == "\U0001f7e6" for r in rows): + md += ( + "\n> \U0001f7e6 = intended structural change accepted via " + "`EXPECTED_CHANGES` (exact pinned base→current values; see " + "`tests/benchmark_compare.py`).\n" + ) return md, has_blocker diff --git a/tests/benchmark_compare_test.py b/tests/benchmark_compare_test.py new file mode 100644 index 00000000..d0ff7b32 --- /dev/null +++ b/tests/benchmark_compare_test.py @@ -0,0 +1,157 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT License. + +"""Tests for the deterministic-metric regression comparator. + +Focuses on the ``EXPECTED_CHANGES`` allowlist used to accept intended +structural node-count changes (e.g. the PR #328 static-cache phase-split) +without masking accidental regressions. +""" + +from __future__ import annotations + +import json +from pathlib import Path + +import pytest + +import tests.benchmark_compare as bc + +# A blue square marks an accepted intended structural change. +_ACCEPTED = "\U0001f7e6" +# A red circle marks a blocking regression. +_BLOCKER = "\U0001f534" + + +def _write(path: Path, models: dict[str, dict[str, int]]) -> str: + path.write_text(json.dumps({"_metadata": {"commit": "deadbeef"}, "models": models})) + return str(path) + + +def _run(tmp_path: Path, base: dict, curr: dict) -> tuple[str, bool]: + baseline = _write(tmp_path / "baseline.json", base) + current = _write(tmp_path / "current.json", curr) + return bc.compare(current, baseline) + + +def test_exact_expected_static_cache_delta_is_waived(tmp_path: Path): + """The exact intended +10 static-cache node delta must not block.""" + md, has_blocker = _run( + tmp_path, + base={"llama (static-cache)": {"num_nodes": 58, "model_size_bytes": 1000}}, + curr={"llama (static-cache)": {"num_nodes": 68, "model_size_bytes": 1000}}, + ) + assert has_blocker is False + assert _ACCEPTED in md + assert _BLOCKER not in md + + +@pytest.mark.parametrize( + "model", + ["llama (static-cache)", "qwen2 (static-cache)", "phi3 (static-cache)"], +) +def test_all_allowlisted_models_waive_their_change(tmp_path: Path, model: str): + base_nodes, curr_nodes = bc.EXPECTED_CHANGES[model]["num_nodes"] + md, has_blocker = _run( + tmp_path, + base={model: {"num_nodes": base_nodes}}, + curr={model: {"num_nodes": curr_nodes}}, + ) + assert has_blocker is False + assert _ACCEPTED in md + + +def test_delta_larger_than_expected_still_blocks(tmp_path: Path): + """An extra unexpected node beyond the intended delta must still block.""" + md, has_blocker = _run( + tmp_path, + base={"llama (static-cache)": {"num_nodes": 58}}, + curr={"llama (static-cache)": {"num_nodes": 69}}, # +11, not the intended +10 + ) + assert has_blocker is True + assert _BLOCKER in md + + +def test_non_allowlisted_model_still_blocks(tmp_path: Path): + """A model outside the allowlist gets no waiver.""" + md, has_blocker = _run( + tmp_path, + base={"llama": {"num_nodes": 58}}, + curr={"llama": {"num_nodes": 68}}, # +17%, over the 10% block threshold + ) + assert has_blocker is True + assert _BLOCKER in md + + +def test_smaller_than_expected_delta_is_not_waived(tmp_path: Path): + """Only the exact delta is waived; a smaller change uses normal thresholds.""" + # +5 on 58 = +8.6%, between warn (5%) and block (10%): a warning, not waived. + md, has_blocker = _run( + tmp_path, + base={"llama (static-cache)": {"num_nodes": 58}}, + curr={"llama (static-cache)": {"num_nodes": 63}}, + ) + assert has_blocker is False + assert _ACCEPTED not in md + assert "\u26a0\ufe0f" in md # warning + + +def test_post_merge_zero_delta_is_inert(tmp_path: Path): + """Once merged, base==head -> 0 delta: no blocker, no accepted-marker.""" + md, has_blocker = _run( + tmp_path, + base={"llama (static-cache)": {"num_nodes": 68}}, + curr={"llama (static-cache)": {"num_nodes": 68}}, + ) + assert has_blocker is False + assert _ACCEPTED not in md + assert _BLOCKER not in md + + +def test_post_merge_repeat_delta_still_blocks(tmp_path: Path): + """Absolute pinning closes the stale-waiver hole. + + After the +10 phase-split merges (baseline becomes 68), a *future* +10 + regression (68 -> 78) must NOT match the (58, 68) waiver and must still block. + """ + md, has_blocker = _run( + tmp_path, + base={"llama (static-cache)": {"num_nodes": 68}}, + curr={"llama (static-cache)": {"num_nodes": 78}}, + ) + assert has_blocker is True + assert _BLOCKER in md + assert _ACCEPTED not in md + + +def test_non_waived_metric_on_allowlisted_model_still_blocks(tmp_path: Path): + """A non-waived metric on an allowlisted model still blocks. + + The waiver is scoped per-metric: an allowlisted model whose num_nodes + matches the pinned (58, 68) must still block on a DIFFERENT metric + (model_size_bytes +30%), guarding against over-broadening the waiver. + """ + md, has_blocker = _run( + tmp_path, + base={"llama (static-cache)": {"num_nodes": 58, "model_size_bytes": 1000}}, + curr={"llama (static-cache)": {"num_nodes": 68, "model_size_bytes": 1300}}, + ) + assert has_blocker is True + assert _BLOCKER in md + assert _ACCEPTED in md # the num_nodes row is still waived + + +def test_expected_changes_keys_are_real_model_display_keys(): + """Guard against a silent waiver no-op if model keys drift. + + A model_type/task rename would move the comparator's allowlist keys away + from the benchmarked display keys, silently re-REDing the intended +10. + """ + from tests.benchmark_build import BENCHMARK_MODELS, _display_key + + valid_keys = {_display_key(e.model_type, e.task_name) for e in BENCHMARK_MODELS} + for model_key in bc.EXPECTED_CHANGES: + assert model_key in valid_keys, ( + f"EXPECTED_CHANGES key {model_key!r} is not a benchmarked model " + f"display key; the waiver would silently never fire." + ) diff --git a/tests/build_graph_test.py b/tests/build_graph_test.py index d6d2f9c8..e6c5f7bc 100644 --- a/tests/build_graph_test.py +++ b/tests/build_graph_test.py @@ -4531,12 +4531,18 @@ def test_static_cache_graph_outputs(self): ) def test_static_cache_has_tensorscatter_and_attention(self): - """Verify graph contains TensorScatter and Attention ops.""" + """Verify graph contains TensorScatter and Attention ops. + + TensorScatter lives at the top level (shared by both phases); + Attention lives inside the phase-split If branches, so the search + recurses into subgraphs. + """ model, _ = self._build_static_cache_model() - op_types = {n.op_type for n in model.graph} + op_types = {n.op_type for n in model.graph.all_nodes()} assert "TensorScatter" in op_types, "Static cache graph should use TensorScatter" assert "Attention" in op_types, "Static cache graph should use Attention" + assert "If" in op_types, "Static cache attention should be phase-split behind an If" def test_static_cache_has_initializers(self): """Verify the graph has model parameters.""" @@ -4555,36 +4561,118 @@ def test_static_cache_graph_validates(self): assert len(proto.SerializeToString()) > 0 def test_static_cache_attention_is_causal(self): - """Verify Attention ops use is_causal=1 in static cache mode.""" + """Verify every Attention op uses 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). Both phase-split branches therefore + use is_causal=0; causality comes from an explicit mask (prefill) or + nonpad_kv_seqlen alone (decode). + """ 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 + attention_nodes = [n for n in model.graph.all_nodes() if n.op_type == "Attention"] + # Two branches (prefill + decode) per layer. + assert len(attention_nodes) == 2 * config.num_hidden_layers for node in attention_nodes: is_causal = node.attributes.get("is_causal") 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 or " + f"nonpad_kv_seqlen)" ) - 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_phase_split_mask_presence(self): + """Each layer is an If whose two branches differ in mask presence. + + Fail-closed structural guard for the phase split. It locates the + ``If`` nodes directly and inspects *both* subgraphs, so it cannot pass + vacuously: a regression to a single unconditional Attention (no If) + fails the If-count assertion, and an inverted or single-sided mask + fails the per-branch checks. + + Invariant (If condition is ``Greater(Shape(query)[1], 1)``): + * ``then_branch`` (multi-token / prefill) Attention HAS an explicit + causal ``attn_mask`` produced by ``GreaterOrEqual`` → memory- + efficient path. + * ``else_branch`` (single-token / decode) Attention OMITS ``attn_mask`` + (input[3] None/empty) → Flash-eligible (ORT disables Flash whenever + attn_mask is present, by pointer not content). + """ 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 + if_nodes = [n for n in model.graph.all_nodes() if n.op_type == "If"] + # (a) Fail-closed: the phase-split If must exist, one per layer. + assert len(if_nodes) == config.num_hidden_layers, ( + f"static-cache attention must phase-split via If " + f"(expected {config.num_hidden_layers}, got {len(if_nodes)}); a " + f"single unconditional-mask Attention would force decode off Flash" + ) - for node in attention_nodes: - # Input 3 (0-indexed) is attn_mask — should be empty/None - 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}" + for if_node in if_nodes: + then_branch = if_node.attributes["then_branch"].as_graph() + else_branch = if_node.attributes["else_branch"].as_graph() + then_attn = self._single_attention(then_branch) + else_attn = self._single_attention(else_branch) + + # (b) then-branch (prefill) MUST carry the GreaterOrEqual mask. + then_mask = then_attn.inputs[3] if len(then_attn.inputs) > 3 else None + assert then_mask is not None and then_mask.name != "", ( + "prefill (then) branch Attention must carry an attn_mask" + ) + producer = then_mask.producer() + assert producer is not None and producer.op_type == "GreaterOrEqual", ( + f"prefill attn_mask should come from the causal-mask " + f"GreaterOrEqual, got " + f"{None if producer is None else producer.op_type}" + ) + + # (b) else-branch (decode) MUST omit the mask to stay Flash-eligible. + else_mask = else_attn.inputs[3] if len(else_attn.inputs) > 3 else None + assert else_mask is None or else_mask.name == "", ( + "decode (else) branch Attention must OMIT attn_mask so ORT " + "keeps it on Flash; a mask here forces the slower memory-" + "efficient path" ) + @staticmethod + def _single_attention(graph): + """Return the sole ``Attention`` node in an If branch subgraph.""" + attention_nodes = [n for n in graph if n.op_type == "Attention"] + assert len(attention_nodes) == 1, ( + f"expected exactly one Attention per If branch, got {len(attention_nodes)}" + ) + return attention_nodes[0] + + def test_static_cache_is_phase_split_behind_if(self): + """Static cache attention must be a per-layer If over masked/maskless. + + Guards against regressing to (a) is_causal=1 (ORT-rejected) or + (b) a single always-masked Attention that would force decode off + Flash. Expect one If per layer plus the causal-mask ops (Range + + GreaterOrEqual) used only in the prefill branch. + """ + model, config = self._build_static_cache_model() + op_counts: dict[str, int] = {} + for node in model.graph.all_nodes(): + op_counts[node.op_type] = op_counts.get(node.op_type, 0) + 1 + + assert op_counts.get("If", 0) == config.num_hidden_layers, ( + "Each layer's static-cache attention should be phase-split via If" + ) + # Causal mask (prefill branch only): one Range + GreaterOrEqual/layer. + assert op_counts.get("Range", 0) >= config.num_hidden_layers, ( + "Prefill branch should build query positions with Range" + ) + assert op_counts.get("GreaterOrEqual", 0) == config.num_hidden_layers, ( + "Causal mask (GreaterOrEqual) should appear once per layer " + "(prefill branch only — decode is maskless)" + ) + def test_static_cache_moe_graph_builds(self): """Build a MoE model (qwen2_moe) with static cache.""" model, _config = self._build_static_cache_model( @@ -4604,8 +4692,9 @@ def test_static_cache_moe_graph_builds(self): assert "position_ids" in input_names assert "attention_mask" not in input_names - # Verify TensorScatter and Attention ops are present - op_types = {n.op_type for n in model.graph} + # Verify TensorScatter (top level) and Attention (inside the + # phase-split If branches) ops are present. + op_types = {n.op_type for n in model.graph.all_nodes()} assert "TensorScatter" in op_types assert "Attention" in op_types diff --git a/tests/static_cache_decode_test.py b/tests/static_cache_decode_test.py new file mode 100644 index 00000000..7312e89a --- /dev/null +++ b/tests/static_cache_decode_test.py @@ -0,0 +1,564 @@ +# 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``. + +Three CUDA regression tests here guard against (a) the original +``is_causal=1`` ``NOT_IMPLEMENTED`` regression and (b) a mask wired onto the +decode branch that would silently push decode off Flash: + +* :func:`test_static_cache_prefill_and_decode_run_on_cuda` (fp32) — both + phases run without ``NOT_IMPLEMENTED``. +* :func:`test_static_cache_decode_runs_maskless_on_cuda` (fp16) — the + executed decode ``Attention`` carries no ``attn_mask`` input (the + structural Flash-eligibility precondition), robust to log-format changes. +* :func:`test_static_cache_decode_selects_flash_kernel_on_cuda` (fp16) — the + direct proof: ORT's VERBOSE kernel-selection log shows decode on **Flash** + and prefill on **Memory-Efficient**. + +The fix sets ``is_causal=0`` and phase-splits the attention behind an ``If`` +keyed on ``Shape(query)[1] > 1``: the multi-token (prefill) branch supplies an +explicit causal mask (:func:`mobius.components._common.create_static_cache_causal_mask`, +memory-efficient path), while the single-token decode branch omits the mask so +ORT keeps it on Flash/XQA — the same kernel the GQA variant uses, so the +profiling comparison stays apples-to-apples. These tests exercise the actual +ONNX Runtime kernel for both phases so the regression cannot silently come +back. They require the CUDA Execution Provider because ``TensorScatter`` and +the external-cache ``Attention`` path are CUDA-only. + +The runnability test (:func:`test_static_cache_prefill_and_decode_run_on_cuda`) +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. The Flash-eligibility guard +(:func:`test_static_cache_decode_runs_maskless_on_cuda`) builds fp16 — the +production precision — and asserts the decode ``If`` branch executes without an +``attn_mask`` input (the structural precondition for ORT to keep decode on +Flash). 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 contextlib +import dataclasses +import json +import os +import re +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._builder import build_from_module +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 + + +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, + *, + ir_dtype: ir.DataType = ir.DataType.FLOAT, + enable_profiling: bool = False, + config_overrides: dict | None = None, +) -> tuple[ort.InferenceSession, object]: + """Build a tiny static-cache qwen2 graph and load it on CUDA. + + Uses the full ``build_from_module`` export path (not bare ``task.build``) + so the phase-split ``If`` subgraphs are exercised through the real + ``optimize_model`` pipeline — that is where a structural regression in + the static-cache attention would surface. + + Args: + tmp_dir: Directory for the saved ONNX model + external data. + ir_dtype: Weight/activation precision. Defaults to fp32 (the + ``is_causal`` guard fires before kernel dtype dispatch, so fp32 + exercises the same external-cache path while avoiding cos/sin + cache casting). Pass ``ir.DataType.FLOAT16`` to build the + production-precision graph used for the Flash-eligibility guard. + enable_profiling: Turn on ORT op-level profiling so callers can + inspect which ``If`` branch executed and with which inputs. + config_overrides: Optional ``_base_config`` field overrides. Used to + build at the production ``head_dim`` (Phi-3.5 = 96) so the + Flash-eligibility guard is dispositive for the real model's head + dimension, not just the tiny default (``head_dim=16``). + """ + config = _base_config(**(config_overrides or {})) + config = dataclasses.replace(config, dtype=ir_dtype) + module = registry.get(_MODEL_TYPE)(config) + task = CausalLMTask(static_cache=True, max_seq_len=_MAX_SEQ_LEN) + package = build_from_module(module, config, task=task, execution_provider="default") + model = package["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_options = ort.SessionOptions() + if enable_profiling: + session_options.enable_profiling = True + session_options.profile_file_prefix = str(Path(tmp_dir) / "prof") + session = ort.InferenceSession( + model_path, + session_options, + 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, np_dtype: np.dtype = _CACHE_DTYPE +) -> 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=np_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 _executed_attention_events(profile_path: str) -> list[dict]: + """Op-level ``Attention`` events from an ORT profiling JSON. + + Each returned dict has ``name`` (carries the ``static_cache_prefill`` / + ``static_cache_decode`` branch tag) and ``has_mask`` — True when the + executed node received a rank-4 ``attn_mask`` input. ORT 1.27's Python + profiler emits only op-level ``Node`` events (no CUDA ``Kernel`` events), + so the *internal* attention kernel (Flash vs memory-efficient) is not + observable here; the mask-input signature is, and the presence of an + ``attn_mask`` is exactly what makes ORT ineligible for Flash. + """ + with open(profile_path) as handle: + events = json.load(handle) + + attention_events: list[dict] = [] + for event in events: + args = event.get("args", {}) + if args.get("op_name") != "Attention": + continue + input_shapes = args.get("input_type_shape", []) + has_mask = any(len(next(iter(shape.values()))) == 4 for shape in input_shapes) + attention_events.append({"name": event["name"], "has_mask": has_mask}) + return attention_events + + +def test_static_cache_decode_runs_maskless_on_cuda(): + """The decode branch executes maskless at runtime (Flash-eligible). + + Build-time tests assert the *graph* phase-splits the static-cache + attention behind an ``If`` (decode branch omits ``attn_mask``). This + test closes the loop at runtime on the production fp16 path: it profiles + a single-token decode and a multi-token prefill on CUDA and asserts the + ``If`` routed correctly and that the executed decode ``Attention`` carries + **no** ``attn_mask`` input. + + Why this matters: ORT disables Flash whenever ``attn_mask`` is present + (by pointer, not content), so a regression that wired the mask onto the + decode branch would silently push the hot decode path onto the slower + memory-efficient kernel and invalidate the Attention-vs-GQA decode + comparison. ORT 1.27's Python profiler does not surface the internal + kernel name, so we assert the structural precondition (mask absence) + that deterministically governs Flash eligibility. + """ + with tempfile.TemporaryDirectory() as tmp_dir: + session, config = _build_static_cache_session( + tmp_dir, ir_dtype=ir.DataType.FLOAT16, enable_profiling=True + ) + 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(2) + + # Single-token decode (S_q = 1): the If must take the maskless branch. + decode_feeds: dict[str, np.ndarray] = { + "input_ids": rng.integers(0, vocab, size=(1, 1), dtype=np.int64), + "position_ids": np.array([[4]], dtype=np.int64), + "write_indices": np.array([4], dtype=np.int64), + "nonpad_kv_seqlen": np.array([5], dtype=np.int64), + } + decode_feeds.update(_empty_caches(num_layers, kv_hidden, np.float16)) + session.run(output_names, decode_feeds) + + # Multi-token prefill (S_q = 4): the If must take the masked branch. + 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, np.float16)) + session.run(output_names, prefill_feeds) + + events = _executed_attention_events(session.end_profiling()) + + decode_events = [e for e in events if "static_cache_decode" in e["name"]] + prefill_events = [e for e in events if "static_cache_prefill" in e["name"]] + + # Decode took the maskless branch on every layer (Flash-eligible). + assert len(decode_events) == num_layers, ( + f"expected {num_layers} decode-branch Attention executions, got {len(decode_events)}" + ) + assert all(not e["has_mask"] for e in decode_events), ( + "decode-branch Attention must run WITHOUT an attn_mask input so ORT " + "keeps it on Flash; a mask here forces the slower memory-efficient " + "path and breaks the decode-latency comparison" + ) + + # Prefill took the masked branch on every layer (memory-efficient path). + assert len(prefill_events) == num_layers, ( + f"expected {num_layers} prefill-branch Attention executions, got {len(prefill_events)}" + ) + assert all(e["has_mask"] for e in prefill_events), ( + "prefill-branch Attention must carry the explicit causal mask" + ) + + +# Matches the opset-24 LLM Attention kernel-selection log line emitted at +# VERBOSE by onnxruntime/core/providers/cuda/llm/attention.cc, e.g. +# "ONNX Attention: using Flash Attention (batch=1, q_seq=1, total_seq=16, ...)" +# "ONNX Attention: using Memory Efficient Attention (batch=1, q_seq=4, ...)" +_ATTENTION_KERNEL_LINE = re.compile( + r"ONNX Attention: using (?P.+?) \(batch=\d+, q_seq=(?P\d+)" +) + + +@contextlib.contextmanager +def _capture_attention_kernel_log(): + """Capture ORT's per-op attention kernel-selection log lines. + + The opset-24 LLM ``Attention`` CUDA kernel logs which kernel it selected + (Flash / Memory-Efficient / unfused) at VERBOSE through the *default* + (process-global) logger, written to the C++ ``stderr`` (fd 2). ORT's + Python profiler does not surface this, so to read it we raise the default + logger severity to VERBOSE and redirect fd 2 around the run. Callers must + read the yielded file *inside* the ``with`` block (it is closed on exit); + fd 2 and the logger severity are restored before exit so assertions made + after the block still report normally. + + Not safe under *in-process* parallelism: it mutates process-global state + (the default logger severity and fd 2). That is fine here — pytest runs + tests sequentially in-process, and ``pytest-xdist`` isolates workers in + separate processes — but do not call it from threads sharing this process. + """ + ort.set_default_logger_severity(0) + saved_stderr_fd = os.dup(2) + try: + with tempfile.TemporaryFile(mode="w+b") as capture_file: + os.dup2(capture_file.fileno(), 2) + try: + yield capture_file + finally: + capture_file.flush() + os.dup2(saved_stderr_fd, 2) + finally: + os.close(saved_stderr_fd) + ort.set_default_logger_severity(2) # back to ORT's default (WARNING) + + +def _selected_attention_kernels(capture_file) -> list[tuple[str, int]]: + """Parse ``(kernel_name, q_seq)`` pairs from a captured verbose log.""" + capture_file.seek(0) + text = capture_file.read().decode("utf-8", "replace") + kernels: list[tuple[str, int]] = [] + for line in text.splitlines(): + match = _ATTENTION_KERNEL_LINE.search(line) + if match is not None: + kernels.append((match.group("kernel").strip(), int(match.group("q_seq")))) + return kernels + + +# Phi-3.5's real attention head dimension (32 heads x 96 = 3072 hidden). The +# tiny default config uses head_dim=16, which exercises the phase-split wiring +# but NOT whether ORT's Flash kernel accepts the production head_dim on this +# GPU. Building at head_dim=96 (heads/hidden scaled down to stay tiny) makes +# the Flash-eligibility assertion dispositive for the model we actually ship. +_PHI35_HEAD_DIM_OVERRIDES = { + "hidden_size": 384, + "num_attention_heads": 4, + "num_key_value_heads": 2, + "head_dim": 96, +} + + +# ORT version whose VERBOSE attention.cc kernel-selection log strings +# (``_ATTENTION_KERNEL_LINE``) this guard was validated against. The log text +# is an internal, unstable contract; on an ORT bump the kernel-name proof must +# be re-validated (or demoted to the profiling harness) rather than fail with a +# confusing regex-miss assertion. Gating on this version makes that churn +# self-announce as a skip at the exact moment the assumption is invalidated. +_VALIDATED_ORT_VERSION = "1.27" + + +@pytest.mark.parametrize( + "config_overrides", + [None, _PHI35_HEAD_DIM_OVERRIDES], + ids=["tiny-head-dim-16", "phi35-head-dim-96"], +) +def test_static_cache_decode_selects_flash_kernel_on_cuda(config_overrides): + """Decode actually selects Flash; prefill selects Memory-Efficient. + + Structural maskless-ness (the test above) is necessary, but the *proof* + that the phase split achieves its purpose is the kernel ORT actually runs. + The opset-24 LLM Attention kernel logs its choice at VERBOSE; this test + captures that log and asserts the single-token decode runs **Flash** + (the same external-cache kernel the GQA variant uses, so the decode- + latency comparison is apples-to-apples) while the multi-token prefill + runs **Memory-Efficient** (mask present, the cheap amortized path). + + This is the regression guard the reviewers required: a change that wired + the causal mask onto the decode branch would flip its kernel from Flash + to Memory-Efficient and fail here, even though finiteness / scatter / If- + count assertions would all still pass. + + Parametrized over head dimension: the tiny default (``head_dim=16``) and + Phi-3.5's production ``head_dim=96``. The latter makes the assertion + dispositive — it empirically confirms ORT's Flash kernel accepts the real + model's head dimension (fp16) on this GPU, rather than silently routing to + Memory-Efficient, which would invalidate the decode-on-Flash premise. + + Skips on an unvalidated ORT version: this proof reads ORT's internal + VERBOSE kernel-selection log strings, so a version bump must re-validate + those strings (the deterministic ``test_static_cache_decode_runs_maskless`` + structural guard remains the version-robust backstop and is not gated). + """ + if not ort.__version__.startswith(_VALIDATED_ORT_VERSION): + pytest.skip( + f"decode-on-Flash kernel proof reads ORT's VERBOSE attention.cc " + f"selection log, validated only on ORT {_VALIDATED_ORT_VERSION}.x " + f"(running {ort.__version__}). Re-validate the log strings in " + f"_ATTENTION_KERNEL_LINE and bump _VALIDATED_ORT_VERSION, or demote " + f"this proof to the profiling harness." + ) + with tempfile.TemporaryDirectory() as tmp_dir: + session, config = _build_static_cache_session( + tmp_dir, + ir_dtype=ir.DataType.FLOAT16, + config_overrides=config_overrides, + ) + 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(3) + + decode_feeds: dict[str, np.ndarray] = { + "input_ids": rng.integers(0, vocab, size=(1, 1), dtype=np.int64), + "position_ids": np.array([[4]], dtype=np.int64), + "write_indices": np.array([4], dtype=np.int64), + "nonpad_kv_seqlen": np.array([5], dtype=np.int64), + } + decode_feeds.update(_empty_caches(num_layers, kv_hidden, np.float16)) + + 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, np.float16)) + + with _capture_attention_kernel_log() as capture_file: + session.run(output_names, decode_feeds) + session.run(output_names, prefill_feeds) + kernels = _selected_attention_kernels(capture_file) + + # The kernel log distinguishes phases by query length (q_seq). + decode_kernels = [name for name, q_seq in kernels if q_seq == 1] + prefill_kernels = [name for name, q_seq in kernels if q_seq == prefill_len] + + assert len(decode_kernels) == num_layers, ( + f"expected {num_layers} decode (q_seq=1) Attention kernel-selection " + f"log lines, got {len(decode_kernels)} (all parsed: {kernels}). If " + f"empty, ORT's verbose attention-kernel log format may have changed." + ) + assert all("Flash" in name for name in decode_kernels), ( + f"decode MUST select Flash Attention to stay apples-to-apples with " + f"GQA's decode kernel; got {decode_kernels}. A mask on the decode " + f"branch flips this to Memory-Efficient." + ) + + assert len(prefill_kernels) == num_layers, ( + f"expected {num_layers} prefill (q_seq={prefill_len}) kernel-selection " + f"log lines, got {len(prefill_kernels)} (all parsed: {kernels})" + ) + assert all("Memory Efficient" in name for name in prefill_kernels), ( + f"prefill should select Memory-Efficient Attention (causal mask " + f"present); got {prefill_kernels}" + ) + + +def test_static_cache_decode_ignores_keys_beyond_nonpad_on_cuda(): + """Out-of-bound cache slots cannot change a decode's output. + + The decode branch runs maskless and relies solely on ``nonpad_kv_seqlen`` + to bound attention to the valid keys. The other e2e tests prove decode + *runs* and stays on Flash, but not that the bound is semantically applied. + This test closes that gap: it compares a decode over a clean carried cache + against a decode over the *same* cache with every slot at or beyond + ``nonpad`` overwritten with large garbage. If the bound is honored the two + decodes produce bit-identical logits; if a regression let the kernel read + the whole pre-allocated cache, the garbage would perturb the softmax and + the logits would diverge. + """ + with tempfile.TemporaryDirectory() as tmp_dir: + session, config = _build_static_cache_session(tmp_dir, ir_dtype=ir.DataType.FLOAT16) + 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(4) + + # 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, np.float16)) + prefill_out = dict(zip(output_names, session.run(output_names, prefill_feeds))) + + # Decode one token into slot 4; valid keys are slots 0..4 (nonpad=5). + 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))) + + # Poison every cache slot at or beyond ``nonpad`` with large garbage; + # those positions must never be attended during decode. Slot 4 (the + # decode write target, within nonpad) is left untouched. + poisoned_caches = _carry_caches(prefill_out, num_layers) + for layer in range(num_layers): + for name in (f"key_cache.{layer}", f"value_cache.{layer}"): + buf = poisoned_caches[name].copy() + buf[:, nonpad:, :] = np.float16(50.0) + poisoned_caches[name] = buf + poisoned_feeds = {**decode_inputs, **poisoned_caches} + perturbed = dict(zip(output_names, session.run(output_names, poisoned_feeds))) + + assert np.array_equal(baseline["logits"], perturbed["logits"]), ( + "decode logits changed when cache slots beyond nonpad_kv_seqlen were " + "poisoned — the nonpad bound is not being honored, so decode is " + "attending to invalid (out-of-range) keys" + )