diff --git a/.github/workflows/main.yml b/.github/workflows/main.yml index 671f259f4..2b4717c8b 100644 --- a/.github/workflows/main.yml +++ b/.github/workflows/main.yml @@ -15,6 +15,68 @@ concurrency: cancel-in-progress: true jobs: + onnx-genai-metadata: + name: ONNX GenAI metadata + runs-on: ubuntu-latest + steps: + - uses: actions/checkout@v7 + - uses: actions/checkout@v7 + with: + repository: justinchuby/onnx-genai + # Pinned by SHA so the producer/runtime contract under test is + # reproducible. PR #828 is maintained by merging main rather than + # rebasing, so this commit remains reachable from its branch history. + # When bumping, verify reachability with `git merge-base --is-ancestor + # `. + ref: 509cd4e9c4471f4cbc59fe44b47168f0ae128fe3 + path: validation/onnx-genai + - uses: actions/setup-python@v7 + with: + python-version: "3.12" + - name: Install dependencies + run: | + pip install torch --index-url https://download.pytorch.org/whl/cpu + pip install -r requirements/ci/requirements.txt + pip install onnxruntime==1.28.0 + pip install -e '.[testing]' + - name: Generate representative packages + run: | + PYTHONPATH=src python tests/generate_onnx_genai_validation_packages.py \ + validation/generated + PYTHONPATH=src python tests/compare_onnx_genai_validation_packages.py \ + tests/fixtures/onnx_genai_workflows validation/generated + - name: Validate package semantics + run: | + # Report every invalid package, not just the first: the job runs under + # `bash -e`, so an unguarded failure inside the loop would abort it and + # hide the remaining results. + failed="" + # The generated tree, not the committed one: validation resolves each + # component's artifact from disk, and the graphs are generated rather + # than committed. The step above has already proven the two trees + # carry the same metadata. + for package in validation/generated/*; do + [ -f "$package/inference_metadata.yaml" ] || continue + cargo run --quiet \ + --manifest-path validation/onnx-genai/Cargo.toml \ + -p onnx-genai-metadata --bin validate_metadata -- "$package" \ + || failed="$failed $package" + done + if [ -n "$failed" ]; then + echo "Invalid packages:$failed" + exit 1 + fi + - name: Execute all workflow packages + run: | + cp tests/onnx_genai_workflow_conformance.rs \ + validation/onnx-genai/crates/onnx-genai-engine/tests/mobius_workflow_conformance.rs + ORT_LIB="$(python -c \ + 'import onnxruntime, pathlib; print(next((pathlib.Path(onnxruntime.__file__).parent / "capi").glob("libonnxruntime.so*")))')" + MOBIUS_WORKFLOW_CONFORMANCE_DIR="$PWD/validation/generated" \ + ONNX_GENAI_ORT_LIB="$ORT_LIB" \ + cargo test --manifest-path validation/onnx-genai/Cargo.toml \ + -p onnx-genai-engine --test mobius_workflow_conformance -- --nocapture + lint: name: Lint runs-on: ubuntu-latest diff --git a/.gitignore b/.gitignore index efcf9d8cd..ad9983b71 100644 --- a/.gitignore +++ b/.gitignore @@ -213,6 +213,19 @@ __marimo__/ *.onnx.data *.gguf +# Generated conformance-package artifacts. The metadata beside them is the +# contract under review and is committed; the graphs and adapter weights are a +# deterministic function of +# tests/generate_onnx_genai_validation_packages.py, so committing them would +# store megabytes of bytes no reviewer reads. CI regenerates and validates +# them; tests get them from the materialized_workflow_packages fixture. +tests/fixtures/onnx_genai_workflows/**/*.safetensors + +# Where CI and the test fixture materialize those packages in full. Nothing +# here is reviewable input: it is the regenerated output the validate and +# conformance steps consume. +validation/ + # Common test dirs output/** cache_dir/** diff --git a/CHANGELOG.md b/CHANGELOG.md index dc6f8a782..f89ee93ff 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -7,6 +7,144 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ## [Unreleased] +### Fixed + +- An exported graph is no longer transcribed into its own workflow component. + A component backed by a shipped `.onnx` file now declares only `ports.roles`; + the artifact answers every question about which ports exist and what shape + they have, and the runtime resolves them against the live session. The + removed block was a second statement of the same ABI with nothing keeping the + two in agreement — the same defect as `model.io`, one level down. Package + validation and runtime conformance both stay at 11/11 against the pinned ONNX + GenAI branch, and the contract tests now resolve every role, invocation + binding and state pair against the graph itself rather than against the + metadata's agreement with a copy of itself, which is a strictly stronger + check. + + Policy graphs keep their contracts, and that boundary was measured rather + than assumed: a workflow value inherits its dtype, rank and request axis from + the port that produced it, so dropping them left row-wise emits untyped and + made 4 of the 11 packages invalid. Those contracts type the workflow's own + dataflow; they do not describe an external interface. + +- Deep decoders are now covered where the cache layer annotation actually + matters. Every metadata package built in the test suite had two layers, and + below ten a cell label sorts identically whether ordered lexicographically, + numerically or by insertion — so a `layer` derived from a cell's position + rather than parsed from its port name would have passed every assertion while + transposing caches on any real model. `canonical_workflow_contract_test` now + builds twelve-layer dynamic, static-cache and hybrid decoders and pins that + the declared layer restates the port name, that ordering by it recovers the + buffer lists, and that a hybrid's alternating groups own layers their cells' + positions never equal. A producer that dropped the parse leaves the 101 + pre-existing assertions green and fails seven of these. + +- Every component a task builds now declares an optimization role. Qwen3-TTS's + four loop-wiring graphs (`code_predictor_prefill`, + `code_predictor_step_embedder`, `code_predictor_indices`, `talker_text_step`) + were absent from `TTSTask.model_roles`, so `build_from_module` fell back to + the `"decoder"` role and offered them the GQA / QKV-packing passes meant for + attention stacks, and `inspect_components` under-reported the package by four + components. They now declare a `"glue"` role: a parameter-free graph that + reads every tensor it uses from a graph input. `arch_validation_test` fails + any task that builds a component it does not declare, and a network-free unit + test pins the same invariant for Qwen3-TTS. + +### One canonical serialized representation + +#### Changed + +- **`pipeline.workflow` is now the only place a package describes its graph + ABI.** No export emits `model.io`, including a bare single-file decoder: that + case is a one-component workflow, not a different kind of document. `model` + keeps package-wide geometry and capabilities and nothing else. Two writable + statements of one fact are a defect whatever they contain — nothing forces + them to agree, and a reader of either never learns the other exists — so a + runtime that wants an optimized single-graph path derives it by lowering the + workflow instead. Verified end to end: the ONNX GenAI runtime executes the + fixed-capacity decode path from the workflow alone, with no `model.io` in the + package. + +#### Added + +- An ONNX component that ships an artifact declares no port contracts. The + `.onnx` file travels inside the package and is authoritative for which ports + exist and what each one's dtype, rank and shape is, so transcribing that into + YAML would be a second writable statement of one fact — the very thing this + section removes — sitting one level below `model.io` rather than beside it. + The runtime resolves ports against the live session, which catches a name the + graph does not expose instead of agreeing with a stale echo of it. A + producer-synthesized policy graph is the exception and states its contracts, + because a workflow value takes its dtype, rank and request axis from the port + that produced it: those contracts are the dataflow's type annotations, not a + description of an external interface. +- Every ONNX component declares `ports.roles`: what it *does* with a value bound + to a port. An invocation records which SSA value reaches a port, not whether + that port is tokens, a mask or logits. Mobius mints these port names in its own + task builders, so it states the mapping (`input_ids`→`token_ids`, + `inputs_embeds`, `attention_mask`, `position_ids`, `logits`, + `last_hidden_state`→`hidden_states`, `encoder_hidden_states`, + `audio_features`) rather than inferring it. A port outside that vocabulary + carries no role. +- State port aliases declare `role` (`key`/`value`) and `layer`. A layer's key + and value buffers are the same dtype and shape, and a cell's label sorts + lexicographically so `cache_10` precedes `cache_2` — pairing per-layer buffers + positionally would silently transpose two layers' caches. Both fields are + emitted together or not at all, so a recurrent or convolution cache is never + given a fabricated index. +- `IndexedScatter.kv_length_ports` names the port carrying the graph-visible + valid length, beside the existing `write_indices_ports`. The two control + vectors are both rank-1 integers and are therefore indistinguishable by shape; + with both named, the whole fixed-capacity ABI is recoverable from the workflow. +- `tests/canonical_workflow_contract_test.py` pins the invariant. It asks one + set of shape-agnostic questions of dynamic, static-cache, FP8, heterogeneous + and composite packages — and of all 11 checked-in fixtures — so a future + feature cannot grow its own top-level block while every feature-specific test + keeps passing. + +### Fixed-capacity (static) KV cache and FP8 KV cache metadata + +#### Added + +- `--features static-cache` now produces onnx-genai metadata instead of being + refused. The producer publishes the write cursor (`write_indices`), the valid + length (`nonpad_kv_seqlen`), the fixed-capacity buffer contracts, the + per-layer input/output pairs, and an `indexed_scatter` state-service update + discipline naming the cursor, the capacity and the per-component port that + carries it. The buffers are declared as `recurrence: {kind: invariant}` loop + cells and the capacity as a `package.cache_capacity` literal workflow input. + Nothing dispatches on model name; the ports are read from the graph. +- Heterogeneous caches keep their own disciplines. Gemma 4's sliding layers stay + on a growing rank-4 BNSH cache while its full-attention layers use rank-3 + fixed-capacity buffers, and only the layers that own a buffer bind ports in a + state group — its KV-shared suffix owns none. +- A `static_cache` package joined the checked-in onnx-genai conformance + fixtures, so the engine exercises the fixed-capacity carry and the write + cursor rather than only the growing-tensor path. + +#### Fixed + +- Gemma 4's shared-KV fallback pinned a 4-D BNSH shape onto *any* borrowed KV + tensor whose rank was not 4. A static-cache source hands over a fully known + rank-3 `[batch, capacity, kv_hidden]` buffer, so this overwrote a correct + shape with a wrong one — corrupting the declared shape of + `updated_key_cache.N` and defeating the rank-3 static-source test further + down, which would then have transposed a rank-3 tensor as BNSH. The fallback + now only supplies a shape when there is none. +- Gemma 4's vision-language decoder dropped `attention_mask` whenever the export + was static, but a Gemma 4 decoder is only *partly* static: its sliding layers + keep a dynamic cache and build their bias from that mask. The hybrid decoder + therefore lost all padding information. Both builders now apply one rule — a + mask exists exactly when some layer still has a dynamic cache — so a fully + static decoder carries no unused port and a hybrid one keeps its mask. +- `--features fp8-kv-cache` no longer silently produces a float16 cache. The + gate only tested whether GQA fusion was *expected*; the pass now reports how + many caches it converted and the build fails when the answer is zero, naming + the reason: FP8 KV storage needs an attention operator with `k_scale`/ + `v_scale` inputs, which a `TensorScatter` + `ai.onnx` `Attention` static-cache + graph does not have. + + ### Qwen3.5/3.6-MoE mixed float/quantized decoder (Olive checkpoints) #### Fixed diff --git a/benchmarks/README.md b/benchmarks/README.md new file mode 100644 index 000000000..4aadee135 --- /dev/null +++ b/benchmarks/README.md @@ -0,0 +1,36 @@ +# Muse workflow benchmark + +`muse_workflow_h200.json` is the shared workload for a paired native and +metadata-workflow benchmark of the published Muse Glimmer INT4 package. Both +paths use the exact 68 token IDs in `muse_prompt_ids.json`, greedy parameters, +token budget, warmups, and steady-decode window. Passing token IDs avoids a +tokenizer adding another beginning-of-text token to the already rendered prompt. +Both runners set `ORT_ENABLE_CUDNN_FLASH_ATTENTION=0`; this matches the native +baseline and avoids comparing different GQA attention kernels. +`request_max_length` is prompt tokens plus new tokens (68 + 128 = 196); +`model_max_context` is the independent 131072-token artifact admission ceiling. +For the text-only VLM path, omit `request.image`. The runtime initializes +`request.image_present=false`, and the false branch supplies empty image features. + +Run the native ORT GenAI path: + +```bash +python scripts/benchmark_muse_native.py \ + --model artifacts/muse-int4-package \ + --output artifacts/muse-int4-package/native-benchmark.json +``` + +Run the ONNX GenAI workflow path with the `profile_native` binary built from +the schema/runtime revision named in the JSON config: + +```bash +python scripts/benchmark_muse_workflow.py \ + --model artifacts/muse-int4-package \ + --runner path/to/profile_native \ + --output artifacts/muse-int4-package/workflow-benchmark.json +``` + +The workflow runner must support `--pipeline --backend ort --ep cuda`, +`--prompt-ids`, and an optional `--image` request binding. The text-only workload +intentionally leaves the image unset so results remain comparable to the +published 61.76 tok/s baseline. diff --git a/benchmarks/muse_568457_release_evidence.json b/benchmarks/muse_568457_release_evidence.json new file mode 100644 index 000000000..b0a37de3d --- /dev/null +++ b/benchmarks/muse_568457_release_evidence.json @@ -0,0 +1,58 @@ +{ + "kind": "workflow_release_diagnostic", + "producer_commit": "80acfc069f0959f9e6580785fad3172bcc4cc0aa", + "runtime_source_head": "568457045ea003e787816995a380bcd9d7443169", + "runner_sha256": "fad10c7c8735c56e42acdfc092c4c6c327090c50de823efc45d4f3309b2d4aa9", + "overlay_manifest_sha256": "fd77589ffb9333d69efe78909b7f4a489da1071a35affb49b243e87a76330c6a", + "metadata_sha256": "8727cb6ab44ff65a4e74b6281c5af9eadba336483f5fa1e6033996c3b99fcced", + "decoder_state_initializer_sha256": "77e4eae0e35cc618214476e0776af3c1b7930862924f431edb62724cadb97af7", + "workload": { + "prompt_tokens": 68, + "generated_tokens": 128, + "warmups": 1, + "runs": 3, + "decode_skip": 8, + "stop_on_eos": false, + "image": null + }, + "metrics": { + "ttft_ms": 67.512, + "decode_ms_per_token": 29.168, + "throughput_tok_s": 34.28, + "native_ttft_ms": 49.02172100264579, + "native_throughput_tok_s": 63.36527373646786, + "throughput_ratio": 0.5409903244885887 + }, + "diagnostics": { + "captures": 1, + "replays": 510, + "synchronizations": 2580, + "h2d_calls": 1556, + "h2d_bytes": 413802656, + "d2h_calls": 512, + "d2h_bytes": 4096, + "d2d_calls": 3068, + "d2d_bytes": 17376, + "island_elapsed_ms": 2974.922 + }, + "token_parity": { + "exact": false, + "first_difference": { + "index": 38, + "runtime": 4243, + "native": 33386 + }, + "runtime_count": 128, + "native_count": 128 + }, + "gate": { + "required_throughput_ratio": 0.99, + "correctness": false, + "paired_performance": false, + "upload_allowed": false + }, + "notes": [ + "The decoder remained outside the captured policy island.", + "This exact prompt-ID runtime still materialized decoder logits through host memory." + ] +} diff --git a/benchmarks/muse_9e575_landed_bridge_evidence.json b/benchmarks/muse_9e575_landed_bridge_evidence.json new file mode 100644 index 000000000..a311b3786 --- /dev/null +++ b/benchmarks/muse_9e575_landed_bridge_evidence.json @@ -0,0 +1,72 @@ +{ + "kind": "landed_bridge_diagnostic", + "producer_commit": "80acfc069f0959f9e6580785fad3172bcc4cc0aa", + "metadata_sha256": "8727cb6ab44ff65a4e74b6281c5af9eadba336483f5fa1e6033996c3b99fcced", + "decoder_state_initializer_sha256": "77e4eae0e35cc618214476e0776af3c1b7930862924f431edb62724cadb97af7", + "measurement_runtime": { + "base_head": "b4422f5709fc41ea2545e1de46b849343ace9824", + "dirty": true, + "log_mtime": "2026-08-14T02:09:47Z", + "log_sha256": "d73e93f03086ff9408c12bb10c191678ed19a0f1650ecb738972911f0361d217", + "manifest_sha256": "940ccc975d87fe1920666e7ac86dad8d9ad4aafd03fa51ff36f066da3d200915", + "five_file_binary_patch_sha256": "53d5643611d760f7a685fa037e7ba7a99faeac35a719af46dbccec8ca6a9a03b" + }, + "landed_runtime": { + "head": "9e5757196b98542390ce11f4ff966a58ab3ef578", + "intermediate_correctness_commit": "3cf21d15", + "binary_patch_sha256": "94845cf46fa8c4c4881d72591881b978b7370ea5a00a96f661857145b5a88449", + "clean_runner_sha256": "4b8b092d5fa104f0c6df67c5ba75de7d761c09ffffc50edf07acaab6b6e450bf" + }, + "workload": { + "prompt_tokens": 68, + "generated_tokens": 128, + "warmups": 1, + "runs": 3, + "decode_skip": 8, + "stop_on_eos": false, + "image": null + }, + "environment": { + "onnxruntime": "1.28.0", + "onnxruntime_api_version": 27, + "onnxruntime_library_sha256": "018d05f615e96f6e0cfb95d2865f63915681d5d2688dd187f599822b245d27bc", + "prompt_ids_sha256": "eab7c0ee593ff2a6331e1678cefadb0969b65cce6f2fad5b75652b8c155399ee", + "ONNX_GENAI_CUDA_GRAPH": "1", + "ORT_ENABLE_CUDNN_FLASH_ATTENTION": "0" + }, + "metrics": { + "ttft_ms": 60.6, + "decode_ms_per_token": 24.682, + "throughput_tok_s": 40.51, + "native_ttft_ms": 49.02172100264579, + "native_throughput_tok_s": 63.36527373646786, + "throughput_ratio": 0.639 + }, + "diagnostics": { + "captures": 1, + "replays": 510, + "synchronizations": 1558, + "h2d_calls": 1044, + "h2d_bytes": 8352, + "d2h_calls": 512, + "d2h_bytes": 4096, + "d2d_calls": 3068, + "d2d_bytes": 17376, + "island_elapsed_ms": 2484.875 + }, + "token_parity": { + "exact": true, + "runtime_count": 128, + "native_count": 128 + }, + "gate": { + "required_throughput_ratio": 0.99, + "real_multitoken_e2e": true, + "paired_performance": false, + "upload_allowed": false + }, + "notes": [ + "The measured dirty bridge is fully represented by the landed runtime diff.", + "A clean landed-runtime rerun is still required for final release performance evidence." + ] +} diff --git a/benchmarks/muse_c885_release_evidence.json b/benchmarks/muse_c885_release_evidence.json new file mode 100644 index 000000000..3da2c22cd --- /dev/null +++ b/benchmarks/muse_c885_release_evidence.json @@ -0,0 +1,58 @@ +{ + "kind": "workflow_release_diagnostic", + "producer_commit": "80acfc069f0959f9e6580785fad3172bcc4cc0aa", + "runtime_source_head": "c885b71b3813fd652690e5ef5154bfc5535e5c1c", + "runner_sha256": "69f22720f35a01b9df23e7e4c688e8a8e8d40b7113f2537e51dd2c3a13c1b015", + "overlay_manifest_sha256": "fd77589ffb9333d69efe78909b7f4a489da1071a35affb49b243e87a76330c6a", + "metadata_sha256": "8727cb6ab44ff65a4e74b6281c5af9eadba336483f5fa1e6033996c3b99fcced", + "decoder_state_initializer_sha256": "77e4eae0e35cc618214476e0776af3c1b7930862924f431edb62724cadb97af7", + "workload": { + "prompt_tokens": 68, + "generated_tokens": 128, + "warmups": 1, + "runs": 3, + "decode_skip": 8, + "stop_on_eos": false, + "image": null + }, + "metrics": { + "ttft_ms": 79.234, + "decode_ms_per_token": 28.913, + "throughput_tok_s": 34.59, + "native_ttft_ms": 49.02172100264579, + "native_throughput_tok_s": 63.36527373646786, + "throughput_ratio": 0.5458825940507668 + }, + "diagnostics": { + "captures": 1, + "replays": 510, + "synchronizations": 2580, + "h2d_calls": 1556, + "h2d_bytes": 413802656, + "d2h_calls": 512, + "d2h_bytes": 4096, + "d2d_calls": 3068, + "d2d_bytes": 17376, + "island_elapsed_ms": 2904.668 + }, + "token_parity": { + "exact": false, + "first_difference": { + "index": 38, + "runtime": 4243, + "native": 33386 + }, + "runtime_count": 128, + "native_count": 128 + }, + "gate": { + "required_throughput_ratio": 0.99, + "correctness": false, + "paired_performance": false, + "upload_allowed": false + }, + "notes": [ + "The decoder remained outside the captured policy island.", + "This is clean-runtime diagnostic evidence, not publication eligibility evidence." + ] +} diff --git a/benchmarks/muse_prompt_ids.json b/benchmarks/muse_prompt_ids.json new file mode 100644 index 000000000..32c712dc1 --- /dev/null +++ b/benchmarks/muse_prompt_ids.json @@ -0,0 +1,70 @@ +[ + 200000, + 200022, + 15651, + 200023, + 4662, + 583, + 262, + 19933, + 12133, + 43910, + 335, + 111198, + 67059, + 38, + 220, + 837, + 34, + 25, + 1761, + 25, + 2468, + 335, + 8048, + 4282, + 38, + 220, + 837, + 34, + 25, + 2834, + 25, + 974, + 1574, + 34956, + 300, + 9762, + 38, + 2244, + 1574, + 15, + 14757, + 73965, + 38, + 392, + 2540, + 706, + 392, + 1556, + 4205, + 200008, + 200022, + 1556, + 200023, + 105195, + 469, + 11341, + 5959, + 290, + 21752, + 14774, + 10950, + 3658, + 290, + 3364, + 26, + 200008, + 200022, + 140680 +] diff --git a/benchmarks/muse_workflow_h200.json b/benchmarks/muse_workflow_h200.json new file mode 100644 index 000000000..a60b61e79 --- /dev/null +++ b/benchmarks/muse_workflow_h200.json @@ -0,0 +1,55 @@ +{ + "scenario": "muse-glimmer-30b-int4-workflow-vs-native", + "package": { + "repository": "justinchuby/Muse-Glimmer-30B-ONNX-INT4-CUDA", + "weights_revision": "bf36a94a4519e14e3c48ad005c6ff1972ab44ccb", + "schema_head": "9e5757196b98542390ce11f4ff966a58ab3ef578", + "artifacts": { + "decoder": "decoder/model.onnx", + "embedding": "embedding/model.onnx", + "vision_encoder": "vision_encoder/model.onnx", + "metadata": "inference_metadata.yaml" + } + }, + "runtime": { + "onnxruntime": "1.28.0", + "onnxruntime_genai_commit": "9e5757196b98542390ce11f4ff966a58ab3ef578", + "execution_provider": "CUDAExecutionProvider", + "cuda_graph": true, + "cudnn_flash_attention": false, + "shared_kv": true + }, + "hardware": { + "gpu": "NVIDIA H200", + "batch_size": 1 + }, + "workload": { + "prompt": "Briefly explain why the sky appears blue during the day.", + "rendered_prompt": "<|begin_of_text|><|start|>system<|message|>You are a helpful AI assistant.\nKnowledge cutoff: 2026-01-04.\nCurrent date: 2026-08-13.\n\nReasoning strength: high.\n\n# Valid recipients: \"self\", \"user\".<|eot|><|start|>user<|message|>Briefly explain why the sky appears blue during the day.<|eot|><|start|>assistant", + "prompt_ids_file": "benchmarks/muse_prompt_ids.json", + "prompt_tokens": 68, + "image": null, + "workflow_media_binding": "omit_optional_request_image", + "max_new_tokens": 128, + "request_max_length": 196, + "model_max_context": 131072, + "stop_on_eos": false, + "warmups": 1, + "runs": 3, + "decode_skip": 8 + }, + "sampling": { + "algorithm": "greedy", + "do_sample": false, + "temperature": 1.0, + "top_k": 1, + "top_p": 1.0, + "seed": 0 + }, + "release_gate": { + "clean_landed_runtime": true, + "exact_token_parity": true, + "cuda_graph_required": true, + "min_throughput_ratio": 0.99 + } +} diff --git a/docs/onnx-genai-performance-conformance.md b/docs/onnx-genai-performance-conformance.md new file mode 100644 index 000000000..833135707 --- /dev/null +++ b/docs/onnx-genai-performance-conformance.md @@ -0,0 +1,268 @@ +# ONNX GenAI workflow performance conformance + +Functional workflow conformance is necessary but not sufficient for release. A metadata-driven +workflow passes only when it matches or improves on the equivalent native implementation under +identical conditions. + +## Controlled comparison + +Every workflow/native pair must record and exactly match: + +- model artifact hash and weights; +- execution provider, device, precision, provider options, and runtime build; +- batch size and every input/state shape; +- prompt and generated-token count; +- sampling algorithm and request parameters, including seed and RNG offset; +- dense/shared/paged KV mode and capacity; +- graph-capture enablement and shape specialization; and +- warmup count, measured iterations, and synchronization timing points. + +Do not compare different quantization, kernels, KV layouts, sampling math, or capture settings. +Report p50/p95 and raw samples; the release gate uses the median after warmup. + +## Required scenarios + +1. **Decoder min-p:** decoder, last-token logits, request-parameterized min-p sampler, and + termination. The steady-state policy path must form and replay one same-device execution + island. +2. **Speculative accept:** proposer, verifier, acceptance/prefix policy, state correction, and + termination with a mostly accepted deterministic proposal fixture. +3. **Speculative reject:** the same artifacts and shapes with a deterministic rejection fixture. + Rollback/truncation tensor math must remain device-resident. +4. **Grammar boundary:** repeat decoder and speculative cases with a stateful grammar adapter. + Record the expected adapter island boundary and verify that pure ONNX work on each side is + still fused and captured independently. + +Run at batch 1 and one representative batched shape. Include fixed decode shapes and a bounded +shape transition case. Use the same accepted/evaluated token history for adaptive-K comparisons. + +## Required instrumentation + +The runtime benchmark record consumed by +`mobius.integrations.onnx_genai.performance.compare_performance` contains: + +- throughput and unit (`tokens_per_second` or `steps_per_second`); +- TTFT for generation workflows; +- peak device memory; +- host-to-device and device-to-host copy counts and bytes; +- explicit device synchronization count; +- ORT/composite session boundary count; +- kernel launch count; +- per-island component list, device, eligibility, capture count, replay count, and fallback reason. + +ORT profiling supplies node placement, kernel launches, session boundaries, and memcpy events. +Allocator/provider telemetry supplies peak device memory and stable-address failures. Island +diagnostics must be sampled after warmup and after measured replay; a merely eligible island is +not evidence of capture. + +The comparison gate rejects mismatched identities before evaluating performance. By default it +allows at most a 5% throughput, TTFT, or memory regression and permits no additional host/device +copies, explicit synchronizations, or session boundaries. Projects may tighten this threshold but +must not silently loosen it for individual model families. + +## Failure reporting + +Never mark a workflow ready from functional E2E alone. For every failed scenario, retain the +native and workflow records and report: + +1. the first divergent metric; +2. the responsible execution island and component boundary; +3. profiler evidence such as CPU placement, memcpy, synchronization, allocation, or unsupported + capture kernel; +4. whether the cause is producer structure, planner lowering, provider support, or runtime memory + management; and +5. the ordinary-execution fallback result. + +As of ONNX GenAI `8bacf8c`, an application-overridable sampler is rejected by island formation +before override resolution. This prevents the required decoder/sampler/termination capture +demonstration even when the selected implementation is pure same-device ONNX. Performance +acceptance remains blocked until override selection precedes island partitioning and the resolved +implementation is evaluated for purity and placement. + +## Fixed-capacity cache and FP8 runtime evidence + +Measured on an NVIDIA H200 with `onnxruntime-gpu==1.29.0`, CUDA execution +provider, against a real `--features static-cache` Qwen2 export. + +### Static cache — executes as specified + +| Check | Result | +| --- | --- | +| B=2 prefill vs per-row B=1 | max abs delta 0.0 | +| Scatter confined to `[0, nonpad_kv_seqlen)` | all rows; zero energy outside | +| Decode, divergent per-row cursors (row 0 advancing, row 1 finished), 3 steps | max abs delta 3.1e-6 | +| Compaction: the finished row reclaims its own last slot each step | verified, no writes past the reclaimed slot | + +Per-row cursor divergence during decode therefore works. Ragged *prefill* is +not claimed and was not measured: ONNX leaves the region between +`nonpad_kv_seqlen` and the query length undefined. + +### FP8 KV cache — not executable with the shipped kernel, and not an IMA + +Session creation fails during graph partitioning; no kernel is ever launched, +so this is not an illegal memory access: + +``` +transformer_memcpy.cc:253 IsNodeCompatibleWithProvider — +Provider type for GroupQueryAttention node 'node_GroupQueryAttention_9' is not set +``` + +ORT reports an unassigned node as an initialization exception rather than a type +error, so the cause was isolated by a three-way comparison of the same graph: + +| Variant | Result | +| --- | --- | +| No FP8 pass | loads and runs on CUDA EP | +| 14-input GQA with `k_scale`/`v_scale`, quantization attributes, **FLOAT** KV | loads and runs on CUDA EP | +| Same node with **FLOAT8E4M3FN** KV | node unassigned, session init fails | + +The rejection is the KV *type constraint* — `tensor(float8e4m3fn)` is not in the +CUDA `GroupQueryAttention` past/present type list in 1.29.0 — not the scale +input arity and not the attributes. The exported graph and its metadata are +well-formed and validate; only the local kernel is missing. + +### Canonical representation — lowering verified + +Against ONNX GenAI `52339e10`, with no `model.io` in any package and no port +contracts on any component that ships an artifact: + +| Check | Result | +| --- | --- | +| `validate_metadata` over the generated packages | 11/11 valid | +| `mobius_workflow_conformance` (engine executes each package) | 11/11 passed | +| `mobius_static_cache_workflow_executes` specifically | passed | + +The last row is the one that matters. The fixed-capacity decode path needs the +write cursor, the valid length and the per-layer buffer pairs; it resolved all +of them by lowering the workflow, which is what makes removing the second copy +safe rather than merely tidy. + +ONNX GenAI later moved the document-level invariants onto +`load_metadata_package`, so the package loader — the path the +`validate_metadata` binary and every on-disk consumer take — now enforces rules +that previously only ran for callers holding a parsed document. That is a strictness increase applied to the exact entry +point our fixtures go through, and all 11 were re-checked against it rather +than assumed to be unaffected. They pass because they carry no `model:` block +at all: a package with one serialized ABI has nothing for a coexistence rule to +find. + +### Why this document names one SHA and not several + +The ONNX GenAI branch we validate against is rebased as its own base moves, so +its commits are rewritten and the old hashes stop being reachable from any ref. +A citation to one of them does not merely go stale: it becomes unresolvable +once the unreferenced object is collected. Only `.github/workflows/main.yml` +pins a hash, because only CI needs to fetch an exact tree, and a bump there is +gated on re-running both suites. Everywhere else, changes are described by what +they do. + +This is not hypothetical. The pin was verified against the branch head, that +head was later rebased, and the pinned commit became reachable from no ref. +`actions/checkout` would have kept succeeding until the object was collected +and then failed with `reference is not a tree` on a commit no one could +inspect. + +Two checks are easy to confuse, and only one of them answers the question: + +- `git ls-remote | grep ` asks whether the commit is a ref *tip*. A + perfectly healthy pin fails this test the moment one more commit lands on + the branch, so a zero here means nothing on its own. +- `git merge-base --is-ancestor ` asks whether the commit is + *reachable*, which is what decides both whether `checkout` resolves it and + whether it can be collected. This is the check that matters. +- `gh api .../commits/` answers 200 for unreferenced objects, so it fails + open and must not be used for either question. + +Because that branch is rebased on someone else's schedule, reachability is not +ours to rely on. The tag `mobius-pr478-pin` in `justinchuby/onnx-genai` anchors +the exact commit this workflow pins, so the object survives any rebase of the +branch it came from. The workflow still pins the SHA rather than the tag: the +SHA is what makes the run reproducible, and the tag exists only to keep it +alive. Move the tag when the pin moves, and delete it when nothing pins it. + +### What the fixtures commit, and what they do not + +Only the metadata is committed. The graphs and the adapter weight file are a +deterministic function of `tests/generate_onnx_genai_validation_packages.py`, +and committing them stored 146 files and roughly 14 MB — including a 5 MB +weight blob — restating in unreadable bytes what the script already says. CI +already regenerated the whole tree and compared it, so the committed copies +were never the thing under test. + +Textproto was the other option and is the right one for ONNX GenAI's own +fixtures, which are 26 KB synthetic graphs. It is the wrong one here: these +carry real weight blobs, a text encoding would grow rather than shrink them, +and this repository does not use the protobuf APIs a textproto writer needs. + +Regeneration takes about four seconds for all eleven packages, so tests that +need a graph build them once per session through the +`materialized_workflow_packages` fixture. + +The change has a consequence worth stating plainly, because it decides where +CI must point: a package whose artifacts are absent does not validate. A +checkout of the committed tree alone fails with `component +'cache_length_update' artifact ... cannot be opened`. That is correct — a +workflow claims to be a complete description of something executable, and it +should not pass when the executable part is missing. Validation and +conformance therefore run against the regenerated tree, and the comparison +step additionally asserts that every artifact the committed metadata names was +really produced, so a generator that quietly stopped emitting one would be +caught rather than leaving an assertion with nothing to check. + +### The role a fixed-capacity cache cannot do without + +A component named in an `indexed_scatter` group's `write_indices_ports` or +`kv_length_ports` must declare a sequence role. Those two control ports are +consumed by exactly one thing — the driver that writes into a preallocated +buffer at an index — and that driver binds from the resolved decode ABI, every +field of which is found by role. A component handed the cursor while declaring +no role is therefore unresolvable as a decoder, and the package degrades to +inferring ports from shapes. + +Nothing upstream rejects that combination, and the reason is worth recording: +identifying the decoder *requires* a sequence role, so a component that omits +one is invisible to the check that would have caught it. A validator's +sole-decoder guard is additionally scoped to workflows with a single ONNX +component, which no package with policy graphs ever is — ours ship ten. The +contradiction is only visible to the producer, so +`test_a_scatter_bound_component_is_always_resolvable_as_a_decoder` asserts it on +both the built packages and the shipped fixtures. Measured: dropping the role +from `static_cache` still validates upstream, and fails here. + +### Where the transcription boundary actually falls + +A component backed by a shipped `.onnx` declares `ports.roles` and nothing else: +the artifact is authoritative for its ports, and `pipeline_admission` checks +declarations against the live session, so a YAML copy can only drift. + +Policy graphs are the exception, and the boundary was measured rather than +assumed. A workflow SSA value inherits its dtype, rank and request axis from the +port that produced it (`validation.rs` binds `value_contracts` from +`ports.outputs`), so a validator — which reads metadata without the artifacts — +has no other source. Dropping policy contracts made 4 of the 11 packages invalid +with `.when is row-wise but .value declares no request_aligned +batch_layout`. Those contracts type the workflow's own dataflow; they are not a +description of an external interface, and they stay. + +## Current measured baseline + +ONNX GenAI `8bacf8c` reports paired five-sample synthetic native/composite measurements over +100 iterations. These establish instrumentation, not Mobius producer readiness: + +| Device | Path | Workflow/native throughput | Warm TTFT workflow/native | Result | +| --- | --- | ---: | ---: | --- | +| CPU | decoder policy | 1.032 | not reported | within 5% | +| CPU | min-p policy | 0.962 | not reported | within 5% | +| H200, ORT 1.28 CUDA | decoder policy | 0.903 | 3.03/17.93 ms | throughput fails | +| H200, ORT 1.28 CUDA | min-p policy | 0.957 | 4.36/3.64 ms | TTFT fails | + +The CUDA islands captured once and replayed 503 times; the speculative verifier/policy fixture +also captured and replayed. Cold workflow startup remained 467/49 ms for decoder and 231/18 ms +for min-p because the first workflow run discovers output extents and constructs stable bindings. +The remaining steady-state decoder throughput gap requires ORT kernel/provider profiling; it is +not explained away by functional parity. Min-p throughput is within the default bar, but its warm +TTFT is not. + +These measurements do not cover the real Mobius package, KV service mode, per-row serving, or the +application-overridable sampler path. Those cases remain blocked/not demonstrated and must produce +records accepted by the comparison gate before PR readiness. diff --git a/docs/onnx-genai-workflows.md b/docs/onnx-genai-workflows.md new file mode 100644 index 000000000..a3fadf438 --- /dev/null +++ b/docs/onnx-genai-workflows.md @@ -0,0 +1,463 @@ +# ONNX GenAI workflow metadata + +Mobius emits the concise public workflow source form. The only control primitives are +`sequence`/`steps`, `invoke`, `loop`, `branch`, and `emit`. + +## One representation + +`pipeline.workflow` is where a package describes itself, and it is the only +place. That is true of a three-graph vision-language package, and it is equally +true of a bare single-file decoder — the single-file case is a one-component +workflow, not a different kind of document with its own top-level keys. `model` +carries package-wide geometry and capabilities; it never carries a port ABI. + +The reason is not tidiness. Two writable statements of the same fact are a +defect whatever they contain, because nothing forces them to agree and a reader +of one never learns that the other said something else. A runtime that wants an +optimized single-graph path gets it by *lowering* the one-component workflow, +which is a derivation and cannot disagree with its source. + +For that to work the workflow carries what the ONNX artifact cannot say. +Every ONNX component declares: + +* `ports.roles` — what the component *does* with a value bound to a port. An + invocation records which SSA value reaches a port, not whether that port is + tokens, a mask or logits, and recovering the difference from spelling is the + name-guessing this format refuses everywhere else. Mobius mints these port + names in its own task builders, so it states the mapping rather than infers + it: `input_ids`→`token_ids`, `inputs_embeds`→`inputs_embeds`, + `attention_mask`→`attention_mask`, `position_ids`→`position_ids`, + `logits`→`logits`, `last_hidden_state`→`hidden_states`, + `encoder_hidden_states`→`encoder_hidden_states`, + `audio_features`→`audio_features`. A port outside that vocabulary carries no + role, because a workflow that guesses is worse than one that stays silent. + +The artifact remains authoritative for port names, dtype, rank, shape, and +opset imports. An ONNX-backed component therefore omits duplicated +`ports.inputs` / `ports.outputs` contracts. Native adapters and policy +components still declare contracts when those contracts type workflow SSA +values and no ONNX artifact exists to supply them. + +State ports need no role entry — the group that carries them already names each +`(input, output)` pair — but they do carry two facts nothing else can recover: + +* `role` (`key` / `value` / `combined`) — a layer's key buffer and its value + buffer are semantically distinct even when their names, dtype, or shape do + not reveal which is which. +* `layer` — a cell's label is producer-chosen and sorts lexicographically, so + `cache_10` precedes `cache_2`; pairing per-layer buffers positionally would + silently transpose two layers' caches. + +When a group has several aliases of one role, every alias declares `layer`; +the same layer number on its key and value aliases pairs them. Geometry remains +per-port: different layers may have different KV head counts, and a layer may +have different K and V head counts. The runtime reads each actual ONNX port +independently; `layer` orders aliases but never asserts equal shapes. A +recurrent or convolution cache has no key/value halves and does not invent +those roles. + +The top-level `schema_version` also versions workflow syntax. The workflow +manifest contains only non-artifact facts such as adapter ABI versions and +capabilities: it does not duplicate an `ir_version`, and it does not copy a +package-wide ONNX opset map when every component artifact already carries its +own exact imports. + +`tests/canonical_workflow_contract_test.py` holds this invariant: it asks the +same shape-agnostic questions of dynamic, static-cache, FP8, heterogeneous and +composite packages, and of every checked-in fixture, so a new feature cannot +grow a private top-level block unnoticed. + +## Structural execution frequency + +- Root `steps` run once per invocation. +- `loop.setup` runs once whenever the loop is entered. +- `loop.steps` run once per iteration. +- Root suffix steps run once after the loop. +- A session-state initializer runs once when that session cell is created. +- Artifact loading and session restoration are runtime lifecycle operations. + +There are no phases, strategies, `run_once`, or execution-frequency flags. + +## Surface and lowered forms + +Serialized metadata uses logical names and concise carries: + +```yaml +state: + cache: + contract: { dtype: float16, rank: 4, shape: [batch, heads, cache_sequence, width] } + scope: invocation + initializer: cache.initial + recurrence: { kind: bounded, axis: 2, max: max_context } +steps: + - kind: loop + setup: [] + steps: + - kind: invoke + component: decoder + inputs: { past_key: cache } + outputs: { present_key: cache.next } + condition: continue + max_iterations: max_output_tokens + carried: [{ cell: cache, next: cache.next }] +``` + +The loader compiles this source form to lexical SSA, branch phi values, loop read/write +versions, and linear effect tokens. Compiler-generated names are not serialized. + +Pure ONNX components have no effect declaration. RNG, caches, counters, and policy state +are ordinary explicit tensors. Effects are reserved for external mutation such as streams, +session mutation, telemetry, and stateful adapter ABIs. The loader threads and joins those +effects through structured control flow. + +ONNX artifacts are authoritative for component input/output names, dtypes, ranks, and +shapes. Mobius declares ports only for adapters or constraints not represented by ONNX. +Resource placement normally determines transfers; workflows do not author transfer steps. + +Versioned component contracts add semantic role mappings without changing execution: + +```yaml +contract: + id: onnx-genai.token-sampler + version: "1" + bindings: { logits: logits, token: token_ids } + parameters: { mode: greedy } +``` + +## Execution islands and graph capture + +Serialized component boundaries preserve policy modularity; they do not require separate ORT +sessions, kernel launches, host round trips, or CUDA Graphs. After artifact loading, override +selection, validation, and SSA lowering, the generic planner may link adjacent pure ONNX invokes +on the same device into an execution island (or equivalent linked composite session). Intermediate +SSA values stay device-resident and optimizer-visible. Decoder logits processors, sampler, +state-update math, and termination predicates can therefore execute as one island. + +Island formation is derived rather than serialized. It ends at structured host control, device +changes, explicit external effects, or stateful host adapters such as grammar clone/commit. +Application overrides are resolved before planning; a selected pure same-device ONNX replacement +remains eligible under the same rules. + +CUDA Graph capture is evaluated per island and concrete shape signature. Eligibility requires: + +- static or bounded runtime-specialized shapes; +- stable device-resident input, output, and state addresses; +- no host data-dependent allocation or control; +- execution-provider support for every selected kernel; and +- explicit tensor-threaded counter RNG seed/offset state. + +The runtime should warm up bindings, capture an eligible equal-shape execution, and replay later +matches. Shape changes, unsupported kernels, allocator instability, or capture errors fall back +to ordinary island execution without changing workflow semantics. Diagnostics should identify +the island's component list and device, eligibility decision, capture/replay counters, and exact +fallback reason. + +Release benchmarking and the native-equivalence acceptance gate are defined in +[`onnx-genai-performance-conformance.md`](onnx-genai-performance-conformance.md). + +### Decoder KV boundary + +Normal decode carries each decoder `present` tensor directly to the corresponding +next-iteration `past` input. Mobius does not emit a generic `kv_update.onnx`. Physical +shared/paged allocation, slots, append, compaction, and in-place mutation belong to the +runtime's model-agnostic KV service. An ONNX state-update component is used only when +semantic tensor math is required, such as accepted-prefix truncation, dense gather, +rollback, or format conversion. + +### Request-parameterized sampling + +The stochastic sampler ABI accepts `temperature`, `top_k`, `top_p`, `min_p`, `seed`, RNG +offset, and a grammar-allow mask as typed inputs. They are request/workflow values rather +than artifact constants, so ordinary option changes do not rebuild the sampler. After grammar +masking and temperature scaling, positive min-p is evaluated as +`scaled_logit >= max_scaled_logit + log(min_p)`; a non-positive value disables the filter. +This preserves fixed `[B,V]` shapes. The generated component is marked +`application_overridable`; an application may +replace it with another implementation of the same versioned port ABI for fundamentally +custom sampling. + +## Fixed-capacity (static) KV cache + +A static-cache export does not grow its KV tensors. Each layer owns a +preallocated `[batch, capacity, kv_hidden]` buffer, and every step writes into +it at a per-row cursor with `TensorScatter` on axis 1. Two integer control +ports drive that write: + +| port | shape | meaning | +| --- | --- | --- | +| `write_indices` | `[batch]` int64 | first slot this step writes, per row | +| `nonpad_kv_seqlen` | `[batch]` int64 | number of valid slots **after** the write | + +The buffers themselves are `key_cache.{layer}` / `value_cache.{layer}` in and +`updated_key_cache.{layer}` / `updated_value_cache.{layer}` out. None of this is +published in a second place: the workflow declares it operationally and that is +the only declaration. The buffers are loop cells with +`recurrence: {kind: invariant}` (they do not grow), the capacity is a +`package.cache_capacity` literal workflow input, and the state service publishes +an `indexed_scatter` update discipline naming the write cursor, the capacity, +the per-component port that carries the cursor (`write_indices_ports`) and the +per-component port that carries the valid length (`kv_length_ports`). The +component those ports belong to declares both, so a consumer resolves the whole +scatter ABI without opening the artifact. + +Because the write cursor and the logical length are the same quantity, both name +the single carried `cache_lengths` cell rather than introducing a second +never-consumed carry. A finished row's length stops advancing, so the slot it +last wrote falls outside its valid prefix and is reclaimed by its next write. +That is deliberate: it is what makes a fixed-capacity buffer safe to keep +serving a batch in which rows finish at different times. + +**Not claimed:** ragged prefill. ONNX leaves the region between +`nonpad_kv_seqlen` and the query length undefined, and the workflow scatters one +same-length chunk per row, so a prefill in which rows have different prompt +lengths is outside the contract. Per-row cursor divergence during *decode* is +fully supported. + +### Heterogeneous caches + +A model may mix disciplines. Gemma 4 keeps its sliding-window layers on a +growing rank-4 BNSH cache while its full-attention layers use a fixed-capacity +rank-3 buffer, and its KV-shared suffix owns no buffer at all. These surface as +separate state-service groups with their own sequence axis, layout, aliasing +rule and update discipline, and only the layers that actually own a buffer bind +ports in a group. Collapsing them into one group would invite a runtime to apply +sliding-window eviction to the global layers, or to allocate caches for layers +that borrow one. + +## FP8 KV cache + +FP8 KV storage is a property of the *attention operator*, not of the cache +tensor alone: the scales that dequantize the cache on read are node inputs +(`k_scale`/`v_scale` at `GroupQueryAttention` slots 12 and 13). The published +contracts therefore repeat whatever dtype the graph declares — `float8_e4m3fn` +— because a runtime sizes the buffers from them, and reporting the model's +compute dtype instead would allocate twice the bytes the model reads. + +Two consequences: + +* `--features static-cache,fp8-kv-cache` is refused at build time. A + static-cache graph scatters into buffers read by `ai.onnx` `Attention`, which + has no scale inputs, so there is no operator that could dequantize an FP8 + buffer. Emitting one anyway would declare FP8 over bytes read as float16. +* A package whose cache is FP8 is valid even where no local kernel can execute + it. The dtype describes the exported graph; kernel availability is a property + of the runtime that happens to be installed. + +## Compact examples + +### Encoder embeddings + +A bidirectional encoder — BERT, ESM-2, ProtBert — is not generative. It reads +the whole sequence once and returns one hidden vector per position, so its +workflow has no loop, no carried state, no KV cache and no sampler. + +```yaml +profiles: + embedding: + kind: embedding + outputs: { last_hidden_state: last_hidden_state } + pooling: { kind: mean, source: last_hidden_state, + mask: request.attention_mask, time_axis: 1, feature_axis: 2 } + batch_invariance: row_independent +steps: + - kind: invoke + component: encoder + inputs: { input_ids: request.input_ids, attention_mask: request.attention_mask } + outputs: { last_hidden_state: encoder.last_hidden_state } + - kind: emit + value: encoder.last_hidden_state + output: last_hidden_state + mode: replace +``` + +The declared inputs are read from the artifact, not from the task signature: +ESM-2 has no token-type embedding, so its graph exposes only `input_ids` and +`attention_mask`, while ProtBert's also exposes `token_type_ids`. + +`batch_invariance: row_independent` is claimed only when the graph consumes an +attention mask, because that is what makes a row's values independent of the +width the batch happened to be padded to. It is also what makes +`pooling.kind: mean` well defined — a reader reduces each row over its own +valid region rather than over the padded extent. + +### Decoder + +```yaml +steps: + - kind: invoke + component: initialize_decoder + inputs: { tokens: prompt_tokens } + outputs: { cache: cache.initial, mask: mask.initial } + - kind: loop + setup: [] + steps: + - kind: invoke + component: decoder + inputs: { tokens: token, cache: cache, attention_mask: mask } + outputs: { logits: logits, present: cache.next } + - kind: invoke + component: sampler + inputs: { logits: logits } + outputs: { token: token.next } + - kind: invoke + component: update_mask + inputs: { current: mask } + outputs: { next: mask.next } + - kind: emit + value: token.next + output: tokens + mode: append + condition: continue + max_iterations: max_output_tokens + carried: + - { cell: cache, next: cache.next } + - { cell: mask, next: mask.next } + - { cell: token, next: token.next } +``` + +### Vision-language + +```yaml +steps: + - kind: invoke + component: image_preprocess + inputs: { encoded: image } + outputs: { pixel_values: pixels, grid: grid } + - kind: invoke + component: vision_encoder + inputs: { pixel_values: pixels, grid: grid } + outputs: { features: image_features } + - kind: invoke + component: embedding + inputs: { tokens: prompt_tokens, image_features: image_features } + outputs: { embeddings: prompt_embeddings } + - kind: loop + setup: [] + steps: + - kind: invoke + component: decoder + inputs: { embeddings: prompt_embeddings, cache: cache } + outputs: { logits: logits, present: cache.next } + condition: continue + max_iterations: max_output_tokens + carried: [{ cell: cache, next: cache.next }] +``` + +Preprocessing, vision, and initial embedding are root-prefix steps and therefore run once. +Only the decoder body runs per generated token. + +### Diffusion + +```yaml +steps: + - kind: loop + setup: + # Schedule, timesteps, and scalar scales are ONNX constant components, so + # the solver reads its sigmas from the graph instead of runtime config. + - { kind: invoke, component: diffusion_schedule, inputs: {}, + outputs: { schedule: diffusion.schedule } } + - { kind: invoke, component: diffusion_timesteps, inputs: {}, + outputs: { schedule: diffusion.timesteps } } + - { kind: invoke, component: latent_row_shape, inputs: {}, + outputs: { shape: diffusion.latent_row_shape } } + # Counter-based RNG: one private stream per row, seeded by the request. + - { kind: invoke, component: latent_noise, + inputs: { seed: request.seed, offset: package.rng_offset, + row_shape: diffusion.latent_row_shape }, + outputs: { noise: diffusion.noise, next_offset: diffusion.rng_offset } } + - { kind: invoke, component: text_encoder, + inputs: { input_ids: request.input_ids }, + outputs: { encoder_hidden_states: conditioning.hidden_states } } + - { kind: invoke, component: text_encoder, + inputs: { input_ids: request.negative_input_ids }, + outputs: { encoder_hidden_states: conditioning.unconditional } } + - { kind: invoke, component: history_initializer, + inputs: { reference: diffusion.noise }, + outputs: { zeros: diffusion.initial_history } } + steps: + - { kind: invoke, component: schedule_lookup, + inputs: { schedule: diffusion.timesteps, step: loop.iteration }, + outputs: { timestep: diffusion.timestep } } + # Classifier-free guidance is two denoiser invocations, not a doubled batch. + - { kind: invoke, component: denoiser, + inputs: { sample: latent_state, timestep: diffusion.timestep, + encoder_hidden_states: conditioning.unconditional }, + outputs: { noise_pred: denoiser.unconditional } } + - { kind: invoke, component: denoiser, + inputs: { sample: latent_state, timestep: diffusion.timestep, + encoder_hidden_states: conditioning.hidden_states }, + outputs: { noise_pred: denoiser.conditional } } + - { kind: invoke, component: guidance_combine, + inputs: { unconditional: denoiser.unconditional, + conditional: denoiser.conditional, + scale: request.guidance_scale }, + outputs: { estimate: denoiser.estimate } } + - { kind: invoke, component: solver_step, + inputs: { sample: latent_state, step: loop.iteration, + schedule: diffusion.schedule, estimate: denoiser.estimate, + history: history }, + outputs: { next_state: latent.body, next_history: history.body } } + - { kind: emit, value: denoiser.estimate, output: noise_estimate, mode: append } + - { kind: emit, value: latent.body, output: latent_trajectory, mode: append } + max_iterations: request.max_iterations + iteration: + value: loop.iteration + contract: { dtype: int64, rank: 1, shape: [batch] } + carried: + - { cell: latent_state, next: latent.body } + - { cell: history, next: history.body } + - { kind: invoke, component: tensor_scale, + inputs: { tensor: latent_state, scale: diffusion.decoder_scale }, + outputs: { scaled: diffusion.decoder_input } } + - { kind: invoke, component: vae_decoder, + inputs: { latent: diffusion.decoder_input }, + outputs: { image: vae.image } } + - { kind: emit, value: latent_state, output: latent, mode: replace } + - { kind: emit, value: vae.image, output: image, mode: replace } + - { kind: emit, value: diffusion.rng_offset, output: rng_offset, mode: replace } +``` + +Conditioning, latent sampling, and VAE decoding are structural: they sit outside the +loop body and therefore run once. Only the denoiser, guidance, and solver run per step. + +Carried state is declared under `state`, where each cell names the value that +initializes it (`latent_state` ← `diffusion.noise`, `history` ← +`diffusion.initial_history`). Nothing about the loop is expressed as a phase or a +strategy; frequency falls out of where a step is nested. + +#### Policy components + +The producer emits the non-network parts of the pipeline as real ONNX graphs so the +runtime never has to reimplement scheduler math: + +| Component | Ports | Purpose | +|---|---|---| +| `diffusion_schedule` / `diffusion_timesteps` | → `schedule` | Sigma schedule and timestep table baked in as initializers, so the solver reads its constants from the graph rather than from runtime config. | +| `decoder_input_scale` | → `value` | Build-time scalar such as `1 / vae_scaling_factor`. | +| `latent_row_shape` | → `shape` | Per-row latent shape consumed by the noise sampler. | +| `schedule_lookup` | `schedule`, `step` → `timestep` | Gathers this step's timestep from the table. | +| `latent_noise` (`onnx-genai.counter-rng@1`) | `seed`, `offset`, `row_shape` → `noise`, `next_offset` | Counter-based Box–Muller normals. Each row draws only from its own seed, so a row's noise does not depend on its batch position, and the advanced counter is returned so the caller can persist it. | +| `history_initializer` | `reference` → `zeros` | Shape-following zero initializer for solver history. | +| `guidance_combine` (`onnx-genai.guidance-combine@1`) | `unconditional`, `conditional`, `scale` → `estimate` | `uncond + scale * (cond - uncond)` with a per-row scale. | +| `tensor_scale` | `tensor`, `scale` → `scaled` | Applies `init_noise_sigma` or the VAE scale. Emitted only when the factor is not 1.0, so no identity multiplies appear in the graph. | +| `solver_step` (`onnx-genai.solver-step@1`) | `sample`, `estimate`, `history`, `step`, `schedule` → `next_state`, `next_history` | The scheduler update. Euler binds `estimate` to its `derivative` port and leaves the history ports unbound; DPM++ 2M carries the previous `x0` estimate as ordinary loop state and masks itself down to first order on the first and last steps. | + +Only components with a stable cross-runtime meaning carry a `contract` id; the constant +and reshaping helpers are plain ONNX graphs identified structurally by their ports. + +#### Why guidance is two invocations + +Classifier-free guidance is conventionally implemented by concatenating the +conditional and unconditional batches and running the denoiser once. The workflow IR +declares batch semantics through `batch_layout`, which has no kind meaning "k rows per +request", so a doubled batch cannot be described truthfully — the runtime would be +unable to attribute a row to a request. Two invocations of the same component plus an +explicit `guidance_combine` keeps every tensor `request_aligned` and leaves batching +decisions to the runtime. + +#### Append-mode outputs + +`mode: append` concatenates chunks along the last axis, so an accumulating output must +name that axis with a symbol of its own (`noise_estimate_width`, not `width`). +Reusing the per-step symbol would bind the same symbol to both the chunk width and the +accumulated width and fail validation. diff --git a/pyproject.toml b/pyproject.toml index 756b03aa5..82c0a8f72 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -19,6 +19,7 @@ dependencies = [ "onnx_ir>=1.0.0", "onnx-shape-inference>=0.3.1", "onnxscript>=0.7.1", + "rfc8785", "safetensors", "torch>=2.1.0", "tqdm", diff --git a/scripts/benchmark_muse_native.py b/scripts/benchmark_muse_native.py new file mode 100644 index 000000000..6e4c550f4 --- /dev/null +++ b/scripts/benchmark_muse_native.py @@ -0,0 +1,207 @@ +#!/usr/bin/env python3 +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT License. + +"""Benchmark the native ORT GenAI path for the paired Muse workflow test.""" + +from __future__ import annotations + +import argparse +import ctypes +import hashlib +import json +import os +import statistics +import time +from pathlib import Path +from typing import Any + +os.environ["ORT_ENABLE_CUDNN_FLASH_ATTENTION"] = "0" + +import onnxruntime as ort +import onnxruntime_genai as og + +_CUDART = ctypes.CDLL("libcudart.so") + + +def _synchronize_cuda() -> None: + status = _CUDART.cudaDeviceSynchronize() + if status: + raise RuntimeError(f"cudaDeviceSynchronize failed with status {status}") + + +def _sha256(path: Path) -> str: + digest = hashlib.sha256() + with path.open("rb") as handle: + for block in iter(lambda: handle.read(1024 * 1024), b""): + digest.update(block) + return digest.hexdigest() + + +def _prompt(tokenizer: Any, text: str, *, include_image: bool) -> str: + content: Any = [{"type": "image"}, {"type": "text", "text": text}] + if not include_image: + content = text + messages = [{"role": "user", "content": content}] + return tokenizer.apply_chat_template(json.dumps(messages), add_generation_prompt=True) + + +def _run_once( + model: Any, + processor: Any, + prompt: str, + image: Any, + *, + prompt_tokens: int, + max_new_tokens: int, + decode_skip: int, + sampling: dict[str, Any], +) -> dict[str, Any]: + inputs = processor(prompt, images=image) if image is not None else processor(prompt) + started = time.perf_counter() + params = og.GeneratorParams(model) + params.set_search_options( + max_length=prompt_tokens + max_new_tokens, + min_length=prompt_tokens + max_new_tokens, + do_sample=bool(sampling["do_sample"]), + temperature=float(sampling["temperature"]), + top_k=int(sampling["top_k"]), + top_p=float(sampling["top_p"]), + ) + generator = og.Generator(model, params) + generator.set_inputs(inputs) + times: list[float] = [] + tokens: list[int] = [] + for _ in range(max_new_tokens): + generator.generate_next_token() + _synchronize_cuda() + times.append(time.perf_counter() - started) + tokens.append(int(generator.get_next_tokens()[0])) + if len(times) <= decode_skip: + raise ValueError("max_new_tokens must exceed decode_skip") + decode_seconds = times[-1] - times[decode_skip - 1] + decode_tokens = len(times) - decode_skip + return { + "token_ids": tokens, + "ttft_ms": times[0] * 1000, + "decode_tokens": decode_tokens, + "decode_seconds": decode_seconds, + "throughput_tok_s": decode_tokens / decode_seconds, + } + + +def main() -> int: + parser = argparse.ArgumentParser() + parser.add_argument("--model", required=True, type=Path) + parser.add_argument( + "--config", + type=Path, + default=Path("benchmarks/muse_workflow_h200.json"), + ) + parser.add_argument("--output", required=True, type=Path) + args = parser.parse_args() + config = json.loads(args.config.read_text()) + workload = config["workload"] + sampling = config["sampling"] + if config["release_gate"] != { + "clean_landed_runtime": True, + "exact_token_parity": True, + "cuda_graph_required": True, + "min_throughput_ratio": 0.99, + }: + raise ValueError("paired Muse benchmark requires the 0.99x hard release gate") + if config["runtime"]["cudnn_flash_attention"]: + raise ValueError("paired Muse benchmark requires cuDNN Flash Attention disabled") + + if ort.__version__ != config["runtime"]["onnxruntime"]: + raise RuntimeError( + f"ORT version mismatch: {ort.__version__} != {config['runtime']['onnxruntime']}" + ) + providers = ort.get_available_providers() + if "CUDAExecutionProvider" not in providers: + raise RuntimeError(f"CUDAExecutionProvider unavailable: {providers}") + + model = og.Model(str(args.model)) + processor = model.create_multimodal_processor() + tokenizer = og.Tokenizer(model) + image_path = workload["image"] + image = og.Images.open(image_path) if image_path else None + prompt = ( + _prompt(tokenizer, workload["prompt"], include_image=True) + if image is not None + else workload["rendered_prompt"] + ) + encoded_prompt = tokenizer.encode(prompt) + prompt_ids = json.loads(Path(workload["prompt_ids_file"]).read_text()) + if encoded_prompt != prompt_ids: + raise RuntimeError("native tokenizer output differs from the canonical prompt IDs") + prompt_tokens = len(encoded_prompt) + if prompt_tokens != int(workload["prompt_tokens"]): + raise RuntimeError( + f"prompt token count mismatch: {prompt_tokens} != {workload['prompt_tokens']}" + ) + request_max_length = prompt_tokens + int(workload["max_new_tokens"]) + if request_max_length != int(workload["request_max_length"]): + raise RuntimeError( + "request max length mismatch: " + f"{request_max_length} != {workload['request_max_length']}" + ) + + for _ in range(int(workload["warmups"])): + _run_once( + model, + processor, + prompt, + image, + prompt_tokens=prompt_tokens, + max_new_tokens=int(workload["max_new_tokens"]), + decode_skip=int(workload["decode_skip"]), + sampling=sampling, + ) + + runs = [ + _run_once( + model, + processor, + prompt, + image, + prompt_tokens=prompt_tokens, + max_new_tokens=int(workload["max_new_tokens"]), + decode_skip=int(workload["decode_skip"]), + sampling=sampling, + ) + for _ in range(int(workload["runs"])) + ] + reference = runs[0]["token_ids"] + if any(run["token_ids"] != reference for run in runs[1:]): + raise RuntimeError("native greedy token output changed across measured runs") + + record = { + "kind": "native", + "config": config, + "environment": { + "onnxruntime": ort.__version__, + "onnxruntime_genai": getattr(og, "__version__", "unknown"), + "providers": providers, + }, + "package": { + "metadata_sha256": _sha256(args.model / "inference_metadata.yaml"), + "genai_config_sha256": _sha256(args.model / "genai_config.json"), + }, + "prompt_tokens": prompt_tokens, + "metrics": { + "ttft_ms": statistics.median(run["ttft_ms"] for run in runs), + "throughput_tok_s": statistics.median(run["throughput_tok_s"] for run in runs), + "ttft_ms_runs": [run["ttft_ms"] for run in runs], + "throughput_tok_s_runs": [run["throughput_tok_s"] for run in runs], + }, + "token_ids": reference, + } + args.output.parent.mkdir(parents=True, exist_ok=True) + args.output.write_text(json.dumps(record, indent=2) + "\n") + print(json.dumps(record["metrics"], indent=2)) + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/scripts/benchmark_muse_workflow.py b/scripts/benchmark_muse_workflow.py new file mode 100644 index 000000000..ac3898c33 --- /dev/null +++ b/scripts/benchmark_muse_workflow.py @@ -0,0 +1,216 @@ +#!/usr/bin/env python3 +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT License. + +"""Run the paired ONNX GenAI workflow benchmark for Muse Glimmer.""" + +from __future__ import annotations + +import argparse +import hashlib +import json +import os +import re +import subprocess +from pathlib import Path +from typing import Any + + +def _sha256(path: Path) -> str: + digest = hashlib.sha256() + with path.open("rb") as handle: + for block in iter(lambda: handle.read(1024 * 1024), b""): + digest.update(block) + return digest.hexdigest() + + +def _runtime_worktree(repo: Path) -> dict[str, Any]: + status = subprocess.run( + ["git", "-C", str(repo), "status", "--porcelain=v1", "--untracked-files=all"], + check=True, + text=True, + capture_output=True, + ).stdout + diff = subprocess.run( + ["git", "-C", str(repo), "diff", "--binary", "HEAD"], + check=True, + capture_output=True, + ).stdout + files = {} + for line in status.splitlines(): + relative = line[3:].split(" -> ")[-1] + path = repo / relative + if path.is_file(): + files[relative] = _sha256(path) + return { + "dirty": bool(status), + "status": status.splitlines(), + "diff_sha256": hashlib.sha256(diff).hexdigest(), + "file_sha256": files, + } + + +def main() -> int: + parser = argparse.ArgumentParser() + parser.add_argument("--model", required=True, type=Path) + parser.add_argument("--runner", required=True, type=Path) + parser.add_argument( + "--runtime-repo", + type=Path, + default=Path(".contract-schema-latest"), + ) + parser.add_argument("--allow-dirty-runtime", action="store_true") + parser.add_argument( + "--config", + type=Path, + default=Path("benchmarks/muse_workflow_h200.json"), + ) + parser.add_argument("--output", required=True, type=Path) + args = parser.parse_args() + + config: dict[str, Any] = json.loads(args.config.read_text()) + workload = config["workload"] + sampling = config["sampling"] + if config["release_gate"] != { + "clean_landed_runtime": True, + "exact_token_parity": True, + "cuda_graph_required": True, + "min_throughput_ratio": 0.99, + }: + raise ValueError("paired Muse benchmark requires the 0.99x hard release gate") + runtime_head = subprocess.run( + ["git", "-C", str(args.runtime_repo), "rev-parse", "HEAD"], + check=True, + text=True, + capture_output=True, + ).stdout.strip() + expected_head = config["runtime"]["onnxruntime_genai_commit"] + if runtime_head != expected_head: + raise ValueError( + f"runtime source is {runtime_head}, paired benchmark requires {expected_head}" + ) + runtime_worktree = _runtime_worktree(args.runtime_repo) + if runtime_worktree["dirty"] and not args.allow_dirty_runtime: + raise ValueError( + "runtime source has uncommitted changes; commit them or pass " + "--allow-dirty-runtime for diagnostic-only evidence" + ) + git_index = Path( + subprocess.run( + [ + "git", + "-C", + str(args.runtime_repo), + "rev-parse", + "--path-format=absolute", + "--git-path", + "index", + ], + check=True, + text=True, + capture_output=True, + ).stdout.strip() + ) + if args.runner.stat().st_mtime < git_index.stat().st_mtime: + raise ValueError( + "runner predates the pinned runtime checkout; rebuild it before benchmarking" + ) + if config["runtime"]["cudnn_flash_attention"]: + raise ValueError("paired Muse benchmark requires cuDNN Flash Attention disabled") + prompt_ids_path = Path(workload["prompt_ids_file"]) + prompt_ids = json.loads(prompt_ids_path.read_text()) + if len(prompt_ids) != int(workload["prompt_tokens"]): + raise ValueError("prompt_ids_file length must equal prompt_tokens") + if int(workload["prompt_tokens"]) + int(workload["max_new_tokens"]) != int( + workload["request_max_length"] + ): + raise ValueError("request_max_length must equal prompt_tokens + max_new_tokens") + if workload["stop_on_eos"] is not False: + raise ValueError("paired Muse benchmark requires stop_on_eos=false") + if sampling != { + "algorithm": "greedy", + "do_sample": False, + "temperature": 1.0, + "top_k": 1, + "top_p": 1.0, + "seed": 0, + }: + raise ValueError("workflow runner currently requires the paired greedy policy") + + command = [ + str(args.runner), + "--model", + str(args.model), + "--pipeline", + "--backend", + "ort", + "--ep", + "cuda", + "--steady", + "--tokens", + str(workload["max_new_tokens"]), + "--warmups", + str(workload["warmups"]), + "--runs", + str(workload["runs"]), + "--decode-skip", + str(workload["decode_skip"]), + "--prompt-ids", + str(prompt_ids_path), + ] + if workload["image"]: + command.extend(["--image", workload["image"]]) + environment = os.environ.copy() + environment["ONNX_GENAI_CUDA_GRAPH"] = "1" + environment["ORT_ENABLE_CUDNN_FLASH_ATTENTION"] = "0" + completed = subprocess.run( + command, + check=True, + text=True, + capture_output=True, + env=environment, + ) + output = completed.stdout + completed.stderr + median = re.search( + r"steady_median: prefill=([0-9.]+) ms decode=([0-9.]+) ms/token " + r"throughput=([0-9.]+) tok/s", + output, + ) + tokens = re.search(r"generated_token_ids: (\[[^\n]+\])", output) + if median is None or tokens is None: + raise RuntimeError(f"workflow runner output is incomplete:\n{output}") + + record = { + "kind": "workflow", + "config": config, + "package": { + "metadata_sha256": _sha256(args.model / "inference_metadata.yaml"), + "genai_config_sha256": _sha256(args.model / "genai_config.json"), + }, + "runtime": { + "source_head": runtime_head, + "runner": str(args.runner.resolve()), + "runner_sha256": _sha256(args.runner), + "worktree": runtime_worktree, + "environment": { + "ONNX_GENAI_CUDA_GRAPH": "1", + "ORT_ENABLE_CUDNN_FLASH_ATTENTION": "0", + }, + }, + "metrics": { + "ttft_ms": float(median.group(1)), + "decode_ms_per_token": float(median.group(2)), + "throughput_tok_s": float(median.group(3)), + }, + "token_ids": json.loads(tokens.group(1)), + "command": command, + "runner_output": output, + } + args.output.parent.mkdir(parents=True, exist_ok=True) + args.output.write_text(json.dumps(record, indent=2) + "\n") + print(json.dumps(record["metrics"], indent=2)) + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/scripts/compare_workflow_performance.py b/scripts/compare_workflow_performance.py new file mode 100755 index 000000000..f45d9a56d --- /dev/null +++ b/scripts/compare_workflow_performance.py @@ -0,0 +1,43 @@ +#!/usr/bin/env python3 +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT License. + +"""Compare equivalent native and metadata-driven workflow benchmark records.""" + +from __future__ import annotations + +import argparse +import json +from pathlib import Path + +from mobius.integrations.onnx_genai.performance import compare_performance + + +def main() -> int: + parser = argparse.ArgumentParser() + parser.add_argument("--workflow", required=True, type=Path) + parser.add_argument("--native", required=True, type=Path) + parser.add_argument("--max-regression-percent", type=float, default=5.0) + args = parser.parse_args() + + with args.workflow.open(encoding="utf-8") as handle: + workflow = json.load(handle) + with args.native.open(encoding="utf-8") as handle: + native = json.load(handle) + result = compare_performance( + workflow, + native, + max_regression_fraction=args.max_regression_percent / 100, + ) + for observation in result.observations: + print(f"MEASURED: {observation}") + for failure in result.failures: + print(f"FAIL: {failure}") + if result.passed: + print("PASS: workflow performance is competitive with native") + return 0 + return 1 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/src/mobius/__init__.py b/src/mobius/__init__.py index 9e3438edd..edaec18a4 100644 --- a/src/mobius/__init__.py +++ b/src/mobius/__init__.py @@ -5,6 +5,18 @@ __all__ = [ "ArchitectureConfig", + "AdapterApplication", + "AdapterArtifact", + "AdapterBatchSelection", + "AdapterSlotSelection", + "AdapterSelectionTensors", + "AdapterServiceOptions", + "AdapterSource", + "AdapterTarget", + "AdapterTargetDescriptor", + "AdapterTargetManifest", + "AdapterTargetSlice", + "AdapterWeights", "AudioConfig", "BaseModelConfig", "CausalLMConfig", @@ -38,15 +50,20 @@ "WorldModelTask", "YolosConfig", "apply_weights", + "adapter_source_from_onnx_adapter", "build", "build_context", "build_diffusers_pipeline", "build_from_gguf", "build_from_module", "build_from_nemo", + "compose_adapter_deltas", "components", "ep_capabilities", "ep_registry", + "fingerprint_model_weights", + "load_peft_adapter", + "generation", "get_build_dtype", "get_ep", "inspect_components", @@ -59,7 +76,7 @@ __version__ = "0.1.0" -from mobius import components, models, tasks +from mobius import components, generation, models, tasks from mobius._build_context import build_context, ep_capabilities, get_build_dtype from mobius._builder import build_from_module from mobius._configs import ( @@ -97,6 +114,23 @@ ModelRegistry, registry, ) +from mobius.adapter_io import adapter_source_from_onnx_adapter, load_peft_adapter +from mobius.adapters import ( + AdapterApplication, + AdapterArtifact, + AdapterBatchSelection, + AdapterSelectionTensors, + AdapterServiceOptions, + AdapterSlotSelection, + AdapterSource, + AdapterTarget, + AdapterTargetDescriptor, + AdapterTargetManifest, + AdapterTargetSlice, + AdapterWeights, + compose_adapter_deltas, + fingerprint_model_weights, +) from mobius.integrations._weight_loading import apply_weights from mobius.integrations.diffusers import build_diffusers_pipeline from mobius.integrations.gguf import build_from_gguf diff --git a/src/mobius/__main__.py b/src/mobius/__main__.py index 19159163a..ce5e78117 100644 --- a/src/mobius/__main__.py +++ b/src/mobius/__main__.py @@ -157,7 +157,19 @@ def _cmd_build(args: argparse.Namespace) -> None: from mobius.tasks import CausalLMTask, ModelTask def _resolve_static_cache_task(model_type: str) -> ModelTask: - """Create the correct static cache task for the given model type.""" + """Create the correct static cache task for the given model type. + + ``--features text-only`` makes :func:`build` swap the checkpoint's + multimodal ``model_type`` for its text-only registry sibling, so the + task must be resolved against the *same* substituted type. Resolving + against the raw checkpoint type instead pairs a text-only module with a + multimodal task, which then fails looking for sub-modules (a vision + tower, a separate decoder) that a text-only module does not have. + """ + if args.text_only: + from mobius._registry import _TEXT_ONLY_MODEL_TYPE + + model_type = _TEXT_ONLY_MODEL_TYPE.get(model_type, model_type) if model_type == "gemma4": from mobius.tasks._gemma4 import Gemma4Task @@ -305,9 +317,21 @@ def _resolve_static_cache_task(model_type: str) -> ModelTask: import transformers config_path = args.config - hf_config = transformers.AutoConfig.from_pretrained( - config_path, trust_remote_code=trust_remote_code - ) + try: + hf_config = transformers.AutoConfig.from_pretrained( + config_path, trust_remote_code=trust_remote_code + ) + except (ValueError, KeyError, OSError): + # A checkpoint predating the mandatory ``model_type`` key still + # names its architecture; resolve it the same way the HF-id path + # does rather than refusing a directory Mobius can build. + from mobius.integrations.transformers._config_resolver import ( + _try_load_config_json, + ) + + hf_config = _try_load_config_json(config_path) + if hf_config is None: + raise model_type = hf_config.model_type parent_config = hf_config if hasattr(hf_config, "text_config"): @@ -435,35 +459,19 @@ def _save_package( print(f" {name}: {path}") elif runtime == "onnx-genai": from mobius.integrations.onnx_genai import write_onnx_genai_config - from mobius.integrations.onnx_genai.inference_metadata import ( - is_native_vlm_package, - write_native_vlm_package_metadata, - ) config = getattr(pkg, "config", None) source = getattr(args, "config", None) or getattr(args, "model", None) - revision_kwargs = ( - {"revision": args.revision} if getattr(args, "revision", None) is not None else {} - ) - if is_native_vlm_package(pkg): - try: - artifacts = write_native_vlm_package_metadata( - pkg, - output_dir, - config=config, - source=source, - **revision_kwargs, - ) - except ValueError as error: - raise SystemExit(f"Error: {error}") from error - else: + try: artifacts = write_onnx_genai_config( pkg, output_dir, config=config, source=source, - **revision_kwargs, + revision=getattr(args, "revision", None), ) + except ValueError as error: + raise SystemExit(f"Error: {error}") from error for name, path in artifacts.items(): print(f" {name}: {path}") diff --git a/src/mobius/_configs/_base.py b/src/mobius/_configs/_base.py index 0dd76977c..91b96448f 100644 --- a/src/mobius/_configs/_base.py +++ b/src/mobius/_configs/_base.py @@ -2848,6 +2848,11 @@ def __post_init__(self): f"({self.encoder_input_channels}) must equal num_mel_bins " f"({self.num_mel_bins})." ) + # The decoder's learned position table is the model's context bound; + # Whisper spells it `max_target_positions`, so mirror it onto the + # architecture-wide field consumers read. + if self.max_position_embeddings == DEFAULT_INT: + self.max_position_embeddings = self.max_target_positions @classmethod def from_transformers(cls, config, parent_config=None) -> WhisperConfig: @@ -2880,6 +2885,8 @@ def from_transformers(cls, config, parent_config=None) -> WhisperConfig: max_target_positions=getattr(config, "max_target_positions", 448), scale_embedding=getattr(config, "scale_embedding", False), decoder_start_token_id=getattr(config, "decoder_start_token_id", None), + bos_token_id=getattr(config, "bos_token_id", None), + eos_token_id=getattr(config, "eos_token_id", None), ) # Model dtype @@ -2965,6 +2972,30 @@ def from_transformers(cls, config, parent_config=None) -> MoonshineConfig: return cls(**options) +def _conv_widths(config, defaults, hidden_size: int) -> tuple[int, ...]: + """Per-layer channel widths of a wav2vec2-family convolutional feature encoder. + + ``conv_dim``, ``conv_kernel`` and ``conv_stride`` describe one conv stack, so + the depth they imply has to agree. M-CTC-T states its single subsampling + convolution through ``conv_kernel``/``conv_stride`` alone and never publishes + ``conv_dim``; inheriting wav2vec2's seven-layer default there would describe a + stack the checkpoint does not have. Size the widths to the declared depth + instead, preferring the checkpoint's own ``conv_channels`` when present. + """ + declared = getattr(config, "conv_dim", None) + if declared: + return tuple(declared) + depth = len(tuple(getattr(config, "conv_kernel", None) or defaults.conv_kernel)) + if depth == len(defaults.conv_dim): + return defaults.conv_dim + channels = getattr(config, "conv_channels", None) + if channels is None: + return (hidden_size,) * depth + if isinstance(channels, (list, tuple)): + return tuple(channels) + return (int(channels),) * depth + + @dataclasses.dataclass class MMSConfig(ArchitectureConfig): """Configuration for MMS (Massively Multilingual Speech) CTC models. @@ -2976,6 +3007,13 @@ class MMSConfig(ArchitectureConfig): weights into the model. HuggingFace class: ``Wav2Vec2ForCTC`` with ``config.model_type == "wav2vec2"`` + + The convolutional feature-encoder geometry (``conv_dim``/``conv_kernel``/ + ``conv_stride``) is *not* boilerplate: it fixes the waveform-to-frame + downsampling ratio, so it must come from the checkpoint rather than from a + hard-coded default. ``facebook/wav2vec2-base-960h`` downsamples by 320 + (``prod(conv_stride)``) and disables conv bias, while + ``facebook/mms-1b-all`` enables it. """ add_adapter: bool = False @@ -2984,24 +3022,78 @@ class MMSConfig(ArchitectureConfig): adapter_stride: int = 2 num_adapter_layers: int = 3 + # Convolutional feature encoder (raw waveform → frames). + conv_dim: tuple[int, ...] = (512, 512, 512, 512, 512, 512, 512) + conv_kernel: tuple[int, ...] = (10, 3, 3, 3, 3, 2, 2) + conv_stride: tuple[int, ...] = (5, 2, 2, 2, 2, 2, 2) + conv_bias: bool = False + feat_extract_norm: str = "group" + + # Transformer encoder shape. + do_stable_layer_norm: bool = False + num_conv_pos_embeddings: int = 128 + num_conv_pos_embedding_groups: int = 16 + layer_norm_eps: float = 1e-5 + def __post_init__(self): if self.output_hidden_size == 0: self.output_hidden_size = self.hidden_size + # Normalize sequence fields so downstream code can index them freely + # regardless of whether the checkpoint used a list or a tuple. + self.conv_dim = tuple(self.conv_dim) + self.conv_kernel = tuple(self.conv_kernel) + self.conv_stride = tuple(self.conv_stride) + if not (len(self.conv_dim) == len(self.conv_kernel) == len(self.conv_stride)): + raise ValueError( + "conv_dim, conv_kernel and conv_stride must have equal length; got " + f"{len(self.conv_dim)}, {len(self.conv_kernel)}, {len(self.conv_stride)}" + ) + if self.feat_extract_norm not in ("group", "layer"): + raise ValueError( + f"feat_extract_norm must be 'group' or 'layer', got {self.feat_extract_norm!r}" + ) + + def feature_extract_output_length(self, num_samples: int) -> int: + """Return the frame count the conv stack emits for *num_samples*. + + Mirrors ``Wav2Vec2PreTrainedModel._get_feat_extract_output_lengths``: + each conv applies ``floor((L - kernel) / stride) + 1``. Callers use it + to segment a padded batch back into per-row transcripts. + """ + length = num_samples + for kernel, stride in zip(self.conv_kernel, self.conv_stride): + length = (length - kernel) // stride + 1 + if self.add_adapter: + for _ in range(self.num_adapter_layers): + length = (length - 1) // self.adapter_stride + 1 + return length @classmethod def from_transformers(cls, config, parent_config=None) -> MMSConfig: """Extract MMSConfig from a HuggingFace Wav2Vec2Config.""" base = ArchitectureConfig.from_transformers(config, parent_config=parent_config) base_fields = _shallow_fields(base) + defaults = cls(hidden_size=1) return cls( **base_fields, - add_adapter=getattr(config, "add_adapter", False), + add_adapter=getattr(config, "add_adapter", False) or False, output_hidden_size=getattr( config, "output_hidden_size", base_fields["hidden_size"] ), adapter_kernel_size=getattr(config, "adapter_kernel_size", 3), adapter_stride=getattr(config, "adapter_stride", 2), num_adapter_layers=getattr(config, "num_adapter_layers", 3), + conv_dim=_conv_widths(config, defaults, base_fields["hidden_size"]), + conv_kernel=tuple(getattr(config, "conv_kernel", None) or defaults.conv_kernel), + conv_stride=tuple(getattr(config, "conv_stride", None) or defaults.conv_stride), + conv_bias=bool(getattr(config, "conv_bias", False)), + feat_extract_norm=getattr(config, "feat_extract_norm", None) or "group", + do_stable_layer_norm=bool(getattr(config, "do_stable_layer_norm", False)), + num_conv_pos_embeddings=getattr(config, "num_conv_pos_embeddings", None) or 128, + num_conv_pos_embedding_groups=( + getattr(config, "num_conv_pos_embedding_groups", None) or 16 + ), + layer_norm_eps=getattr(config, "layer_norm_eps", None) or 1e-5, ) diff --git a/src/mobius/_constants.py b/src/mobius/_constants.py index 5e420cf14..c80cc4780 100644 --- a/src/mobius/_constants.py +++ b/src/mobius/_constants.py @@ -9,3 +9,25 @@ # Separated into its own module to avoid circular imports between # tasks, components, and the top-level package. OPSET_VERSION = 24 + +# --------------------------------------------------------------------------- +# Static KV cache ABI +# --------------------------------------------------------------------------- +# A static-cache export scatters each step's keys and values into pre-allocated, +# fixed-capacity buffers instead of concatenating a growing cache. The two +# control ports below are plain integer vectors, so they are *shape-indistin- +# guishable* from one another and from every other per-row integer input: no +# consumer can recover their roles from the graph. They are therefore a declared +# ABI, minted here once and read back by the metadata producers, rather than +# names any consumer is expected to guess. +STATIC_CACHE_WRITE_INDICES = "write_indices" +"""Per-row destination of this step's scatter, along the cache sequence axis.""" + +STATIC_CACHE_KV_SEQUENCE_LENGTH = "nonpad_kv_seqlen" +"""Per-row count of valid cache entries *after* this step's scatter.""" + +STATIC_CACHE_SEQUENCE_AXIS = 1 +"""Cache axis the scatter addresses: buffers are ``[batch, capacity, kv_hidden]``.""" + +STATIC_CACHE_LAYOUT = "bsh" +"""Element layout of a static cache buffer: batch, sequence slot, packed KV hidden.""" diff --git a/src/mobius/_inspect.py b/src/mobius/_inspect.py index 5ff82f516..abe7ae0f1 100644 --- a/src/mobius/_inspect.py +++ b/src/mobius/_inspect.py @@ -32,9 +32,11 @@ class ComponentInfo: name: Component name. This is the ``ModelPackage`` key mobius produces (and the subfolder name a multi-component export is saved under). role: Optimization role of the component, e.g. ``"decoder"``, - ``"encoder"``, or ``"embedding"``. Mobius uses this to gate fusion - passes (only ``"decoder"`` receives GQA / QKV-packing). It is the - value declared in the task's ``model_roles``. + ``"encoder"``, ``"embedding"``, or ``"glue"``. Mobius uses this to + gate fusion passes (only ``"decoder"`` receives GQA / QKV-packing). + ``"glue"`` marks a parameter-free graph that only wires a + generation loop — it carries no weights and no fusion applies. It + is the value declared in the task's ``model_roles``. source_paths: Runtime ``named_modules()`` paths that make up this component inside the full HuggingFace model. These are not checkpoint/state-dict key prefixes. A single component may map to diff --git a/src/mobius/_inspect_test.py b/src/mobius/_inspect_test.py index cbdc5e006..2da97bde1 100644 --- a/src/mobius/_inspect_test.py +++ b/src/mobius/_inspect_test.py @@ -322,7 +322,15 @@ def test_unresolvable_config_raises(monkeypatch): "qwen3_tts", "fastconformer_rnnt", ] -_EMPTY_SOURCE_PATHS = {("trocr", "encoder")} +_EMPTY_SOURCE_PATHS = { + ("trocr", "encoder"), + # Loop-wiring graphs mobius authors itself — no HuggingFace sub-module + # backs them. + ("qwen3_tts", "code_predictor_prefill"), + ("qwen3_tts", "code_predictor_step_embedder"), + ("qwen3_tts", "code_predictor_indices"), + ("qwen3_tts", "talker_text_step"), +} @pytest.mark.parametrize("model_type", _MULTI_COMPONENT_MODEL_TYPES) diff --git a/src/mobius/_model_package.py b/src/mobius/_model_package.py index 7740e240a..7f386ae52 100644 --- a/src/mobius/_model_package.py +++ b/src/mobius/_model_package.py @@ -20,24 +20,61 @@ __all__ = ["ModelPackage"] +import hashlib import inspect import logging import os +import shutil import threading from collections import UserDict -from collections.abc import Callable +from collections.abc import Callable, Iterator +from contextlib import contextmanager from typing import Any import onnx_ir as ir +import rfc8785 import torch import tqdm from mobius._optimizations import fold_initializers_after_weights +from mobius.adapters import ( + AdapterArtifact, + AdapterServiceOptions, + AdapterTargetManifest, +) +from mobius.generation import PolicyComponent from mobius.integrations._weight_loading import _assign_weight logger = logging.getLogger(__name__) +def _adapter_dtype_name(dtype: ir.DataType) -> str: + names = { + ir.DataType.FLOAT16: "float16", + ir.DataType.FLOAT: "float32", + ir.DataType.BFLOAT16: "bfloat16", + } + try: + return names[dtype] + except KeyError as error: + raise ValueError( + f"adapter dtype {dtype.name} must be a floating-point adapter dtype" + ) from error + + +def _adapter_source_file(artifact: AdapterArtifact) -> str: + if artifact.source.path is None: + raise ValueError( + f"adapter {artifact.name!r} cannot preserve an in-memory source format" + ) + if artifact.source.format == "onnx_adapter": + return artifact.source.path + raise ValueError( + f"adapter {artifact.name!r} source format {artifact.source.format!r} " + "cannot be preserved" + ) + + class ModelPackage(UserDict[str, ir.Model]): """A dict-like collection of named ``ir.Model`` objects. @@ -50,9 +87,26 @@ def __init__( self, models: dict[str, ir.Model] | None = None, config: object | None = None, + policy_components: dict[str, PolicyComponent] | None = None, + adapter_artifacts: dict[str, AdapterArtifact] | None = None, + adapter_target_manifest: AdapterTargetManifest | None = None, + adapter_service_options: AdapterServiceOptions | None = None, ) -> None: super().__init__(models or {}) self.config = config + self.policy_components = dict(policy_components or {}) + self.adapter_target_manifest = adapter_target_manifest + self.adapter_service_options = adapter_service_options or AdapterServiceOptions() + if adapter_target_manifest is not None: + adapter_target_manifest.validate(self.data) + self.adapter_artifacts: dict[str, AdapterArtifact] = {} + for name, artifact in (adapter_artifacts or {}).items(): + if name != artifact.name: + raise ValueError( + f"adapter catalog key {name!r} does not match artifact name " + f"{artifact.name!r}" + ) + self.add_adapter_artifact(artifact) def __repr__(self) -> str: names = ", ".join(repr(k) for k in self.data) @@ -78,6 +132,8 @@ def save( components: Callable[[str], bool] | None = None, progress_bar: bool = True, check_weights: bool = True, + include_policy_components: bool = True, + include_adapter_artifacts: bool = True, ) -> None: """Save all component models to a directory. @@ -129,6 +185,10 @@ def save( check_weights: Whether to verify that all initializers have weight data before saving. Defaults to ``True``. Set to ``False`` when saving skeleton models without weights. + include_policy_components: Save attached generation-policy ONNX + components under ``policies/``. Defaults to ``True``. + include_adapter_artifacts: Save attached parameter-adapter bundles + under ``adapters/``. Defaults to ``True``. Raises: ValueError: If *external_data* is not ``"onnx"`` or @@ -161,22 +221,335 @@ def save( else: model_dir = directory path = os.path.join(model_dir, "model.onnx") - if external_data == "safetensors": - ir.save_safetensors( - model, - path, - max_shard_size_bytes=max_shard_size_bytes, - callback=callback, + with _namespaced_symbolic_dimensions(model, f"component.{name}") as saved_model: + if external_data == "safetensors": + ir.save_safetensors( + saved_model, + path, + max_shard_size_bytes=max_shard_size_bytes, + callback=callback, + ) + else: + save_kwargs: dict[str, Any] = { + "external_data": "model.onnx.data", + "max_shard_size_bytes": max_shard_size_bytes, + "callback": callback, + } + if "max_workers" in inspect.signature(ir.save).parameters: + save_kwargs["max_workers"] = max_workers + ir.save(saved_model, path, **save_kwargs) + + if include_policy_components: + self.save_policy_components(directory, check_weights=check_weights) + if include_adapter_artifacts: + self.save_adapter_artifacts(directory) + + def add_policy_component(self, name: str, component: PolicyComponent) -> None: + """Attach a reusable generation-policy graph to this package.""" + if not name or "/" in name or "\\" in name: + raise ValueError("Policy component name must be a non-empty path segment") + self.policy_components[name] = component + + def add_adapter_artifact( + self, artifact: AdapterArtifact, *, validate_base: bool = True + ) -> None: + """Attach a model-agnostic adapter artifact to this package. + + Persistence and ONNX GenAI metadata emission intentionally remain separate + until the runtime artifact/schema contract is finalized. + """ + if artifact.name in self.adapter_artifacts: + raise ValueError(f"adapter artifact {artifact.name!r} is already attached") + if validate_base: + artifact.validate_base( + self.data, + fingerprint_targets=( + self.adapter_target_manifest.targets + if self.adapter_target_manifest is not None + else None + ), + ) + if self.adapter_target_manifest is not None: + if artifact.base_fingerprint != self.adapter_target_manifest.base_fingerprint: + raise ValueError( + f"adapter {artifact.name!r} base fingerprint does not match " + "the authoritative target manifest" ) - else: - save_kwargs: dict[str, Any] = { - "external_data": "model.onnx.data", - "max_shard_size_bytes": max_shard_size_bytes, - "callback": callback, + missing = artifact.target_bindings - self.adapter_target_manifest.bindings + if missing: + raise ValueError( + f"adapter artifact {artifact.name!r} contains targets outside " + f"the authoritative manifest: {sorted(map(str, missing))}" + ) + self.adapter_artifacts[artifact.name] = artifact + + def save_adapter_artifacts(self, directory: str) -> dict[str, dict[str, object]]: + """Save exact adapter bundles and return ONNX GenAI catalog entries.""" + if not self.adapter_artifacts: + return {} + if self.adapter_target_manifest is None: + raise ValueError( + "adapter artifacts require an authoritative adapter target manifest" + ) + self.adapter_target_manifest.validate(self.data) + adapter_dir = os.path.join(directory, "adapters") + os.makedirs(adapter_dir, exist_ok=True) + descriptors = { + descriptor.target: descriptor + for descriptor in self.adapter_target_manifest.targets + } + manifest_targets = { + target["id"]: target + for target in self.adapter_target_manifest_metadata()["targets"] + } + catalog: dict[str, dict[str, object]] = {} + identities: set[tuple[str, str]] = set() + for artifact_index, (alias, artifact) in enumerate( + sorted(self.adapter_artifacts.items()) + ): + identity_version = (artifact.stable_identity, artifact.version) + if identity_version in identities: + raise ValueError( + f"adapter identity/version {identity_version[0]}@" + f"{identity_version[1]} must be unique" + ) + identities.add(identity_version) + ordered_weights = sorted( + artifact.weights, + key=lambda item: ( + item.target.component, + item.target.parameter, + item.target_id or "", + ), + ) + dtypes = {weight.dtype for weight in artifact.weights} + if len(dtypes) != 1: + raise ValueError( + f"adapter {alias!r} has heterogeneous target dtypes, " + "which the ONNX GenAI artifact contract cannot represent" + ) + rank = ordered_weights[0].rank + alpha = ordered_weights[0].alpha + dtype = _adapter_dtype_name(dtypes.pop()) + bindings: list[dict[str, object]] = [] + portable_targets: dict[str, dict[str, list[float]]] = {} + for weight in ordered_weights: + descriptor = descriptors[weight.target] + target_id = weight.target_id or descriptor.semantic_name + if target_id not in manifest_targets: + raise ValueError( + f"adapter {alias!r} references target ID {target_id!r} " + "outside the authoritative manifest" + ) + target_policy = manifest_targets[target_id] + if target_policy.get("rank", weight.rank) != weight.rank: + raise ValueError( + f"adapter {alias!r} target {target_id!r} rank {weight.rank} " + f"violates manifest policy {target_policy['rank']}" + ) + if target_policy.get("alpha", weight.alpha) != weight.alpha: + raise ValueError( + f"adapter {alias!r} target {target_id!r} alpha {weight.alpha} " + f"violates manifest policy {target_policy['alpha']}" + ) + slice_policy = target_policy.get("output_slice") + if isinstance(slice_policy, dict): + if slice_policy.get("rank", weight.rank) != weight.rank: + raise ValueError( + f"adapter {alias!r} target {target_id!r} rank {weight.rank} " + f"violates output-slice policy {slice_policy['rank']}" + ) + if slice_policy.get("alpha", weight.alpha) != weight.alpha: + raise ValueError( + f"adapter {alias!r} target {target_id!r} alpha {weight.alpha} " + f"violates output-slice policy {slice_policy['alpha']}" + ) + weight_key = weight.weight_key or descriptor.semantic_name + binding: dict[str, object] = { + "target": target_id, + "weight_key": weight_key, + } + if weight.rank != rank: + binding["rank"] = weight.rank + if weight.alpha != alpha: + binding["alpha"] = weight.alpha + bindings.append(binding) + portable_targets[weight_key] = { + "a": weight.a.numpy().reshape(-1).astype("float32").tolist(), + "b": weight.b.numpy().reshape(-1).astype("float32").tolist(), + } + + artifact_dir = os.path.join(directory, "adapters", alias) + os.makedirs(artifact_dir, exist_ok=True) + weight_artifacts: list[dict[str, object]] = [] + if self.adapter_service_options.portable_fallback: + relative_location = f"adapters/{alias}/adapter.json" + destination = os.path.join(directory, relative_location) + payload = rfc8785.dumps({"targets": portable_targets}) + with open(destination, "wb") as handle: + handle.write(payload) + weight_artifacts.append( + { + "location": relative_location, + "loader_capability": "onnx-genai.adapters.json@1", + "sha256": hashlib.sha256(payload).hexdigest(), + "scale_encoding": "alpha_over_rank", + "format": "json", + } + ) + if ( + self.adapter_service_options.preserve_source_format + and artifact.source.format != "in_memory" + ): + if artifact.source.format == "onnx_adapter": + source_path = _adapter_source_file(artifact) + relative_location = f"adapters/{alias}/adapter.onnx_adapter" + destination = os.path.join(directory, relative_location) + shutil.copyfile(source_path, destination) + with open(destination, "rb") as handle: + payload = handle.read() + weight_artifacts.append( + { + "location": relative_location, + "loader_capability": "onnxruntime.lora-adapter@1", + "sha256": hashlib.sha256(payload).hexdigest(), + "scale_encoding": "baked", + "format": "ort_genai", + } + ) + elif artifact.source.format == "peft_safetensors": + if artifact.source.path is None: + raise ValueError(f"adapter {alias!r} PEFT source path is absent") + source_dir = artifact.source.path + source_weights = os.path.join(source_dir, "adapter_model.safetensors") + source_config = os.path.join(source_dir, "adapter_config.json") + relative_location = f"adapters/{alias}/adapter_model.safetensors" + relative_config = f"adapters/{alias}/adapter_config.json" + destination = os.path.join(directory, relative_location) + config_destination = os.path.join(directory, relative_config) + shutil.copyfile(source_weights, destination) + shutil.copyfile(source_config, config_destination) + with open(destination, "rb") as handle: + payload = handle.read() + with open(config_destination, "rb") as handle: + config_payload = handle.read() + weight_artifacts.append( + { + "location": relative_location, + "loader_capability": "onnx-genai.adapters.hf-peft@1", + "sha256": hashlib.sha256(payload).hexdigest(), + "config_location": relative_config, + "config_sha256": hashlib.sha256(config_payload).hexdigest(), + "scale_encoding": "alpha_over_rank", + "format": "hf_peft", + } + ) + else: + raise ValueError( + f"adapter {alias!r} source format {artifact.source.format!r} " + "cannot be preserved" + ) + if not weight_artifacts: + raise ValueError( + f"adapter {alias!r} must emit a portable or preserved source artifact" + ) + provenance = {"producer": artifact.source.producer} + if artifact.source.base_model: + provenance["source"] = artifact.source.base_model + if artifact.source.revision: + provenance["revision"] = artifact.source.revision + catalog[alias] = { + "index": artifact_index, + "identity": artifact.stable_identity, + "version": artifact.version, + "base_model_fingerprint": artifact.base_fingerprint, + "rank": rank, + "alpha": alpha, + "dtype": dtype, + "provenance": provenance, + "weights": weight_artifacts, + "bindings": bindings, + } + return catalog + + def adapter_target_manifest_metadata(self) -> dict[str, object]: + """Serialize the authoritative generic LoRA target manifest.""" + if self.adapter_target_manifest is None: + raise ValueError("adapter target manifest is not attached") + targets: list[dict[str, object]] = [] + for descriptor in sorted( + self.adapter_target_manifest.targets, + key=lambda item: item.semantic_name, + ): + initializer = self.data[descriptor.target.component].graph.initializers[ + descriptor.target.parameter + ] + activation_dtype = _adapter_dtype_name( + descriptor.activation_dtype or initializer.dtype + ) + base: dict[str, object] = { + "id": descriptor.semantic_name, + "component": descriptor.target.component, + "initializer": descriptor.target.parameter, + "node_name": descriptor.node_name, + "output_name": descriptor.output_name, + "activation_dtype": activation_dtype, + "input_features": descriptor.input_size, + "output_features": descriptor.output_size, + } + if descriptor.layer_index is not None: + base["layer_index"] = descriptor.layer_index + if descriptor.rank is not None: + base["rank"] = descriptor.rank + if descriptor.alpha is not None: + base["alpha"] = descriptor.alpha + if descriptor.graph_input_a is not None: + graph_inputs = { + "a": descriptor.graph_input_a, + "b": descriptor.graph_input_b, } - if "max_workers" in inspect.signature(ir.save).parameters: - save_kwargs["max_workers"] = max_workers - ir.save(model, path, **save_kwargs) + if descriptor.graph_input_scale is not None: + graph_inputs["scale"] = descriptor.graph_input_scale + base["graph_inputs"] = graph_inputs + targets.append(base) + for target_slice in descriptor.slices: + sliced = dict(base) + sliced["id"] = f"{descriptor.semantic_name}.{target_slice.role}" + output_slice: dict[str, object] = { + "role": target_slice.role, + "offset": target_slice.offset, + "width": target_slice.width, + } + if target_slice.rank is not None: + output_slice["rank"] = target_slice.rank + if target_slice.alpha is not None: + output_slice["alpha"] = target_slice.alpha + sliced["output_slice"] = output_slice + targets.append(sliced) + return {"targets": targets} + + def save_policy_components( + self, + directory: str, + *, + check_weights: bool = True, + ) -> dict[str, str]: + """Save attached policy graphs and return package-relative artifact paths.""" + if not self.policy_components: + return {} + policy_dir = os.path.join(directory, "policies") + os.makedirs(policy_dir, exist_ok=True) + artifacts: dict[str, str] = {} + for name, component in self.policy_components.items(): + if check_weights: + _check_weights(name, component.model) + relative_path = f"policies/{name}.onnx" + # Policy components are separate ONNX artifacts with public ABI + # dimension names such as ``batch``. Namespacing those dimensions + # makes an otherwise exact semantic contract ineligible for runtime fusion. + ir.save(component.model, os.path.join(directory, relative_path)) + artifacts[name] = relative_path + return artifacts @classmethod def load(cls, directory: str) -> ModelPackage: @@ -203,13 +576,29 @@ def load(cls, directory: str) -> ModelPackage: if os.path.isdir(subdir) and os.path.isfile(model_path): models[entry] = ir.load(model_path) if models: - return cls(models) + package = cls(models) + package._load_policy_components(directory) + return package # Fall back to flat layout for filename in sorted(os.listdir(directory)): if filename.endswith(".onnx"): name = filename.removesuffix(".onnx") models[name] = ir.load(os.path.join(directory, filename)) - return cls(models) + package = cls(models) + package._load_policy_components(directory) + return package + + def _load_policy_components(self, directory: str) -> None: + policy_dir = os.path.join(directory, "policies") + if not os.path.isdir(policy_dir): + return + for filename in sorted(os.listdir(policy_dir)): + if not filename.endswith(".onnx"): + continue + model = ir.load(os.path.join(policy_dir, filename)) + self.policy_components[filename.removesuffix(".onnx")] = ( + PolicyComponent.from_model(model) + ) # -- Weight application ------------------------------------------------ @@ -285,6 +674,58 @@ def apply_weights( fold_initializers_after_weights(model) +@contextmanager +def _namespaced_symbolic_dimensions( + model: ir.Model, + namespace: str, +) -> Iterator[ir.Model]: + """Namespace interface symbols and discard non-contractual intermediate aliases.""" + interface_values = {id(value) for value in (*model.graph.inputs, *model.graph.outputs)} + values: dict[int, ir.Value] = {} + graphs: list[ir.GraphProtocol] = [] + nodes = ir.traversal.RecursiveGraphIterator(model.graph, enter_graph=graphs.append) + for node in nodes: + for value in (*node.inputs, *node.outputs): + if value is not None: + values[id(value)] = value + for graph in graphs: + for value in (*graph.inputs, *graph.outputs, *graph.initializers.values()): + values[id(value)] = value + + originals: list[tuple[ir.Value, ir.Shape]] = [] + symbols: dict[str, str] = {} + try: + for value in values.values(): + if value.shape is None: + continue + dimensions: list[int | str | ir.SymbolicDim] = [] + changed = False + for dimension in value.shape: + if isinstance(dimension, int): + dimensions.append(dimension) + continue + if dimension.value is None: + dimensions.append(dimension) + continue + if id(value) not in interface_values: + dimensions.append(ir.SymbolicDim(None)) + changed = True + continue + text = str(dimension) + dimensions.append(symbols.setdefault(text, f"{namespace}.{text}")) + changed = True + if changed: + originals.append((value, value.shape)) + denotations = [ + value.shape.get_denotation(index) for index in range(len(value.shape)) + ] + value.shape = ir.Shape(dimensions, denotations) + yield model + finally: + for value, shape in originals: + value.shape = shape + + def _make_progress_callback(): """Create a thread-safe tqdm progress-bar callback for ``ir.save``. diff --git a/src/mobius/_model_package_test.py b/src/mobius/_model_package_test.py index 50f2b9780..ada31473d 100644 --- a/src/mobius/_model_package_test.py +++ b/src/mobius/_model_package_test.py @@ -15,8 +15,13 @@ from mobius._builder import build_from_module from mobius._configs import VisionConfig -from mobius._model_package import ModelPackage, _make_progress_callback +from mobius._model_package import ( + ModelPackage, + _make_progress_callback, + _namespaced_symbolic_dimensions, +) from mobius._testing import make_config +from mobius.generation import build_greedy_sampler from mobius.models.base import CausalLMModel from mobius.models.gemma3 import Gemma3MultiModalModel from mobius.tasks import CausalLMTask, VisionLanguageTask @@ -365,6 +370,125 @@ def test_save_creates_directory(self, tmp_path): pkg.save(str(outdir)) assert (outdir / "model.onnx").exists() + def test_policy_components_roundtrip(self, tmp_path): + pkg = ModelPackage({"model": _make_simple_model()}) + pkg.add_policy_component("sample", build_greedy_sampler()) + pkg.save(str(tmp_path)) + + assert (tmp_path / "policies" / "sample.onnx").exists() + loaded = ModelPackage.load(str(tmp_path)) + assert loaded.policy_components["sample"].contract_id == "onnx-genai.token-sampler@1" + + def test_save_namespaces_component_symbols_without_mutating_package(self, tmp_path): + input_value = ir.Value( + name="input", + type=ir.TensorType(ir.DataType.FLOAT), + shape=ir.Shape(["batch", "sequence"]), + ) + input_value.shape.set_denotation(0, "DATA_BATCH") + intermediate = ir.Value( + name="intermediate", + type=ir.TensorType(ir.DataType.FLOAT), + shape=ir.Shape(["batch", "sequence"]), + ) + output_value = ir.Value( + name="output", + type=ir.TensorType(ir.DataType.FLOAT), + shape=ir.Shape(["batch", "sequence"]), + ) + model = ir.Model( + ir.Graph( + [input_value], + [output_value], + nodes=[ + ir.Node("", "Identity", [input_value], outputs=[intermediate]), + ir.Node("", "Identity", [intermediate], outputs=[output_value]), + ], + name="symbolic", + ), + ir_version=10, + ) + pkg = ModelPackage({"decoder": model}) + + pkg.save(str(tmp_path)) + + saved = ir.load(tmp_path / "model.onnx") + assert [str(dimension) for dimension in saved.graph.inputs[0].shape] == [ + "component.decoder.batch", + "component.decoder.sequence", + ] + assert saved.graph.inputs[0].shape.get_denotation(0) == "DATA_BATCH" + assert all( + dimension.value is None for dimension in next(iter(saved.graph)).outputs[0].shape + ) + assert [str(dimension) for dimension in model.graph.inputs[0].shape] == [ + "batch", + "sequence", + ] + + def test_save_anonymizes_nested_graph_intermediate_symbols(self): + condition = ir.Value( + name="condition", + type=ir.TensorType(ir.DataType.BOOL), + shape=ir.Shape([]), + ) + data = ir.Value( + name="data", + type=ir.TensorType(ir.DataType.FLOAT), + shape=ir.Shape(["batch", "sequence"]), + ) + branch_output = ir.Value( + name="branch_output", + type=ir.TensorType(ir.DataType.FLOAT), + shape=ir.Shape(["batch", "branch_sequence"]), + ) + branch = ir.Graph( + [], + [branch_output], + nodes=[ir.Node("", "Identity", [data], outputs=[branch_output])], + ) + output = ir.Value( + name="output", + type=ir.TensorType(ir.DataType.FLOAT), + shape=ir.Shape(["batch", "sequence"]), + ) + model = ir.Model( + ir.Graph( + [condition, data], + [output], + nodes=[ + ir.Node( + "", + "If", + [condition], + attributes={ + "then_branch": ir.AttrGraph("then_branch", branch), + "else_branch": ir.AttrGraph("else_branch", branch), + }, + outputs=[output], + ) + ], + ), + ir_version=10, + ) + + with _namespaced_symbolic_dimensions(model, "component.decoder"): + assert all(dimension.value is None for dimension in branch_output.shape) + + assert str(branch_output.shape[1]) == "branch_sequence" + + def test_save_preserves_public_policy_symbols(self, tmp_path): + sampler = build_greedy_sampler() + pkg = ModelPackage({"model": _make_simple_model()}) + pkg.add_policy_component("sample", sampler) + + pkg.save(str(tmp_path)) + + saved = ir.load(tmp_path / "policies" / "sample.onnx") + assert str(saved.graph.inputs[0].shape[0]) == "batch" + assert str(saved.graph.inputs[0].shape[1]) == "vocabulary" + assert str(sampler.model.graph.inputs[0].shape[0]) == "batch" + class TestModelPackageApplyWeights: def test_single_component(self): diff --git a/src/mobius/_optimizations.py b/src/mobius/_optimizations.py index af5af8f93..defbdc8d8 100644 --- a/src/mobius/_optimizations.py +++ b/src/mobius/_optimizations.py @@ -62,6 +62,7 @@ ) from mobius.functions import register_function_bodies from mobius.rewrite_rules import ( + clip_to_min_max_rules, decompose_attention_pass, decompose_rope_rules, gelu_fusion_rules, @@ -355,6 +356,13 @@ def _get_optimization_passes( if caps.requires_graph_capture_rewrite: lower.append(("StaticEmptyKV", list(static_empty_kv_rules()))) + # --- bfloat16 Clip lowering (all EPs) --- + # ORT has no bfloat16 Clip kernel and expands the op's ONNX function body + # into Less/Where, which also lack bfloat16 kernels — the unassigned nodes + # abort session creation. Min/Max are kernel-backed and exactly equivalent. + if dtype == ir.DataType.BFLOAT16: + lower.append(("ClipToMinMax", list(clip_to_min_max_rules()))) + return fuse, lower @@ -593,7 +601,28 @@ def _should_inline(func: ir.Function) -> bool: stacklevel=4, ) else: - Fp8KvCachePass(kv_cache_scales)(model) + fp8_pass = Fp8KvCachePass(kv_cache_scales) + fp8_pass(model) + if fp8_pass.converted == 0: + # The gate above tests the *intent* to fuse GQA; only the pass + # knows the outcome. FP8 KV storage is a property of the + # attention operator, because the scales that dequantize the + # cache on read are node inputs: GroupQueryAttention has + # ``k_scale``/``v_scale``, ai.onnx ``Attention`` has no such + # slot. Retyping a cache that no operator can dequantize would + # declare FP8 over data that is read as fp16 — silently wrong + # numerics — and quietly leaving the cache at the model dtype + # would hand back a package that does not do what was asked. + raise ValueError( + "fp8_kv_cache=True was requested but the optimized graph " + "exposes no GroupQueryAttention KV cache to convert. FP8 KV " + "storage needs an attention operator with k_scale/v_scale " + "inputs to dequantize the cache on read; a static-cache " + "export scatters into fixed buffers read by ai.onnx " + "Attention, which has no such inputs. Build without " + "--features fp8-kv-cache, or without --features static-cache " + "so the decoder fuses to GroupQueryAttention." + ) def fold_initializers_after_weights(model: ir.Model) -> None: diff --git a/src/mobius/_passes/_fp8_kv_cache.py b/src/mobius/_passes/_fp8_kv_cache.py index 9da74f6e5..e991579be 100644 --- a/src/mobius/_passes/_fp8_kv_cache.py +++ b/src/mobius/_passes/_fp8_kv_cache.py @@ -140,6 +140,10 @@ class Fp8KvCachePass(ir.passes.InPlacePass): def __init__(self, scales: dict[int, tuple[float, float]] | None = None) -> None: super().__init__() self._scales = _validate_scales(scales or {}) + #: Number of KV caches retyped by the last :meth:`call`. A caller that + #: asked for FP8 needs to know it actually happened; a graph whose + #: attention operator carries no KV-scale inputs leaves this at zero. + self.converted = 0 def call(self, model: ir.Model) -> ir.passes.PassResult: graph = model.graph @@ -211,6 +215,7 @@ def call(self, model: ir.Model) -> ir.passes.PassResult: modified = True converted += 1 + self.converted = converted if converted: logger.info( "Fp8KvCachePass: converted %d GroupQueryAttention KV cache(s) to FP8", diff --git a/src/mobius/_registry.py b/src/mobius/_registry.py index b49cc0788..6eacfdbb3 100644 --- a/src/mobius/_registry.py +++ b/src/mobius/_registry.py @@ -125,6 +125,7 @@ from mobius.models.ctrl import CTRLCausalLMModel from mobius.models.depth_anything import DepthAnythingForDepthEstimation from mobius.models.distilbert import DistilBertModel +from mobius.models.esm import EsmConfig, EsmModel from mobius.models.falcon import ( BloomCausalLMModel, FalconCausalLMModel, @@ -719,7 +720,7 @@ def _detect_fallback_registration(hf_config) -> ModelRegistration | None: "electra": ModelRegistration(BertModel, task="feature-extraction"), "ernie": ModelRegistration(BertModel, task="feature-extraction"), "ernie_m": ModelRegistration(BertModel, task="feature-extraction"), - "esm": ModelRegistration(BertModel, task="feature-extraction"), + "esm": ModelRegistration(EsmModel, task="feature-extraction", config_class=EsmConfig), "flaubert": ModelRegistration(BertModel, task="feature-extraction"), "ibert": ModelRegistration(BertModel, task="feature-extraction"), "layoutlm": ModelRegistration(BertModel, task="feature-extraction"), @@ -872,6 +873,12 @@ def _create_default_registry() -> ModelRegistry: "muse_glimmer_text": "muse_glimmer_text", "gemma3n": "gemma3n_text", "gemma3n_text": "gemma3n_text", + # Shipped Gemma 4 multimodal checkpoints (e.g. ``google/gemma-4-E2B-it``) + # declare ``model_type="gemma4"`` with a nested ``text_config`` whose own + # ``model_type`` is ``gemma4_text``. Both resolve to the same + # ``Gemma4CausalLMModel`` backbone, so ``text_only=True`` is supported. + "gemma4": "gemma4_text", + "gemma4_text": "gemma4_text", "gemma4_unified": "gemma4_unified_text", "gemma4_unified_text": "gemma4_unified_text", # Qwen3.5-MoE-VL (Qwen3.6-35B-A3B): export just the hybrid MoE text diff --git a/src/mobius/_registry_test.py b/src/mobius/_registry_test.py index 2d92ac64b..a4aab61a7 100644 --- a/src/mobius/_registry_test.py +++ b/src/mobius/_registry_test.py @@ -8,6 +8,7 @@ import pytest from mobius._registry import ( + _TEXT_ONLY_MODEL_TYPE, ModelRegistration, _detect_fallback_registration, registry, @@ -174,3 +175,51 @@ def test_registry_no_suggestion_for_random_string(self): """Very different strings don't produce suggestions.""" with pytest.raises(KeyError, match=r"Use registry\.register"): registry.get("zzzzzzzzz_not_a_model") + + +class TestTextOnlyModelTypeOverrides: + """``build(..., text_only=True)`` resolution table. + + Every multimodal ``model_type`` whose text backbone can be exported alone + needs an entry here; the CLI surfaces this as ``--features text-only``. + """ + + def test_every_target_is_registered(self): + for source, target in _TEXT_ONLY_MODEL_TYPE.items(): + assert target in registry, ( + f"text-only override {source!r} -> {target!r} names an unregistered model_type" + ) + + def test_mapping_is_idempotent(self): + """Applying the override twice must be a no-op.""" + for target in set(_TEXT_ONLY_MODEL_TYPE.values()): + assert _TEXT_ONLY_MODEL_TYPE.get(target) == target, ( + f"text-only target {target!r} is missing a self-mapping, so " + "text_only=True fails once the type is already text-only" + ) + + @pytest.mark.parametrize( + "model_type", + # Shipped Gemma 4 checkpoints declare model_type="gemma4" with a + # nested text_config of model_type="gemma4_text"; both must resolve. + ["gemma4", "gemma4_text"], + ) + def test_gemma4_resolves_to_text_backbone(self, model_type): + assert _TEXT_ONLY_MODEL_TYPE[model_type] == "gemma4_text" + + def test_nested_text_config_type_is_mapped(self): + """A VL type and its nested ``text_config`` type must agree. + + ``build(text_only=True)`` may resolve either the outer multimodal + ``model_type`` or the nested ``text_config.model_type`` depending on how + the checkpoint is loaded, so both spellings must land on the same + backbone. + """ + for source, target in _TEXT_ONLY_MODEL_TYPE.items(): + if not source.endswith("_text"): + continue + outer = source[: -len("_text")] + if outer in _TEXT_ONLY_MODEL_TYPE: + assert _TEXT_ONLY_MODEL_TYPE[outer] == target, ( + f"{outer!r} and {source!r} resolve to different backbones" + ) diff --git a/src/mobius/adapter_io.py b/src/mobius/adapter_io.py new file mode 100644 index 000000000..19a621885 --- /dev/null +++ b/src/mobius/adapter_io.py @@ -0,0 +1,157 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT License. + +"""Load standard adapter sources into the model-agnostic Mobius representation.""" + +from __future__ import annotations + +__all__ = ["adapter_source_from_onnx_adapter", "load_peft_adapter"] + +import hashlib +import json +import math +import re +from collections.abc import Mapping +from pathlib import Path + +import numpy as np +import onnx_ir as ir +from safetensors.numpy import load_file + +from mobius.adapters import ( + AdapterArtifact, + AdapterSource, + AdapterTarget, + AdapterWeights, +) + +_PEFT_CONFIG = "adapter_config.json" +_PEFT_WEIGHTS = "adapter_model.safetensors" +_FACTOR_PATTERN = re.compile(r"^(?P.+)\.lora_(?P[AB])(?:\.[^.]+)?\.weight$") + + +def _source_checksum(paths: list[Path]) -> str: + digest = hashlib.sha256() + for path in paths: + digest.update(path.name.encode()) + digest.update(b"\0") + digest.update(path.read_bytes()) + digest.update(b"\0") + return f"sha256:{digest.hexdigest()}" + + +def _pattern_value(patterns: Mapping[str, object], module_key: str, default: object) -> object: + matches = [(key, value) for key, value in patterns.items() if module_key.endswith(key)] + if not matches: + return default + return max(matches, key=lambda item: len(item[0]))[1] + + +def _resolve_target( + module_key: str, targets: Mapping[str, AdapterTarget] +) -> tuple[str, AdapterTarget]: + if module_key in targets: + return module_key, targets[module_key] + matches = [(name, target) for name, target in targets.items() if module_key.endswith(name)] + if len(matches) != 1: + raise ValueError( + f"PEFT module {module_key!r} resolves to {len(matches)} producer targets; " + "provide one exact or unique suffix binding" + ) + return matches[0] + + +def load_peft_adapter( + directory: str | Path, + *, + target_bindings: Mapping[str, AdapterTarget], + base_fingerprint: str, + name: str | None = None, +) -> AdapterArtifact: + """Load PEFT config/safetensors using producer-declared, model-agnostic bindings.""" + directory = Path(directory) + config_path = directory / _PEFT_CONFIG + weights_path = directory / _PEFT_WEIGHTS + if not config_path.is_file() or not weights_path.is_file(): + raise ValueError( + f"PEFT adapter directory {directory} must contain " + f"{_PEFT_CONFIG} and {_PEFT_WEIGHTS}" + ) + config = json.loads(config_path.read_text()) + tensors = load_file(weights_path) + target_modules = tuple(config.get("target_modules", ())) + if not target_modules: + raise ValueError("PEFT adapter target_modules must not be empty") + default_rank = int(config.get("r", 0)) + default_alpha = float(config.get("lora_alpha", default_rank)) + rank_pattern = config.get("rank_pattern", {}) + alpha_pattern = config.get("alpha_pattern", {}) + + pending: dict[str, dict[str, np.ndarray]] = {} + for key, tensor in tensors.items(): + match = _FACTOR_PATTERN.match(key) + if match is None: + continue + module_key = match.group("module") + if not any(module_key.endswith(target) for target in target_modules): + raise ValueError(f"PEFT tensor {key!r} is not covered by target_modules") + pending.setdefault(module_key, {})[match.group("factor")] = tensor + if not pending: + raise ValueError("PEFT adapter contains no LoRA A/B tensors") + + loaded_weights: list[AdapterWeights] = [] + for module_key, factors in sorted(pending.items()): + if set(factors) != {"A", "B"}: + raise ValueError(f"PEFT module {module_key!r} must contain paired A/B factors") + a = np.ascontiguousarray(factors["A"]) + b = np.ascontiguousarray(factors["B"]) + rank = int(_pattern_value(rank_pattern, module_key, default_rank)) + alpha = float(_pattern_value(alpha_pattern, module_key, default_alpha)) + if rank <= 0 or not math.isfinite(alpha): + raise ValueError( + f"PEFT module {module_key!r} has invalid rank/alpha {rank}/{alpha}" + ) + if a.ndim != 2 or b.ndim != 2 or a.shape[0] != rank or b.shape[1] != rank: + raise ValueError( + f"PEFT module {module_key!r} factors must have shapes " + f"[rank,K]/[N,rank] for rank {rank}, got {a.shape}/{b.shape}" + ) + weight_key, target = _resolve_target(module_key, target_bindings) + loaded_weights.append( + AdapterWeights( + target, + ir.tensor(a), + ir.tensor(b), + alpha, + weight_key=weight_key, + ) + ) + + source = AdapterSource( + "peft_safetensors", + path=str(directory), + checksum=_source_checksum([config_path, weights_path]), + base_model=config.get("base_model_name_or_path"), + revision=config.get("revision"), + ) + return AdapterArtifact( + name=name or directory.name, + base_fingerprint=base_fingerprint, + weights=tuple(loaded_weights), + source=source, + ) + + +def adapter_source_from_onnx_adapter( + path: str | Path, +) -> AdapterSource: + """Declare an ORT FlatBuffers adapter source without making it mandatory.""" + path = Path(path) + payload = path.read_bytes() + if len(payload) < 8 or payload[4:8] != b"TORT": + raise ValueError(f"ONNX adapter {path} does not contain the TORT identifier") + return AdapterSource( + "onnx_adapter", + path=str(path), + checksum=f"sha256:{hashlib.sha256(payload).hexdigest()}", + ) diff --git a/src/mobius/adapters.py b/src/mobius/adapters.py new file mode 100644 index 000000000..f3ca6b449 --- /dev/null +++ b/src/mobius/adapters.py @@ -0,0 +1,770 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT License. + +"""Model-agnostic low-rank adapter artifacts and per-request selection state.""" + +from __future__ import annotations + +__all__ = [ + "AdapterApplication", + "AdapterArtifact", + "AdapterBatchSelection", + "AdapterSlotSelection", + "AdapterSelectionTensors", + "AdapterServiceOptions", + "AdapterSource", + "AdapterTarget", + "AdapterTargetDescriptor", + "AdapterTargetManifest", + "AdapterTargetSlice", + "AdapterWeights", + "compose_adapter_deltas", + "fingerprint_model_weights", +] + +import dataclasses +import hashlib +import itertools +import json +import math +import sys +from collections.abc import Mapping, Sequence +from typing import Any, Literal + +import numpy as np +import onnx_ir as ir +import rfc8785 + + +def _validate_identifier(value: str, description: str) -> None: + if not value or "/" in value or "\\" in value: + raise ValueError(f"{description} must be a non-empty path segment") + + +def _tensor_bytes(tensor: ir.Tensor) -> bytes: + array = np.ascontiguousarray(tensor.numpy()) + if array.dtype.byteorder == ">" or ( + array.dtype.byteorder == "=" and sys.byteorder == "big" + ): + array = array.byteswap().view(array.dtype.newbyteorder("<")) + return array.tobytes(order="C") + + +def _update_tensor_hash(digest: Any, tensor: ir.Tensor) -> None: + shape = [int(dimension) for dimension in tensor.shape] + digest.update( + json.dumps( + {"dtype": tensor.dtype.name, "shape": shape}, + sort_keys=True, + separators=(",", ":"), + ).encode() + ) + digest.update(b"\0") + digest.update(_tensor_bytes(tensor)) + + +def _canonical_attribute_value(attribute: ir.Attr) -> object: + value = attribute.value + if attribute.type in { + ir.AttributeType.FLOAT, + ir.AttributeType.INT, + ir.AttributeType.STRING, + }: + return value + if attribute.type in { + ir.AttributeType.FLOATS, + ir.AttributeType.INTS, + ir.AttributeType.STRINGS, + }: + return list(value) + if attribute.type == ir.AttributeType.TENSOR: + return { + "dtype": int(value.dtype), + "shape": [int(dimension) for dimension in value.shape], + "sha256": hashlib.sha256(_tensor_bytes(value)).hexdigest(), + } + if attribute.type == ir.AttributeType.TENSORS: + return [ + { + "dtype": int(tensor.dtype), + "shape": [int(dimension) for dimension in tensor.shape], + "sha256": hashlib.sha256(_tensor_bytes(tensor)).hexdigest(), + } + for tensor in value + ] + raise ValueError( + f"cannot canonicalize adapter target consumer attribute " + f"{attribute.name!r} of type {attribute.type.name}" + ) + + +def _target_fingerprint_record( + models: Mapping[str, ir.Model], + resolved_target: AdapterTarget | AdapterTargetDescriptor, +) -> dict[str, object]: + target = ( + resolved_target.target + if isinstance(resolved_target, AdapterTargetDescriptor) + else resolved_target + ) + model = models.get(target.component) + if model is None: + raise ValueError(f"cannot fingerprint unknown component {target.component!r}") + initializer = model.graph.initializers.get(target.parameter) + if initializer is None: + raise ValueError( + f"cannot fingerprint unknown parameter {target.component!r}/{target.parameter!r}" + ) + if initializer.const_value is None: + raise ValueError( + f"cannot fingerprint unloaded initializer " + f"{target.component!r}/{target.parameter!r}" + ) + consumers: list[dict[str, object]] = [] + for node_ordinal, node in enumerate(model.graph): + for input_ordinal, node_input in enumerate(node.inputs): + if node_input is not initializer: + continue + consumers.append( + { + "attributes": { + name: _canonical_attribute_value(attribute) + for name, attribute in sorted(node.attributes.items()) + }, + "domain": node.domain or "ai.onnx", + "input_ordinal": input_ordinal, + "node_ordinal": node_ordinal, + "op_type": node.op_type, + } + ) + record: dict[str, object] = { + "component": target.component, + "consumers": consumers, + "dtype": int(initializer.dtype), + "initializer": target.parameter, + "shape": [int(dimension) for dimension in initializer.shape], + "tensor_sha256": hashlib.sha256(_tensor_bytes(initializer.const_value)).hexdigest(), + } + if isinstance(resolved_target, AdapterTargetDescriptor): + record["id"] = resolved_target.semantic_name + record["node_name"] = resolved_target.node_name + record["output_name"] = resolved_target.output_name + record["activation_dtype"] = int(resolved_target.activation_dtype or initializer.dtype) + if resolved_target.graph_input_a is not None: + record["graph_inputs"] = { + "a": resolved_target.graph_input_a, + "b": resolved_target.graph_input_b, + **( + {"scale": resolved_target.graph_input_scale} + if resolved_target.graph_input_scale is not None + else {} + ), + } + return record + + +def fingerprint_model_weights( + models: Mapping[str, ir.Model], + targets: Sequence[AdapterTarget | AdapterTargetDescriptor] | None = None, +) -> str: + """Fingerprint the exact immutable base parameters targeted by adapters.""" + if targets is None: + targets = tuple( + AdapterTarget(component, parameter) + for component, model in sorted(models.items()) + for parameter in sorted(model.graph.initializers) + ) + unique_targets = sorted( + targets, + key=lambda item: ( + item.target.component + if isinstance(item, AdapterTargetDescriptor) + else item.component, + item.target.parameter + if isinstance(item, AdapterTargetDescriptor) + else item.parameter, + item.semantic_name if isinstance(item, AdapterTargetDescriptor) else "", + ), + ) + if not unique_targets: + raise ValueError("adapter base fingerprint requires at least one target") + records: list[dict[str, object]] = [] + for target in unique_targets: + record = _target_fingerprint_record(models, target) + records.append(record) + if isinstance(target, AdapterTargetDescriptor): + for target_slice in sorted(target.slices, key=lambda item: item.role): + sliced = dict(record) + sliced["id"] = f"{target.semantic_name}.{target_slice.role}" + sliced["output_slice"] = { + "role": target_slice.role, + "offset": target_slice.offset, + "width": target_slice.width, + } + records.append(sliced) + canonical = rfc8785.dumps( + { + "schema": "onnx-genai-targeted-base-v1", + "targets": records, + } + ) + return f"onnx-genai-targeted-base-v1:sha256:{hashlib.sha256(canonical).hexdigest()}" + + +@dataclasses.dataclass(frozen=True) +class AdapterTarget: + """A base-model parameter addressed without architecture-specific aliases.""" + + component: str + parameter: str + + def __post_init__(self) -> None: + _validate_identifier(self.component, "adapter target component") + if not self.parameter: + raise ValueError("adapter target parameter must be non-empty") + + +@dataclasses.dataclass(frozen=True) +class AdapterTargetSlice: + """One semantic child occupying a contiguous fused-projection output slice.""" + + role: str + offset: int + width: int + rank: int | None = None + alpha: float | None = None + + def __post_init__(self) -> None: + if not self.role: + raise ValueError("adapter target slice role must be non-empty") + if self.offset < 0 or self.width <= 0: + raise ValueError("adapter target slice offset/width must be non-negative/positive") + if self.rank is not None and self.rank <= 0: + raise ValueError("adapter target slice rank must be positive") + if self.alpha is not None and not math.isfinite(self.alpha): + raise ValueError("adapter target slice alpha must be finite") + + +@dataclasses.dataclass(frozen=True) +class AdapterTargetDescriptor: + """Authoritative producer binding from a semantic target to an exact ONNX value.""" + + target: AdapterTarget + semantic_name: str + node_name: str + output_name: str + input_size: int + output_size: int + layer_index: int | None = None + rank: int | None = None + alpha: float | None = None + slices: tuple[AdapterTargetSlice, ...] = () + activation_dtype: ir.DataType | None = None + graph_input_a: str | None = None + graph_input_b: str | None = None + graph_input_scale: str | None = None + + def __post_init__(self) -> None: + if not self.semantic_name or not self.node_name or not self.output_name: + raise ValueError("adapter semantic, node, and output names must be non-empty") + if self.input_size <= 0 or self.output_size <= 0: + raise ValueError("adapter target input/output dimensions must be positive") + if self.layer_index is not None and self.layer_index < 0: + raise ValueError("adapter target layer index must be non-negative") + if self.rank is not None and self.rank <= 0: + raise ValueError("adapter target rank must be positive") + if self.alpha is not None and not math.isfinite(self.alpha): + raise ValueError("adapter target alpha must be finite") + roles = [item.role for item in self.slices] + if len(roles) != len(set(roles)): + raise ValueError("adapter target slice roles must be unique") + ordered = sorted(self.slices, key=lambda item: item.offset) + for previous, current in itertools.pairwise(ordered): + if previous.offset + previous.width > current.offset: + raise ValueError("adapter target slices must not overlap") + if ordered and ordered[-1].offset + ordered[-1].width > self.output_size: + raise ValueError("adapter target slice exceeds the projection output dimension") + if (self.graph_input_a is None) != (self.graph_input_b is None): + raise ValueError("adapter graph inputs must declare paired A/B names") + graph_inputs = [ + name + for name in (self.graph_input_a, self.graph_input_b, self.graph_input_scale) + if name is not None + ] + if any(not name for name in graph_inputs) or len(graph_inputs) != len( + set(graph_inputs) + ): + raise ValueError( + "adapter graph inputs must contain distinct non-empty A/B/scale names" + ) + + +@dataclasses.dataclass(frozen=True) +class AdapterTargetManifest: + """Authoritative model-export target map; runtimes need no family discovery.""" + + base_fingerprint: str + targets: tuple[AdapterTargetDescriptor, ...] + + def __post_init__(self) -> None: + if not self.base_fingerprint: + raise ValueError("adapter target manifest base fingerprint must be non-empty") + if not self.targets: + raise ValueError("adapter target manifest must contain at least one target") + bindings = [descriptor.target for descriptor in self.targets] + if len(bindings) != len(set(bindings)): + raise ValueError("adapter target manifest contains duplicate bindings") + semantics = [descriptor.semantic_name for descriptor in self.targets] + if len(semantics) != len(set(semantics)): + raise ValueError("adapter target manifest contains duplicate semantic names") + + def validate(self, models: Mapping[str, ir.Model]) -> None: + for descriptor in self.targets: + model = models.get(descriptor.target.component) + if model is None: + raise ValueError( + f"adapter manifest targets unknown component " + f"{descriptor.target.component!r}" + ) + initializer = model.graph.initializers.get(descriptor.target.parameter) + if initializer is None: + raise ValueError( + f"adapter manifest targets unknown parameter " + f"{descriptor.target.component!r}/{descriptor.target.parameter!r}" + ) + shape = [int(dimension) for dimension in initializer.shape] + expected_shape = [descriptor.output_size, descriptor.input_size] + if shape != expected_shape: + raise ValueError( + f"adapter manifest parameter {descriptor.target.parameter!r} " + f"has shape {shape}, expected {expected_shape}" + ) + nodes = [node for node in model.graph if node.name == descriptor.node_name] + if len(nodes) != 1: + raise ValueError( + f"adapter manifest node {descriptor.node_name!r} resolved " + f"{len(nodes)} times" + ) + if initializer not in nodes[0].inputs: + raise ValueError( + f"adapter manifest node {descriptor.node_name!r} does not consume " + f"parameter {descriptor.target.parameter!r}" + ) + if descriptor.output_name not in { + output.name for output in nodes[0].outputs if output.name is not None + }: + raise ValueError( + f"adapter manifest node {descriptor.node_name!r} does not produce " + f"{descriptor.output_name!r}" + ) + actual_fingerprint = fingerprint_model_weights(models, self.targets) + if actual_fingerprint != self.base_fingerprint: + raise ValueError( + "adapter target manifest base fingerprint mismatch: " + f"expected {self.base_fingerprint}, got {actual_fingerprint}" + ) + + @property + def bindings(self) -> frozenset[AdapterTarget]: + """Exact base parameters covered by the manifest.""" + return frozenset(descriptor.target for descriptor in self.targets) + + +@dataclasses.dataclass(frozen=True) +class AdapterSource: + """Artifact provenance retained independently from runtime container choice.""" + + format: Literal["in_memory", "peft_safetensors", "onnx_adapter"] + path: str | None = None + checksum: str | None = None + base_model: str | None = None + revision: str | None = None + producer: str = "mobius" + + def __post_init__(self) -> None: + if not self.producer: + raise ValueError("adapter source producer must be non-empty") + if self.format != "in_memory" and not self.path: + raise ValueError(f"{self.format} adapter source requires a path") + if self.checksum is not None and not self.checksum.startswith("sha256:"): + raise ValueError("adapter source checksum must use sha256") + + +@dataclasses.dataclass(frozen=True) +class AdapterWeights: + """LoRA factors for one target, computing ``B @ A * alpha / rank``.""" + + target: AdapterTarget + a: ir.Tensor + b: ir.Tensor + alpha: float + weight_key: str | None = None + target_id: str | None = None + + def __post_init__(self) -> None: + if len(self.a.shape) != 2 or len(self.b.shape) != 2: + raise ValueError("adapter A and B factors must both be rank-2 tensors") + if int(self.a.shape[0]) <= 0: + raise ValueError("adapter rank must be positive") + if int(self.b.shape[1]) != int(self.a.shape[0]): + raise ValueError( + "adapter B input dimension must equal adapter A rank " + f"({self.b.shape[1]} != {self.a.shape[0]})" + ) + if self.a.dtype != self.b.dtype: + raise ValueError("adapter A and B factors must have the same dtype") + if not math.isfinite(self.alpha) or self.alpha <= 0.0: + raise ValueError("adapter alpha must be finite and greater than zero") + if self.weight_key is not None and not self.weight_key: + raise ValueError("adapter weight key must be non-empty") + if self.target_id is not None and not self.target_id: + raise ValueError("adapter target ID must be non-empty") + + @property + def rank(self) -> int: + return int(self.a.shape[0]) + + @property + def dtype(self) -> ir.DataType: + return self.a.dtype + + def delta(self) -> np.ndarray: + """Materialize the reference LoRA update for validation and parity tests.""" + return (self.b.numpy() @ self.a.numpy()) * (self.alpha / self.rank) + + +@dataclasses.dataclass(frozen=True) +class AdapterArtifact: + """Immutable low-rank adapter data bound to one exact base-model fingerprint.""" + + name: str + base_fingerprint: str + weights: tuple[AdapterWeights, ...] + source: AdapterSource = dataclasses.field( + default_factory=lambda: AdapterSource("in_memory") + ) + identity: str | None = None + version: str = "1" + + def __post_init__(self) -> None: + _validate_identifier(self.name, "adapter name") + if not self.base_fingerprint: + raise ValueError("adapter base fingerprint must be non-empty") + if self.identity is not None and not self.identity: + raise ValueError("adapter identity must be non-empty") + if not self.version: + raise ValueError("adapter version must be non-empty") + if not self.weights: + raise ValueError("adapter artifact must contain at least one target") + targets = [weight.target for weight in self.weights] + if len(targets) != len(set(targets)): + raise ValueError("adapter artifact contains duplicate targets") + + @property + def checksum(self) -> str: + """Deterministic checksum covering bindings, metadata, and tensors.""" + digest = hashlib.sha256() + digest.update(self.name.encode()) + digest.update(b"\0") + digest.update(self.base_fingerprint.encode()) + for weight in sorted( + self.weights, key=lambda item: (item.target.component, item.target.parameter) + ): + metadata = { + "alpha": weight.alpha, + "component": weight.target.component, + "parameter": weight.target.parameter, + "rank": weight.rank, + } + digest.update(json.dumps(metadata, sort_keys=True, separators=(",", ":")).encode()) + digest.update(b"\0") + _update_tensor_hash(digest, weight.a) + _update_tensor_hash(digest, weight.b) + return f"sha256:{digest.hexdigest()}" + + @property + def nbytes(self) -> int: + """Resident tensor bytes required by this artifact's factor pages.""" + return sum( + weight.a.numpy().nbytes + weight.b.numpy().nbytes for weight in self.weights + ) + + def validate_base( + self, + models: Mapping[str, ir.Model], + *, + fingerprint_targets: Sequence[AdapterTarget | AdapterTargetDescriptor] | None = None, + ) -> None: + """Validate fingerprint, target existence, dtype, and matrix dimensions.""" + target_shapes: dict[str, tuple[AdapterTarget, int, int]] = {} + for resolved in fingerprint_targets or (): + if not isinstance(resolved, AdapterTargetDescriptor): + continue + target_shapes[resolved.semantic_name] = ( + resolved.target, + resolved.input_size, + resolved.output_size, + ) + for target_slice in resolved.slices: + target_shapes[f"{resolved.semantic_name}.{target_slice.role}"] = ( + resolved.target, + resolved.input_size, + target_slice.width, + ) + for weight in self.weights: + model = models.get(weight.target.component) + if model is None: + raise ValueError( + f"adapter {self.name!r} targets unknown component " + f"{weight.target.component!r}" + ) + initializer = model.graph.initializers.get(weight.target.parameter) + if initializer is None: + raise ValueError( + f"adapter {self.name!r} targets unknown parameter " + f"{weight.target.component!r}/{weight.target.parameter!r}" + ) + expected_shape = [int(weight.b.shape[0]), int(weight.a.shape[1])] + actual_shape = [int(dimension) for dimension in initializer.shape] + resolved_shape = target_shapes.get(weight.target_id or "") + if weight.target_id is not None and resolved_shape is None: + raise ValueError( + f"adapter target ID {weight.target_id!r} is absent from " + "the authoritative target manifest" + ) + if resolved_shape is None and actual_shape != expected_shape: + raise ValueError( + f"adapter target {weight.target.component!r}/" + f"{weight.target.parameter!r} has shape {actual_shape}, " + f"but B @ A has shape {expected_shape}" + ) + if resolved_shape is not None: + resolved_target, input_features, output_features = resolved_shape + if resolved_target != weight.target: + raise ValueError( + f"adapter target ID {weight.target_id!r} resolves to " + f"{resolved_target}, not {weight.target}" + ) + descriptor = next( + item + for item in fingerprint_targets or () + if isinstance(item, AdapterTargetDescriptor) + and item.target == weight.target + ) + base_shape = [descriptor.output_size, descriptor.input_size] + if actual_shape != base_shape or expected_shape != [ + output_features, + input_features, + ]: + raise ValueError( + f"adapter target ID {weight.target_id!r} expects base/factor " + f"shapes {base_shape}/[{output_features}, {input_features}], " + f"got {actual_shape}/{expected_shape}" + ) + if initializer.dtype != weight.dtype: + raise ValueError( + f"adapter target {weight.target.component!r}/" + f"{weight.target.parameter!r} has dtype {initializer.dtype.name}, " + f"but adapter factors have dtype {weight.dtype.name}" + ) + actual_fingerprint = fingerprint_model_weights( + models, + fingerprint_targets or tuple(weight.target for weight in self.weights), + ) + if actual_fingerprint != self.base_fingerprint: + raise ValueError( + f"adapter {self.name!r} base fingerprint mismatch: " + f"expected {self.base_fingerprint}, got {actual_fingerprint}" + ) + + @property + def target_bindings(self) -> frozenset[AdapterTarget]: + """Exact base parameters modified by this artifact.""" + return frozenset(weight.target for weight in self.weights) + + @property + def stable_identity(self) -> str: + """Stable identity used independently from the package catalog alias.""" + return self.identity or self.name + + def validate_checksum(self, expected: str) -> None: + """Reject corrupted or substituted adapter tensor data.""" + if self.checksum != expected: + raise ValueError( + f"adapter {self.name!r} checksum mismatch: " + f"expected {expected}, got {self.checksum}" + ) + + +@dataclasses.dataclass(frozen=True) +class AdapterApplication: + """One adapter and its request-local scale in composition order.""" + + adapter: str + scale: float = 1.0 + + def __post_init__(self) -> None: + _validate_identifier(self.adapter, "adapter application name") + if not math.isfinite(self.scale) or not -16.0 <= self.scale <= 16.0: + raise ValueError("adapter application scale must be finite and within [-16, 16]") + + +@dataclasses.dataclass(frozen=True) +class AdapterServiceOptions: + """Producer-neutral runtime lifecycle, planning, and artifact format options.""" + + segments: str = "request.adapter_segments" + adapter_counts: str = "request.adapter_counts" + scales: str = "request.adapter_scales" + active: str | None = None + max_adapters: int = 4 + application_capability: str = "onnx-genai.adapters@1" + portable_fallback: bool = True + cache_max_entries: int = 16 + bucket_by_adapter_set: bool = True + stable_buffers: bool = True + invalidate_capture_on_eviction: bool = True + preserve_source_format: bool = False + discovery_fallback: Literal["disabled", "tooling_only"] = "disabled" + + def __post_init__(self) -> None: + if not self.application_capability: + raise ValueError("adapter application capability must be non-empty") + if self.cache_max_entries <= 0: + raise ValueError("adapter cache max_entries must be greater than zero") + if self.max_adapters <= 0: + raise ValueError("adapter max_adapters must be greater than zero") + if self.discovery_fallback not in {"disabled", "tooling_only"}: + raise ValueError("adapter discovery fallback must be disabled or tooling_only") + if not self.portable_fallback and not self.preserve_source_format: + raise ValueError( + "adapter service must emit a portable fallback or preserved source format" + ) + + +@dataclasses.dataclass(frozen=True) +class AdapterSlotSelection: + """Adapter composition for one stable semantic serving slot.""" + + slot_id: int + request_epoch: int + adapters: tuple[AdapterApplication, ...] = () + + def __post_init__(self) -> None: + if self.request_epoch < 0: + raise ValueError("adapter request epoch must be non-negative") + names = [application.adapter for application in self.adapters] + if len(names) != len(set(names)): + raise ValueError("adapter slot contains duplicate adapter") + + +@dataclasses.dataclass(frozen=True) +class AdapterBatchSelection: + """Fixed-shape, compaction-safe adapter state for a heterogeneous batch.""" + + slots: tuple[AdapterSlotSelection, ...] + + def __post_init__(self) -> None: + slot_ids = [slot.slot_id for slot in self.slots] + if len(slot_ids) != len(set(slot_ids)): + raise ValueError("adapter batch slot IDs must be unique") + + def validate_catalog(self, artifacts: Mapping[str, AdapterArtifact]) -> None: + for slot in self.slots: + for application in slot.adapters: + if application.adapter not in artifacts: + raise ValueError( + f"slot {slot.slot_id} selects unknown adapter {application.adapter!r}" + ) + + def compact(self, permutation: Sequence[int]) -> AdapterBatchSelection: + """Apply the same physical-slot permutation used for all workflow state.""" + if sorted(permutation) != list(range(len(self.slots))): + raise ValueError("adapter compaction must be a permutation of all batch slots") + return AdapterBatchSelection(tuple(self.slots[index] for index in permutation)) + + def to_tensors( + self, + artifacts: Mapping[str, AdapterArtifact], + *, + max_adapters: int, + active: Sequence[bool] | None = None, + ) -> AdapterSelectionTensors: + """Lower aliases to fixed-shape request tensors without serializing numeric IDs.""" + if max_adapters <= 0: + raise ValueError("adapter max_adapters must be greater than zero") + if active is None: + active = [True] * len(self.slots) + if len(active) != len(self.slots): + raise ValueError("adapter active slots must match the selection batch size") + self.validate_catalog(artifacts) + aliases = tuple(sorted(artifacts)) + indices = {alias: index for index, alias in enumerate(aliases)} + segments = np.full((len(self.slots), max_adapters), -1, dtype=np.int64) + scales = np.zeros((len(self.slots), max_adapters), dtype=np.float32) + counts = np.zeros((len(self.slots),), dtype=np.int64) + for slot_index, (slot, is_active) in enumerate(zip(self.slots, active)): + if not is_active: + continue + if len(slot.adapters) > max_adapters: + raise ValueError( + f"adapter slot {slot.slot_id} selects {len(slot.adapters)} adapters, " + f"exceeding max_adapters {max_adapters}" + ) + counts[slot_index] = len(slot.adapters) + for composition_index, application in enumerate(slot.adapters): + segments[slot_index, composition_index] = indices[application.adapter] + scales[slot_index, composition_index] = application.scale + return AdapterSelectionTensors( + slot_ids=np.asarray([slot.slot_id for slot in self.slots], dtype=np.int64), + request_epochs=np.asarray( + [slot.request_epoch for slot in self.slots], dtype=np.int64 + ), + segments=segments, + adapter_counts=counts, + scales=scales, + active=np.asarray(active, dtype=np.bool_), + aliases=aliases, + ) + + @property + def referenced_adapters(self) -> frozenset[str]: + """Live adapter set that a paged runtime must pin against eviction.""" + return frozenset( + application.adapter for slot in self.slots for application in slot.adapters + ) + + +@dataclasses.dataclass(frozen=True) +class AdapterSelectionTensors: + """Fixed-shape SSA request buffers for the ``onnx-genai.adapters@1`` ABI.""" + + slot_ids: np.ndarray + request_epochs: np.ndarray + segments: np.ndarray + adapter_counts: np.ndarray + scales: np.ndarray + active: np.ndarray + aliases: tuple[str, ...] + + +def compose_adapter_deltas( + slot: AdapterSlotSelection, + artifacts: Mapping[str, AdapterArtifact], +) -> dict[AdapterTarget, np.ndarray]: + """Compose a slot's selected adapters into reference parameter updates.""" + deltas: dict[AdapterTarget, np.ndarray] = {} + for application in slot.adapters: + try: + artifact = artifacts[application.adapter] + except KeyError as error: + raise ValueError( + f"slot {slot.slot_id} selects unknown adapter {application.adapter!r}" + ) from error + for weight in artifact.weights: + update = weight.delta() * application.scale + if weight.target in deltas: + deltas[weight.target] = deltas[weight.target] + update + else: + deltas[weight.target] = update + return deltas diff --git a/src/mobius/adapters_test.py b/src/mobius/adapters_test.py new file mode 100644 index 000000000..edba87d6b --- /dev/null +++ b/src/mobius/adapters_test.py @@ -0,0 +1,874 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT License. + +"""Tests for generic low-rank adapter artifacts and request state.""" + +from __future__ import annotations + +import hashlib +import json +import shutil +import uuid +from pathlib import Path + +import numpy as np +import onnx_ir as ir +import pytest +from safetensors.numpy import load_file, save_file + +from mobius import ( + AdapterApplication, + AdapterArtifact, + AdapterBatchSelection, + AdapterServiceOptions, + AdapterSlotSelection, + AdapterSource, + AdapterTarget, + AdapterTargetDescriptor, + AdapterTargetManifest, + AdapterTargetSlice, + AdapterWeights, + ModelPackage, + adapter_source_from_onnx_adapter, + compose_adapter_deltas, + fingerprint_model_weights, + load_peft_adapter, +) +from mobius.integrations.onnx_genai.inference_metadata import ( + add_adapter_service_to_metadata, +) + + +def _model(weight: np.ndarray | None = None) -> ir.Model: + values = np.arange(12, dtype=np.float32).reshape(3, 4) if weight is None else weight + initializer = ir.Value( + name="projection.weight", + const_value=ir.tensor(values), + type=ir.TensorType(ir.DataType.FLOAT), + shape=ir.Shape(values.shape), + ) + x = ir.val( + "hidden_states", + type=ir.TensorType(ir.DataType.FLOAT), + shape=ir.Shape([1, 4]), + ) + projection = ir.Node("", "MatMul", [x, initializer], name="projection") + output = projection.outputs[0] + output.name = "projection.output" + graph = ir.Graph( + [x], + [output], + nodes=[projection], + initializers=[initializer], + name="adapter_test_model", + ) + return ir.Model(graph, ir_version=10) + + +def _weights( + *, + component: str = "decoder", + parameter: str = "projection.weight", + a: np.ndarray | None = None, + b: np.ndarray | None = None, + alpha: float = 4.0, + weight_key: str | None = None, + target_id: str | None = None, +) -> AdapterWeights: + a_values = np.arange(8, dtype=np.float32).reshape(2, 4) / 10 if a is None else a + b_values = np.arange(6, dtype=np.float32).reshape(3, 2) / 10 if b is None else b + return AdapterWeights( + AdapterTarget(component, parameter), + ir.tensor(a_values), + ir.tensor(b_values), + alpha, + weight_key=weight_key, + target_id=target_id, + ) + + +def _artifact(model: ir.Model, name: str = "style") -> AdapterArtifact: + models = {"decoder": model} + weights = _weights() + return AdapterArtifact( + name=name, + base_fingerprint=fingerprint_model_weights(models, (weights.target,)), + weights=(weights,), + ) + + +def _manifest(model: ir.Model, *, include_second: bool = False) -> AdapterTargetManifest: + targets = [ + AdapterTargetDescriptor( + AdapterTarget("decoder", "projection.weight"), + semantic_name="layers.0.self_attn.q_proj", + node_name="projection", + output_name="projection.output", + input_size=4, + output_size=3, + layer_index=0, + slices=(AdapterTargetSlice("q", 0, 3, rank=2, alpha=4.0),), + activation_dtype=ir.DataType.FLOAT, + graph_input_a="lora.q_proj.a", + graph_input_b="lora.q_proj.b", + graph_input_scale="lora.q_proj.scale", + ) + ] + if include_second: + targets.append( + AdapterTargetDescriptor( + AdapterTarget("decoder", "other.weight"), + semantic_name="layers.0.self_attn.v_proj", + node_name="other", + output_name="other.output", + input_size=4, + output_size=3, + ) + ) + return AdapterTargetManifest( + fingerprint_model_weights({"decoder": model}, tuple(targets)), + tuple(targets), + ) + + +def test_artifact_validates_base_target_shape_dtype_and_checksum() -> None: + model = _model() + artifact = _artifact(model) + + artifact.validate_base({"decoder": model}) + assert artifact.checksum.startswith("sha256:") + assert artifact.checksum == _artifact(model).checksum + + changed = _model(np.ones((3, 4), dtype=np.float32)) + with pytest.raises(ValueError, match="base fingerprint mismatch"): + artifact.validate_base({"decoder": changed}) + + +def test_targeted_fingerprint_excludes_unrelated_weights_and_includes_consumers() -> None: + model = _model() + target = AdapterTarget("decoder", "projection.weight") + fingerprint = fingerprint_model_weights({"decoder": model}, (target,)) + assert fingerprint.startswith("onnx-genai-targeted-base-v1:sha256:") + + unrelated = ir.Value( + name="unrelated.weight", + const_value=ir.tensor(np.ones((1,), dtype=np.float32)), + type=ir.TensorType(ir.DataType.FLOAT), + shape=ir.Shape([1]), + ) + model.graph.initializers.add(unrelated) + assert fingerprint_model_weights({"decoder": model}, (target,)) == fingerprint + + next(iter(model.graph)).attributes["producer_contract"] = ir.AttrInt64( + "producer_contract", 1 + ) + assert fingerprint_model_weights({"decoder": model}, (target,)) != fingerprint + + +def test_authoritative_target_manifest_validates_exact_graph_binding() -> None: + model = _model() + manifest = _manifest(model) + manifest.validate({"decoder": model}) + assert manifest.bindings == {AdapterTarget("decoder", "projection.weight")} + + stale = AdapterTargetManifest( + manifest.base_fingerprint, + ( + AdapterTargetDescriptor( + AdapterTarget("decoder", "projection.weight"), + "layers.0.self_attn.q_proj", + "projection", + "stale.output", + 4, + 3, + ), + ), + ) + with pytest.raises(ValueError, match="does not produce"): + stale.validate({"decoder": model}) + + +@pytest.mark.parametrize( + ("weights", "message"), + [ + (_weights(parameter="missing.weight"), "unknown parameter"), + ( + _weights(a=np.ones((2, 5), dtype=np.float32)), + "B @ A has shape", + ), + ], +) +def test_artifact_rejects_invalid_targets(weights: AdapterWeights, message: str) -> None: + model = _model() + artifact = AdapterArtifact( + name="invalid", + base_fingerprint=fingerprint_model_weights({"decoder": model}), + weights=(weights,), + ) + with pytest.raises(ValueError, match=message): + artifact.validate_base({"decoder": model}) + + +def test_adapter_weights_reject_mismatched_factor_dtype() -> None: + with pytest.raises(ValueError, match="same dtype"): + _weights(a=np.ones((2, 4), dtype=np.float16)) + + +def test_lora_delta_matches_reference_math() -> None: + weights = _weights(alpha=6.0) + expected = weights.b.numpy() @ weights.a.numpy() * 3.0 + np.testing.assert_allclose(weights.delta(), expected, rtol=1e-6, atol=1e-6) + + +def test_composed_delta_matches_scaled_sum() -> None: + model = _model() + style = _artifact(model, "style") + speaker = AdapterArtifact( + "speaker", + style.base_fingerprint, + (_weights(alpha=2.0),), + ) + row = AdapterSlotSelection( + 100, + 1, + ( + AdapterApplication("style", 0.25), + AdapterApplication("speaker", 1.5), + ), + ) + actual = compose_adapter_deltas(row, {"style": style, "speaker": speaker}) + target = AdapterTarget("decoder", "projection.weight") + expected = style.weights[0].delta() * 0.25 + speaker.weights[0].delta() * 1.5 + np.testing.assert_allclose(actual[target], expected, rtol=1e-6, atol=1e-6) + + +def test_zero_one_and_composed_per_row_adapters() -> None: + batch = AdapterBatchSelection( + ( + AdapterSlotSelection(slot_id=100, request_epoch=4), + AdapterSlotSelection( + slot_id=101, + request_epoch=7, + adapters=(AdapterApplication("style", 0.5),), + ), + AdapterSlotSelection( + slot_id=102, + request_epoch=2, + adapters=( + AdapterApplication("style", 0.25), + AdapterApplication("speaker", 1.5), + ), + ), + ) + ) + model = _model() + catalog = { + "style": _artifact(model, "style"), + "speaker": _artifact(model, "speaker"), + } + batch.validate_catalog(catalog) + + assert batch.slots[0].adapters == () + assert [item.adapter for item in batch.slots[2].adapters] == ["style", "speaker"] + + +def test_compaction_preserves_semantic_rows_and_slot_reuse_uses_epoch() -> None: + original = AdapterBatchSelection( + ( + AdapterSlotSelection(100, 1, (AdapterApplication("style"),)), + AdapterSlotSelection(101, 5, (AdapterApplication("speaker", 0.25),)), + ) + ) + compacted = original.compact([1, 0]) + assert [slot.slot_id for slot in compacted.slots] == [101, 100] + assert compacted.slots[0].request_epoch == 5 + assert compacted.compact([1, 0]) == original + assert compacted.referenced_adapters == {"style", "speaker"} + + reused_slot = AdapterSlotSelection( + slot_id=200, + request_epoch=2, + adapters=(AdapterApplication("speaker"),), + ) + assert reused_slot != original.slots[0] + + +def test_selection_lowers_to_stable_fixed_shape_request_tensors() -> None: + model = _model() + artifact = _artifact(model) + batch = AdapterBatchSelection( + ( + AdapterSlotSelection( + 100, + 4, + ( + AdapterApplication("red", 0.5), + AdapterApplication("blue", -0.25), + ), + ), + AdapterSlotSelection(101, 5, (AdapterApplication("blue", 1.0),)), + ) + ) + tensors = batch.to_tensors( + {"blue": artifact, "red": artifact}, + max_adapters=3, + active=[True, False], + ) + assert tensors.aliases == ("blue", "red") + np.testing.assert_array_equal(tensors.slot_ids, [100, 101]) + np.testing.assert_array_equal(tensors.request_epochs, [4, 5]) + np.testing.assert_array_equal(tensors.segments, [[1, 0, -1], [-1, -1, -1]]) + np.testing.assert_array_equal(tensors.adapter_counts, [2, 0]) + np.testing.assert_array_equal(tensors.scales, [[0.5, -0.25, 0.0], [0.0, 0.0, 0.0]]) + assert tensors.segments.dtype == np.int64 + assert tensors.scales.dtype == np.float32 + + compacted = batch.compact([1, 0]).to_tensors( + {"blue": artifact, "red": artifact}, + max_adapters=3, + ) + np.testing.assert_array_equal(compacted.slot_ids, [101, 100]) + np.testing.assert_array_equal(compacted.request_epochs, [5, 4]) + with pytest.raises(ValueError, match="exceeding max_adapters"): + batch.to_tensors({"blue": artifact, "red": artifact}, max_adapters=1) + + +def test_model_package_catalog_validates_and_rejects_duplicates() -> None: + model = _model() + manifest = _manifest(model) + artifact = AdapterArtifact("style", manifest.base_fingerprint, (_weights(),)) + package = ModelPackage({"decoder": model}, adapter_target_manifest=manifest) + package.add_adapter_artifact(artifact) + assert package.adapter_artifacts["style"] is artifact + assert artifact.nbytes == 56 + + with pytest.raises(ValueError, match="already attached"): + package.add_adapter_artifact(artifact) + + with pytest.raises(ValueError, match="checksum mismatch"): + artifact.validate_checksum("sha256:" + "0" * 64) + + +def test_model_package_allows_distinct_manifest_targets_per_adapter() -> None: + model = _model() + other = ir.Value( + name="other.weight", + const_value=ir.tensor(np.ones((3, 4), dtype=np.float32)), + type=ir.TensorType(ir.DataType.FLOAT), + shape=ir.Shape([3, 4]), + ) + x = model.graph.inputs[0] + node = ir.Node("", "MatMul", [x, other], name="other") + node.outputs[0].name = "other.output" + model.graph.append(node) + model.graph.initializers.add(other) + manifest = _manifest(model, include_second=True) + package = ModelPackage({"decoder": model}, adapter_target_manifest=manifest) + package.add_adapter_artifact( + AdapterArtifact("style", manifest.base_fingerprint, (_weights(),)) + ) + second = AdapterArtifact( + "speaker", + manifest.base_fingerprint, + (_weights(parameter="other.weight"),), + ) + package.add_adapter_artifact(second) + assert package.adapter_artifacts["style"].target_bindings == { + AdapterTarget("decoder", "projection.weight") + } + assert package.adapter_artifacts["speaker"].target_bindings == { + AdapterTarget("decoder", "other.weight") + } + + +def test_peft_migration_source_preserves_rank_alpha_and_provenance() -> None: + directory = Path("artifacts") / f"adapter-peft-test-{uuid.uuid4().hex}" + directory.mkdir(parents=True) + try: + config = { + "base_model_name_or_path": "synthetic/base", + "revision": "producer-fixture", + "r": 4, + "lora_alpha": 8.0, + "target_modules": ["q_proj"], + "rank_pattern": {"layers.0.self_attn.q_proj": 2}, + "alpha_pattern": {"self_attn.q_proj": 6.0}, + } + (directory / "adapter_config.json").write_text(json.dumps(config)) + module = "base_model.model.layers.0.self_attn.q_proj" + a = np.arange(8, dtype=np.float32).reshape(2, 4) + b = np.arange(6, dtype=np.float32).reshape(3, 2) + save_file( + { + f"{module}.lora_A.weight": a, + f"{module}.lora_B.weight": b, + }, + directory / "adapter_model.safetensors", + ) + model = _model() + manifest = _manifest(model) + artifact = load_peft_adapter( + directory, + name="peft-style", + base_fingerprint=manifest.base_fingerprint, + target_bindings={ + "layers.0.self_attn.q_proj": AdapterTarget("decoder", "projection.weight") + }, + ) + artifact.validate_base({"decoder": model}, fingerprint_targets=manifest.targets) + assert artifact.weights[0].rank == 2 + assert artifact.weights[0].alpha == pytest.approx(6.0) + assert artifact.source.format == "peft_safetensors" + assert artifact.source.base_model == "synthetic/base" + assert artifact.source.revision == "producer-fixture" + package = ModelPackage( + {"decoder": model}, + adapter_target_manifest=manifest, + adapter_service_options=AdapterServiceOptions( + portable_fallback=False, + preserve_source_format=True, + ), + ) + package.add_adapter_artifact(artifact) + output = directory / "package" + catalog = package.save_adapter_artifacts(str(output)) + declaration = catalog["peft-style"]["weights"][0] + assert declaration["format"] == "hf_peft" + assert declaration["loader_capability"] == "onnx-genai.adapters.hf-peft@1" + assert declaration["scale_encoding"] == "alpha_over_rank" + assert catalog["peft-style"]["provenance"] == { + "producer": "mobius", + "source": "synthetic/base", + "revision": "producer-fixture", + } + saved = load_file(output / declaration["location"]) + assert set(saved) == { + f"{module}.lora_A.weight", + f"{module}.lora_B.weight", + } + assert (output / declaration["config_location"]).read_bytes() == ( + directory / "adapter_config.json" + ).read_bytes() + np.testing.assert_allclose( + artifact.weights[0].delta(), b @ a * 3.0, rtol=1e-6, atol=1e-6 + ) + finally: + shutil.rmtree(directory) + + +def test_peft_rank_and_alpha_patterns_emit_heterogeneous_binding_overrides() -> None: + directory = Path("artifacts") / f"adapter-peft-pattern-test-{uuid.uuid4().hex}" + source = directory / "source" + source.mkdir(parents=True) + try: + model = _model() + other = ir.Value( + name="other.weight", + const_value=ir.tensor(np.ones((3, 4), dtype=np.float32)), + type=ir.TensorType(ir.DataType.FLOAT), + shape=ir.Shape([3, 4]), + ) + node = ir.Node("", "MatMul", [model.graph.inputs[0], other], name="other") + node.outputs[0].name = "other.output" + model.graph.append(node) + model.graph.initializers.add(other) + manifest = _manifest(model, include_second=True) + + config = { + "r": 1, + "lora_alpha": 2.0, + "target_modules": ["q_proj", "v_proj"], + "rank_pattern": {"layers.0.self_attn.q_proj": 2}, + "alpha_pattern": {"layers.0.self_attn.q_proj": 6.0}, + } + (source / "adapter_config.json").write_text(json.dumps(config)) + save_file( + { + "base_model.layers.0.self_attn.q_proj.lora_A.weight": np.ones( + (2, 4), dtype=np.float32 + ), + "base_model.layers.0.self_attn.q_proj.lora_B.weight": np.ones( + (3, 2), dtype=np.float32 + ), + "base_model.layers.0.self_attn.v_proj.lora_A.weight": np.ones( + (1, 4), dtype=np.float32 + ), + "base_model.layers.0.self_attn.v_proj.lora_B.weight": np.ones( + (3, 1), dtype=np.float32 + ), + }, + source / "adapter_model.safetensors", + ) + artifact = load_peft_adapter( + source, + name="heterogeneous", + base_fingerprint=manifest.base_fingerprint, + target_bindings={ + "layers.0.self_attn.q_proj": AdapterTarget("decoder", "projection.weight"), + "layers.0.self_attn.v_proj": AdapterTarget("decoder", "other.weight"), + }, + ) + package = ModelPackage( + {"decoder": model}, + adapter_target_manifest=manifest, + adapter_service_options=AdapterServiceOptions( + portable_fallback=False, + preserve_source_format=True, + ), + ) + package.add_adapter_artifact(artifact) + catalog = package.save_adapter_artifacts(str(directory / "package")) + + assert set(catalog) == {"heterogeneous"} + declaration = catalog["heterogeneous"] + assert declaration["rank"] == 1 + assert declaration["alpha"] == pytest.approx(2.0) + bindings = {binding["target"]: binding for binding in declaration["bindings"]} + assert bindings["layers.0.self_attn.q_proj"] == { + "target": "layers.0.self_attn.q_proj", + "weight_key": "layers.0.self_attn.q_proj", + "rank": 2, + "alpha": 6.0, + } + assert bindings["layers.0.self_attn.v_proj"] == { + "target": "layers.0.self_attn.v_proj", + "weight_key": "layers.0.self_attn.v_proj", + } + assert len(declaration["weights"]) == 1 + assert declaration["weights"][0]["format"] == "hf_peft" + finally: + shutil.rmtree(directory) + + +def test_onnx_adapter_migration_source_is_optional_and_checksummed() -> None: + directory = Path("artifacts") / f"adapter-ort-test-{uuid.uuid4().hex}" + directory.mkdir(parents=True) + try: + path = directory / "style.onnx_adapter" + path.write_bytes(b"\x00\x00\x00\x00TORTsynthetic") + source = adapter_source_from_onnx_adapter(path) + assert source.format == "onnx_adapter" + assert source.checksum is not None + artifact = AdapterArtifact( + "style", + _artifact(_model()).base_fingerprint, + (_weights(),), + source=source, + ) + assert artifact.source.path == str(path) + finally: + shutil.rmtree(directory) + + +def test_onnx_adapter_source_can_be_declared_for_native_capability() -> None: + directory = Path("artifacts") / f"adapter-native-test-{uuid.uuid4().hex}" + source_directory = directory / "source" + output_directory = directory / "package" + source_directory.mkdir(parents=True) + output_directory.mkdir() + try: + source_path = source_directory / "style.onnx_adapter" + source_path.write_bytes(b"\x00\x00\x00\x00TORTsynthetic") + model = _model() + manifest = _manifest(model) + package = ModelPackage( + {"decoder": model}, + adapter_target_manifest=manifest, + adapter_service_options=AdapterServiceOptions( + portable_fallback=False, + preserve_source_format=True, + ), + ) + package.add_adapter_artifact( + AdapterArtifact( + "style", + manifest.base_fingerprint, + (_weights(alpha=2.0),), + source=adapter_source_from_onnx_adapter(source_path), + ) + ) + catalog = package.save_adapter_artifacts(str(output_directory)) + declared = catalog["style"]["weights"][0] + assert declared["format"] == "ort_genai" + assert declared["loader_capability"] == "onnxruntime.lora-adapter@1" + assert declared["scale_encoding"] == "baked" + assert declared["location"] == "adapters/style/adapter.onnx_adapter" + assert catalog["style"]["bindings"] == [ + { + "target": "layers.0.self_attn.q_proj", + "weight_key": "layers.0.self_attn.q_proj", + } + ] + copied = output_directory / declared["location"] + assert copied.read_bytes() == source_path.read_bytes() + assert declared["sha256"] == hashlib.sha256(copied.read_bytes()).hexdigest() + finally: + shutil.rmtree(directory) + + +def test_adapter_source_rejects_unprovenanced_external_artifact() -> None: + with pytest.raises(ValueError, match="requires a path"): + AdapterSource("onnx_adapter") + + +def test_exact_onnx_genai_catalog_and_portable_bundle_serialization() -> None: + directory = Path("artifacts") / f"adapter-export-test-{uuid.uuid4().hex}" + directory.mkdir(parents=True) + try: + model = _model() + manifest = _manifest(model) + package = ModelPackage( + {"decoder": model}, + adapter_target_manifest=manifest, + adapter_service_options=AdapterServiceOptions( + active="request.active", + max_adapters=2, + cache_max_entries=2, + ), + ) + package.add_adapter_artifact( + AdapterArtifact( + "red", + manifest.base_fingerprint, + (_weights(alpha=2.0),), + identity="style-red", + version="2026.08", + ) + ) + metadata = { + "pipeline": { + "workflow": { + "manifest": {"capabilities": []}, + "inputs": { + "request.slot_ids": { + "contract": { + "dtype": "int64", + "rank": 1, + "shape": ["batch"], + }, + "role": {"kind": "opaque"}, + "source": { + "kind": "application", + "name": "serving.slot_ids", + }, + }, + "request.active": { + "contract": { + "dtype": "bool", + "rank": 1, + "shape": ["batch"], + }, + "role": { + "kind": "runtime", + "version": "1.0", + "role": "adapter_active", + }, + "source": {"kind": "request"}, + }, + }, + "components": {"decoder": {"implementation": {"kind": "binding"}}}, + "steps": [], + } + } + } + add_adapter_service_to_metadata(metadata, package, str(directory)) + service = metadata["adapters"] + assert "adapters" not in metadata["pipeline"]["workflow"] + assert service["base_model_fingerprint"].startswith( + "onnx-genai-targeted-base-v1:sha256:" + ) + assert service["selection"] == { + "segments": "request.adapter_segments", + "adapter_counts": "request.adapter_counts", + "scales": "request.adapter_scales", + "active": "request.active", + "max_adapters": 2, + } + assert service["application_capability"] == "onnx-genai.adapters@1" + assert service["discovery_fallback"] == "disabled" + target = service["target_manifest"]["targets"][0] + assert target == { + "id": "layers.0.self_attn.q_proj", + "component": "decoder", + "initializer": "projection.weight", + "layer_index": 0, + "node_name": "projection", + "output_name": "projection.output", + "activation_dtype": "float32", + "input_features": 4, + "output_features": 3, + "graph_inputs": { + "a": "lora.q_proj.a", + "b": "lora.q_proj.b", + "scale": "lora.q_proj.scale", + }, + } + assert service["target_manifest"]["targets"][1]["output_slice"] == { + "role": "q", + "offset": 0, + "width": 3, + "rank": 2, + "alpha": 4.0, + } + assert service["cache"] == {"max_entries": 2, "eviction": "lru"} + assert service["planning"] == { + "bucket_by_adapter_set": True, + "stable_buffers": True, + "invalidate_capture_on_eviction": True, + } + assert not any( + "request_epochs" in name for name in metadata["pipeline"]["workflow"]["inputs"] + ) + artifact = service["artifacts"]["red"] + assert artifact["index"] == 0 + assert artifact["identity"] == "style-red" + assert artifact["version"] == "2026.08" + assert artifact["rank"] == 2 + assert artifact["alpha"] == pytest.approx(2.0) + assert artifact["dtype"] == "float32" + assert artifact["provenance"] == {"producer": "mobius"} + assert artifact["bindings"] == [ + { + "target": "layers.0.self_attn.q_proj", + "weight_key": "layers.0.self_attn.q_proj", + } + ] + weight = artifact["weights"][0] + payload = (directory / weight["location"]).read_bytes() + assert weight["format"] == "json" + assert weight["loader_capability"] == "onnx-genai.adapters.json@1" + assert weight["scale_encoding"] == "alpha_over_rank" + assert weight["location"] == "adapters/red/adapter.json" + assert len(weight["sha256"]) == 64 + assert weight["sha256"] == hashlib.sha256(payload).hexdigest() + bundle = json.loads(payload) + assert set(bundle["targets"]) == {"layers.0.self_attn.q_proj"} + assert len(bundle["targets"]["layers.0.self_attn.q_proj"]["a"]) == 8 + assert len(bundle["targets"]["layers.0.self_attn.q_proj"]["b"]) == 6 + finally: + shutil.rmtree(directory) + + +def test_top_level_adapter_metadata_supports_bare_model_package() -> None: + directory = Path("artifacts") / f"adapter-bare-test-{uuid.uuid4().hex}" + directory.mkdir(parents=True) + try: + model = _model() + target = AdapterTarget("model", "projection.weight") + descriptor = AdapterTargetDescriptor( + target, + semantic_name="projection", + node_name="projection", + output_name="projection.output", + input_size=4, + output_size=3, + activation_dtype=ir.DataType.FLOAT, + ) + fingerprint = fingerprint_model_weights({"model": model}, (descriptor,)) + package = ModelPackage( + {"model": model}, + adapter_target_manifest=AdapterTargetManifest(fingerprint, (descriptor,)), + ) + package.add_adapter_artifact( + AdapterArtifact( + "style", + fingerprint, + (_weights(component="model"),), + ) + ) + metadata: dict[str, object] = {"schema_version": "v1"} + add_adapter_service_to_metadata(metadata, package, str(directory)) + assert metadata["adapters"]["target_manifest"]["targets"][0]["component"] == "model" + assert "pipeline" not in metadata + finally: + shutil.rmtree(directory) + + +def test_wire_contract_emits_per_target_rank_override() -> None: + model = _model() + other = ir.Value( + name="other.weight", + const_value=ir.tensor(np.ones((3, 4), dtype=np.float32)), + type=ir.TensorType(ir.DataType.FLOAT), + shape=ir.Shape([3, 4]), + ) + node = ir.Node("", "MatMul", [model.graph.inputs[0], other], name="other") + node.outputs[0].name = "other.output" + model.graph.append(node) + model.graph.initializers.add(other) + manifest = _manifest(model, include_second=True) + package = ModelPackage({"decoder": model}, adapter_target_manifest=manifest) + package.add_adapter_artifact( + AdapterArtifact( + "mixed-rank", + manifest.base_fingerprint, + ( + _weights(target_id="layers.0.self_attn.q_proj.q"), + _weights( + parameter="other.weight", + a=np.ones((1, 4), dtype=np.float32), + b=np.ones((3, 1), dtype=np.float32), + ), + ), + ) + ) + directory = Path("artifacts") / f"adapter-rank-test-{uuid.uuid4().hex}" + directory.mkdir(parents=True) + try: + artifact = package.save_adapter_artifacts(str(directory))["mixed-rank"] + assert artifact["rank"] == 1 + bindings = {binding["target"]: binding for binding in artifact["bindings"]} + assert bindings["layers.0.self_attn.q_proj.q"]["rank"] == 2 + assert "rank" not in bindings["layers.0.self_attn.v_proj"] + finally: + shutil.rmtree(directory) + + +def test_wire_contract_rejects_manifest_rank_policy_violation() -> None: + model = _model() + descriptor = AdapterTargetDescriptor( + AdapterTarget("decoder", "projection.weight"), + semantic_name="projection", + node_name="projection", + output_name="projection.output", + input_size=4, + output_size=3, + rank=1, + alpha=4.0, + ) + manifest = AdapterTargetManifest( + fingerprint_model_weights({"decoder": model}, (descriptor,)), + (descriptor,), + ) + package = ModelPackage({"decoder": model}, adapter_target_manifest=manifest) + package.add_adapter_artifact( + AdapterArtifact("style", manifest.base_fingerprint, (_weights(),)) + ) + directory = Path("artifacts") / f"adapter-policy-test-{uuid.uuid4().hex}" + directory.mkdir(parents=True) + try: + with pytest.raises(ValueError, match="violates manifest policy"): + package.save_adapter_artifacts(str(directory)) + finally: + shutil.rmtree(directory) + + +def test_application_scale_matches_runtime_bound() -> None: + with pytest.raises(ValueError, match=r"within \[-16, 16\]"): + AdapterApplication("style", 16.1) + + +def test_selection_rejects_duplicate_adapter() -> None: + application = AdapterApplication("style") + with pytest.raises(ValueError, match="contains duplicate adapter"): + AdapterSlotSelection(100, 0, (application, application)) + + +def test_selection_rejects_unknown_adapter_and_invalid_permutation() -> None: + batch = AdapterBatchSelection( + (AdapterSlotSelection(100, 0, (AdapterApplication("missing"),)),) + ) + with pytest.raises(ValueError, match="unknown adapter"): + batch.validate_catalog({}) + with pytest.raises(ValueError, match="permutation"): + batch.compact([1]) diff --git a/src/mobius/generation/__init__.py b/src/mobius/generation/__init__.py new file mode 100644 index 000000000..a9c7951a3 --- /dev/null +++ b/src/mobius/generation/__init__.py @@ -0,0 +1,158 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT License. + +"""Reusable ONNX generation-policy components.""" + +from __future__ import annotations + +from mobius.generation._policy_components import ( + SOLVER_BUILDERS, + PolicyCapabilities, + PolicyComponent, + attach_policy_components, + build_adaptive_k_policy, + build_batch_minimum, + build_boolean_not, + build_code_frame_update, + build_code_history_append, + build_codec_layout_transpose, + build_counter_rng_normal, + build_ddim_solver_step, + build_decoder_state_initializer, + build_decoder_step_update, + build_duplex_agent_frame_select, + build_duplex_cell_to_frame, + build_duplex_frame_assemble, + build_duplex_frame_commit, + build_duplex_stream_append, + build_duplex_stream_tail, + build_duplex_teacher_select, + build_duplex_user_stream_merge, + build_duplex_waveform_append, + build_effectful_identity, + build_empty_features, + build_eos_termination, + build_euler_model_input, + build_euler_solver_step, + build_flow_match_solver_step, + build_grammar_logits_processor, + build_greedy_sampler, + build_guidance_combine, + build_identity_model_input, + build_integer_add, + build_integer_minimum, + build_integer_row_broadcast, + build_iteration_cast, + build_last_token_logits, + build_masked_token_update, + build_model_token_cast, + build_multistep_solver_step, + build_pack_latents_2x2, + build_proposal_metrics, + build_scalar_constant, + build_scalar_integer_add, + build_schedule_constant, + build_schedule_history_append, + build_schedule_lookup, + build_seeded_categorical_sampler, + build_selective_integer_add, + build_sequence_concat, + build_sequence_length, + build_shape_constant, + build_speculative_acceptance, + build_speculative_state_rollback, + build_tensor_scale, + build_termination_batch_initializer, + build_token_block_identity, + build_token_state_update, + build_token_to_slot, + build_true_cfg, + build_tts_decoder_state_initializer, + build_tts_decoder_step_update, + build_tts_state_initializer, + build_unpack_latents_2x2, + build_video_conv_cache_initializer, + build_video_decode_chunk, + build_video_decode_chunk_count, + build_video_latent_initializer, + build_video_latent_permute, + build_video_latent_unscale, + build_zeros_like, + rotary_axis_count, +) + +__all__ = [ + "SOLVER_BUILDERS", + "PolicyCapabilities", + "PolicyComponent", + "attach_policy_components", + "build_adaptive_k_policy", + "build_batch_minimum", + "build_boolean_not", + "build_code_frame_update", + "build_code_history_append", + "build_codec_layout_transpose", + "build_counter_rng_normal", + "build_ddim_solver_step", + "build_decoder_state_initializer", + "build_decoder_step_update", + "build_duplex_agent_frame_select", + "build_duplex_cell_to_frame", + "build_duplex_frame_assemble", + "build_duplex_frame_commit", + "build_duplex_stream_append", + "build_duplex_stream_tail", + "build_duplex_teacher_select", + "build_duplex_user_stream_merge", + "build_duplex_waveform_append", + "build_effectful_identity", + "build_empty_features", + "build_eos_termination", + "build_euler_model_input", + "build_euler_solver_step", + "build_flow_match_solver_step", + "build_grammar_logits_processor", + "build_greedy_sampler", + "build_guidance_combine", + "build_identity_model_input", + "build_integer_add", + "build_integer_minimum", + "build_integer_row_broadcast", + "build_iteration_cast", + "build_last_token_logits", + "build_masked_token_update", + "build_model_token_cast", + "build_multistep_solver_step", + "build_pack_latents_2x2", + "build_proposal_metrics", + "build_scalar_constant", + "build_scalar_integer_add", + "build_schedule_constant", + "build_schedule_history_append", + "build_schedule_lookup", + "build_seeded_categorical_sampler", + "build_selective_integer_add", + "build_sequence_concat", + "build_sequence_length", + "build_shape_constant", + "build_speculative_acceptance", + "build_speculative_state_rollback", + "build_tensor_scale", + "build_termination_batch_initializer", + "build_token_block_identity", + "build_token_state_update", + "build_token_to_slot", + "build_true_cfg", + "build_tts_decoder_state_initializer", + "build_tts_decoder_step_update", + "build_tts_state_initializer", + "build_unpack_latents_2x2", + "build_video_conv_cache_initializer", + "build_video_decode_chunk", + "build_video_decode_chunk_count", + "build_video_latent_initializer", + "build_video_latent_permute", + "build_video_latent_unscale", + "build_zeros_like", + "rotary_axis_count", +] diff --git a/src/mobius/generation/_duplex_policy_test.py b/src/mobius/generation/_duplex_policy_test.py new file mode 100644 index 000000000..2c4ae482b --- /dev/null +++ b/src/mobius/generation/_duplex_policy_test.py @@ -0,0 +1,239 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT License. + +"""Tests for the full-duplex delay-cache policy components. + +The reference implementations below are a direct transcription of the upstream +``moshi.models.lm.LMGen.prepare_step_input`` / ``_step`` delay bookkeeping used by +Moshi-family full-duplex models (Moshi, PersonaPlex). The ONNX components must +reproduce them exactly, because a single misplaced ring slot silently corrupts +the interleaved text/agent/user token streams. +""" + +from __future__ import annotations + +import numpy as np +import onnx_ir as ir +import onnxruntime as ort +import pytest + +from mobius.generation import ( + build_duplex_frame_assemble, + build_duplex_frame_commit, + build_duplex_stream_append, + build_duplex_stream_tail, + build_duplex_teacher_select, + build_duplex_waveform_append, +) + +# PersonaPlex / Moshi channel layout: text, 8 agent streams, 8 user streams. +DELAYS = [0, 0, 1, 1, 1, 1, 1, 1, 1, 0, 1, 1, 1, 1, 1, 1, 1] +CHANNELS = len(DELAYS) +MAX_DELAY = max(DELAYS) +CACHE_T = MAX_DELAY + 3 +TEXT_INITIAL = 32000 +AUDIO_INITIAL = 2048 +INITIAL = [TEXT_INITIAL] + [AUDIO_INITIAL] * (CHANNELS - 1) + + +def _session(component) -> ort.InferenceSession: + proto = ir.to_proto(component.model) + return ort.InferenceSession(proto.SerializeToString(), providers=["CPUExecutionProvider"]) + + +def _reference_assemble(cache, provided, offset, stream_tokens): + """Upstream ``prepare_step_input`` for ``offset >= 1``.""" + cache = cache.copy() + provided = provided.copy() + for k, delay in enumerate(DELAYS): + token = int(stream_tokens[0, k]) + if token >= 0: + pos = (offset + delay) % CACHE_T + cache[0, k, pos] = token + provided[0, k, pos] = True + for k, delay in enumerate(DELAYS): + if offset <= delay: + cache[0, k, offset % CACHE_T] = INITIAL[k] + provided[0, k, offset % CACHE_T] = True + input_pos = (offset - 1) % CACHE_T + target_pos = offset % CACHE_T + return ( + cache, + provided, + cache[:, :, input_pos : input_pos + 1].copy(), + cache[:, :, target_pos : target_pos + 1].copy(), + provided[:, :, target_pos : target_pos + 1].copy(), + ) + + +def _reference_commit(cache, provided, offset, frame): + """Upstream cache commit plus delay-compensated read-out.""" + cache = cache.copy() + provided = provided.copy() + input_pos = (offset - 1) % CACHE_T + target_pos = offset % CACHE_T + provided[0, :, input_pos] = False + for k in range(CHANNELS): + if not provided[0, k, target_pos]: + cache[0, k, target_pos] = int(frame[0, k]) + out = np.array( + [cache[0, k, (offset - MAX_DELAY + DELAYS[k]) % CACHE_T] for k in range(CHANNELS)], + np.int64, + )[None] + return cache, provided, out, offset + 1, offset > MAX_DELAY + + +def _random_state(rng): + cache = rng.integers(0, 2048, size=(1, CHANNELS, CACHE_T)).astype(np.int64) + provided = rng.random((1, CHANNELS, CACHE_T)) < 0.5 + return cache, provided + + +@pytest.mark.parametrize("offset", [1, 2, 3, 4, 5, 8, 123]) +def test_duplex_frame_assemble_matches_reference(offset: int) -> None: + rng = np.random.default_rng(offset) + session = _session(build_duplex_frame_assemble(channels=CHANNELS, cache_length=CACHE_T)) + for trial in range(4): + cache, provided = _random_state(rng) + stream_tokens = rng.integers(-1, 2048, size=(1, CHANNELS)).astype(np.int64) + if trial == 0: # live phase: only the user streams carry tokens + stream_tokens[:, :9] = -1 + got = session.run( + None, + { + "token_cache": cache, + "token_provided": provided, + "offset": np.array(offset, np.int64), + "stream_tokens": stream_tokens, + "delays": np.array(DELAYS, np.int64), + "initial_tokens": np.array(INITIAL, np.int64), + }, + ) + want = _reference_assemble(cache, provided, offset, stream_tokens) + for index, (actual, expected) in enumerate(zip(got, want)): + np.testing.assert_array_equal(actual, expected, err_msg=f"output {index}") + + +@pytest.mark.parametrize("offset", [1, 2, 3, 4, 7, 122]) +def test_duplex_frame_commit_matches_reference(offset: int) -> None: + rng = np.random.default_rng(1000 + offset) + session = _session( + build_duplex_frame_commit(channels=CHANNELS, cache_length=CACHE_T, max_delay=MAX_DELAY) + ) + for _ in range(4): + cache, provided = _random_state(rng) + frame = rng.integers(0, 2048, size=(1, CHANNELS)).astype(np.int64) + got = session.run( + None, + { + "token_cache": cache, + "token_provided": provided, + "offset": np.array(offset, np.int64), + "frame": frame, + "delays": np.array(DELAYS, np.int64), + }, + ) + want = _reference_commit(cache, provided, offset, frame) + np.testing.assert_array_equal(got[0], want[0]) + np.testing.assert_array_equal(got[1], want[1]) + np.testing.assert_array_equal(got[2], want[2]) + assert int(got[3]) == want[3] + assert bool(got[4]) == want[4] + + +def test_duplex_assemble_commit_round_trip_preserves_streams() -> None: + """A full delayed round trip returns the tokens that were fed in. + + Streams with delay 1 are emitted one frame late, so feeding a known user + stream and reading it back through the ring must reproduce it exactly. + """ + assemble = _session(build_duplex_frame_assemble(channels=CHANNELS, cache_length=CACHE_T)) + commit = _session( + build_duplex_frame_commit(channels=CHANNELS, cache_length=CACHE_T, max_delay=MAX_DELAY) + ) + rng = np.random.default_rng(7) + cache = np.full((1, CHANNELS, CACHE_T), -1, np.int64) + provided = np.zeros((1, CHANNELS, CACHE_T), bool) + cache[0, :, 0] = INITIAL + offset = 1 + fed: list[list[int]] = [] + emitted: list[list[int]] = [] + for _ in range(12): + user = rng.integers(0, 2048, size=8).astype(np.int64) + fed.append(user.tolist()) + stream_tokens = np.full((1, CHANNELS), -1, np.int64) + stream_tokens[0, 9:] = user + cache, provided, _, _target, _target_provided = assemble.run( + None, + { + "token_cache": cache, + "token_provided": provided, + "offset": np.array(offset, np.int64), + "stream_tokens": stream_tokens, + "delays": np.array(DELAYS, np.int64), + "initial_tokens": np.array(INITIAL, np.int64), + }, + ) + frame = rng.integers(0, 2048, size=(1, CHANNELS)).astype(np.int64) + cache, provided, out, next_offset, emit = commit.run( + None, + { + "token_cache": cache, + "token_provided": provided, + "offset": np.array(offset, np.int64), + "frame": frame, + "delays": np.array(DELAYS, np.int64), + }, + ) + if bool(emit): + emitted.append(out[0, 9:].tolist()) + offset = int(next_offset) + # user streams have delay 1 and the read-out subtracts max_delay 1, so the + # emitted user stream is exactly what was fed on the same frame. + assert emitted == fed[: len(emitted)] + assert len(emitted) == 11 + + +def test_duplex_teacher_select_prefers_supplied_tokens() -> None: + session = _session(build_duplex_teacher_select(channels=CHANNELS)) + target = np.arange(CHANNELS, dtype=np.int64).reshape(1, CHANNELS, 1) + 100 + provided = np.zeros((1, CHANNELS, 1), bool) + provided[0, 3, 0] = True + for index in (0, 3, 16): + token = session.run( + None, + { + "target": target, + "target_provided": provided, + "sampled": np.array([7], np.int64), + "index": np.array(index, np.int64), + }, + )[0] + assert token.tolist() == ([103] if index == 3 else [7]) + + +def test_duplex_stream_append_and_tail() -> None: + append = _session(build_duplex_stream_append(streams=8)) + tail = _session(build_duplex_stream_tail(streams=8)) + prefix = np.zeros((1, 8, 0), np.int64) + frames = [] + for step in range(5): + frame = np.arange(8, dtype=np.int64).reshape(1, 8) + step * 8 + frames.append(frame) + prefix = append.run(None, {"prefix": prefix, "frame": frame})[0] + assert prefix.shape == (1, 8, 5) + np.testing.assert_array_equal(prefix[:, :, -1], frames[-1]) + got = tail.run(None, {"prefix": prefix, "count": np.array(2, np.int64)})[0] + np.testing.assert_array_equal(got, prefix[:, :, -2:]) + + +def test_duplex_waveform_append_grows_packed_audio() -> None: + append = _session(build_duplex_waveform_append()) + tail = _session(build_duplex_stream_tail(streams=1, dtype=ir.DataType.FLOAT)) + prefix = np.zeros((1, 1, 0), np.float32) + for step in range(3): + chunk = np.full((1, 1, 1920), float(step), np.float32) + prefix = append.run(None, {"prefix": prefix, "chunk": chunk})[0] + assert prefix.shape == (1, 1, 5760) + got = tail.run(None, {"prefix": prefix, "count": np.array(1920, np.int64)})[0] + np.testing.assert_allclose(got, np.full((1, 1, 1920), 2.0, np.float32)) diff --git a/src/mobius/generation/_policy_components.py b/src/mobius/generation/_policy_components.py new file mode 100644 index 000000000..a261710e5 --- /dev/null +++ b/src/mobius/generation/_policy_components.py @@ -0,0 +1,3019 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT License. + +"""Small, model-agnostic ONNX graphs for generation policy and state math. + +These components deliberately receive policy parameters as tensor inputs. They +contain no model-family dispatch and can therefore be invoked from a generic +workflow IR just like neural ONNX components. +""" + +from __future__ import annotations + +import json +from collections.abc import Sequence +from dataclasses import dataclass +from typing import Any, Protocol + +import onnx_ir as ir +from onnxscript import GraphBuilder + +from mobius._constants import OPSET_VERSION + +_POLICY_CONTRACT_ID_METADATA = "mobius.generation.policy_contract_id" +_POLICY_CONTRACT_METADATA = "mobius.generation.policy_contract" +_POLICY_EFFECTS_METADATA = "mobius.generation.policy_effects" + + +@dataclass(frozen=True) +class PolicyComponent: + """A versioned semantic contract and its executable ONNX model.""" + + contract_id: str + model: ir.Model + contract: dict[str, object] + effects: tuple[str, ...] + + def __post_init__(self) -> None: + if "@" not in self.contract_id: + raise ValueError("policy contract_id must include a version") + self.model.graph.metadata_props[_POLICY_CONTRACT_ID_METADATA] = self.contract_id + self.model.graph.metadata_props[_POLICY_CONTRACT_METADATA] = json.dumps(self.contract) + self.model.graph.metadata_props[_POLICY_EFFECTS_METADATA] = json.dumps(self.effects) + + @classmethod + def from_model(cls, model: ir.Model) -> PolicyComponent: + """Restore a component from contract metadata embedded in its ONNX graph.""" + contract_id = model.graph.metadata_props.get(_POLICY_CONTRACT_ID_METADATA) + if contract_id is None: + raise ValueError("ONNX policy component is missing its versioned contract ID") + contract = json.loads(model.graph.metadata_props[_POLICY_CONTRACT_METADATA]) + effects = tuple(json.loads(model.graph.metadata_props[_POLICY_EFFECTS_METADATA])) + return cls(contract_id, model, contract, effects) + + +@dataclass(frozen=True) +class PolicyCapabilities: + """Data-driven declaration of policy math required by a package.""" + + sampler: str | None = None + eos_termination: bool = False + solver: str | None = None + masked_update: bool = False + speculative_acceptance: bool = False + grammar_guidance: bool = False + adaptive_k_max: int | None = None + token_state_update: bool = False + + +class _PolicyPackage(Protocol): + def add_policy_component(self, name: str, component: PolicyComponent) -> None: ... + + +def attach_policy_components( + pkg: _PolicyPackage, + capabilities: PolicyCapabilities, +) -> dict[str, str]: + """Attach exactly the policy artifacts selected by declared capabilities.""" + builders = { + "greedy": build_greedy_sampler, + "seeded_categorical": build_seeded_categorical_sampler, + } + solvers = SOLVER_BUILDERS + if capabilities.sampler not in {None, *builders}: + raise ValueError(f"Unsupported sampler policy {capabilities.sampler!r}") + if capabilities.solver not in {None, *solvers}: + raise ValueError(f"Unsupported solver policy {capabilities.solver!r}") + + selected: list[tuple[str, PolicyComponent]] = [] + if capabilities.sampler is not None: + selected.append(("token_sampler", builders[capabilities.sampler]())) + if capabilities.eos_termination: + selected.append(("termination", build_eos_termination())) + if capabilities.solver is not None: + selected.append(("solver_step", solvers[capabilities.solver]())) + if capabilities.masked_update: + selected.append(("masked_update", build_masked_token_update())) + if capabilities.speculative_acceptance: + selected.append(("speculative_acceptance", build_speculative_acceptance())) + if capabilities.grammar_guidance: + selected.append(("grammar_guidance", build_grammar_logits_processor())) + if capabilities.adaptive_k_max is not None: + selected.append( + ( + "adaptive_k", + build_adaptive_k_policy(max_k=capabilities.adaptive_k_max), + ) + ) + if capabilities.token_state_update: + selected.append(("token_state_update", build_token_state_update())) + + for name, component in selected: + pkg.add_policy_component(name, component) + return {name: f"policies/{name}.onnx" for name, _ in selected} + + +#: Element types the ONNX ``ConstantOfShape`` kernel can produce. The narrow +#: float types are absent, so a policy graph that must materialize an fp8 or +#: fp4 tensor has to fill a supported dtype and cast the result. +_CONSTANT_OF_SHAPE_DTYPES = frozenset( + { + ir.DataType.FLOAT, + ir.DataType.FLOAT16, + ir.DataType.BFLOAT16, + ir.DataType.DOUBLE, + ir.DataType.INT8, + ir.DataType.INT16, + ir.DataType.INT32, + ir.DataType.INT64, + ir.DataType.UINT8, + ir.DataType.UINT16, + ir.DataType.UINT32, + ir.DataType.UINT64, + ir.DataType.BOOL, + } +) + + +def _component( + contract_id: str, + graph: ir.Graph, + contract: dict[str, object], + *_effects: str, +) -> PolicyComponent: + # ONNX policy components are pure: RNG and state are explicit tensor data. + model = ir.Model(graph, ir_version=11) + model.producer_name = "mobius" + return PolicyComponent(contract_id, model, contract, ()) + + +def _make_graph(name: str) -> tuple[ir.Graph, GraphBuilder]: + graph = ir.Graph( + [], + [], + nodes=[], + name=name, + opset_imports={"": OPSET_VERSION}, + ) + return graph, GraphBuilder(graph) + + +def _set_public_shape(value: ir.Value, shape: list[str | int]) -> None: + """Set ABI dimensions without GraphBuilder's graph-name qualification.""" + value.shape = ir.Shape(shape) + + +def build_greedy_sampler( + *, + effect: str = "sample", + row_selective: bool = False, +) -> PolicyComponent: + """Build row-selective greedy sampling over the final axis.""" + graph, builder = _make_graph("greedy_sampler") + op = builder.op + logits = builder.input( + "logits", + dtype=ir.DataType.FLOAT, + shape=["batch", "vocabulary"], + ) + sampled = op.ArgMax(logits, axis=-1, keepdims=0) + if row_selective: + active = builder.input("active", dtype=ir.DataType.BOOL, shape=["batch"]) + done = builder.input("done", dtype=ir.DataType.BOOL, shape=["batch"]) + enabled = op.And(active, op.Not(done)) + token_ids = op.Where(enabled, sampled, op.Constant(value_int=-1)) + else: + token_ids = sampled + builder.add_output(token_ids, "token") + return _component( + "onnx-genai.token-sampler@2" if row_selective else "onnx-genai.token-sampler@1", + graph, + { + "role": "token_sampler", + "mode": "greedy", + **({"batching": "per_row", "inactive_rows": "preserve"} if row_selective else {}), + "logits": "logits", + **({"active": "active", "done": "done"} if row_selective else {}), + "token": "token", + "effect": effect, + }, + effect, + ) + + +def build_last_token_logits( + input_dtype: ir.DataType = ir.DataType.FLOAT, +) -> PolicyComponent: + """Build ``[B,T,V] -> [B,V]`` selection and normalize logits to float32.""" + graph, builder = _make_graph("last_token_logits") + logits = builder.input( + "logits", + dtype=input_dtype, + shape=["batch", "sequence", "vocabulary"], + ) + selected = builder.op.Gather(logits, builder.op.Constant(value_int=-1), axis=1) + if input_dtype != ir.DataType.FLOAT: + selected = builder.op.Cast(selected, to=ir.DataType.FLOAT) + selected.shape = ir.Shape(["batch", "vocabulary"]) + builder.add_output(selected, "last_logits") + return _component("mobius.policy.auxiliary@1", graph, {}) + + +def build_boolean_not() -> PolicyComponent: + """Build one synchronized ``continue = Not(Any(done))`` predicate.""" + graph, builder = _make_graph("boolean_not") + done = builder.input("done", dtype=ir.DataType.BOOL, shape=["batch"]) + any_done = builder.op.ReduceMax( + builder.op.Cast(done, to=ir.DataType.INT64), + keepdims=1, + ) + continued = builder.op.Equal(any_done, builder.op.Constant(value_int=0)) + continued.shape = ir.Shape([1]) + builder.add_output(continued, "continue") + return _component("mobius.policy.auxiliary@1", graph, {}) + + +def build_integer_minimum() -> PolicyComponent: + """Compute the per-batch minimum of two integer lengths.""" + graph, builder = _make_graph("integer_minimum") + left = builder.input("left", ir.DataType.INT64, ["batch"]) + right = builder.input("right", ir.DataType.INT64, ["batch"]) + minimum = builder.op.Min(left, right) + minimum.shape = ir.Shape(["batch"]) + builder.add_output(minimum, "minimum") + return _component("mobius.policy.auxiliary@1", graph, {}) + + +def build_integer_add() -> PolicyComponent: + """Add two per-row integer state values.""" + graph, builder = _make_graph("integer_add") + left = builder.input("left", ir.DataType.INT64, ["batch"]) + right = builder.input("right", ir.DataType.INT64, ["batch"]) + total = builder.op.Add(left, right) + total.shape = ir.Shape(["batch"]) + builder.add_output(total, "total") + return _component("mobius.policy.auxiliary@1", graph, {}) + + +def build_selective_integer_add() -> PolicyComponent: + """Add per-row integer state only for active, unfinished rows.""" + graph, builder = _make_graph("selective_integer_add") + op = builder.op + left = builder.input("left", ir.DataType.INT64, ["batch"]) + right = builder.input("right", ir.DataType.INT64, ["batch"]) + active = builder.input("active", ir.DataType.BOOL, ["batch"]) + done = builder.input("done", ir.DataType.BOOL, ["batch"]) + enabled = op.And(active, op.Not(done)) + total = op.Where(enabled, op.Add(left, right), left) + total.shape = ir.Shape(["batch"]) + builder.add_output(total, "total") + return _component("mobius.policy.auxiliary@1", graph, {}) + + +def build_integer_row_broadcast() -> PolicyComponent: + """Broadcast one scalar loop control to the current dynamic batch.""" + graph, builder = _make_graph("integer_row_broadcast") + op = builder.op + value = builder.input("value", ir.DataType.INT64, [1]) + active = builder.input("active", ir.DataType.BOOL, ["batch"]) + rows = op.Shape(active) + result = op.Expand(value, rows) + result.shape = ir.Shape(["batch"]) + builder.add_output(result, "rows") + return _component("mobius.policy.auxiliary@1", graph, {}) + + +def build_termination_batch_initializer() -> PolicyComponent: + """Normalize explicit ragged EOS sets and per-row generation limits.""" + graph, builder = _make_graph("termination_batch_initializer") + op = builder.op + eos_ids = builder.input("input_eos_ids", ir.DataType.INT64, ["batch", "num_eos"]) + eos_lengths = builder.input("input_eos_lengths", ir.DataType.INT64, ["batch"]) + max_iterations = builder.input("input_max_iterations", ir.DataType.INT64, ["batch"]) + fallback_max_iterations = builder.input("fallback_max_iterations", ir.DataType.INT64, [1]) + active = builder.input("active", ir.DataType.BOOL, ["batch"]) + batch_shape = op.Shape(active) + row_max_iterations = op.Where( + op.Greater(max_iterations, op.Constant(value_int=0)), + max_iterations, + op.Expand(fallback_max_iterations, batch_shape), + ) + row_eos_ids = op.Identity(eos_ids) + eos_count = op.Identity(eos_lengths) + row_eos_ids.shape = ir.Shape(["batch", "num_eos"]) + eos_count.shape = ir.Shape(["batch"]) + row_max_iterations.shape = ir.Shape(["batch"]) + builder.add_output(row_eos_ids, "row_eos_ids") + builder.add_output(eos_count, "eos_lengths") + builder.add_output(row_max_iterations, "max_iterations") + return _component("mobius.policy.auxiliary@1", graph, {}) + + +def build_batch_minimum() -> PolicyComponent: + """Synchronize a per-batch integer length to one conservative scalar.""" + graph, builder = _make_graph("batch_minimum") + values = builder.input("values", ir.DataType.INT64, ["batch"]) + minimum = builder.op.ReduceMin(values, keepdims=1) + minimum.shape = ir.Shape([1]) + builder.add_output(minimum, "minimum") + return _component("mobius.policy.auxiliary@1", graph, {}) + + +def build_proposal_metrics() -> PolicyComponent: + """Derive evaluated width and budget fullness from a dense proposal.""" + graph, builder = _make_graph("proposal_metrics") + op = builder.op + tokens = builder.input("proposed_tokens", ir.DataType.INT64, ["batch", "proposal"]) + requested_k = builder.input("requested_k", ir.DataType.INT64, ["batch"]) + batch = op.Shape(tokens, start=0, end=1) + length = op.Expand(op.Shape(tokens, start=1, end=2), batch) + filled = op.Equal(length, requested_k) + length.shape = ir.Shape(["batch"]) + filled.shape = ir.Shape(["batch"]) + builder.add_output(length, "evaluated") + builder.add_output(filled, "filled_proposal_budget") + return _component("mobius.policy.auxiliary@1", graph, {}) + + +def build_sequence_length() -> PolicyComponent: + """Expand a dense rank-2 token block's sequence width per batch.""" + graph, builder = _make_graph("sequence_length") + op = builder.op + tokens = builder.input("tokens", ir.DataType.INT64, ["batch", "sequence"]) + length = op.Expand( + op.Shape(tokens, start=1, end=2), + op.Shape(tokens, start=0, end=1), + ) + length.shape = ir.Shape(["batch"]) + builder.add_output(length, "length") + return _component("mobius.policy.auxiliary@1", graph, {}) + + +def build_iteration_cast(dtype: ir.DataType) -> PolicyComponent: + """Cast the generic int64 loop induction value for a model timestep port.""" + graph, builder = _make_graph("iteration_cast") + iteration = builder.input("iteration", dtype=ir.DataType.INT64, shape=["batch"]) + timestep = builder.op.Cast(iteration, to=dtype) + timestep.shape = ir.Shape(["batch"]) + builder.add_output(timestep, "timestep") + return _component("mobius.policy.auxiliary@1", graph, {}) + + +def build_schedule_constant(values: list[float]) -> PolicyComponent: + """Materialize a producer-selected diffusion schedule inside ONNX.""" + if len(values) < 2: + raise ValueError("a diffusion schedule requires at least two values") + graph, builder = _make_graph("diffusion_schedule") + schedule = builder.op.Constant( + value=ir.tensor(values, dtype=ir.DataType.FLOAT), + ) + schedule.shape = ir.Shape([len(values)]) + builder.add_output(schedule, "schedule") + return _component("mobius.policy.auxiliary@1", graph, {}) + + +def build_schedule_lookup(dtype: ir.DataType) -> PolicyComponent: + """Gather the current schedule value and cast it for a model timestep port.""" + graph, builder = _make_graph("schedule_lookup") + op = builder.op + schedule = builder.input("schedule", ir.DataType.FLOAT, ["schedule_length"]) + step = builder.input("step", ir.DataType.INT64, ["batch"]) + timestep = op.Cast(op.Gather(schedule, step, axis=0), to=dtype) + timestep.shape = ir.Shape(["batch"]) + builder.add_output(timestep, "timestep") + return _component("mobius.policy.auxiliary@1", graph, {}) + + +def build_tts_state_initializer(num_code_groups: int) -> PolicyComponent: + """Create an empty codec history and zeroed current frame from prompt batch.""" + if num_code_groups < 1: + raise ValueError("num_code_groups must be positive") + graph, builder = _make_graph("tts_state_initializer") + op = builder.op + prompt = builder.input( + "prompt_tokens", dtype=ir.DataType.INT64, shape=["batch", "sequence"] + ) + batch = op.Shape(prompt, start=0, end=1) + frame_shape = op.Concat(batch, op.Constant(value_ints=[num_code_groups]), axis=0) + history_shape = op.Concat(batch, op.Constant(value_ints=[0, num_code_groups]), axis=0) + frame = op.ConstantOfShape(frame_shape, value=ir.tensor([0], dtype=ir.DataType.INT64)) + token_slot = op.ConstantOfShape( + op.Concat(batch, op.Constant(value_ints=[1]), axis=0), + value=ir.tensor([0], dtype=ir.DataType.INT64), + ) + history = op.ConstantOfShape(history_shape, value=ir.tensor([0], dtype=ir.DataType.INT64)) + frame.shape = ir.Shape(["batch", num_code_groups]) + token_slot.shape = ir.Shape(["batch", 1]) + history.shape = ir.Shape(["batch", 0, num_code_groups]) + builder.add_output(frame, "frame_codes") + builder.add_output(token_slot, "token_slot") + builder.add_output(history, "code_history") + return _component("mobius.policy.auxiliary@1", graph, {}) + + +def build_tts_decoder_state_initializer( + decoder: ir.Model, + *, + graph_name: str, + embedding_input: str, + attention_mask_input: str, + position_ids_input: str, + cache_inputs: list[str], +) -> PolicyComponent: + """Initialize masks, positions, and empty KV state from prefill embeddings.""" + graph, builder = _make_graph(graph_name) + op = builder.op + inputs = {value.name: value for value in decoder.graph.inputs} + embedding = inputs[embedding_input] + prompt = builder.input( + "prefill_embeds", + embedding.dtype, + ["batch", "prefill_sequence", list(embedding.shape)[-1]], + ) + batch_shape = op.Shape(prompt, start=0, end=1) + sequence_shape = op.Shape(prompt, start=1, end=2) + mask_shape = op.Concat(batch_shape, sequence_shape, axis=0) + attention_value = inputs[attention_mask_input] + attention = op.Cast( + op.ConstantOfShape(mask_shape, value=ir.tensor([1])), + to=attention_value.dtype, + ) + attention.shape = attention_value.shape + body_attention = op.Concat( + attention, + op.Cast( + op.ConstantOfShape( + op.Concat(batch_shape, op.Constant(value_ints=[1]), axis=0), + value=ir.tensor([1]), + ), + to=attention_value.dtype, + ), + axis=1, + ) + body_attention.shape = ir.Shape(["batch", "prefill_sequence + 1"]) + + position_value = inputs[position_ids_input] + position_rank = len(position_value.shape or []) + position_range = op.Range( + op.Constant(value_int=0), + op.Squeeze(sequence_shape, [0]), + op.Constant(value_int=1), + ) + if position_rank == 2: + position_shape = mask_shape + positions = op.Expand(op.Unsqueeze(position_range, [0]), position_shape) + body_position = op.Expand( + op.Cast(sequence_shape, to=position_value.dtype), + op.Concat(batch_shape, op.Constant(value_ints=[1]), axis=0), + ) + body_position.shape = ir.Shape(["batch", 1]) + elif position_rank == 3: + position_shape = op.Concat( + op.Constant(value_ints=[3]), batch_shape, sequence_shape, axis=0 + ) + positions = op.Expand(op.Unsqueeze(position_range, [0, 1]), position_shape) + body_position = op.Expand( + op.Reshape( + op.Cast(sequence_shape, to=position_value.dtype), + op.Constant(value_ints=[1, 1, 1]), + ), + op.Concat( + op.Constant(value_ints=[3]), + batch_shape, + op.Constant(value_ints=[1]), + axis=0, + ), + ) + body_position.shape = ir.Shape([3, "batch", 1]) + else: + raise ValueError("TTS decoder position_ids must be rank 2 or 3") + positions = op.Cast(positions, to=position_value.dtype) + positions.shape = position_value.shape + + builder.add_output(attention, attention_mask_input) + builder.add_output(positions, position_ids_input) + builder.add_output(body_attention, "body_attention_mask") + builder.add_output(body_position, "body_position_ids") + for name in cache_inputs: + value = inputs[name] + dimensions = list(value.shape or []) + shape_parts = [] + for axis, dimension in enumerate(dimensions): + text = str(getattr(dimension, "value", dimension)) + if axis == 0: + shape_parts.append(batch_shape) + elif "sequence" in text: + shape_parts.append(op.Constant(value_ints=[0])) + elif isinstance(dimension, int): + shape_parts.append(op.Constant(value_ints=[dimension])) + else: + raise ValueError(f"cache input {name!r} has unsupported dimension {text!r}") + empty = op.ConstantOfShape( + op.Concat(*shape_parts, axis=0), + value=ir.tensor( + [0.0 if value.dtype.is_floating_point else 0], + dtype=value.dtype, + ), + ) + empty.shape = value.shape + builder.add_output(empty, name) + return _component("mobius.policy.auxiliary@1", graph, {}) + + +def build_tts_decoder_step_update( + *, + graph_name: str, + attention_dtype: ir.DataType, + position_dtype: ir.DataType, + position_rank: int, +) -> PolicyComponent: + """Append one decoder mask slot and increment rank-2 or rank-3 positions.""" + graph, builder = _make_graph(graph_name) + op = builder.op + attention = builder.input("attention_mask", attention_dtype, ["batch", "context"]) + one_shape = op.Concat( + op.Shape(attention, start=0, end=1), + op.Constant(value_ints=[1]), + axis=0, + ) + one = op.CastLike(op.ConstantOfShape(one_shape, value=ir.tensor([1])), attention) + next_attention = op.Concat(attention, one, axis=1) + next_attention.shape = ir.Shape(["batch", "context + 1"]) + if position_rank == 2: + position_shape: list[int | str] = ["batch", 1] + elif position_rank == 3: + position_shape = [3, "batch", 1] + else: + raise ValueError("TTS decoder position_ids must be rank 2 or 3") + position = builder.input("position_ids", position_dtype, position_shape) + next_position = op.Add(position, op.CastLike(op.Constant(value_int=1), position)) + next_position.shape = position.shape + builder.add_output(next_attention, "next_attention_mask") + builder.add_output(next_position, "next_position_ids") + return _component("mobius.policy.auxiliary@1", graph, {}) + + +def build_code_frame_update( + num_code_groups: int, *, scalar_index: bool = False +) -> PolicyComponent: + """Scatter one predicted code into the current codec frame.""" + graph, builder = _make_graph("code_frame_update") + op = builder.op + frame = builder.input( + "frame_codes", + dtype=ir.DataType.INT64, + shape=["batch", num_code_groups], + ) + token = builder.input("token", dtype=ir.DataType.INT64, shape=["batch"]) + index = builder.input( + "index", + dtype=ir.DataType.INT64, + shape=[] if scalar_index else ["batch"], + ) + if scalar_index: + index = op.Expand(index, op.Shape(token)) + updated = op.ScatterElements( + frame, + op.Unsqueeze(index, op.Constant(value_ints=[-1])), + op.Unsqueeze(token, op.Constant(value_ints=[-1])), + axis=1, + ) + updated.shape = frame.shape + builder.add_output(updated, "next_frame") + return _component("mobius.policy.auxiliary@1", graph, {}) + + +def build_code_history_append(num_code_groups: int) -> PolicyComponent: + """Append a completed codec frame to a growing frame history.""" + graph, builder = _make_graph("code_history_append") + op = builder.op + history = builder.input( + "history", + dtype=ir.DataType.INT64, + shape=["batch", "frames", num_code_groups], + ) + frame = builder.input("frame", dtype=ir.DataType.INT64, shape=["batch", num_code_groups]) + next_history = op.Concat( + history, + op.Unsqueeze(frame, op.Constant(value_ints=[1])), + axis=1, + ) + next_history.shape = ir.Shape(["batch", "frames + 1", num_code_groups]) + builder.add_output(next_history, "next_history") + return _component("mobius.policy.auxiliary@1", graph, {}) + + +def build_codec_layout_transpose(num_code_groups: int) -> PolicyComponent: + """Convert frame-major ``[B,F,G]`` history to codec-major ``[B,G,F]``.""" + graph, builder = _make_graph("codec_layout_transpose") + history = builder.input( + "history", + dtype=ir.DataType.INT64, + shape=["batch", "frames", num_code_groups], + ) + codes = builder.op.Transpose(history, perm=[0, 2, 1]) + codes.shape = ir.Shape(["batch", num_code_groups, "frames"]) + builder.add_output(codes, "codes") + return _component("mobius.policy.auxiliary@1", graph, {}) + + +def build_model_token_cast(dtype: ir.DataType) -> PolicyComponent: + """Cast the canonical int64 token state to a decoder's integer dtype.""" + graph, builder = _make_graph("model_token_cast") + token = builder.input("token", dtype=ir.DataType.INT64, shape=["batch", 1]) + model_token = builder.op.Cast(token, to=dtype) + model_token.shape = ir.Shape(["batch", 1]) + builder.add_output(model_token, "model_token") + return _component("mobius.policy.auxiliary@1", graph, {}) + + +def build_empty_features(dtype: ir.DataType, feature_size: int) -> PolicyComponent: + """Build the empty feature matrix used by multimodal text-only requests.""" + graph, builder = _make_graph("empty_features") + features = builder.op.ConstantOfShape( + builder.op.Constant(value_ints=[0, feature_size]), + value=ir.tensor([0.0], dtype=dtype), + ) + features.shape = ir.Shape([0, feature_size]) + builder.add_output(features, "features") + return _component("mobius.policy.auxiliary@1", graph, {}) + + +def rotary_axis_count(position_value: ir.Value) -> int | None: + """Number of rotary axes a decoder's ``position_ids`` input carries. + + A plain decoder reads ``(batch, sequence)`` positions and gets ``None``. A + decoder with multi-axis rotary embeddings reads + ``(sections, batch, sequence)``, where the leading axis is a fixed count of + rotary axes and must therefore be a static dimension: it is a property of + the exported graph, not of the request. + """ + shape = position_value.shape + if shape is None or len(shape) != 3: + return None + leading = shape[0] + sections = getattr(leading, "value", leading) + if not isinstance(sections, int): + raise TypeError( + f"position input {position_value.name!r} declares a rank-3 shape " + f"whose leading rotary-axis count {sections!r} is symbolic. The " + "number of rotary axes is fixed by the exported graph, so it must " + "be a static dimension." + ) + return sections + + +def build_decoder_state_initializer( + decoder: ir.Model, + *, + token_input: str | None, + prompt_dtype: ir.DataType | None = None, + attention_mask_input: str | None, + position_ids_input: str | None, + cache_inputs: list[str], + fixed_capacity: bool = False, + ragged: bool = False, + write_indices_output: str | None = None, +) -> PolicyComponent: + """Build prompt-derived decoder state, optionally with capture-stable storage. + + ``write_indices_output`` names the graph port of a static (indexed-scatter) + KV cache. Prefill writes the whole prompt chunk starting at slot zero, so the + initial destinations are zeros and the resulting logical length is the prompt + length; both are emitted here so no consumer has to infer them. + """ + if fixed_capacity and attention_mask_input is None: + raise ValueError( + "fixed-capacity decoder state requires an attention-mask input to carry " + "each row's logical length" + ) + graph, builder = _make_graph("decoder_state_initializer") + op = builder.op + decoder_inputs = {value.name: value for value in decoder.graph.inputs} + if prompt_dtype is None: + if token_input is None: + raise ValueError("token_input or prompt_dtype is required") + prompt_dtype = decoder_inputs[token_input].dtype + prompt = builder.input( + "prompt_tokens", + dtype=prompt_dtype, + shape=["batch", "prompt_sequence"], + ) + prompt_shape = op.Shape(prompt) + batch_shape = op.Shape(prompt, start=0, end=1) + sequence_shape = op.Shape(prompt, start=1, end=2) + sequence_length = op.Squeeze(sequence_shape, op.Constant(value_ints=[0])) + if ragged: + provided_prompt_lengths = builder.input( + "prompt_lengths", + dtype=ir.DataType.INT64, + shape=["batch"], + ) + full_prompt_lengths = op.Expand( + op.Unsqueeze(sequence_length, [0]), + batch_shape, + ) + prompt_lengths = op.Where( + op.Greater(provided_prompt_lengths, op.Constant(value_int=0)), + provided_prompt_lengths, + full_prompt_lengths, + ) + else: + prompt_lengths = op.Expand(op.Unsqueeze(sequence_length, [0]), batch_shape) + capacity = None + if fixed_capacity: + max_iterations = builder.input( + "max_iterations", + dtype=ir.DataType.INT64, + shape=[1], + ) + capacity = op.Add( + sequence_length, + op.Squeeze(max_iterations, [0]), + ) + attention_shape = op.Concat( + batch_shape, + op.Unsqueeze(capacity, op.Constant(value_ints=[0])), + axis=0, + ) + offsets = op.Range( + op.Constant(value_int=0), + capacity, + op.Constant(value_int=1), + ) + offsets = op.Expand( + op.Unsqueeze(offsets, op.Constant(value_ints=[0])), + attention_shape, + ) + + attention = None + body_attention = None + if attention_mask_input is not None: + attention_value = decoder_inputs[attention_mask_input] + if fixed_capacity: + # Native ORT GenAI binds one persistent full-capacity mask for prefill + # and decode, then enables the next logical slot before each decode. + attention = op.Cast( + op.Less(offsets, op.Unsqueeze(prompt_lengths, [-1])), + to=attention_value.dtype, + ) + attention.shape = ir.Shape(["batch", "capacity"]) + body_attention = op.Identity(attention) + body_attention.shape = attention.shape + else: + offsets = op.Range( + op.Constant(value_int=0), + sequence_length, + op.Constant(value_int=1), + ) + offsets = op.Expand(op.Unsqueeze(offsets, [0]), prompt_shape) + attention = op.Cast( + op.Less(offsets, op.Unsqueeze(prompt_lengths, [-1])), + to=attention_value.dtype, + ) + attention.shape = attention_value.shape + body_attention = op.Concat( + attention, + op.Cast( + op.ConstantOfShape( + op.Concat(batch_shape, op.Constant(value_ints=[1]), axis=0), + value=ir.tensor([1]), + ), + to=attention_value.dtype, + ), + axis=1, + ) + body_attention.shape = ir.Shape(["batch", "prompt_sequence + 1"]) + token_slot = op.ConstantOfShape( + op.Concat(batch_shape, op.Constant(value_ints=[1]), axis=0), + value=ir.tensor([0], dtype=ir.DataType.INT64), + ) + token_slot.shape = ir.Shape(["batch", 1]) + + positions = None + body_position = None + if position_ids_input is not None: + position_value = decoder_inputs[position_ids_input] + sections = rotary_axis_count(position_value) + positions = op.Range( + op.Constant(value_int=0), + op.Squeeze(sequence_shape, op.Constant(value_ints=[0])), + op.Constant(value_int=1), + ) + positions = op.Expand( + op.Unsqueeze(positions, op.Constant(value_ints=[0])), prompt_shape + ) + positions = op.Cast(positions, to=position_value.dtype) + # (batch, prompt_sequence) + positions.shape = ir.Shape(["batch", "prompt_sequence"]) + body_position = op.Unsqueeze( + op.Cast(prompt_lengths, to=position_value.dtype), + [-1], + ) + # (batch, 1) + body_position.shape = ir.Shape(["batch", 1]) + if sections is not None: + # A decoder with multi-axis rotary positions reads + # (sections, batch, sequence): one position row per rotary axis. + # Every axis carries the same sequential position here, which is + # what the axes agree on for a pure token stream. A component that + # lays media out differently across axes states that layout itself; + # this initializer never invents one. + positions = op.Expand( + op.Unsqueeze(positions, op.Constant(value_ints=[0])), + op.Concat(op.Constant(value_ints=[sections]), prompt_shape, axis=0), + ) + positions.shape = ir.Shape([sections, "batch", "prompt_sequence"]) + body_position = op.Expand( + op.Unsqueeze(body_position, op.Constant(value_ints=[0])), + op.Concat( + op.Constant(value_ints=[sections]), + batch_shape, + op.Constant(value_ints=[1]), + axis=0, + ), + ) + body_position.shape = ir.Shape([sections, "batch", 1]) + if attention_mask_input is not None: + assert attention is not None and body_attention is not None + builder.add_output(attention, attention_mask_input) + if position_ids_input is not None: + assert positions is not None and body_position is not None + builder.add_output(positions, position_ids_input) + if attention_mask_input is not None: + assert body_attention is not None + builder.add_output(body_attention, "body_attention_mask") + if position_ids_input is not None: + builder.add_output(body_position, "body_position_ids") + builder.add_output(token_slot, "token_slot") + if ragged: + generated_lengths = op.ConstantOfShape( + batch_shape, + value=ir.tensor([0], dtype=ir.DataType.INT64), + ) + generated_lengths.shape = ir.Shape(["batch"]) + builder.add_output(generated_lengths, "generated_lengths") + if fixed_capacity or write_indices_output is not None: + cache_lengths = op.Identity(prompt_lengths) + cache_lengths.shape = ir.Shape(["batch"]) + builder.add_output(cache_lengths, "cache_lengths") + if write_indices_output is not None: + # Prefill scatters the whole prompt chunk from slot zero for every row; + # the per-row cursor only diverges once decode advances rows separately. + write_indices = op.ConstantOfShape( + batch_shape, + value=ir.tensor([0], dtype=ir.DataType.INT64), + ) + write_indices.shape = ir.Shape(["batch"]) + builder.add_output(write_indices, write_indices_output) + + for name in cache_inputs: + value = decoder_inputs[name] + if value.shape is None: + raise ValueError(f"cache input {name!r} must declare a shape") + dimensions = list(value.shape) + shape_parts = [] + for axis, dimension in enumerate(dimensions): + dimension_text = str(getattr(dimension, "value", dimension)) + if axis == 0: + shape_parts.append(batch_shape) + elif "sequence" in dimension_text: + if fixed_capacity: + assert capacity is not None + shape_parts.append(op.Unsqueeze(capacity, op.Constant(value_ints=[0]))) + else: + shape_parts.append(op.Constant(value_ints=[0])) + elif isinstance(dimension, int): + shape_parts.append(op.Constant(value_ints=[dimension])) + else: + raise ValueError( + f"cache input {name!r} has unsupported symbolic " + f"dimension {dimension_text!r}" + ) + cache_shape = op.Concat(*shape_parts, axis=0) + # ``ConstantOfShape`` kernels do not implement the narrow float types, so + # an fp8 KV cache must be materialized in a supported dtype and cast. + fill_dtype = ( + value.dtype if value.dtype in _CONSTANT_OF_SHAPE_DTYPES else ir.DataType.FLOAT + ) + zero = 0.0 if fill_dtype.is_floating_point else 0 + empty = op.ConstantOfShape( + cache_shape, + value=ir.tensor([zero], dtype=fill_dtype), + ) + if fill_dtype != value.dtype: + empty = op.Cast(empty, to=value.dtype) + empty.shape = ( + ir.Shape( + [ + "capacity" if "sequence" in str(getattr(d, "value", d)) else d + for d in dimensions + ] + ) + if fixed_capacity + else value.shape + ) + builder.add_output(empty, name) + + return _component("mobius.policy.auxiliary@1", graph, {}) + + +def build_decoder_step_update( + *, + attention_dtype: ir.DataType | None, + position_dtype: ir.DataType | None, + fixed_capacity: bool = False, + position_sections: int | None = None, +) -> PolicyComponent: + """Build one-token attention-mask and position update.""" + if attention_dtype is None and position_dtype is None: + raise ValueError("decoder step update requires an attention mask or position ids") + if attention_dtype is None and fixed_capacity: + raise ValueError("fixed-capacity decoder step update requires an attention mask") + graph, builder = _make_graph("decoder_step_update") + op = builder.op + if attention_dtype is not None: + attention = builder.input( + "attention_mask", + dtype=attention_dtype, + shape=["batch", "context"], + ) + if fixed_capacity: + logical_length = builder.input( + "logical_length", + dtype=ir.DataType.INT64, + shape=["batch"], + ) + offsets = op.Range( + op.Constant(value_int=0), + op.Squeeze(op.Shape(attention, start=1, end=2), [0]), + op.Constant(value_int=1), + ) + slots = op.Equal( + op.Unsqueeze(offsets, [0]), + op.Unsqueeze(logical_length, [1]), + ) + next_attention = op.Where( + slots, + op.CastLike(op.Constant(value_int=1), attention), + attention, + ) + next_attention.shape = ir.Shape(["batch", "context"]) + else: + batch_shape = op.Shape(attention, start=0, end=1) + one_shape = op.Concat(batch_shape, op.Constant(value_ints=[1]), axis=0) + one = op.CastLike(op.ConstantOfShape(one_shape, value=ir.tensor([1])), attention) + next_attention = op.Concat(attention, one, axis=1) + next_attention.shape = ir.Shape(["batch", "context + 1"]) + builder.add_output(next_attention, "next_attention_mask") + if position_dtype is not None: + # A multi-axis rotary decoder carries one position row per rotary axis; + # advancing one token advances every axis, so the update is the same + # `+1` at either rank and only the declared shape differs. + position_shape: list[int | str] = ( + ["batch", 1] if position_sections is None else [position_sections, "batch", 1] + ) + position = builder.input( + "position_ids", + dtype=position_dtype, + shape=position_shape, + ) + next_position = op.Add(position, op.CastLike(op.Constant(value_int=1), position)) + next_position.shape = ir.Shape(position_shape) + builder.add_output(next_position, "next_position_ids") + return _component("mobius.policy.auxiliary@1", graph, {}) + + +def build_grammar_logits_processor() -> PolicyComponent: + """Apply a grammar adapter's mask and return a forced or sampled token.""" + graph, builder = _make_graph("grammar_guided_sampler") + op = builder.op + logits = builder.input("logits", ir.DataType.FLOAT, ["batch", "vocabulary"]) + logits_mask = builder.input("logits_mask", ir.DataType.BOOL, ["batch", "vocabulary"]) + forced_tokens = builder.input("forced_tokens", ir.DataType.INT64, ["batch", 1]) + forced_length = builder.input("forced_length", ir.DataType.INT64, ["batch"]) + blocked = op.CastLike(op.Constant(value_float=-3.4028235e38), logits) + masked_logits = op.Where(logits_mask, logits, blocked) + sampled = op.ArgMax(masked_logits, axis=-1, keepdims=1) + token = op.Where( + op.Unsqueeze(op.Greater(forced_length, op.Constant(value_int=0)), [-1]), + forced_tokens, + sampled, + ) + token.shape = ir.Shape(["batch", 1]) + builder.add_output(token, "token") + return _component("onnx-genai.grammar-guidance@1", graph, {}) + + +def build_adaptive_k_policy(*, max_k: int = 16, min_k: int = 1) -> PolicyComponent: + """Build the advisory adjacent-probe adaptive-K controller from ORT GenAI.""" + if not 1 <= min_k <= max_k: + raise ValueError("adaptive K requires 1 <= min_k <= max_k") + graph, builder = _make_graph("adaptive_k_policy") + op = builder.op + k_slots = max_k + 1 + estimate_slots = 4 * k_slots + 4 + + current_k = builder.input("current_k", ir.DataType.INT64, ["batch"]) + accepted = builder.input("accepted", ir.DataType.INT64, ["batch"]) + evaluated = builder.input("evaluated", ir.DataType.INT64, ["batch"]) + committed_tokens = builder.input("committed_tokens", ir.DataType.INT64, ["batch"]) + filled_proposal_budget = builder.input( + "filled_proposal_budget", ir.DataType.BOOL, ["batch"] + ) + draft_ms = builder.input("draft_ms", ir.DataType.FLOAT, ["batch"]) + target_ms = builder.input("target_ms", ir.DataType.FLOAT, ["batch"]) + estimates = builder.input("estimates", ir.DataType.FLOAT, ["batch", estimate_slots]) + + def section(start: int, end: int): + return op.Slice( + estimates, + op.Constant(value_ints=[start]), + op.Constant(value_ints=[end]), + op.Constant(value_ints=[1]), + ) + + token_estimates = section(0, k_slots) + millisecond_estimates = section(k_slots, 2 * k_slots) + acceptance_estimates = section(2 * k_slots, 3 * k_slots) + sample_counts = section(3 * k_slots, 4 * k_slots) + controller = section(4 * k_slots, estimate_slots) + probe_origin_k = op.Clip( + op.Cast(op.Slice(controller, [0], [1], [1]), to=ir.DataType.INT64), + op.Constant(value_int=0), + op.Constant(value_int=max_k), + ) + probe_observations = op.Cast(op.Slice(controller, [1], [2], [1]), to=ir.DataType.INT64) + stable_observations = op.Cast(op.Slice(controller, [2], [3], [1]), to=ir.DataType.INT64) + probe_cooldown = op.Cast(op.Slice(controller, [3], [4], [1]), to=ir.DataType.INT64) + + index = op.Unsqueeze(current_k, [-1]) + zero_i = op.Constant(value_int=0) + one_i = op.Constant(value_int=1) + two_i = op.Constant(value_int=2) + total_ms = op.Add(draft_ms, target_ms) + finite_time = op.Not(op.Or(op.IsNaN(total_ms), op.IsInf(total_ms))) + valid = op.And( + op.Greater(evaluated, zero_i), + op.And( + op.Greater(committed_tokens, zero_i), + op.And( + filled_proposal_budget, + op.And(op.Greater(total_ms, op.Constant(value_float=0.0)), finite_time), + ), + ), + ) + valid_col = op.Unsqueeze(valid, [-1]) + + old_tokens = op.GatherElements(token_estimates, index, axis=1) + old_ms = op.GatherElements(millisecond_estimates, index, axis=1) + old_acceptance = op.GatherElements(acceptance_estimates, index, axis=1) + old_samples = op.GatherElements(sample_counts, index, axis=1) + sample_tokens = op.Unsqueeze(op.Cast(committed_tokens, to=ir.DataType.FLOAT), [-1]) + sample_acceptance = op.Unsqueeze( + op.Div( + op.Cast(accepted, to=ir.DataType.FLOAT), + op.Cast(evaluated, to=ir.DataType.FLOAT), + ), + [-1], + ) + sample_ms = op.Unsqueeze(total_ms, [-1]) + first_sample = op.Equal(old_samples, op.Constant(value_float=0.0)) + alpha = op.Constant(value_float=0.25) + + def ewma(old, sample): + return op.Where( + first_sample, + sample, + op.Add(old, op.Mul(alpha, op.Sub(sample, old))), + ) + + updated_tokens = op.Where(valid_col, ewma(old_tokens, sample_tokens), old_tokens) + updated_ms = op.Where(valid_col, ewma(old_ms, sample_ms), old_ms) + updated_acceptance = op.Where( + valid_col, ewma(old_acceptance, sample_acceptance), old_acceptance + ) + updated_samples = op.Where( + valid_col, + op.Add(old_samples, op.Constant(value_float=1.0)), + old_samples, + ) + next_tokens = op.ScatterElements(token_estimates, index, updated_tokens, axis=1) + next_ms = op.ScatterElements(millisecond_estimates, index, updated_ms, axis=1) + next_acceptance = op.ScatterElements( + acceptance_estimates, index, updated_acceptance, axis=1 + ) + next_samples = op.ScatterElements(sample_counts, index, updated_samples, axis=1) + + current_throughput = op.Div( + updated_tokens, + op.Max(updated_ms, op.Constant(value_float=1e-12)), + ) + origin_index = probe_origin_k + origin_tokens = op.GatherElements(token_estimates, origin_index, axis=1) + origin_ms = op.GatherElements(millisecond_estimates, origin_index, axis=1) + origin_throughput = op.Div( + origin_tokens, + op.Max(origin_ms, op.Constant(value_float=1e-12)), + ) + current_col = index + in_probe = op.Greater(probe_origin_k, zero_i) + next_probe_count = op.Add(probe_observations, one_i) + probing_up = op.Greater(current_col, probe_origin_k) + severe_regression = op.And( + valid_col, + op.And( + in_probe, + op.And( + op.Greater(origin_throughput, op.Constant(value_float=0.0)), + op.Less( + current_throughput, + op.Mul(origin_throughput, op.Constant(value_float=0.8)), + ), + ), + ), + ) + probe_ready = op.And( + valid_col, + op.And( + in_probe, + op.And( + op.Not(severe_regression), + op.GreaterOrEqual(next_probe_count, two_i), + ), + ), + ) + throughput_safe = op.Or( + op.LessOrEqual(origin_throughput, op.Constant(value_float=0.0)), + op.GreaterOrEqual( + current_throughput, + op.Mul(origin_throughput, op.Constant(value_float=0.97)), + ), + ) + acceptance_safe = op.Or( + op.Not(probing_up), + op.GreaterOrEqual(updated_acceptance, op.Constant(value_float=0.75)), + ) + keep_probe = op.And(probe_ready, op.And(throughput_safe, acceptance_safe)) + finish_probe = op.Or(severe_regression, probe_ready) + probe_k = op.Where( + finish_probe, + op.Where(keep_probe, current_col, probe_origin_k), + current_col, + ) + probe_origin = op.Where(finish_probe, zero_i, probe_origin_k) + probe_count = op.Where( + finish_probe, + zero_i, + op.Where(valid_col, next_probe_count, probe_observations), + ) + probe_stable = op.Where(finish_probe, zero_i, stable_observations) + probe_next_cooldown = op.Where( + finish_probe, + op.Where(keep_probe, one_i, op.Constant(value_int=6)), + probe_cooldown, + ) + + stable_count = op.Where(valid_col, op.Add(stable_observations, one_i), stable_observations) + cooling = op.And(valid_col, op.Greater(probe_cooldown, zero_i)) + stable_cooldown = op.Where(cooling, op.Sub(probe_cooldown, one_i), probe_cooldown) + can_probe = op.And( + valid_col, + op.And( + op.Equal(probe_cooldown, zero_i), + op.GreaterOrEqual(updated_samples, op.Constant(value_float=2.0)), + ), + ) + probe_down = op.And( + can_probe, + op.And( + op.Less(updated_acceptance, op.Constant(value_float=0.5)), + op.Greater(current_col, op.Constant(value_int=min_k)), + ), + ) + probe_up = op.And( + can_probe, + op.And( + op.GreaterOrEqual(updated_acceptance, op.Constant(value_float=0.75)), + op.And( + op.Less(current_col, op.Constant(value_int=max_k)), + op.GreaterOrEqual(stable_count, two_i), + ), + ), + ) + start_probe = op.Or(probe_down, probe_up) + candidate_k = op.Where( + probe_down, + op.Sub(current_col, one_i), + op.Add(current_col, one_i), + ) + stable_k = op.Where(start_probe, candidate_k, current_col) + stable_origin = op.Where(start_probe, current_col, probe_origin_k) + stable_count = op.Where(start_probe, zero_i, stable_count) + + computed_k = op.Squeeze(op.Where(in_probe, probe_k, stable_k), [-1]) + next_probe_origin = op.Where(in_probe, probe_origin, stable_origin) + next_probe_observations = op.Where( + in_probe, + probe_count, + op.Where(start_probe, zero_i, probe_observations), + ) + next_stable_observations = op.Where(in_probe, probe_stable, stable_count) + next_probe_cooldown = op.Where(in_probe, probe_next_cooldown, stable_cooldown) + next_controller = op.Cast( + op.Concat( + next_probe_origin, + next_probe_observations, + next_stable_observations, + next_probe_cooldown, + axis=1, + ), + to=ir.DataType.FLOAT, + ) + computed_estimates = op.Concat( + next_tokens, + next_ms, + next_acceptance, + next_samples, + next_controller, + axis=1, + ) + next_k = op.Where(valid, computed_k, current_k) + next_estimates = op.Where(valid_col, computed_estimates, estimates) + next_k.shape = ir.Shape(["batch"]) + next_estimates.shape = ir.Shape(["batch", estimate_slots]) + builder.add_output(next_k, "next_k") + builder.add_output(next_estimates, "next_estimates") + return _component( + "onnx-genai.adaptive-proposal-budget@1", + graph, + { + "role": "adaptive_proposal_budget", + "current_k": "current_k", + "accepted": "accepted", + "evaluated": "evaluated", + "committed_tokens": "committed_tokens", + "filled_proposal_budget": "filled_proposal_budget", + "draft_ms": "draft_ms", + "target_ms": "target_ms", + "estimates": "estimates", + "next_k": "next_k", + "next_estimates": "next_estimates", + "effect": "adaptive", + }, + "adaptive", + ) + + +def build_seeded_categorical_sampler() -> PolicyComponent: + """Build request-parameterized categorical sampling with explicit RNG state. + + Threefry is counter based: identical tensor inputs produce the same token. + Temperature, top-k, top-p, and min-p remain request inputs; + changing ordinary generation options never regenerates this artifact. + """ + graph, builder = _make_graph("seeded_categorical_sampler") + op = builder.op + logits = builder.input("logits", ir.DataType.FLOAT, ["batch", "vocabulary"]) + temperature = builder.input("temperature", ir.DataType.FLOAT, ["batch"]) + top_k = builder.input("top_k", ir.DataType.INT64, ["batch"]) + top_p = builder.input("top_p", ir.DataType.FLOAT, ["batch"]) + min_p = builder.input("min_p", ir.DataType.FLOAT, ["batch"]) + seed = builder.input("seed", ir.DataType.INT64, ["batch"]) + counter = builder.input("counter", ir.DataType.INT64, ["batch"]) + active = builder.input("active", ir.DataType.BOOL, ["batch"]) + done = builder.input("done", ir.DataType.BOOL, ["batch"]) + enabled = op.And(active, op.Not(done)) + + # Threefry2x64: a counter-based Random123 generator with no hidden state. + # Unsigned arithmetic gives the specified modulo-2^64 round behavior. + k0 = op.Cast(seed, to=ir.DataType.UINT64) + k1 = op.Constant(value_int=0) + k1 = op.Cast(k1, to=ir.DataType.UINT64) + parity = op.Cast(op.Constant(value_int=0x1BD11BDAA9FC1A22), to=ir.DataType.UINT64) + k2 = op.BitwiseXor(op.BitwiseXor(k0, k1), parity) + keys = [k0, k1, k2] + x0 = op.Add(op.Cast(counter, to=ir.DataType.UINT64), k0) + x1 = op.Add(k1, op.Cast(op.Constant(value_int=0), to=ir.DataType.UINT64)) + rotations = [16, 42, 12, 31, 16, 32, 24, 21] + for round_index in range(20): + x0 = op.Add(x0, x1) + rotation = rotations[round_index % len(rotations)] + left = op.BitShift( + x1, + op.Cast(op.Constant(value_int=rotation), to=ir.DataType.UINT64), + direction="LEFT", + ) + right = op.BitShift( + x1, + op.Cast(op.Constant(value_int=64 - rotation), to=ir.DataType.UINT64), + direction="RIGHT", + ) + x1 = op.BitwiseXor(op.BitwiseOr(left, right), x0) + if (round_index + 1) % 4 == 0: + injection = (round_index + 1) // 4 + x0 = op.Add(x0, keys[injection % 3]) + x1 = op.Add( + op.Add(x1, keys[(injection + 1) % 3]), + op.Cast(op.Constant(value_int=injection), to=ir.DataType.UINT64), + ) + mantissa = op.BitShift( + x0, + op.Cast(op.Constant(value_int=11), to=ir.DataType.UINT64), + direction="RIGHT", + ) + uniform = op.Div( + op.Cast(mantissa, to=ir.DataType.DOUBLE), + op.Cast( + op.Constant(value_float=9_007_199_254_740_992.0), + to=ir.DataType.DOUBLE, + ), + ) + uniform = op.Cast(uniform, to=ir.DataType.FLOAT) + + blocked = op.CastLike(op.Constant(value_float=-3.4028235e38), logits) + safe_temperature = op.Unsqueeze( + op.Max(temperature, op.Constant(value_float=1e-6)), + [-1], + ) + scaled_logits = op.Div(logits, safe_temperature) + safe_min_p = op.Unsqueeze( + op.Clip( + min_p, + op.Constant(value_float=1e-20), + op.Constant(value_float=1.0), + ), + [-1], + ) + min_p_threshold = op.Add( + op.ReduceMax(scaled_logits, axes=[-1], keepdims=1), + op.Log(safe_min_p), + ) + min_p_mask = op.Or( + op.Unsqueeze(op.LessOrEqual(min_p, op.Constant(value_float=0.0)), [-1]), + op.GreaterOrEqual(scaled_logits, min_p_threshold), + ) + scaled_logits = op.Where(min_p_mask, scaled_logits, blocked) + + vocabulary = op.Shape(logits, start=1, end=2) + requested_k = op.Where( + op.Greater(top_k, op.Constant(value_int=0)), + top_k, + vocabulary, + ) + effective_k = op.Min(op.Max(requested_k, op.Constant(value_int=1)), vocabulary) + _, top_indices = op.TopK( + scaled_logits, + vocabulary, + axis=-1, + largest=1, + sorted=1, + _outputs=2, + ) + ranks = op.Range( + op.Constant(value_int=0), + op.Squeeze(vocabulary, [0]), + op.Constant(value_int=1), + ) + ranks = op.Expand(op.Unsqueeze(ranks, [0]), op.Shape(logits)) + keep_top_k = op.Less(ranks, op.Unsqueeze(effective_k, [-1])) + top_k_mask = op.Greater( + op.ScatterElements( + op.ConstantOfShape( + op.Shape(logits), + value=ir.tensor([0], dtype=ir.DataType.INT64), + ), + top_indices, + op.Cast(keep_top_k, to=ir.DataType.INT64), + axis=1, + ), + op.Constant(value_int=0), + ) + top_k_logits = op.Where(top_k_mask, scaled_logits, blocked) + probabilities = op.Softmax(top_k_logits, axis=-1) + + _, sorted_indices = op.TopK( + probabilities, + vocabulary, + axis=-1, + largest=1, + sorted=1, + _outputs=2, + ) + sorted_probabilities = op.GatherElements(probabilities, sorted_indices, axis=1) + cumulative_probabilities = op.CumSum( + sorted_probabilities, + op.Constant(value_int=1), + ) + safe_top_p = op.Unsqueeze( + op.Clip( + top_p, + op.Constant(value_float=1e-6), + op.Constant(value_float=1.0), + ), + [-1], + ) + keep_sorted = op.Less( + op.Sub(cumulative_probabilities, sorted_probabilities), + safe_top_p, + ) + top_p_mask = op.Greater( + op.ScatterElements( + op.ConstantOfShape( + op.Shape(logits), + value=ir.tensor([0], dtype=ir.DataType.INT64), + ), + sorted_indices, + op.Cast(keep_sorted, to=ir.DataType.INT64), + axis=1, + ), + op.Constant(value_int=0), + ) + probabilities = op.Where( + top_p_mask, + probabilities, + op.CastLike(op.Constant(value_float=0.0), probabilities), + ) + probabilities = op.Div( + probabilities, + op.Max( + op.ReduceSum(probabilities, axes=[-1], keepdims=1), + op.Constant(value_float=1e-20), + ), + ) + axis = op.Constant(value_int=-1) + cumulative = op.CumSum(probabilities, axis) + uniform = op.Unsqueeze(uniform, op.Constant(value_ints=[-1])) + candidates = op.GreaterOrEqual(cumulative, uniform) + token_ids = op.ArgMax( + op.Cast(candidates, to=ir.DataType.INT64), + axis=-1, + keepdims=0, + ) + token_ids = op.Where(enabled, token_ids, op.Constant(value_int=-1)) + next_counter = op.Where( + enabled, + op.Add(counter, op.Constant(value_int=1)), + counter, + ) + _set_public_shape(logits, ["batch", "vocabulary"]) + for value in (temperature, top_k, top_p, min_p, seed, counter, active, done): + _set_public_shape(value, ["batch"]) + _set_public_shape(token_ids, ["batch"]) + _set_public_shape(next_counter, ["batch"]) + builder.add_output(token_ids, "token") + builder.add_output(next_counter, "next_counter") + return _component( + "onnx-genai.token-sampler@2", + graph, + { + "role": "token_sampler", + "mode": "seeded_stochastic", + "batching": "per_row", + "inactive_rows": "preserve", + "logits": "logits", + "token": "token", + "temperature": "temperature", + "top_k": "top_k", + "top_p": "top_p", + "min_p": "min_p", + "active": "active", + "done": "done", + "seed": "seed", + "counter": "counter", + "next_counter": "next_counter", + "effect": "rng", + }, + "rng", + ) + + +def build_eos_termination(*, row_selective: bool = False) -> PolicyComponent: + """Build an EOS predicate for batched current tokens and an EOS-id set.""" + graph, builder = _make_graph("eos_termination") + op = builder.op + token_ids = builder.input("tokens", ir.DataType.INT64, ["batch"]) + eos_ids = builder.input( + "eos_ids", + ir.DataType.INT64, + ["batch", "num_eos"] if row_selective else ["num_eos"], + ) + eos_lengths = ( + builder.input("eos_lengths", ir.DataType.INT64, ["batch"]) if row_selective else None + ) + iteration = builder.input( + "iteration", + ir.DataType.INT64, + [1] if row_selective else ["batch"], + ) + max_iterations = builder.input( + "max_iterations", + ir.DataType.INT64, + ["batch"], + ) + tokens = op.Unsqueeze(token_ids, op.Constant(value_ints=[-1])) + eos = eos_ids if row_selective else op.Unsqueeze(eos_ids, [0]) + matches = op.Equal(tokens, eos) + if row_selective: + assert eos_lengths is not None + eos_positions = op.Range( + op.Constant(value_int=0), + op.Squeeze(op.Shape(eos_ids, start=1, end=2), [0]), + op.Constant(value_int=1), + ) + valid_eos = op.Less( + op.Unsqueeze(eos_positions, [0]), + op.Unsqueeze(eos_lengths, [-1]), + ) + matches = op.And(matches, valid_eos) + match_count = op.ReduceSum( + op.Cast(matches, to=ir.DataType.INT64), + axes=[-1], + keepdims=0, + ) + hit_eos = op.Greater(match_count, op.Constant(value_int=0)) + hit_limit = op.GreaterOrEqual( + op.Add(iteration, op.Constant(value_int=1)), + max_iterations, + ) + if row_selective: + active = builder.input("active", ir.DataType.BOOL, ["batch"]) + newly_done = op.And(active, op.Or(hit_eos, hit_limit)) + next_active = op.And(active, op.Not(newly_done)) + done = op.Not(next_active) + else: + done = op.Or(hit_eos, hit_limit) + next_active = op.Not(done) + continued = op.Greater( + op.ReduceMax(op.Cast(next_active, to=ir.DataType.INT64), keepdims=1), + op.Constant(value_int=0), + ) + if row_selective: + assert eos_lengths is not None + _set_public_shape(token_ids, ["batch"]) + _set_public_shape(eos_ids, ["batch", "num_eos"]) + _set_public_shape(eos_lengths, ["batch"]) + _set_public_shape(iteration, [1]) + _set_public_shape(max_iterations, ["batch"]) + _set_public_shape(active, ["batch"]) + _set_public_shape(done, ["batch"]) + _set_public_shape(next_active, ["batch"]) + _set_public_shape(continued, [1]) + builder.add_output(done, "done") + if row_selective: + builder.add_output(next_active, "next_active") + builder.add_output(continued, "continue") + return _component( + ( + "onnx-genai.termination-predicate@2" + if row_selective + else "onnx-genai.termination-predicate@1" + ), + graph, + { + "role": "termination_predicate", + "tokens": "tokens", + "eos_ids": "eos_ids", + "iteration": "iteration", + "max_iterations": "max_iterations", + **( + { + "eos_lengths": "eos_lengths", + "active": "active", + "batching": "per_row", + "inactive_rows": "preserve", + } + if row_selective + else {} + ), + "done": "done", + **({"next_active": "next_active"} if row_selective else {}), + "continue": "continue", + "effect": "termination", + }, + "termination", + ) + + +def build_scalar_constant(value: float) -> PolicyComponent: + """Materialize one producer-selected scalar as a rank-1 tensor.""" + graph, builder = _make_graph("scalar_constant") + constant = builder.op.Constant(value=ir.tensor([value], dtype=ir.DataType.FLOAT)) + constant.shape = ir.Shape([1]) + builder.add_output(constant, "value") + return _component("mobius.policy.auxiliary@1", graph, {}) + + +def build_shape_constant(dims: list[int]) -> PolicyComponent: + """Materialize a producer-selected integer shape vector.""" + graph, builder = _make_graph("shape_constant") + constant = builder.op.Constant( + value=ir.tensor([int(dim) for dim in dims], dtype=ir.DataType.INT64) + ) + constant.shape = ir.Shape([len(dims)]) + builder.add_output(constant, "shape") + return _component("mobius.policy.auxiliary@1", graph, {}) + + +def build_tensor_scale(dtype: ir.DataType = ir.DataType.FLOAT) -> PolicyComponent: + """Scale a rank-4 state tensor by a broadcast scalar factor.""" + graph, builder = _make_graph("tensor_scale") + op = builder.op + tensor = builder.input("tensor", dtype, ["batch", "channels", "height", "width"]) + scale = builder.input("scale", ir.DataType.FLOAT, [1]) + scaled = op.Mul(tensor, op.Cast(scale, to=dtype)) + scaled.shape = tensor.shape + builder.add_output(scaled, "scaled") + return _component("mobius.policy.auxiliary@1", graph, {}) + + +def build_zeros_like(dtype: ir.DataType = ir.DataType.FLOAT) -> PolicyComponent: + """Produce a zero tensor shaped like its reference, for state initializers.""" + graph, builder = _make_graph("zeros_like") + op = builder.op + reference = builder.input("reference", dtype, ["batch", "channels", "height", "width"]) + zeros = op.Mul(reference, op.Cast(op.Constant(value_float=0.0), to=dtype)) + zeros.shape = reference.shape + builder.add_output(zeros, "zeros") + return _component("mobius.policy.auxiliary@1", graph, {}) + + +def build_guidance_combine(dtype: ir.DataType = ir.DataType.FLOAT) -> PolicyComponent: + """Combine two conditioned estimates by a per-row guidance scale. + + ``estimate = unconditional + scale * (conditional - unconditional)`` is the + classifier-free-guidance extrapolation. The scale is an ordinary per-row + tensor input, so a caller can vary it per request without a rebuild. + """ + graph, builder = _make_graph("guidance_combine") + op = builder.op + unconditional = builder.input( + "unconditional", dtype, ["batch", "channels", "height", "width"] + ) + conditional = builder.input("conditional", dtype, ["batch", "channels", "height", "width"]) + scale = builder.input("scale", ir.DataType.FLOAT, ["batch"]) + # (batch,) -> (batch, 1, 1, 1) so the row scale broadcasts over the latent. + factor = op.Unsqueeze(op.Cast(scale, to=dtype), op.Constant(value_ints=[1, 2, 3])) + estimate = op.Add(unconditional, op.Mul(factor, op.Sub(conditional, unconditional))) + estimate.shape = unconditional.shape + builder.add_output(estimate, "estimate") + return _component( + "onnx-genai.guidance-combine@1", + graph, + { + "role": "guidance_combine", + "unconditional": "unconditional", + "conditional": "conditional", + "scale": "scale", + "estimate": "estimate", + }, + ) + + +_RNG_MODULUS = 2147483647 +_RNG_MULTIPLIER = 48271 +_RNG_STRIDE = 2654435761 + + +def _counter_uniform(op, key, counter): + """Map a counter-based ``(key, counter)`` pair onto a uniform in ``(0, 1)``. + + The stream is a pure function of its inputs, so a row's noise depends only + on its own seed and offset and never on batch position or iteration order. + """ + modulus = op.Constant(value_int=_RNG_MODULUS) + multiplier = op.Constant(value_int=_RNG_MULTIPLIER) + state = op.Mod( + op.Add( + op.Add(key, op.Constant(value_int=1)), + op.Mul(counter, op.Constant(value_int=_RNG_STRIDE)), + ), + modulus, + fmod=0, + ) + for _ in range(3): + state = op.Mod(op.Mul(state, multiplier), modulus, fmod=0) + # Integer division stands in for a right shift: BitShift is unsigned-only. + state = op.BitwiseXor(state, op.Div(state, op.Constant(value_int=2048))) + state = op.Mod(op.Mul(state, multiplier), modulus, fmod=0) + # Shift off zero so the logarithm in the Box-Muller transform stays finite. + return op.Div( + op.Cast(op.Add(state, op.Constant(value_int=1)), to=ir.DataType.FLOAT), + op.Constant(value_float=float(_RNG_MODULUS + 2)), + ) + + +def build_counter_rng_normal(dtype: ir.DataType = ir.DataType.FLOAT) -> PolicyComponent: + """Draw standard-normal noise from explicit counter RNG state. + + Inputs are a per-row ``seed``, a per-row ``offset`` counter, and the target + ``shape``. The component is pure: it consumes counter state and returns the + advanced counter, so RNG progress is loop-carried workflow state rather than + hidden session state inside an operator. + """ + graph, builder = _make_graph("counter_rng_normal") + op = builder.op + seed = builder.input("seed", ir.DataType.INT64, ["batch"]) + offset = builder.input("offset", ir.DataType.INT64, ["batch"]) + row_shape = builder.input("row_shape", ir.DataType.INT64, ["row_rank"]) + + # The batch extent comes from the per-row seed, so the draw is always + # request-aligned no matter how many rows the runtime batched together. + shape = op.Concat(op.Shape(seed), row_shape, axis=0) + row_elements = op.ReduceProd(row_shape, keepdims=1) + # (row_elements,) counter positions, shared by every row's private stream. + positions = op.Range( + op.Constant(value_int=0), + op.Squeeze(row_elements, op.Constant(value_ints=[0])), + op.Constant(value_int=1), + ) + positions = op.Unsqueeze(positions, op.Constant(value_ints=[0])) + base = op.Mul(op.Unsqueeze(offset, op.Constant(value_ints=[1])), row_elements) + counter = op.Add(base, positions) + key = op.Unsqueeze(seed, op.Constant(value_ints=[1])) + + # Box-Muller over two independent counter blocks keeps the draw stateless. + uniform_radius = _counter_uniform(op, key, counter) + uniform_angle = _counter_uniform( + op, key, op.Add(counter, op.Mul(row_elements, op.Constant(value_int=2))) + ) + radius = op.Sqrt(op.Mul(op.Constant(value_float=-2.0), op.Log(uniform_radius))) + angle = op.Mul(op.Constant(value_float=6.283185307179586), uniform_angle) + flat = op.Mul(radius, op.Cos(angle)) + noise = op.Cast(op.Reshape(flat, shape), to=dtype) + noise.shape = ir.Shape(["batch", "channels", "height", "width"]) + next_offset = op.Add(offset, op.Constant(value_int=1)) + next_offset.shape = ir.Shape(["batch"]) + builder.add_output(noise, "noise") + builder.add_output(next_offset, "next_offset") + return _component( + "onnx-genai.counter-rng@1", + graph, + { + "role": "counter_rng", + "seed": "seed", + "offset": "offset", + "row_shape": "row_shape", + "noise": "noise", + "next_offset": "next_offset", + }, + ) + + +_IMAGE_LATENT_DIMS: tuple[str, ...] = ("batch", "channels", "height", "width") + + +def build_euler_model_input( + dtype: ir.DataType = ir.DataType.FLOAT, + latent_dims: Sequence[str] = _IMAGE_LATENT_DIMS, +) -> PolicyComponent: + """Scale a latent for the Euler denoiser input at the current sigma. + + ``latent_dims`` names the latent axes. The per-row sigma is broadcast over + every axis after the batch, so a video latent that carries a temporal axis + works without a separate component. + """ + graph, builder = _make_graph("euler_model_input") + op = builder.op + sample = builder.input("sample", dtype, list(latent_dims)) + step = builder.input("step", ir.DataType.INT64, ["batch"]) + schedule = builder.input("schedule", ir.DataType.FLOAT, ["schedule_length"]) + sigma = op.Gather(schedule, step, axis=0) + scale = op.Sqrt(op.Add(op.Mul(sigma, sigma), op.Constant(value_float=1.0))) + scale = op.Cast(scale, to=dtype) + scale = op.Unsqueeze(scale, op.Constant(value_ints=list(range(1, len(latent_dims))))) + model_input = op.Div(sample, scale) + model_input.shape = sample.shape + builder.add_output(model_input, "model_input") + return _component("mobius.policy.auxiliary@1", graph, {}) + + +def build_euler_solver_step( + dtype: ir.DataType = ir.DataType.FLOAT, + latent_dims: Sequence[str] = _IMAGE_LATENT_DIMS, +) -> PolicyComponent: + """Build the generic Euler update ``x_next = x + dx * (sigma_next-sigma)``. + + ``latent_dims`` names the latent axes so the same update serves image and + video latents; the sigma delta broadcasts over every axis after the batch. + """ + graph, builder = _make_graph("euler_solver_step") + op = builder.op + sample = builder.input("sample", dtype, list(latent_dims)) + derivative = builder.input("derivative", dtype, list(latent_dims)) + step = builder.input("step", ir.DataType.INT64, ["batch"]) + schedule = builder.input("schedule", ir.DataType.FLOAT, ["schedule_length"]) + final_index = op.Sub(op.Shape(schedule, start=0, end=1), op.Constant(value_ints=[1])) + next_step = op.Min(op.Add(step, op.Constant(value_int=1)), final_index) + sigma = op.Gather(schedule, step, axis=0) + sigma_next = op.Gather(schedule, next_step, axis=0) + delta = op.Sub(sigma_next, sigma) + delta = op.Cast(delta, to=dtype) + delta = op.Unsqueeze(delta, op.Constant(value_ints=list(range(1, len(latent_dims))))) + next_sample = op.Add(sample, op.Mul(derivative, delta)) + builder.add_output(next_sample, "next_state") + return _component( + "onnx-genai.solver-step@1", + graph, + { + "role": "solver_step", + "state": "sample", + "estimate": "derivative", + "step": "step", + "schedule": "schedule", + "next_state": "next_state", + "effect": "solver", + }, + "solver", + ) + + +def build_multistep_solver_step( + dtype: ir.DataType = ir.DataType.FLOAT, +) -> PolicyComponent: + """Build a second-order multistep solver update with explicit history state. + + This is the DPM-Solver++(2M) midpoint update expressed entirely in ONNX. The + solver's memory - the previous data-space estimate - is an ordinary tensor + port, so the workflow carries it as declared state instead of the runtime + holding hidden scheduler attributes. + + ``next = (sig_next/sig_now) * sample - alpha_next * (sig_next/sig_now_ratio - 1) + * (D0 + 0.5 * D1)`` where ``D0`` is the current data estimate and ``D1`` the + finite difference against the previous one. The first and last steps have no + usable history, so ``D1`` is masked to zero there, which reduces the update to + the first-order form exactly as a multistep scheme's warm-up and final step do. + """ + graph, builder = _make_graph("multistep_solver_step") + op = builder.op + sample = builder.input("sample", dtype, ["batch", "channels", "height", "width"]) + estimate = builder.input("estimate", dtype, ["batch", "channels", "height", "width"]) + history = builder.input("history", dtype, ["batch", "channels", "height", "width"]) + step = builder.input("step", ir.DataType.INT64, ["batch"]) + schedule = builder.input("schedule", ir.DataType.FLOAT, ["schedule_length"]) + + one = op.Constant(value_int=1) + zero = op.Constant(value_int=0) + final_index = op.Sub(op.Shape(schedule, start=0, end=1), op.Constant(value_ints=[1])) + next_step = op.Min(op.Add(step, one), final_index) + previous_step = op.Max(op.Sub(step, one), zero) + sigma_now = op.Gather(schedule, step, axis=0) + sigma_next = op.Gather(schedule, next_step, axis=0) + sigma_previous = op.Gather(schedule, previous_step, axis=0) + + def alpha(sigma): + return op.Div( + op.Constant(value_float=1.0), + op.Sqrt(op.Add(op.Mul(sigma, sigma), op.Constant(value_float=1.0))), + ) + + alpha_now = alpha(sigma_now) + alpha_next = alpha(sigma_next) + # The variance-preserving noise level, sigma_t in the DPM-Solver derivation. + noise_now = op.Mul(sigma_now, alpha_now) + noise_next = op.Mul(sigma_next, alpha_next) + + def row_scalar(value): + return op.Unsqueeze(op.Cast(value, to=dtype), op.Constant(value_ints=[1, 2, 3])) + + # Data-space estimate x0 from the epsilon-space model output. + data_estimate = op.Div( + op.Sub(sample, op.Mul(row_scalar(noise_now), estimate)), + row_scalar(alpha_now), + ) + ratio = op.Div(sigma_next, sigma_now) + sample_coefficient = row_scalar(op.Div(noise_next, noise_now)) + data_coefficient = row_scalar( + op.Mul(alpha_next, op.Sub(ratio, op.Constant(value_float=1.0))) + ) + + # Half-log-SNR spacing ratio; lambda(sigma) = -log(sigma) for this schedule. + interval = op.Sub(op.Log(sigma_now), op.Log(sigma_next)) + previous_interval = op.Sub(op.Log(sigma_previous), op.Log(sigma_now)) + difference = op.Mul( + row_scalar(op.Div(interval, previous_interval)), + op.Sub(data_estimate, history), + ) + warm = op.Greater(step, zero) + penultimate = op.Less(op.Add(step, one), final_index) + usable = op.Unsqueeze(op.And(warm, penultimate), op.Constant(value_ints=[1, 2, 3])) + difference = op.Where(usable, difference, op.Cast(op.Constant(value_float=0.0), to=dtype)) + + next_state = op.Sub( + op.Mul(sample_coefficient, sample), + op.Mul( + data_coefficient, + op.Add( + data_estimate, + op.Mul(op.Cast(op.Constant(value_float=0.5), to=dtype), difference), + ), + ), + ) + next_state.shape = sample.shape + data_estimate.shape = sample.shape + builder.add_output(next_state, "next_state") + builder.add_output(data_estimate, "next_history") + return _component( + "onnx-genai.solver-step@1", + graph, + { + "role": "solver_step", + "state": "sample", + "estimate": "estimate", + "history": "history", + "step": "step", + "schedule": "schedule", + "next_state": "next_state", + "next_history": "next_history", + }, + ) + + +def build_flow_match_solver_step( + dtype: ir.DataType = ir.DataType.FLOAT, +) -> PolicyComponent: + """Build the Euler update for rank-3 (packed/patchified) latents. + + Identical arithmetic to :func:`build_euler_solver_step` — ``x_next = x + dx * + (sigma_next - sigma)`` — but broadcasts the per-batch step size over a + ``(batch, sequence, channels)`` latent instead of a rank-4 image latent. + Flow-matching transformers (Qwen Image, Flux, SD3) carry latents in this + packed layout. + """ + graph, builder = _make_graph("flow_match_solver_step") + op = builder.op + sample = builder.input("sample", dtype, ["batch", "sequence", "channels"]) + derivative = builder.input("derivative", dtype, ["batch", "sequence", "channels"]) + step = builder.input("step", ir.DataType.INT64, ["batch"]) + schedule = builder.input("schedule", ir.DataType.FLOAT, ["schedule_length"]) + final_index = op.Sub(op.Shape(schedule, start=0, end=1), op.Constant(value_ints=[1])) + next_step = op.Min(op.Add(step, op.Constant(value_int=1)), final_index) + sigma = op.Gather(schedule, step, axis=0) + sigma_next = op.Gather(schedule, next_step, axis=0) + delta = op.Cast(op.Sub(sigma_next, sigma), to=dtype) + delta = op.Unsqueeze(delta, op.Constant(value_ints=[1, 2])) + next_sample = op.Add(sample, op.Mul(derivative, delta)) + builder.add_output(next_sample, "next_state") + return _component( + "onnx-genai.solver-step@1", + graph, + { + "role": "solver_step", + "state": "sample", + "estimate": "derivative", + "step": "step", + "schedule": "schedule", + "next_state": "next_state", + "effect": "solver", + }, + "solver", + ) + + +def build_ddim_solver_step( + dtype: ir.DataType = ir.DataType.FLOAT, + latent_dims: Sequence[str] = _IMAGE_LATENT_DIMS, + *, + clip_sample_range: float | None = None, +) -> PolicyComponent: + """Build the deterministic DDIM update (``eta = 0``, epsilon prediction). + + ``schedule`` holds the cumulative alpha of every denoising step followed by + the alpha of the final step's predecessor, so entry ``i + 1`` is the + ``alpha_prev`` of step ``i``:: + + pred_x0 = (x - sqrt(1 - a_t) * eps) / sqrt(a_t) + x_prev = sqrt(a_prev) * pred_x0 + sqrt(1 - a_prev) * eps + + ``clip_sample_range`` reproduces schedulers configured with + ``clip_sample=True``, which clamp the predicted clean sample before the + reverse step. ``latent_dims`` names the latent axes, so the same update + serves image and video latents. + """ + graph, builder = _make_graph("ddim_solver_step") + op = builder.op + sample = builder.input("sample", dtype, list(latent_dims)) + derivative = builder.input("derivative", dtype, list(latent_dims)) + step = builder.input("step", ir.DataType.INT64, ["batch"]) + schedule = builder.input("schedule", ir.DataType.FLOAT, ["schedule_length"]) + final_index = op.Sub(op.Shape(schedule, start=0, end=1), op.Constant(value_ints=[1])) + next_step = op.Min(op.Add(step, op.Constant(value_int=1)), final_index) + broadcast_axes = op.Constant(value_ints=list(range(1, len(latent_dims)))) + alpha = op.Unsqueeze(op.Cast(op.Gather(schedule, step, axis=0), to=dtype), broadcast_axes) + alpha_prev = op.Unsqueeze( + op.Cast(op.Gather(schedule, next_step, axis=0), to=dtype), broadcast_axes + ) + one = op.CastLike(op.Constant(value_float=1.0), sample) + # Recover the predicted clean latent, then re-noise it to alpha_prev. + pred_original = op.Div( + op.Sub(sample, op.Mul(op.Sqrt(op.Sub(one, alpha)), derivative)), op.Sqrt(alpha) + ) + if clip_sample_range is not None: + limit = op.CastLike(op.Constant(value_float=float(clip_sample_range)), sample) + pred_original = op.Clip(pred_original, op.Neg(limit), limit) + next_sample = op.Add( + op.Mul(op.Sqrt(alpha_prev), pred_original), + op.Mul(op.Sqrt(op.Sub(one, alpha_prev)), derivative), + ) + _set_public_shape(next_sample, list(latent_dims)) + builder.add_output(next_sample, "next_state") + return _component( + "onnx-genai.solver-step@1", + graph, + { + "role": "solver_step", + "state": "sample", + "estimate": "derivative", + "step": "step", + "schedule": "schedule", + "next_state": "next_state", + "effect": "solver", + }, + "solver", + ) + + +def build_pack_latents_2x2(dtype: ir.DataType = ir.DataType.FLOAT) -> PolicyComponent: + """Patchify a 3D VAE latent into the transformer's packed token layout. + + ``(B, C, T, H, W) -> (B, T*(H/2)*(W/2), C*4)`` by folding each 2x2 spatial + patch into the channel axis, matching ``QwenImagePipeline._pack_latents``. + Shapes are derived from the input at runtime so the component stays valid + for any resolution. + """ + graph, builder = _make_graph("pack_latents") + op = builder.op + latent = builder.input( + "latent_sample", dtype, ["batch", "channels", "frames", "height", "width"] + ) + two = op.Constant(value_ints=[2]) + batch = op.Shape(latent, start=0, end=1) + channels = op.Shape(latent, start=1, end=2) + frames = op.Shape(latent, start=2, end=3) + height = op.Shape(latent, start=3, end=4) + width = op.Shape(latent, start=4, end=5) + half_h = op.Div(height, two) + half_w = op.Div(width, two) + # (B, C, T, H, W) -> (B, C, T, H/2, 2, W/2, 2) + patched = op.Reshape( + latent, + op.Concat(batch, channels, frames, half_h, two, half_w, two, axis=0), + ) + # -> (B, T, H/2, W/2, C, 2, 2) so each 2x2 patch is contiguous per channel + patched = op.Transpose(patched, perm=[0, 2, 3, 5, 1, 4, 6]) + tokens = op.Mul(op.Mul(frames, half_h), half_w) + packed = op.Reshape( + patched, + op.Concat(batch, tokens, op.Mul(channels, op.Constant(value_ints=[4])), axis=0), + ) + _set_public_shape(packed, ["batch", "sequence", "packed_channels"]) + builder.add_output(packed, "packed_latent") + return _component("mobius.policy.auxiliary@1", graph, {}) + + +def build_unpack_latents_2x2(dtype: ir.DataType = ir.DataType.FLOAT) -> PolicyComponent: + """Invert :func:`build_pack_latents_2x2` back to a 3D VAE latent. + + ``(B, S, C*4) -> (B, C, 1, H*2, W*2)`` given the packed latent grid + ``height``/``width`` (in packed tokens). The single frame matches the image + pipelines, whose VAE always encodes one temporal chunk. + """ + graph, builder = _make_graph("unpack_latents") + op = builder.op + packed = builder.input("packed_latent", dtype, ["batch", "sequence", "packed_channels"]) + height = builder.input("height", ir.DataType.INT64, [1]) + width = builder.input("width", ir.DataType.INT64, [1]) + two = op.Constant(value_ints=[2]) + batch = op.Shape(packed, start=0, end=1) + channels = op.Div(op.Shape(packed, start=2, end=3), op.Constant(value_ints=[4])) + # (B, S, C*4) -> (B, H, W, C, 2, 2) + grid = op.Reshape(packed, op.Concat(batch, height, width, channels, two, two, axis=0)) + # -> (B, C, H, 2, W, 2) -> (B, C, 1, H*2, W*2) + grid = op.Transpose(grid, perm=[0, 3, 1, 4, 2, 5]) + latent = op.Reshape( + grid, + op.Concat( + batch, + channels, + op.Constant(value_ints=[1]), + op.Mul(height, two), + op.Mul(width, two), + axis=0, + ), + ) + _set_public_shape(latent, ["batch", "channels", "frames", "height", "width"]) + builder.add_output(latent, "latent_sample") + return _component("mobius.policy.auxiliary@1", graph, {}) + + +def build_sequence_concat(dtype: ir.DataType = ir.DataType.FLOAT) -> PolicyComponent: + """Concatenate two token sequences along the sequence axis. + + Image-editing denoisers attend jointly over the generated tokens and the + source-image tokens, so the model input is ``concat([target, source], 1)`` + and the estimate is sliced back to the target length inside the denoiser. + """ + graph, builder = _make_graph("sequence_concat") + op = builder.op + target = builder.input("target", dtype, ["batch", "target_sequence", "channels"]) + source = builder.input("source", dtype, ["batch", "source_sequence", "channels"]) + joined = op.Concat(target, source, axis=1) + _set_public_shape(joined, ["batch", "sequence", "channels"]) + builder.add_output(joined, "sequence") + return _component("mobius.policy.auxiliary@1", graph, {}) + + +def build_true_cfg( + dtype: ir.DataType = ir.DataType.FLOAT, + *, + guidance_scale: float = 4.0, +) -> PolicyComponent: + """Combine conditional/unconditional estimates with Qwen's true CFG. + + ``combined = uncond + scale * (cond - uncond)``, then rescaled to preserve + the conditional estimate's per-token norm:: + + noise_pred = combined * (||cond||_2 / ||combined||_2) + + The reduction runs in float32: squared activations overflow float16, and + onnxruntime ships no bfloat16 ``ReduceL2`` kernel at all (a bfloat16 graph + containing one cannot be assigned to any provider and fails to load). + """ + graph, builder = _make_graph("true_cfg") + op = builder.op + cond = builder.input("conditional", dtype, ["batch", "sequence", "channels"]) + uncond = builder.input("unconditional", dtype, ["batch", "sequence", "channels"]) + cond_f32 = op.Cast(cond, to=ir.DataType.FLOAT) + uncond_f32 = op.Cast(uncond, to=ir.DataType.FLOAT) + scale = op.Constant(value_float=float(guidance_scale)) + combined = op.Add(uncond_f32, op.Mul(op.Sub(cond_f32, uncond_f32), scale)) + cond_norm = op.ReduceL2(cond_f32, [-1], keepdims=True) + combined_norm = op.Max(op.ReduceL2(combined, [-1], keepdims=True), 1e-12) + guided = op.Cast(op.Mul(combined, op.Div(cond_norm, combined_norm)), to=dtype) + _set_public_shape(guided, ["batch", "sequence", "channels"]) + builder.add_output(guided, "estimate") + return _component("mobius.policy.auxiliary@1", graph, {}) + + +SOLVER_BUILDERS = { + "euler": build_euler_solver_step, + "multistep": build_multistep_solver_step, +} + + +def build_identity_model_input( + dtype: ir.DataType = ir.DataType.FLOAT, + latent_dims: Sequence[str] = _IMAGE_LATENT_DIMS, +) -> PolicyComponent: + """Pass the latent to the denoiser unchanged. + + DDIM-style schedulers define ``scale_model_input`` as the identity. Keeping + the node explicit means the workflow reads the same way for every solver + and the input scaling stays a declared, swappable policy. + """ + graph, builder = _make_graph("identity_model_input") + op = builder.op + sample = builder.input("sample", dtype, list(latent_dims)) + step = builder.input("step", ir.DataType.INT64, ["batch"]) + schedule = builder.input("schedule", ir.DataType.FLOAT, ["schedule_length"]) + model_input = op.Identity(sample) + # The step and schedule are unused by this policy but stay on the signature + # so a package can swap solvers without rewiring the workflow. + _ = op.Gather(schedule, step, axis=0) + _set_public_shape(model_input, list(latent_dims)) + builder.add_output(model_input, "model_input") + return _component("mobius.policy.auxiliary@1", graph, {}) + + +def build_masked_token_update() -> PolicyComponent: + """Build replacement of masked positions with explicit RNG-counter threading.""" + graph, builder = _make_graph("masked_token_update") + op = builder.op + current = builder.input("current_tokens", ir.DataType.INT64, ["batch", "sequence"]) + proposed = builder.input("proposed_tokens", ir.DataType.INT64, ["batch", "sequence"]) + logits = builder.input( + "logits", + ir.DataType.FLOAT, + ["batch", "sequence", "vocabulary"], + ) + masked = builder.input("masked", ir.DataType.BOOL, ["batch", "sequence"]) + step = builder.input("step", ir.DataType.INT64, ["batch"]) + total_steps = builder.input("total_steps", ir.DataType.INT64, ["batch"]) + seed = builder.input("seed", ir.DataType.INT64, ["batch"]) + offset = builder.input("offset", ir.DataType.INT64, ["batch"]) + sequence_length = op.Shape(masked, start=1, end=2) + positions = op.Range( + op.Constant(value_int=0), + op.Squeeze(sequence_length, op.Constant(value_ints=[0])), + op.Constant(value_int=1), + ) + positions = op.Unsqueeze(positions, op.Constant(value_ints=[0])) + stream = op.Add( + positions, + op.Unsqueeze( + op.Add(op.Add(seed, offset), op.Mul(step, op.Constant(value_int=17))), + op.Constant(value_ints=[-1]), + ), + ) + tie_noise = op.Div( + op.Cast(op.Mod(stream, op.Constant(value_int=997), fmod=0), to=ir.DataType.FLOAT), + op.Constant(value_float=997_000_000.0), + ) + probabilities = op.Softmax(logits, axis=-1) + confidence = op.Squeeze( + op.GatherElements( + probabilities, + op.Unsqueeze(proposed, op.Constant(value_ints=[-1])), + axis=-1, + ), + op.Constant(value_ints=[-1]), + ) + confidence = op.Add(confidence, tie_noise) + negative = op.CastLike(op.Constant(value_float=-1.0), confidence) + ranked_scores = op.Where(masked, confidence, negative) + left = op.Unsqueeze(ranked_scores, op.Constant(value_ints=[-1])) + right = op.Unsqueeze(ranked_scores, op.Constant(value_ints=[-2])) + rank = op.ReduceSum( + op.Cast(op.Greater(right, left), to=ir.DataType.INT64), + axes=[-1], + keepdims=0, + ) + remaining_before = op.ReduceSum( + op.Cast(masked, to=ir.DataType.INT64), + axes=[-1], + keepdims=0, + ) + steps_left = op.Max( + op.Constant(value_int=1), + op.Sub(total_steps, step), + ) + quota = op.Div( + op.Add(op.Sub(remaining_before, op.Constant(value_int=1)), steps_left), + steps_left, + ) + committed = op.And( + masked, + op.Less(rank, op.Unsqueeze(quota, op.Constant(value_ints=[-1]))), + ) + updated = op.Where(committed, proposed, current) + updated.shape = ir.Shape(["batch", "sequence"]) + remaining = op.And(masked, op.Not(committed)) + remaining.shape = ir.Shape(["batch", "sequence"]) + remaining_count = op.ReduceSum( + op.Cast(remaining, to=ir.DataType.INT64), + axes=[-1], + keepdims=0, + ) + done = op.Equal(remaining_count, op.Constant(value_int=0)) + done.shape = ir.Shape(["batch"]) + continued = op.Equal( + op.ReduceMax(op.Cast(done, to=ir.DataType.INT64), keepdims=1), + op.Constant(value_int=0), + ) + continued.shape = ir.Shape([1]) + next_offset = op.Add( + op.Add(offset, op.Squeeze(sequence_length, op.Constant(value_ints=[0]))), + op.Mul(seed, op.Constant(value_int=0)), + ) + next_offset.shape = ir.Shape(["batch"]) + builder.add_output(updated, "next_state") + builder.add_output(remaining, "next_mask") + builder.add_output(next_offset, "next_offset") + builder.add_output(done, "done") + builder.add_output(continued, "continue") + return _component( + "onnx-genai.masked-update@1", + graph, + { + "role": "masked_update", + "state": "current_tokens", + "proposal": "proposed_tokens", + "mask": "masked", + "step": "step", + "next_state": "next_state", + "next_mask": "next_mask", + "continue": "continue", + "rng": { + "seed": "seed", + "offset": "offset", + "next_offset": "next_offset", + }, + "effect": "update", + }, + "update", + ) + + +def build_speculative_acceptance() -> PolicyComponent: + """Build per-token speculative acceptance and accepted-prefix length.""" + graph, builder = _make_graph("speculative_acceptance") + op = builder.op + target_scores = builder.input( + "target_scores", ir.DataType.FLOAT, ["batch", "draft_sequence", "vocabulary"] + ) + proposed_tokens = builder.input( + "proposed_tokens", ir.DataType.INT64, ["batch", "draft_sequence"] + ) + seed = builder.input("seed", ir.DataType.INT64, ["batch"]) + offset = builder.input("offset", ir.DataType.INT64, ["batch"]) + target_tokens = op.ArgMax(target_scores, axis=-1, keepdims=0) + accepted = op.Equal(target_tokens, proposed_tokens) + rejected = op.Cast(op.Not(accepted), to=ir.DataType.INT64) + rejection_count = op.CumSum(rejected, op.Constant(value_int=-1)) + prefix = op.Cast( + op.Equal(rejection_count, op.Constant(value_int=0)), + to=ir.DataType.INT64, + ) + verified_count = op.ReduceSum(prefix, axes=[-1], keepdims=0) + first_rejection = op.And( + op.Cast(rejected, to=ir.DataType.BOOL), + op.Equal(rejection_count, op.Constant(value_int=1)), + ) + zeros = op.ConstantOfShape( + op.Shape(proposed_tokens), + value=ir.tensor([0], dtype=ir.DataType.INT64), + ) + # Publish the verified prefix plus the verifier's correction at the first + # mismatch. Trailing slots remain zero and are bounded by accepted_len. + accepted_tokens = op.Where( + op.Cast(prefix, to=ir.DataType.BOOL), + proposed_tokens, + op.Where(first_rejection, target_tokens, zeros), + ) + accepted_tokens.shape = ir.Shape(["batch", "draft_sequence"]) + draft_length = op.Shape(proposed_tokens, start=1, end=2) + done = op.Equal(verified_count, draft_length) + accepted_count = op.Min( + op.Add(verified_count, op.Cast(op.Not(done), to=ir.DataType.INT64)), + draft_length, + ) + continued = op.Not(done) + positions = op.Range( + op.Constant(value_int=0), + op.Squeeze(draft_length, op.Constant(value_ints=[0])), + op.Constant(value_int=1), + ) + valid = op.Less( + op.Unsqueeze(positions, op.Constant(value_ints=[0])), + op.Unsqueeze(accepted_count, op.Constant(value_ints=[-1])), + ) + accepted_tokens = op.Where(valid, accepted_tokens, zeros) + accepted_tokens.shape = ir.Shape(["batch", "draft_sequence"]) + next_offset = op.Add( + offset, + op.Squeeze(draft_length, op.Constant(value_ints=[0])), + ) + next_offset = op.Add(next_offset, op.Mul(seed, op.Constant(value_int=0))) + accepted_count.shape = ir.Shape(["batch"]) + done.shape = ir.Shape(["batch"]) + verified_count.shape = ir.Shape(["batch"]) + continued.shape = ir.Shape(["batch"]) + next_offset.shape = ir.Shape(["batch"]) + builder.add_output(accepted_tokens, "accepted_tokens") + builder.add_output(accepted_count, "accepted_len") + builder.add_output(done, "done") + builder.add_output(next_offset, "next_offset") + builder.add_output(verified_count, "rollback_len") + builder.add_output(continued, "continue") + return _component( + "onnx-genai.speculative-verifier@1", + graph, + { + "role": "speculative_verifier", + "target_scores": "target_scores", + "proposed_tokens": "proposed_tokens", + "accepted_tokens": "accepted_tokens", + "accepted_len": "accepted_len", + "done": "done", + "continue": "continue", + "rng": { + "seed": "seed", + "offset": "offset", + "next_offset": "next_offset", + }, + "effect": "verify", + }, + "verify", + ) + + +def build_speculative_state_rollback( + dtype: ir.DataType, + shape: list[int | str], + *, + sequence_axis: int, + effect: str = "rollback", +) -> PolicyComponent: + """Trim tentative recurrent state to ``past_length + accepted_length``.""" + if not 0 <= sequence_axis < len(shape): + raise ValueError("sequence_axis must index the state shape") + graph, builder = _make_graph("speculative_state_rollback") + op = builder.op + past = builder.input("past_state", dtype, shape) + tentative_shape = list(shape) + tentative_shape[sequence_axis] = "tentative_sequence" + tentative = builder.input("tentative_state", dtype, tentative_shape) + accepted_len = builder.input("accepted_len", ir.DataType.INT64, [1]) + past_len = op.Shape(past, start=sequence_axis, end=sequence_axis + 1) + end = op.Add(past_len, accepted_len) + corrected = op.Slice( + tentative, + op.Constant(value_ints=[0]), + end, + op.Constant(value_ints=[sequence_axis]), + op.Constant(value_ints=[1]), + ) + corrected_shape = list(shape) + corrected_shape[sequence_axis] = "accepted_sequence" + corrected.shape = ir.Shape(corrected_shape) + builder.add_output(corrected, "corrected_state") + return _component("mobius.policy.auxiliary@1", graph, {}, effect) + + +def build_effectful_identity( + name: str, + dtype: ir.DataType, + shape: list[int | str], + *, + effect: str, +) -> PolicyComponent: + """Publish a typed branch-local state value with a linear effect.""" + graph, builder = _make_graph(name) + value = builder.input("value", dtype, shape) + builder.add_output(builder.op.Identity(value), "next_value") + return _component("mobius.policy.auxiliary@1", graph, {}, effect) + + +def build_token_block_identity() -> PolicyComponent: + """Publish a branch-local speculative token block with a linear effect.""" + graph, builder = _make_graph("token_block_identity") + tokens = builder.input("tokens", ir.DataType.INT64, ["batch", "draft_sequence"]) + builder.add_output(builder.op.Identity(tokens), "next_tokens") + return _component("mobius.policy.auxiliary@1", graph, {}, "state") + + +def build_token_state_update(*, row_selective: bool = False) -> PolicyComponent: + """Selectively update one-token state while preserving suppressed rows.""" + graph, builder = _make_graph("token_state_update") + op = builder.op + current = builder.input("current", ir.DataType.INT64, ["batch", 1]) + update = builder.input( + "update", + ir.DataType.INT64, + ["batch", 1] if row_selective else ["batch"], + ) + if row_selective: + active = builder.input("active", ir.DataType.BOOL, ["batch"]) + done = builder.input("done", ir.DataType.BOOL, ["batch"]) + enabled = op.And(active, op.Not(done)) + next_state = op.Where( + op.Unsqueeze(enabled, [-1]), + update, + current, + ) + else: + next_state = op.Unsqueeze(update, [-1]) + if row_selective: + _set_public_shape(current, ["batch", 1]) + _set_public_shape(update, ["batch", 1]) + _set_public_shape(active, ["batch"]) + _set_public_shape(done, ["batch"]) + _set_public_shape(next_state, ["batch", 1]) + builder.add_output(next_state, "next") + return _component( + "onnx-genai.state-update@2" if row_selective else "onnx-genai.state-update@1", + graph, + { + "role": "state_update", + "current": "current", + "update": "update", + **({"batching": "per_row", "inactive_rows": "preserve"} if row_selective else {}), + **({"active": "active", "done": "done"} if row_selective else {}), + "next": "next", + "effect": "state", + }, + "state", + ) + + +def build_token_to_slot() -> PolicyComponent: + """Convert a canonical sampled token vector to a one-token ID tensor.""" + graph, builder = _make_graph("token_to_slot") + token = builder.input("token", ir.DataType.INT64, ["batch"]) + slot = builder.op.Unsqueeze(token, [-1]) + slot.shape = ir.Shape(["batch", 1]) + builder.add_output(slot, "slot") + return _component("mobius.policy.auxiliary@1", graph, {}) + + +def build_video_latent_initializer( + dtype: ir.DataType, + init_noise_sigma: float, + history_dtype: ir.DataType = ir.DataType.INT64, +) -> PolicyComponent: + """Scale request noise into the scheduler's starting video latent. + + The noise carries the temporal axis, so nothing here collapses a clip to a + single frame: the component is rank-agnostic and only applies the + scheduler's ``init_noise_sigma``. + """ + graph, builder = _make_graph("video_latent_initializer") + op = builder.op + noise = builder.input("noise", dtype, ["batch", "frames", "channels", "height", "width"]) + latent = op.Mul(noise, op.CastLike(op.Constant(value_float=init_noise_sigma), noise)) + _set_public_shape(latent, ["batch", "frames", "channels", "height", "width"]) + builder.add_output(latent, "latent") + # An empty scheduler history: the denoise loop appends one timestep per step. + history = op.ConstantOfShape( + op.Concat(op.Shape(noise, start=0, end=1), op.Constant(value_ints=[0]), axis=0), + value=ir.tensor([0], dtype=history_dtype), + ) + _set_public_shape(history, ["batch", 0]) + builder.add_output(history, "history") + return _component("mobius.policy.auxiliary@1", graph, {}) + + +def build_schedule_history_append(dtype: ir.DataType) -> PolicyComponent: + """Append the current timestep to the scheduler's history. + + Multistep video schedulers consume the trajectory of previous timesteps, so + the history is real state rather than telemetry. + """ + graph, builder = _make_graph("schedule_history_append") + op = builder.op + history = builder.input("history", dtype, ["batch", "history"]) + timestep = builder.input("timestep", dtype, ["batch"]) + updated = op.Concat(history, op.Unsqueeze(timestep, op.Constant(value_ints=[1])), axis=1) + _set_public_shape(updated, ["batch", "history"]) + builder.add_output(updated, "next") + return _component("mobius.policy.auxiliary@1", graph, {}) + + +def build_video_decode_chunk_count(latent_frame_axis: int = 2) -> PolicyComponent: + """Number of causal decode chunks a latent clip is split into. + + Mirrors ``AutoencoderKLCogVideoX._decode``: ``max(latent_frames // 2, 1)``. + """ + graph, builder = _make_graph("video_decode_chunk_count") + op = builder.op + latent = builder.input( + "latent", + ir.DataType.FLOAT, + ["batch", "channels", "latent_frames", "height", "width"], + ) + frames = op.Shape(latent, start=latent_frame_axis, end=latent_frame_axis + 1) + count = op.Max( + op.Div(frames, op.Constant(value_ints=[2])), + op.Constant(value_ints=[1]), + ) + _set_public_shape(count, [1]) + builder.add_output(count, "count") + return _component("mobius.policy.auxiliary@1", graph, {}) + + +def build_video_decode_chunk(latent_frame_axis: int = 2) -> PolicyComponent: + """Slice the latent frames belonging to one causal decode chunk. + + Reproduces the reference chunk walk, where the odd frame left over by the + two-frame stride is folded into the first chunk:: + + remaining = latent_frames % 2 + start = 2 * step + (0 if step == 0 else remaining) + end = 2 * (step + 1) + remaining + """ + graph, builder = _make_graph("video_decode_chunk") + op = builder.op + latent = builder.input( + "latent", + ir.DataType.FLOAT, + ["batch", "channels", "latent_frames", "height", "width"], + ) + step = builder.input("step", ir.DataType.INT64, ["batch"]) + + two = op.Constant(value_ints=[2]) + zero = op.Constant(value_ints=[0]) + frames = op.Shape(latent, start=latent_frame_axis, end=latent_frame_axis + 1) + remaining = op.Mod(frames, two) + # The loop induction value is batch-broadcast; every row walks the same clip. + index = op.Slice(step, zero, op.Constant(value_ints=[1]), zero) + offset = op.Where(op.Equal(index, zero), zero, remaining) + start = op.Add(op.Mul(index, two), offset) + end = op.Add(op.Mul(op.Add(index, op.Constant(value_ints=[1])), two), remaining) + chunk = op.Slice(latent, start, end, op.Constant(value_ints=[latent_frame_axis])) + _set_public_shape(chunk, ["batch", "channels", "chunk_frames", "height", "width"]) + builder.add_output(chunk, "chunk") + return _component("mobius.policy.auxiliary@1", graph, {}) + + +def build_video_conv_cache_initializer( + entries: list[tuple[str, int, int]], + dtype: ir.DataType = ir.DataType.FLOAT, +) -> PolicyComponent: + """Zero-length causal convolution caches sized from the latent grid. + + ``entries`` are ``(port, channels, spatial_scale)``. A zero-length temporal + axis is the encoding of "no previous chunk", which makes the first decode + chunk replicate its own first frame exactly as the reference does. + """ + graph, builder = _make_graph("video_conv_cache_initializer") + op = builder.op + latent = builder.input( + "latent", + dtype, + ["batch", "channels", "latent_frames", "height", "width"], + ) + batch = op.Shape(latent, start=0, end=1) + height = op.Shape(latent, start=3, end=4) + width = op.Shape(latent, start=4, end=5) + zero = ir.tensor([0.0], dtype=dtype) + for port, channels, scale in entries: + scaled_height = ( + height if scale == 1 else op.Mul(height, op.Constant(value_ints=[scale])) + ) + scaled_width = width if scale == 1 else op.Mul(width, op.Constant(value_ints=[scale])) + shape = op.Concat( + batch, + op.Constant(value_ints=[channels, 0]), + scaled_height, + scaled_width, + axis=0, + ) + cache = op.ConstantOfShape(shape, value=zero) + _set_public_shape( + cache, + [ + "batch", + channels, + 0, + "height" if scale == 1 else f"{scale}*height", + "width" if scale == 1 else f"{scale}*width", + ], + ) + builder.add_output(cache, port) + return _component("mobius.policy.auxiliary@1", graph, {}) + + +def build_video_latent_permute(perm: list[int]) -> PolicyComponent: + """Reorder a video latent between the denoiser and VAE layouts. + + CogVideoX denoises ``[batch, frames, channels, height, width]`` but decodes + ``[batch, channels, frames, height, width]``; the transposition is part of + the pipeline contract, not an implementation detail of either model. + """ + graph, builder = _make_graph("video_latent_permute") + op = builder.op + source = builder.input( + "latent", ir.DataType.FLOAT, ["batch", "frames", "channels", "height", "width"] + ) + permuted = op.Transpose(source, perm=perm) + _set_public_shape(permuted, ["batch", "channels", "frames", "height", "width"]) + builder.add_output(permuted, "permuted") + return _component("mobius.policy.auxiliary@1", graph, {}) + + +def build_video_latent_unscale(scaling_factor: float) -> PolicyComponent: + """Undo the autoencoder's latent scaling before decoding.""" + graph, builder = _make_graph("video_latent_unscale") + op = builder.op + latent = builder.input( + "latent", ir.DataType.FLOAT, ["batch", "channels", "frames", "height", "width"] + ) + unscaled = op.Div(latent, op.CastLike(op.Constant(value_float=scaling_factor), latent)) + _set_public_shape(unscaled, ["batch", "channels", "frames", "height", "width"]) + builder.add_output(unscaled, "unscaled") + return _component("mobius.policy.auxiliary@1", graph, {}) + + +_INT64_MAX = 2**63 - 1 + + +def build_scalar_integer_add() -> PolicyComponent: + """Add two rank-0 integer control values. + + Loop induction values and substep indices are scalars, not per-row state, so + they need a rank-0 adder rather than the ``[batch]`` :func:`build_integer_add`. + """ + graph, builder = _make_graph("scalar_integer_add") + left = builder.input("left", ir.DataType.INT64, []) + right = builder.input("right", ir.DataType.INT64, []) + total = builder.op.Add(left, right) + total.shape = ir.Shape([]) + builder.add_output(total, "total") + return _component("mobius.policy.auxiliary@1", graph, {}) + + +def _duplex_positions( + op: Any, + offset: Any, + delays: Any, + cache_length: int, + channels: int, + batch_shape: Any, +) -> Any: + """Build ``[batch, channels, 1]`` ring indices for ``(offset + delays) % CT``.""" + positions = op.Mod(op.Add(delays, offset), op.Constant(value_int=cache_length)) + positions = op.Reshape(positions, op.Constant(value_ints=[1, channels, 1])) + target_shape = op.Concat(batch_shape, op.Constant(value_ints=[channels, 1]), axis=0) + return op.Expand(positions, target_shape) + + +def build_duplex_frame_assemble( + *, channels: int = 17, cache_length: int = 4 +) -> PolicyComponent: + """Write per-stream tokens into the delay ring cache and read one model frame. + + Full-duplex codec language models interleave several token streams (one text + stream plus interleaved agent and user acoustic streams) that are each shifted + by a small per-stream delay. The delay compensation is a ring buffer of + ``cache_length`` frames indexed by ``(offset + delay) % cache_length``. + + ``stream_tokens`` carries every externally supplied token for this step; a + negative entry means "this stream has nothing to contribute, the model must + predict it". ``initial_tokens`` primes streams whose delay has not elapsed. + """ + graph, builder = _make_graph("duplex_frame_assemble") + op = builder.op + cache = builder.input("token_cache", ir.DataType.INT64, ["batch", channels, cache_length]) + provided = builder.input( + "token_provided", ir.DataType.BOOL, ["batch", channels, cache_length] + ) + offset = builder.input("offset", ir.DataType.INT64, []) + stream_tokens = builder.input("stream_tokens", ir.DataType.INT64, ["batch", channels]) + delays = builder.input("delays", ir.DataType.INT64, [channels]) + initial_tokens = builder.input("initial_tokens", ir.DataType.INT64, [channels]) + + batch_shape = op.Shape(cache, start=0, end=1) + + # 1. scatter externally supplied tokens at their delayed ring slot. + write_index = _duplex_positions(op, offset, delays, cache_length, channels, batch_shape) + token_update = op.Unsqueeze(stream_tokens, [-1]) + has_token = op.GreaterOrEqual(token_update, op.Constant(value_int=0)) + cache = op.ScatterElements( + cache, + write_index, + op.Where(has_token, token_update, op.GatherElements(cache, write_index, axis=2)), + axis=2, + ) + # ORT has no bool Where kernel; setting a flag is a saturating Or. + provided = op.ScatterElements( + provided, + write_index, + op.Or(op.GatherElements(provided, write_index, axis=2), has_token), + axis=2, + ) + + # 2. prime streams whose delay has not yet elapsed (offset <= delay). + ring = op.Mod(offset, op.Constant(value_int=cache_length)) + target_index = op.Expand( + op.Reshape(ring, op.Constant(value_ints=[1, 1, 1])), + op.Concat(batch_shape, op.Constant(value_ints=[channels, 1]), axis=0), + ) + primed = op.Reshape( + op.LessOrEqual(op.Expand(offset, op.Constant(value_ints=[channels])), delays), + op.Constant(value_ints=[1, channels, 1]), + ) + primed = op.Expand(primed, op.Shape(target_index)) + initial_update = op.Expand( + op.Reshape(initial_tokens, op.Constant(value_ints=[1, channels, 1])), + op.Shape(target_index), + ) + cache = op.ScatterElements( + cache, + target_index, + op.Where(primed, initial_update, op.GatherElements(cache, target_index, axis=2)), + axis=2, + ) + provided = op.ScatterElements( + provided, + target_index, + op.Or(op.GatherElements(provided, target_index, axis=2), primed), + axis=2, + ) + + # 3. read the model input frame (offset - 1) and the teacher-forcing target. + input_index = op.Expand( + op.Reshape( + op.Mod( + op.Add( + op.Sub(offset, op.Constant(value_int=1)), + op.Constant(value_int=cache_length), + ), + op.Constant(value_int=cache_length), + ), + op.Constant(value_ints=[1, 1, 1]), + ), + op.Shape(target_index), + ) + input_frame = op.GatherElements(cache, input_index, axis=2) + target = op.GatherElements(cache, target_index, axis=2) + target_provided = op.GatherElements(provided, target_index, axis=2) + + cache.shape = ir.Shape(["batch", channels, cache_length]) + provided.shape = ir.Shape(["batch", channels, cache_length]) + for value in (input_frame, target, target_provided): + value.shape = ir.Shape(["batch", channels, 1]) + builder.add_output(cache, "next_token_cache") + builder.add_output(provided, "next_token_provided") + builder.add_output(input_frame, "input_frame") + builder.add_output(target, "target") + builder.add_output(target_provided, "target_provided") + return _component("mobius.policy.auxiliary@1", graph, {}) + + +def build_duplex_frame_commit( + *, channels: int = 17, cache_length: int = 4, max_delay: int = 1 +) -> PolicyComponent: + """Commit predicted tokens and read the delay-compensated output frame. + + Predictions only fill ring slots that were not externally supplied, so a + teacher-forced stream keeps its supplied value. The emitted frame undoes the + per-stream delay by reading each stream at ``offset - max_delay + delay``, + which is only well defined once ``offset`` has passed ``max_delay``. + """ + graph, builder = _make_graph("duplex_frame_commit") + op = builder.op + cache = builder.input("token_cache", ir.DataType.INT64, ["batch", channels, cache_length]) + provided = builder.input( + "token_provided", ir.DataType.BOOL, ["batch", channels, cache_length] + ) + offset = builder.input("offset", ir.DataType.INT64, []) + frame = builder.input("frame", ir.DataType.INT64, ["batch", channels]) + delays = builder.input("delays", ir.DataType.INT64, [channels]) + + batch_shape = op.Shape(cache, start=0, end=1) + cell_shape = op.Concat(batch_shape, op.Constant(value_ints=[channels, 1]), axis=0) + false_like = op.ConstantOfShape(cell_shape, value=ir.tensor([False])) + + def ring_index(value: Any) -> Any: + return op.Expand( + op.Reshape( + op.Mod( + op.Add(value, op.Constant(value_int=cache_length)), + op.Constant(value_int=cache_length), + ), + op.Constant(value_ints=[1, 1, 1]), + ), + cell_shape, + ) + + # 1. retire the slot that was just consumed as model input. + input_index = ring_index(op.Sub(offset, op.Constant(value_int=1))) + provided = op.ScatterElements(provided, input_index, false_like, axis=2) + + # 2. fill only the slots the caller did not supply. + target_index = ring_index(offset) + target_provided = op.GatherElements(provided, target_index, axis=2) + existing = op.GatherElements(cache, target_index, axis=2) + cache = op.ScatterElements( + cache, + target_index, + op.Where(target_provided, existing, op.Unsqueeze(frame, [-1])), + axis=2, + ) + + # 3. undo the per-stream delay: stream k is read at offset - max_delay + delay[k]. + read_index = _duplex_positions( + op, + op.Sub(offset, op.Constant(value_int=max_delay)), + delays, + cache_length, + channels, + batch_shape, + ) + out_frame = op.Squeeze(op.GatherElements(cache, read_index, axis=2), [-1]) + next_offset = op.Add(offset, op.Constant(value_int=1)) + emit = op.Greater(offset, op.Constant(value_int=max_delay)) + + cache.shape = ir.Shape(["batch", channels, cache_length]) + provided.shape = ir.Shape(["batch", channels, cache_length]) + out_frame.shape = ir.Shape(["batch", channels]) + next_offset.shape = ir.Shape([]) + emit.shape = ir.Shape([]) + builder.add_output(cache, "next_token_cache") + builder.add_output(provided, "next_token_provided") + builder.add_output(out_frame, "out_frame") + builder.add_output(next_offset, "next_offset") + builder.add_output(emit, "emit") + return _component("mobius.policy.auxiliary@1", graph, {}) + + +def build_duplex_teacher_select(*, channels: int = 17) -> PolicyComponent: + """Choose the supplied token over the sampled token for one stream index.""" + graph, builder = _make_graph("duplex_teacher_select") + op = builder.op + target = builder.input("target", ir.DataType.INT64, ["batch", channels, 1]) + target_provided = builder.input( + "target_provided", ir.DataType.BOOL, ["batch", channels, 1] + ) + sampled = builder.input("sampled", ir.DataType.INT64, ["batch"]) + index = builder.input("index", ir.DataType.INT64, []) + stream = op.Reshape(index, op.Constant(value_ints=[1])) + picked = op.Squeeze(op.Gather(target, stream, axis=1), [1, 2]) + flag = op.Squeeze(op.Gather(target_provided, stream, axis=1), [1, 2]) + token = op.Where(flag, picked, sampled) + token.shape = ir.Shape(["batch"]) + builder.add_output(token, "token") + return _component("mobius.policy.auxiliary@1", graph, {}) + + +def build_duplex_stream_append( + *, streams: int = 8, dtype: ir.DataType = ir.DataType.INT64 +) -> PolicyComponent: + """Append one frame to a growing ``[batch, streams, length]`` prefix.""" + graph, builder = _make_graph("duplex_stream_append") + op = builder.op + prefix = builder.input("prefix", dtype, ["batch", streams, "length"]) + frame = builder.input("frame", dtype, ["batch", streams]) + appended = op.Concat(prefix, op.Unsqueeze(frame, [-1]), axis=2) + appended.shape = ir.Shape(["batch", streams, "length + 1"]) + builder.add_output(appended, "next_prefix") + return _component("mobius.policy.auxiliary@1", graph, {}) + + +def build_duplex_waveform_append(*, dtype: ir.DataType = ir.DataType.FLOAT) -> PolicyComponent: + """Append a packed audio chunk to a growing ``[batch, 1, samples]`` prefix.""" + graph, builder = _make_graph("duplex_waveform_append") + op = builder.op + prefix = builder.input("prefix", dtype, ["batch", 1, "samples"]) + chunk = builder.input("chunk", dtype, ["batch", 1, "chunk"]) + appended = op.Concat(prefix, chunk, axis=2) + appended.shape = ir.Shape(["batch", 1, "samples + chunk"]) + builder.add_output(appended, "next_prefix") + return _component("mobius.policy.auxiliary@1", graph, {}) + + +def build_duplex_stream_tail( + *, streams: int = 8, dtype: ir.DataType = ir.DataType.INT64, rank: int = 3 +) -> PolicyComponent: + """Read the trailing ``count`` positions of a growing prefix. + + Stateless codec graphs are replayed over an accumulated prefix, so only the + newest ``count`` positions belong to the current event. + """ + graph, builder = _make_graph("duplex_stream_tail") + op = builder.op + shape = ["batch", streams, "length"] if rank == 3 else ["batch", "length"] + prefix = builder.input("prefix", dtype, shape) + count = builder.input("count", ir.DataType.INT64, []) + axis = rank - 1 + length = op.Squeeze(op.Shape(prefix, start=axis, end=axis + 1), [0]) + start = op.Reshape(op.Sub(length, count), op.Constant(value_ints=[1])) + tail = op.Slice( + prefix, + start, + op.Constant(value_ints=[_INT64_MAX]), + op.Constant(value_ints=[axis]), + ) + tail.shape = ir.Shape([*shape[:-1], "count"]) + builder.add_output(tail, "tail") + return _component("mobius.policy.auxiliary@1", graph, {}) + + +def build_duplex_user_stream_merge(*, channels: int = 17, streams: int = 8) -> PolicyComponent: + """Overlay freshly encoded user codes onto the supplied stream-token frame. + + The trailing ``streams`` channels of a full-duplex frame carry the incoming + user audio, so the codec output always wins over the request-supplied frame. + """ + graph, builder = _make_graph("duplex_user_stream_merge") + op = builder.op + frame = builder.input("frame_codes", ir.DataType.INT64, ["batch", channels]) + codes = builder.input("codes", ir.DataType.INT64, ["batch", streams, 1]) + head = op.Slice( + frame, + op.Constant(value_ints=[0]), + op.Constant(value_ints=[channels - streams]), + op.Constant(value_ints=[1]), + ) + merged = op.Concat(head, op.Squeeze(codes, [-1]), axis=1) + merged.shape = ir.Shape(["batch", channels]) + builder.add_output(merged, "stream_tokens") + return _component("mobius.policy.auxiliary@1", graph, {}) + + +def build_duplex_cell_to_frame(*, channels: int = 17) -> PolicyComponent: + """Drop the single-position axis of one ring-buffer cell.""" + graph, builder = _make_graph("duplex_cell_to_frame") + target = builder.input("target", ir.DataType.INT64, ["batch", channels, 1]) + frame = builder.op.Squeeze(target, [-1]) + frame.shape = ir.Shape(["batch", channels]) + builder.add_output(frame, "frame") + return _component("mobius.policy.auxiliary@1", graph, {}) + + +def build_duplex_agent_frame_select( + *, channels: int = 17, streams: int = 8 +) -> PolicyComponent: + """Read the agent acoustic streams out of a delay-compensated frame.""" + graph, builder = _make_graph("duplex_agent_frame_select") + op = builder.op + frame = builder.input("frame", ir.DataType.INT64, ["batch", channels]) + codes = op.Slice( + frame, + op.Constant(value_ints=[1]), + op.Constant(value_ints=[1 + streams]), + op.Constant(value_ints=[1]), + ) + codes.shape = ir.Shape(["batch", streams]) + builder.add_output(codes, "codes") + return _component("mobius.policy.auxiliary@1", graph, {}) diff --git a/src/mobius/generation/_policy_components_test.py b/src/mobius/generation/_policy_components_test.py new file mode 100644 index 000000000..feb29153b --- /dev/null +++ b/src/mobius/generation/_policy_components_test.py @@ -0,0 +1,1474 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT License. + +from __future__ import annotations + +import numpy as np +import onnx_ir as ir +import onnxruntime as ort +import pytest + +from mobius._model_package import ModelPackage +from mobius.generation import ( + PolicyCapabilities, + attach_policy_components, + build_adaptive_k_policy, + build_batch_minimum, + build_boolean_not, + build_code_frame_update, + build_counter_rng_normal, + build_ddim_solver_step, + build_decoder_state_initializer, + build_decoder_step_update, + build_empty_features, + build_eos_termination, + build_euler_model_input, + build_euler_solver_step, + build_flow_match_solver_step, + build_grammar_logits_processor, + build_greedy_sampler, + build_guidance_combine, + build_integer_minimum, + build_integer_row_broadcast, + build_last_token_logits, + build_masked_token_update, + build_model_token_cast, + build_multistep_solver_step, + build_pack_latents_2x2, + build_proposal_metrics, + build_scalar_constant, + build_schedule_history_append, + build_seeded_categorical_sampler, + build_sequence_concat, + build_shape_constant, + build_speculative_acceptance, + build_speculative_state_rollback, + build_tensor_scale, + build_termination_batch_initializer, + build_token_state_update, + build_true_cfg, + build_unpack_latents_2x2, + build_video_conv_cache_initializer, + build_video_decode_chunk, + build_video_decode_chunk_count, + build_video_latent_initializer, + build_zeros_like, + rotary_axis_count, +) +from mobius.generation._policy_components import _make_graph + + +def _run(component, tmp_path, feeds): + path = tmp_path / f"{component.model.graph.name}.onnx" + ir.save(component.model, path) + session = ort.InferenceSession(str(path), providers=["CPUExecutionProvider"]) + return session.run(None, feeds) + + +def _run_model(model, tmp_path, feeds): + path = tmp_path / f"{model.graph.name}.onnx" + ir.save(model, path) + session = ort.InferenceSession(str(path), providers=["CPUExecutionProvider"]) + return session.run(None, feeds) + + +def test_grammar_logits_processor_applies_mask_and_forced_tokens(tmp_path): + logits = np.array([[1.0, 4.0, 3.0], [5.0, 2.0, 1.0]], np.float32) + (tokens,) = _run( + build_grammar_logits_processor(), + tmp_path, + { + "logits": logits, + "logits_mask": np.array([[True, False, True], [True, True, True]], np.bool_), + "forced_tokens": np.array([[0], [2]], np.int64), + "forced_length": np.array([0, 1], np.int64), + }, + ) + np.testing.assert_array_equal(tokens, [[2], [2]]) + + +def test_adaptive_k_policy_probes_and_keeps_faster_adjacent_width(tmp_path): + max_k = 4 + k_slots = max_k + 1 + current_k = np.array([2], np.int64) + estimates = np.zeros((1, 4 * k_slots + 4), np.float32) + + def observe(*, evaluated, accepted, committed, draft_ms, target_ms): + nonlocal current_k, estimates + current_k, estimates = _run( + build_adaptive_k_policy(max_k=max_k), + tmp_path, + { + "current_k": current_k, + "accepted": np.array([accepted], np.int64), + "evaluated": np.array([evaluated], np.int64), + "committed_tokens": np.array([committed], np.int64), + "filled_proposal_budget": np.array([True], np.bool_), + "draft_ms": np.array([draft_ms], np.float32), + "target_ms": np.array([target_ms], np.float32), + "estimates": estimates, + }, + ) + + observe(evaluated=2, accepted=2, committed=3, draft_ms=1.0, target_ms=2.0) + observe(evaluated=2, accepted=2, committed=3, draft_ms=1.0, target_ms=2.0) + np.testing.assert_array_equal(current_k, [3]) + np.testing.assert_array_equal(estimates[:, 4 * k_slots], [2.0]) + + observe(evaluated=3, accepted=3, committed=4, draft_ms=1.0, target_ms=2.0) + observe(evaluated=3, accepted=3, committed=4, draft_ms=1.0, target_ms=2.0) + np.testing.assert_array_equal(current_k, [3]) + np.testing.assert_array_equal(estimates[:, 4 * k_slots], [0.0]) + np.testing.assert_array_equal(estimates[:, 4 * k_slots + 3], [1.0]) + assert estimates[0, 3] / estimates[0, k_slots + 3] > 1.0 + + +def test_speculative_guidance_length_and_budget_math(tmp_path): + (minimum,) = _run( + build_integer_minimum(), + tmp_path, + { + "left": np.array([3, 1], np.int64), + "right": np.array([2, 4], np.int64), + }, + ) + evaluated, filled = _run( + build_proposal_metrics(), + tmp_path, + { + "proposed_tokens": np.zeros((2, 3), np.int64), + "requested_k": np.array([3, 2], np.int64), + }, + ) + np.testing.assert_array_equal(minimum, [2, 1]) + np.testing.assert_array_equal(evaluated, [3, 3]) + np.testing.assert_array_equal(filled, [True, False]) + (synchronized,) = _run( + build_batch_minimum(), + tmp_path, + {"values": np.array([2, 1, 3], np.int64)}, + ) + np.testing.assert_array_equal(synchronized, [1]) + + +def test_adaptive_k_ignores_invalid_telemetry_per_batch(tmp_path): + estimates = np.arange(48, dtype=np.float32).reshape(2, 24) + next_k, next_estimates = _run( + build_adaptive_k_policy(max_k=4), + tmp_path, + { + "current_k": np.array([2, 3], np.int64), + "accepted": np.array([2, 2], np.int64), + "evaluated": np.array([0, 3], np.int64), + "committed_tokens": np.array([3, 0], np.int64), + "filled_proposal_budget": np.array([True, True], np.bool_), + "draft_ms": np.array([1.0, np.nan], np.float32), + "target_ms": np.array([2.0, 2.0], np.float32), + "estimates": estimates, + }, + ) + np.testing.assert_array_equal(next_k, [2, 3]) + np.testing.assert_array_equal(next_estimates, estimates) + + +def test_greedy_sampler_runtime(tmp_path): + (tokens,) = _run( + build_greedy_sampler(), + tmp_path, + {"logits": np.array([[0.2, 0.7, 0.1], [2.0, 1.0, 3.0]], np.float32)}, + ) + np.testing.assert_array_equal(tokens, [1, 2]) + + +def test_row_selective_greedy_sampler_suppresses_inactive_rows(tmp_path): + (tokens,) = _run( + build_greedy_sampler(row_selective=True), + tmp_path, + { + "logits": np.array([[0.0, 2.0], [3.0, 1.0], [1.0, 4.0]], np.float32), + "active": np.array([True, False, True], np.bool_), + "done": np.array([False, False, True], np.bool_), + }, + ) + np.testing.assert_array_equal(tokens, [1, -1, -1]) + + +def test_last_token_logits_and_continue_predicate_runtime(tmp_path): + logits = np.arange(24, dtype=np.float32).reshape(2, 3, 4) + (last,) = _run(build_last_token_logits(), tmp_path, {"logits": logits}) + np.testing.assert_array_equal(last, logits[:, -1, :]) + + half_logits = logits.astype(np.float16) + (last_float,) = _run( + build_last_token_logits(ir.DataType.FLOAT16), + tmp_path, + {"logits": half_logits}, + ) + assert last_float.dtype == np.float32 + np.testing.assert_array_equal(last_float, logits[:, -1, :]) + + (continued,) = _run( + build_boolean_not(), + tmp_path, + {"done": np.array([True, False])}, + ) + np.testing.assert_array_equal(continued, [False]) + + +def test_decoder_state_initializer_and_step_update_runtime(tmp_path): + inputs = [ + ir.Value( + name="input_ids", + type=ir.TensorType(ir.DataType.INT64), + shape=ir.Shape(["batch", "sequence"]), + ), + ir.Value( + name="attention_mask", + type=ir.TensorType(ir.DataType.INT64), + shape=ir.Shape(["batch", "past_sequence + sequence"]), + ), + ir.Value( + name="position_ids", + type=ir.TensorType(ir.DataType.INT64), + shape=ir.Shape(["batch", "sequence"]), + ), + ir.Value( + name="past_key_values.0.key", + type=ir.TensorType(ir.DataType.FLOAT), + shape=ir.Shape(["batch", 2, "past_sequence", 4]), + ), + ] + decoder = ir.Model(ir.Graph(inputs, [], nodes=[], name="decoder"), ir_version=11) + initializer = build_decoder_state_initializer( + decoder, + token_input="input_ids", + attention_mask_input="attention_mask", + position_ids_input="position_ids", + cache_inputs=["past_key_values.0.key"], + ) + outputs = _run( + initializer, + tmp_path, + {"prompt_tokens": np.array([[3, 4, 5]], np.int64)}, + ) + np.testing.assert_array_equal(outputs[0], [[1, 1, 1]]) + np.testing.assert_array_equal(outputs[1], [[0, 1, 2]]) + np.testing.assert_array_equal(outputs[2], [[1, 1, 1, 1]]) + np.testing.assert_array_equal(outputs[3], [[3]]) + assert outputs[5].shape == (1, 2, 0, 4) + + updated = _run( + build_decoder_step_update( + attention_dtype=ir.DataType.INT64, + position_dtype=ir.DataType.INT64, + ), + tmp_path, + { + "attention_mask": outputs[2], + "position_ids": outputs[3], + }, + ) + np.testing.assert_array_equal(updated[0], [[1, 1, 1, 1, 1]]) + np.testing.assert_array_equal(updated[1], [[4]]) + (cast_token,) = _run( + build_model_token_cast(ir.DataType.INT32), + tmp_path, + {"token": np.array([[8]], np.int64)}, + ) + assert cast_token.dtype == np.int32 + np.testing.assert_array_equal(cast_token, [[8]]) + + +def test_decoder_fixed_capacity_state_matches_native_capture_layout(tmp_path): + inputs = [ + ir.Value( + name="input_ids", + type=ir.TensorType(ir.DataType.INT64), + shape=ir.Shape(["batch", "sequence"]), + ), + ir.Value( + name="attention_mask", + type=ir.TensorType(ir.DataType.INT64), + shape=ir.Shape(["batch", "past_sequence + sequence"]), + ), + ir.Value( + name="past_key_values.0.key", + type=ir.TensorType(ir.DataType.FLOAT16), + shape=ir.Shape(["batch", 2, "past_sequence", 4]), + ), + ] + decoder = ir.Model(ir.Graph(inputs, [], nodes=[], name="decoder"), ir_version=11) + prompt = np.arange(68, dtype=np.int64).reshape(1, 68) + outputs = _run( + build_decoder_state_initializer( + decoder, + token_input="input_ids", + attention_mask_input="attention_mask", + position_ids_input=None, + cache_inputs=["past_key_values.0.key"], + fixed_capacity=True, + ), + tmp_path, + { + "prompt_tokens": prompt, + "max_iterations": np.array([128], np.int64), + }, + ) + attention, body_attention, token, cache_lengths, cache = outputs + expected_body_attention = np.zeros((1, 196), np.int64) + expected_body_attention[:, :68] = 1 + np.testing.assert_array_equal(attention, expected_body_attention) + np.testing.assert_array_equal(body_attention, expected_body_attention) + np.testing.assert_array_equal(token, [[0]]) + np.testing.assert_array_equal(cache_lengths, [68]) + assert cache.shape == (1, 2, 196, 4) + + (next_attention,) = _run( + build_decoder_step_update( + attention_dtype=ir.DataType.INT64, + position_dtype=None, + fixed_capacity=True, + ), + tmp_path, + { + "attention_mask": body_attention, + "logical_length": np.array([68], np.int64), + }, + ) + expected_body_attention[:, :69] = 1 + np.testing.assert_array_equal(next_attention, expected_body_attention) + + +def test_decoder_ragged_batch_uses_stable_capacity_and_independent_rows(tmp_path): + inputs = [ + ir.Value( + name="input_ids", + type=ir.TensorType(ir.DataType.INT64), + shape=ir.Shape(["batch", "sequence"]), + ), + ir.Value( + name="attention_mask", + type=ir.TensorType(ir.DataType.INT64), + shape=ir.Shape(["batch", "past_sequence + sequence"]), + ), + ir.Value( + name="past_key_values.0.key", + type=ir.TensorType(ir.DataType.FLOAT16), + shape=ir.Shape(["batch", 2, "past_sequence", 4]), + ), + ] + decoder = ir.Model(ir.Graph(inputs, [], nodes=[], name="decoder"), ir_version=11) + component = build_decoder_state_initializer( + decoder, + token_input="input_ids", + attention_mask_input="attention_mask", + position_ids_input=None, + cache_inputs=["past_key_values.0.key"], + fixed_capacity=True, + ragged=True, + ) + feeds = { + "prompt_tokens": np.arange(12, dtype=np.int64).reshape(3, 4), + "prompt_lengths": np.array([4, 2, 1], np.int64), + "max_iterations": np.array([5], np.int64), + } + attention, body_attention, token, generated, cache_lengths, cache = _run( + component, tmp_path, feeds + ) + expected = np.array( + [ + [1, 1, 1, 1, 0, 0, 0, 0, 0], + [1, 1, 0, 0, 0, 0, 0, 0, 0], + [1, 0, 0, 0, 0, 0, 0, 0, 0], + ], + np.int64, + ) + np.testing.assert_array_equal(attention, expected) + np.testing.assert_array_equal(body_attention, expected) + np.testing.assert_array_equal(token, np.zeros((3, 1), np.int64)) + np.testing.assert_array_equal(generated, [0, 0, 0]) + np.testing.assert_array_equal(cache_lengths, [4, 2, 1]) + assert cache.shape == (3, 2, 9, 4) + + for row in range(3): + row_outputs = _run( + component, + tmp_path, + { + "prompt_tokens": feeds["prompt_tokens"][row : row + 1], + "prompt_lengths": feeds["prompt_lengths"][row : row + 1], + "max_iterations": feeds["max_iterations"], + }, + ) + for batched, independent in zip( + (attention, body_attention, token, generated, cache_lengths), + row_outputs[:-1], + strict=True, + ): + np.testing.assert_array_equal(batched[row : row + 1], independent) + assert row_outputs[-1].shape == (1, 2, 9, 4) + + +def test_empty_features_runtime(tmp_path): + (features,) = _run( + build_empty_features(ir.DataType.FLOAT16, 64), + tmp_path, + {}, + ) + assert features.dtype == np.float16 + assert features.shape == (0, 64) + + +def test_decoder_policy_chain_generates_multiple_tokens_from_prompt_only(tmp_path): + graph, builder = _make_graph("decoder_stub") + op = builder.op + input_ids = builder.input("input_ids", ir.DataType.INT64, ["batch", "sequence"]) + builder.input("attention_mask", ir.DataType.INT64, ["batch", "context"]) + builder.input("position_ids", ir.DataType.INT64, ["batch", "sequence"]) + past = builder.input( + "past_key_values.0.key", + ir.DataType.FLOAT, + ["batch", 2, "past_sequence", 4], + ) + logits_shape = op.Concat(op.Shape(input_ids), op.Constant(value_ints=[8]), axis=0) + logits = op.ConstantOfShape( + logits_shape, + value=ir.tensor([0.0], dtype=ir.DataType.FLOAT), + ) + logits.shape = ir.Shape(["batch", "sequence", 8]) + builder.add_output(logits, "logits") + builder.add_output(op.Identity(past), "present.0.key") + decoder = ir.Model(graph, ir_version=11) + + prompt = np.array([[4, 5]], np.int64) + max_output_tokens = np.array([3], np.int64) + initialized = _run( + build_decoder_state_initializer( + decoder, + token_input="input_ids", + attention_mask_input="attention_mask", + position_ids_input="position_ids", + cache_inputs=["past_key_values.0.key"], + ), + tmp_path, + {"prompt_tokens": prompt}, + ) + attention, positions, body_attention, body_position, token, cache = initialized + logits, cache = _run_model( + decoder, + tmp_path, + { + "input_ids": prompt, + "attention_mask": attention, + "position_ids": positions, + "past_key_values.0.key": cache, + }, + ) + + emitted = [] + for iteration in range(3): + (last,) = _run(build_last_token_logits(), tmp_path, {"logits": logits}) + (sample,) = _run(build_greedy_sampler(), tmp_path, {"logits": last}) + emitted.append(int(sample[0])) + done, _continue = _run( + build_eos_termination(), + tmp_path, + { + "tokens": sample, + "eos_ids": np.array([7], np.int64), + "iteration": np.array([iteration], np.int64), + "max_iterations": max_output_tokens, + }, + ) + (token,) = _run( + build_token_state_update(), + tmp_path, + {"current": token, "update": sample}, + ) + logits, cache = _run_model( + decoder, + tmp_path, + { + "input_ids": token, + "attention_mask": body_attention, + "position_ids": body_position, + "past_key_values.0.key": cache, + }, + ) + body_attention, body_position = _run( + build_decoder_step_update( + attention_dtype=ir.DataType.INT64, + position_dtype=ir.DataType.INT64, + ), + tmp_path, + { + "attention_mask": body_attention, + "position_ids": body_position, + }, + ) + if done[0]: + break + + assert emitted == [0, 0, 0] + assert done.tolist() == [True] + + +def test_seeded_sampler_is_counter_based_and_reproducible(tmp_path): + component = build_seeded_categorical_sampler() + feeds = { + "logits": np.array([[0.0, 0.0, 0.0, 0.0]], np.float32), + "temperature": np.array([1.0], np.float32), + "top_k": np.array([0], np.int64), + "top_p": np.array([1.0], np.float32), + "min_p": np.array([0.0], np.float32), + "seed": np.array([7], np.int64), + "counter": np.array([11], np.int64), + "active": np.array([True], np.bool_), + "done": np.array([False], np.bool_), + } + first = _run(component, tmp_path, feeds) + second = _run(component, tmp_path, feeds) + np.testing.assert_array_equal(first[0], second[0]) + np.testing.assert_array_equal(first[1], [12]) + + +def test_batched_policy_v2_ports_use_exact_public_shapes(): + components = { + "sampler": build_seeded_categorical_sampler(), + "termination": build_eos_termination(row_selective=True), + "state": build_token_state_update(row_selective=True), + } + expected = { + "sampler": { + "logits": ["batch", "vocabulary"], + "temperature": ["batch"], + "top_k": ["batch"], + "top_p": ["batch"], + "min_p": ["batch"], + "seed": ["batch"], + "counter": ["batch"], + "active": ["batch"], + "done": ["batch"], + "token": ["batch"], + "next_counter": ["batch"], + }, + "termination": { + "tokens": ["batch"], + "eos_ids": ["batch", "num_eos"], + "eos_lengths": ["batch"], + "iteration": ["1"], + "max_iterations": ["batch"], + "active": ["batch"], + "done": ["batch"], + "next_active": ["batch"], + "continue": ["1"], + }, + "state": { + "current": ["batch", "1"], + "update": ["batch", "1"], + "active": ["batch"], + "done": ["batch"], + "next": ["batch", "1"], + }, + } + for name, component in components.items(): + ports = [*component.model.graph.inputs, *component.model.graph.outputs] + assert {port.name: [str(dim) for dim in port.shape] for port in ports} == expected[ + name + ] + + +def test_seeded_sampler_applies_request_top_k(tmp_path): + (token, _) = _run( + build_seeded_categorical_sampler(), + tmp_path, + { + "logits": np.array([[10.0, 9.0, 8.0, 7.0]], np.float32), + "temperature": np.array([0.7], np.float32), + "top_k": np.array([1], np.int64), + "top_p": np.array([0.5], np.float32), + "min_p": np.array([0.0], np.float32), + "seed": np.array([17], np.int64), + "counter": np.array([0], np.int64), + "active": np.array([True], np.bool_), + "done": np.array([False], np.bool_), + }, + ) + np.testing.assert_array_equal(token, [0]) + + +def test_seeded_sampler_applies_request_min_p_in_logit_space(tmp_path): + (token, next_counter) = _run( + build_seeded_categorical_sampler(), + tmp_path, + { + "logits": np.array([[0.0, -0.1, -10.0]], np.float32), + "temperature": np.array([1.0], np.float32), + "top_k": np.array([0], np.int64), + "top_p": np.array([1.0], np.float32), + "min_p": np.array([0.95], np.float32), + "seed": np.array([23], np.int64), + "counter": np.array([4], np.int64), + "active": np.array([True], np.bool_), + "done": np.array([False], np.bool_), + }, + ) + np.testing.assert_array_equal(token, [0]) + np.testing.assert_array_equal(next_counter, [5]) + + +def test_seeded_sampler_heterogeneous_batch_matches_independent_rows(tmp_path): + component = build_seeded_categorical_sampler() + feeds = { + "logits": np.array( + [ + [2.0, 1.0, 0.5, -1.0, -2.0], + [0.0, 0.1, 0.2, 0.3, 0.4], + [4.0, 3.0, 2.0, 1.0, 0.0], + [-1.0, 0.0, 1.0, 2.0, 3.0], + ], + np.float32, + ), + "temperature": np.array([0.5, 1.5, 0.8, 2.0], np.float32), + "top_k": np.array([1, 3, 0, 2], np.int64), + "top_p": np.array([1.0, 0.8, 0.6, 0.9], np.float32), + "min_p": np.array([0.0, 0.05, 0.2, 0.1], np.float32), + "seed": np.array([3, 7, 11, 13], np.int64), + "counter": np.array([0, 5, 9, 12], np.int64), + "active": np.array([True, True, False, True], np.bool_), + "done": np.array([False, False, False, True], np.bool_), + } + batched_token, batched_offset = _run(component, tmp_path, feeds) + for row in range(4): + row_feeds = {name: value[row : row + 1] for name, value in feeds.items()} + row_token, row_offset = _run(component, tmp_path, row_feeds) + np.testing.assert_array_equal(batched_token[row : row + 1], row_token) + np.testing.assert_array_equal(batched_offset[row : row + 1], row_offset) + np.testing.assert_array_equal(batched_token[2:], [-1, -1]) + np.testing.assert_array_equal(batched_offset, [1, 6, 9, 12]) + + +def test_row_selective_state_and_termination_preserve_inactive_rows(tmp_path): + (next_state,) = _run( + build_token_state_update(row_selective=True), + tmp_path, + { + "current": np.array([[10], [20], [30]], np.int64), + "update": np.array([[11], [21], [31]], np.int64), + "active": np.array([True, False, True], np.bool_), + "done": np.array([False, False, True], np.bool_), + }, + ) + np.testing.assert_array_equal(next_state, [[11], [20], [30]]) + + done, next_active, continued = _run( + build_eos_termination(row_selective=True), + tmp_path, + { + "tokens": np.array([2, 8, 9], np.int64), + "eos_ids": np.array([[2, 9], [2, 9], [2, 9]], np.int64), + "eos_lengths": np.array([2, 1, 2], np.int64), + "iteration": np.array([0], np.int64), + "max_iterations": np.array([5, 5, 5], np.int64), + "active": np.array([True, False, True], np.bool_), + }, + ) + np.testing.assert_array_equal(done, [True, True, True]) + np.testing.assert_array_equal(next_active, [False, False, False]) + np.testing.assert_array_equal(continued, [False]) + + +def test_row_selective_termination_heterogeneous_batch_matches_independent_rows( + tmp_path, +): + component = build_eos_termination(row_selective=True) + feeds = { + "tokens": np.array([2, 9, 5, 7], np.int64), + "eos_ids": np.array([[2, 99], [8, 9], [5, 6], [1, 2]], np.int64), + "eos_lengths": np.array([1, 1, 2, 2], np.int64), + "iteration": np.array([2], np.int64), + "max_iterations": np.array([5, 5, 10, 3], np.int64), + "active": np.array([True, True, True, True], np.bool_), + } + done, next_active, continued = _run(component, tmp_path, feeds) + np.testing.assert_array_equal(done, [True, False, True, True]) + np.testing.assert_array_equal(next_active, [False, True, False, False]) + np.testing.assert_array_equal(continued, [True]) + for row in range(4): + row_outputs = _run( + component, + tmp_path, + { + name: value if name == "iteration" else value[row : row + 1] + for name, value in feeds.items() + }, + ) + np.testing.assert_array_equal(done[row : row + 1], row_outputs[0]) + np.testing.assert_array_equal(next_active[row : row + 1], row_outputs[1]) + + +def test_termination_batch_controls_initialize_dynamic_rows(tmp_path): + eos_ids, eos_lengths, max_iterations = _run( + build_termination_batch_initializer(), + tmp_path, + { + "input_eos_ids": np.array([[2, 3], [7, -1]], np.int64), + "input_eos_lengths": np.array([2, 1], np.int64), + "input_max_iterations": np.array([4, -1], np.int64), + "fallback_max_iterations": np.array([8], np.int64), + "active": np.array([True, False], np.bool_), + }, + ) + np.testing.assert_array_equal(eos_ids, [[2, 3], [7, -1]]) + np.testing.assert_array_equal(eos_lengths, [2, 1]) + np.testing.assert_array_equal(max_iterations, [4, 8]) + + (iteration_rows,) = _run( + build_integer_row_broadcast(), + tmp_path, + { + "value": np.array([3], np.int64), + "active": np.array([True, False], np.bool_), + }, + ) + np.testing.assert_array_equal(iteration_rows, [3, 3]) + + +def test_eos_termination_runtime(tmp_path): + terminated, continued = _run( + build_eos_termination(), + tmp_path, + { + "tokens": np.array([2, 8, 9], np.int64), + "eos_ids": np.array([2, 9], np.int64), + "iteration": np.array([0, 4, 1], np.int64), + "max_iterations": np.array([5, 5, 2], np.int64), + }, + ) + np.testing.assert_array_equal(terminated, [True, True, True]) + np.testing.assert_array_equal(continued, [False]) + + +def test_euler_solver_runtime_parity(tmp_path): + sample = np.arange(8, dtype=np.float32).reshape(1, 2, 2, 2) + derivative = np.full_like(sample, 0.25) + (actual,) = _run( + build_euler_solver_step(), + tmp_path, + { + "sample": sample, + "derivative": derivative, + "step": np.array([0], np.int64), + "schedule": np.array([1.5, 0.5], np.float32), + }, + ) + np.testing.assert_allclose(actual, sample - derivative) + + +def test_ddim_solver_runtime_parity_on_a_video_latent(tmp_path): + # [batch, frames, channels, height, width]: the update has to broadcast the + # per-row alphas over a temporal axis, not just over an image. + dims = ["batch", "frames", "channels", "height", "width"] + sample = np.linspace(-2.0, 2.0, 24, dtype=np.float32).reshape(1, 3, 2, 2, 2) + estimate = np.full_like(sample, 0.25) + schedule = np.array([0.4, 0.9, 1.0], np.float32) + (actual,) = _run( + build_ddim_solver_step(latent_dims=dims), + tmp_path, + { + "sample": sample, + "derivative": estimate, + "step": np.array([0], np.int64), + "schedule": schedule, + }, + ) + alpha, alpha_prev = schedule[0], schedule[1] + expected = ( + np.sqrt(alpha_prev) * ((sample - np.sqrt(1 - alpha) * estimate) / np.sqrt(alpha)) + + np.sqrt(1 - alpha_prev) * estimate + ) + np.testing.assert_allclose(actual, expected, rtol=1e-5, atol=1e-6) + + +def test_ddim_solver_clips_the_predicted_clean_latent(tmp_path): + dims = ["batch", "frames", "channels", "height", "width"] + sample = np.full((1, 1, 1, 1, 1), 8.0, np.float32) + estimate = np.zeros_like(sample) + (actual,) = _run( + build_ddim_solver_step(latent_dims=dims, clip_sample_range=1.0), + tmp_path, + { + "sample": sample, + "derivative": estimate, + "step": np.array([0], np.int64), + "schedule": np.array([0.25, 0.81], np.float32), + }, + ) + # pred_x0 = 8 / sqrt(0.25) = 16, clipped to 1, then renoised to alpha_prev. + np.testing.assert_allclose(actual, np.sqrt(0.81), rtol=1e-6) + + +def test_video_decode_chunk_walk_matches_the_reference(tmp_path): + latent = np.arange(5 * 2, dtype=np.float32).reshape(1, 1, 5, 1, 2) + (count,) = _run(build_video_decode_chunk_count(), tmp_path, {"latent": latent}) + np.testing.assert_array_equal(count, [2]) + chunks = [ + _run( + build_video_decode_chunk(), + tmp_path, + {"latent": latent, "step": np.array([step], np.int64)}, + )[0] + for step in range(int(count[0])) + ] + # Five latent frames split as three then two: the odd frame is folded into + # the first chunk, and the chunks tile the clip without gaps or overlap. + assert [chunk.shape[2] for chunk in chunks] == [3, 2] + np.testing.assert_array_equal(np.concatenate(chunks, axis=2), latent) + + single = np.arange(3 * 2, dtype=np.float32).reshape(1, 1, 3, 1, 2) + (single_count,) = _run(build_video_decode_chunk_count(), tmp_path, {"latent": single}) + np.testing.assert_array_equal(single_count, [1]) + + +def test_video_conv_cache_initializer_sizes_each_resolution(tmp_path): + latent = np.zeros((2, 4, 3, 5, 6), np.float32) + caches = _run( + build_video_conv_cache_initializer( + [("conv_cache.conv_in", 4, 1), ("conv_cache.conv_out", 8, 4)] + ), + tmp_path, + {"latent": latent}, + ) + # Zero frames is how "no previous chunk" is expressed; the spatial extents + # still have to match the resolution each cached convolution runs at. + assert caches[0].shape == (2, 4, 0, 5, 6) + assert caches[1].shape == (2, 8, 0, 20, 24) + + +def test_scheduler_history_starts_empty_and_grows_per_step(tmp_path): + noise = np.zeros((2, 3, 4, 2, 2), np.float32) + latent, history = _run( + build_video_latent_initializer(ir.DataType.FLOAT, 2.0), tmp_path, {"noise": noise} + ) + assert latent.shape == noise.shape + assert history.shape == (2, 0) + assert history.dtype == np.int64 + for step, timestep in enumerate([600, 300, 0]): + (history,) = _run( + build_schedule_history_append(ir.DataType.INT64), + tmp_path, + {"history": history, "timestep": np.array([timestep, timestep], np.int64)}, + ) + assert history.shape == (2, step + 1) + np.testing.assert_array_equal(history, [[600, 300, 0], [600, 300, 0]]) + + +def test_masked_update_runtime_parity(tmp_path): + logits = np.zeros((1, 3, 7), dtype=np.float32) + logits[0, 1, 5] = 1.0 + logits[0, 2, 6] = 4.0 + outputs = _run( + build_masked_token_update(), + tmp_path, + { + "current_tokens": np.array([[1, 99, 99]], np.int64), + "proposed_tokens": np.array([[4, 5, 6]], np.int64), + "logits": logits, + "masked": np.array([[False, True, True]]), + "step": np.array([0], np.int64), + "total_steps": np.array([2], np.int64), + "seed": np.array([7], np.int64), + "offset": np.array([11], np.int64), + }, + ) + np.testing.assert_array_equal(outputs[0], [[1, 99, 6]]) + np.testing.assert_array_equal(outputs[1], [[False, True, False]]) + np.testing.assert_array_equal(outputs[2], [14]) + np.testing.assert_array_equal(outputs[3], [False]) + np.testing.assert_array_equal(outputs[4], [True]) + + final = _run( + build_masked_token_update(), + tmp_path, + { + "current_tokens": outputs[0], + "proposed_tokens": np.array([[4, 5, 6]], np.int64), + "logits": logits, + "masked": outputs[1], + "step": np.array([1], np.int64), + "total_steps": np.array([2], np.int64), + "seed": np.array([7], np.int64), + "offset": outputs[2], + }, + ) + np.testing.assert_array_equal(final[0], [[1, 5, 6]]) + np.testing.assert_array_equal(final[1], [[False, False, False]]) + np.testing.assert_array_equal(final[3], [True]) + np.testing.assert_array_equal(final[4], [False]) + + +def test_speculative_acceptance_prefix_runtime(tmp_path): + ( + accepted_tokens, + count, + done, + next_offset, + rollback_len, + continued, + ) = _run( + build_speculative_acceptance(), + tmp_path, + { + "target_scores": np.array( + [[[0, 1], [1, 0], [0, 1], [1, 0]]], + np.float32, + ), + "proposed_tokens": np.array([[1, 0, 0, 0]], np.int64), + "seed": np.array([3], np.int64), + "offset": np.array([8], np.int64), + }, + ) + np.testing.assert_array_equal(accepted_tokens, [[1, 0, 1, 0]]) + np.testing.assert_array_equal(count, [3]) + np.testing.assert_array_equal(done, [False]) + np.testing.assert_array_equal(next_offset, [12]) + np.testing.assert_array_equal(rollback_len, [2]) + np.testing.assert_array_equal(continued, [True]) + (corrected_cache,) = _run( + build_speculative_state_rollback( + ir.DataType.FLOAT, + ["batch", 1, "past_sequence", 2], + sequence_axis=2, + ), + tmp_path, + { + "past_state": np.zeros((1, 1, 2, 2), np.float32), + "tentative_state": np.zeros((1, 1, 6, 2), np.float32), + "accepted_len": rollback_len, + }, + ) + assert corrected_cache.shape[2] == 4 + + +def test_speculative_acceptance_preserves_per_row_prefixes(tmp_path): + ( + accepted_tokens, + count, + done, + _, + rollback_len, + continued, + ) = _run( + build_speculative_acceptance(), + tmp_path, + { + "target_scores": np.array( + [ + [[0, 1], [1, 0], [0, 1], [1, 0]], + [[0, 1], [1, 0], [0, 1], [1, 0]], + ], + np.float32, + ), + "proposed_tokens": np.array([[1, 1, 0, 0], [1, 0, 1, 0]], np.int64), + "seed": np.array([3, 4], np.int64), + "offset": np.array([0, 0], np.int64), + }, + ) + np.testing.assert_array_equal(count, [2, 4]) + np.testing.assert_array_equal( + accepted_tokens, + [[1, 0, 0, 0], [1, 0, 1, 0]], + ) + np.testing.assert_array_equal(done, [False, True]) + np.testing.assert_array_equal(rollback_len, [1, 4]) + np.testing.assert_array_equal(continued, [True, False]) + + +def test_speculative_state_rollback_trims_tentative_cache(tmp_path): + past = np.arange(4, dtype=np.float32).reshape(1, 1, 2, 2) + tentative = np.arange(12, dtype=np.float32).reshape(1, 1, 6, 2) + (corrected,) = _run( + build_speculative_state_rollback( + ir.DataType.FLOAT, + ["batch", 1, "past_sequence", 2], + sequence_axis=2, + ), + tmp_path, + { + "past_state": past, + "tentative_state": tentative, + "accepted_len": np.array([2], np.int64), + }, + ) + assert corrected.shape == (1, 1, 4, 2) + np.testing.assert_array_equal(corrected, tentative[:, :, :4, :]) + + +def test_euler_model_input_scales_by_sigma(tmp_path): + (scaled,) = _run( + build_euler_model_input(), + tmp_path, + { + "sample": np.full((1, 1, 1, 1), 10.0, np.float32), + "step": np.array([0], np.int64), + "schedule": np.array([2.0, 0.0], np.float32), + }, + ) + np.testing.assert_allclose(scaled, 10.0 / np.sqrt(5.0), rtol=1e-6) + + +def test_guidance_combine_extrapolates_per_row(tmp_path): + # Classifier-free guidance: uncond + scale * (cond - uncond), with the scale + # supplied per request row rather than baked into the graph. + unconditional = np.array([[[[1.0]]], [[[2.0]]]], dtype=np.float32) + conditional = np.array([[[[3.0]]], [[[-2.0]]]], dtype=np.float32) + (guided,) = _run( + build_guidance_combine(), + tmp_path, + { + "unconditional": unconditional, + "conditional": conditional, + "scale": np.array([7.5, 0.0], np.float32), + }, + ) + np.testing.assert_allclose( + guided, + unconditional + + np.array([7.5, 0.0], np.float32).reshape(2, 1, 1, 1) * (conditional - unconditional), + rtol=1e-6, + ) + + +def test_multistep_solver_matches_dpmsolverpp_second_order(tmp_path): + # Reproduces diffusers' DPMSolverMultistepScheduler.step for dpmsolver++ + # midpoint updates, including the first-order fallback on the first step. + schedule = np.array([8.0, 4.0, 1.0, 0.0], np.float32) + sample = np.linspace(-1.0, 1.0, 8, dtype=np.float32).reshape(1, 2, 2, 2) + estimate = np.linspace(0.5, -0.5, 8, dtype=np.float32).reshape(1, 2, 2, 2) + history = np.full_like(sample, 0.25) + + def reference(sample, estimate, history, step, first_order): + sigma, sigma_next = schedule[step], schedule[step + 1] + alpha = 1.0 / np.sqrt(sigma**2 + 1.0) + alpha_next = 1.0 / np.sqrt(sigma_next**2 + 1.0) + noise, noise_next = sigma * alpha, sigma_next * alpha_next + x0 = (sample - noise * estimate) / alpha + ratio = noise_next / noise + if first_order: + return ratio * sample - alpha_next * (sigma_next / sigma - 1.0) * x0 + step_size = np.log(sigma) - np.log(sigma_next) + previous = np.log(schedule[step - 1]) - np.log(sigma) + difference = (step_size / previous) * (x0 - history) + return ratio * sample - alpha_next * (sigma_next / sigma - 1.0) * ( + x0 + 0.5 * difference + ) + + # Step 0 has no usable history, so the solver must fall back to first order. + next_state, next_history = _run( + build_multistep_solver_step(), + tmp_path, + { + "sample": sample, + "estimate": estimate, + "history": history, + "step": np.array([0], np.int64), + "schedule": schedule, + }, + ) + np.testing.assert_allclose( + next_state, reference(sample, estimate, history, 0, True), rtol=1e-5, atol=1e-6 + ) + + # Step 1 uses the carried estimate; the returned history is the new one. + second_state, second_history = _run( + build_multistep_solver_step(), + tmp_path, + { + "sample": sample, + "estimate": estimate, + "history": next_history, + "step": np.array([1], np.int64), + "schedule": schedule, + }, + ) + np.testing.assert_allclose( + second_state, + reference(sample, estimate, next_history, 1, False), + rtol=1e-5, + atol=1e-6, + ) + + # The final step drops back to first order the way lower_order_final does. + final_state, _ = _run( + build_multistep_solver_step(), + tmp_path, + { + "sample": sample, + "estimate": estimate, + "history": second_history, + "step": np.array([2], np.int64), + "schedule": schedule, + }, + ) + np.testing.assert_allclose( + final_state, + reference(sample, estimate, second_history, 2, True), + rtol=1e-5, + atol=1e-6, + ) + + +def test_counter_rng_is_reproducible_and_row_private(tmp_path): + component = build_counter_rng_normal() + feeds = { + "seed": np.array([1234, 4321], np.int64), + "offset": np.array([0, 0], np.int64), + "row_shape": np.array([4, 8, 8], np.int64), + } + noise, next_offset = _run(component, tmp_path, feeds) + assert noise.shape == (2, 4, 8, 8) + np.testing.assert_array_equal(next_offset, [1, 1]) + + repeat, _ = _run(component, tmp_path, feeds) + np.testing.assert_array_equal(noise, repeat) + + # A row's draw depends only on its own seed, not on its batch position. + swapped, _ = _run( + component, + tmp_path, + {**feeds, "seed": np.array([4321, 1234], np.int64)}, + ) + np.testing.assert_array_equal(swapped[0], noise[1]) + np.testing.assert_array_equal(swapped[1], noise[0]) + + # Advancing the counter draws a different, decorrelated block. + advanced, advanced_offset = _run( + component, + tmp_path, + {**feeds, "offset": np.array([1, 1], np.int64)}, + ) + np.testing.assert_array_equal(advanced_offset, [2, 2]) + assert np.abs(advanced - noise).max() > 1e-3 + assert abs(float(np.corrcoef(advanced.ravel(), noise.ravel())[0, 1])) < 0.1 + # Box-Muller output must look standard normal. + assert abs(float(noise.mean())) < 0.1 + assert abs(float(noise.std()) - 1.0) < 0.1 + + +def test_tensor_scale_and_zeros_like_shape_from_their_input(tmp_path): + sample = np.arange(12, dtype=np.float32).reshape(1, 3, 2, 2) + (scaled,) = _run( + build_tensor_scale(), + tmp_path, + {"tensor": sample, "scale": np.array([0.5], np.float32)}, + ) + np.testing.assert_allclose(scaled, sample * 0.5) + (zeros,) = _run(build_zeros_like(), tmp_path, {"reference": sample}) + np.testing.assert_array_equal(zeros, np.zeros_like(sample)) + + +def test_scalar_and_shape_constants_publish_their_values(tmp_path): + (value,) = _run(build_scalar_constant(0.18215), tmp_path, {}) + np.testing.assert_allclose(value, 0.18215, rtol=1e-6) + (shape,) = _run(build_shape_constant([4, 64, 64]), tmp_path, {}) + np.testing.assert_array_equal(shape, [4, 64, 64]) + + +def test_code_frame_update_accepts_scalar_loop_index(tmp_path): + (updated,) = _run( + build_code_frame_update(4, scalar_index=True), + tmp_path, + { + "frame_codes": np.zeros((2, 4), np.int64), + "token": np.array([5, 7], np.int64), + "index": np.array(2, np.int64), + }, + ) + np.testing.assert_array_equal(updated, [[0, 0, 5, 0], [0, 0, 7, 0]]) + + +def test_token_state_update_runtime(tmp_path): + (next_state,) = _run( + build_token_state_update(), + tmp_path, + { + "current": np.array([[1], [3]], np.int64), + "update": np.array([5, 7], np.int64), + }, + ) + np.testing.assert_array_equal(next_state, [[5], [7]]) + + +def test_capability_driven_attachment_is_model_agnostic(): + package = ModelPackage() + artifacts = attach_policy_components( + package, + PolicyCapabilities( + sampler="greedy", + eos_termination=True, + solver="euler", + masked_update=True, + speculative_acceptance=True, + token_state_update=True, + ), + ) + + assert artifacts == { + "token_sampler": "policies/token_sampler.onnx", + "termination": "policies/termination.onnx", + "solver_step": "policies/solver_step.onnx", + "masked_update": "policies/masked_update.onnx", + "speculative_acceptance": "policies/speculative_acceptance.onnx", + "token_state_update": "policies/token_state_update.onnx", + } + assert {component.contract_id for component in package.policy_components.values()} == { + "onnx-genai.token-sampler@1", + "onnx-genai.termination-predicate@1", + "onnx-genai.solver-step@1", + "onnx-genai.masked-update@1", + "onnx-genai.speculative-verifier@1", + "onnx-genai.state-update@1", + } + + +def test_state_initializer_allocates_fp8_cache_through_a_cast(tmp_path): + """An fp8 KV cache must not be materialized by ``ConstantOfShape`` directly. + + ORT's ``ConstantOfShape`` kernel has no fp8 output implementation, so an + fp8 cache emitted that way fails at session initialization with + "Unsupported value attribute datatype: 17". Fill a supported dtype and cast. + """ + inputs = [ + ir.Value( + name="input_ids", + type=ir.TensorType(ir.DataType.INT64), + shape=ir.Shape(["batch", "sequence"]), + ), + ir.Value( + name="attention_mask", + type=ir.TensorType(ir.DataType.INT64), + shape=ir.Shape(["batch", "sequence"]), + ), + ir.Value( + name="past_key_values.0.key", + type=ir.TensorType(ir.DataType.FLOAT8E4M3FN), + shape=ir.Shape(["batch", 2, "past_sequence", 4]), + ), + ] + decoder = ir.Model(ir.Graph(inputs, [], nodes=[], name="decoder"), ir_version=11) + initializer = build_decoder_state_initializer( + decoder, + token_input="input_ids", + attention_mask_input="attention_mask", + position_ids_input=None, + cache_inputs=["past_key_values.0.key"], + ) + + cache = next( + value + for value in initializer.model.graph.outputs + if value.name == "past_key_values.0.key" + ) + assert cache.dtype == ir.DataType.FLOAT8E4M3FN + assert cache.producer().op_type == "Cast" + fill = cache.producer().inputs[0].producer() + assert fill.op_type == "ConstantOfShape" + assert fill.attributes["value"].value.dtype == ir.DataType.FLOAT + + outputs = _run( + initializer, + tmp_path, + {"prompt_tokens": np.array([[3, 4, 5]], np.int64)}, + ) + assert outputs[-1].shape == (1, 2, 0, 4) + + +def _reference_pack(latent: np.ndarray) -> np.ndarray: + """``QwenImagePipeline._pack_latents`` in numpy: (B,C,T,H,W) -> (B,T*H/2*W/2,C*4).""" + batch, channels, frames, height, width = latent.shape + packed = latent.reshape(batch, channels, frames, height // 2, 2, width // 2, 2) + packed = packed.transpose(0, 2, 3, 5, 1, 4, 6) + return packed.reshape(batch, frames * (height // 2) * (width // 2), channels * 4) + + +def test_flow_match_solver_step_runtime_parity(tmp_path): + """Flow matching integrates ``x + (sigma_next - sigma) * v`` on rank-3 tokens.""" + sample = np.arange(12, dtype=np.float32).reshape(1, 3, 4) + derivative = np.full_like(sample, 0.5) + (actual,) = _run( + build_flow_match_solver_step(), + tmp_path, + { + "sample": sample, + "derivative": derivative, + "step": np.array([1], np.int64), + "schedule": np.array([1.0, 0.6, 0.2], np.float32), + }, + ) + # step 1 moves sigma 0.6 -> 0.2, so the update is -0.4 * derivative. + np.testing.assert_allclose(actual, sample - 0.4 * derivative, rtol=1e-6, atol=1e-6) + + +def test_pack_latents_matches_diffusers_patchify(tmp_path): + latent = np.arange(2 * 4 * 1 * 4 * 6, dtype=np.float32).reshape(2, 4, 1, 4, 6) + (packed,) = _run(build_pack_latents_2x2(), tmp_path, {"latent_sample": latent}) + assert packed.shape == (2, 6, 16) + np.testing.assert_array_equal(packed, _reference_pack(latent)) + + +def test_unpack_latents_inverts_pack(tmp_path): + """Round-tripping must be exact: the loop packs once and unpacks once. + + ``height``/``width`` are the *packed* token grid, i.e. half the latent + spatial extent, because each token folds a 2x2 patch into its channels. + """ + latent = np.arange(1 * 4 * 1 * 6 * 4, dtype=np.float32).reshape(1, 4, 1, 6, 4) + (packed,) = _run(build_pack_latents_2x2(), tmp_path, {"latent_sample": latent}) + (restored,) = _run( + build_unpack_latents_2x2(), + tmp_path, + { + "packed_latent": packed, + "height": np.array([3], np.int64), + "width": np.array([2], np.int64), + }, + ) + np.testing.assert_array_equal(restored, latent) + + +def test_sequence_concat_joins_target_then_source(tmp_path): + """Order matters: the denoiser slices its estimate back off the front.""" + target = np.ones((1, 2, 3), np.float32) + source = np.full((1, 4, 3), 2.0, np.float32) + (joined,) = _run(build_sequence_concat(), tmp_path, {"target": target, "source": source}) + assert joined.shape == (1, 6, 3) + np.testing.assert_array_equal(joined, np.concatenate([target, source], axis=1)) + + +def test_true_cfg_matches_diffusers_norm_rescale(tmp_path): + """True CFG rescales the guided estimate back to the conditional norm. + + ``QwenImageEditPlusPipeline`` computes ``comb = neg + s * (cond - neg)`` and + then multiplies by ``||cond||/||comb||`` over the channel axis, so guidance + changes direction without inflating magnitude. + """ + rng = np.random.default_rng(3) + cond = rng.standard_normal((2, 3, 4)).astype(np.float32) + uncond = rng.standard_normal((2, 3, 4)).astype(np.float32) + (actual,) = _run( + build_true_cfg(guidance_scale=4.0), + tmp_path, + {"conditional": cond, "unconditional": uncond}, + ) + comb = uncond + 4.0 * (cond - uncond) + cond_norm = np.linalg.norm(cond, axis=-1, keepdims=True) + comb_norm = np.linalg.norm(comb, axis=-1, keepdims=True) + np.testing.assert_allclose(actual, comb * (cond_norm / comb_norm), rtol=1e-5, atol=1e-5) + + +def test_true_cfg_is_identity_at_unit_guidance(tmp_path): + rng = np.random.default_rng(5) + cond = rng.standard_normal((1, 2, 4)).astype(np.float32) + uncond = rng.standard_normal((1, 2, 4)).astype(np.float32) + (actual,) = _run( + build_true_cfg(guidance_scale=1.0), + tmp_path, + {"conditional": cond, "unconditional": uncond}, + ) + np.testing.assert_allclose(actual, cond, rtol=1e-5, atol=1e-5) + + +def _multi_axis_decoder(sections) -> ir.Model: + """A decoder whose ``position_ids`` carries `sections` rotary axes.""" + inputs = [ + ir.Value( + name="input_ids", + type=ir.TensorType(ir.DataType.INT64), + shape=ir.Shape(["batch", "sequence"]), + ), + ir.Value( + name="attention_mask", + type=ir.TensorType(ir.DataType.INT64), + shape=ir.Shape(["batch", "past_sequence + sequence"]), + ), + ir.Value( + name="position_ids", + type=ir.TensorType(ir.DataType.INT64), + shape=ir.Shape([sections, "batch", "sequence"]), + ), + ir.Value( + name="past_key_values.0.key", + type=ir.TensorType(ir.DataType.FLOAT), + shape=ir.Shape(["batch", 2, "past_sequence", 4]), + ), + ] + return ir.Model(ir.Graph(inputs, [], nodes=[], name="decoder"), ir_version=11) + + +def test_rotary_axis_count_reads_the_declared_leading_dimension(): + decoder = _multi_axis_decoder(3) + positions = {value.name: value for value in decoder.graph.inputs}["position_ids"] + assert rotary_axis_count(positions) == 3 + + +def test_rotary_axis_count_is_none_for_plain_rank_2_positions(): + # A plain decoder has no rotary-axis dimension to broadcast over, so the + # policy graphs keep emitting (batch, sequence) positions. + positions = ir.Value( + name="position_ids", + type=ir.TensorType(ir.DataType.INT64), + shape=ir.Shape(["batch", "sequence"]), + ) + assert rotary_axis_count(positions) is None + + +def test_rotary_axis_count_refuses_a_symbolic_axis_count(): + # The number of rotary axes is fixed by the exported graph. A symbolic + # leading dimension means the export did not state it, and guessing would + # silently mis-shape every position the decoder reads. + positions = ir.Value( + name="position_ids", + type=ir.TensorType(ir.DataType.INT64), + shape=ir.Shape(["sections", "batch", "sequence"]), + ) + with pytest.raises(TypeError, match="symbolic"): + rotary_axis_count(positions) + + +def test_multi_axis_positions_are_broadcast_to_every_rotary_axis(tmp_path): + # A 3D-rotary decoder reads (sections, batch, sequence) positions. The + # prefill range and the per-step increment must be shaped for it, or the + # decoder silently reads a rank-2 tensor as one axis of three. + initializer = build_decoder_state_initializer( + _multi_axis_decoder(3), + token_input="input_ids", + attention_mask_input="attention_mask", + position_ids_input="position_ids", + cache_inputs=["past_key_values.0.key"], + ) + outputs = _run( + initializer, + tmp_path, + {"prompt_tokens": np.array([[3, 4, 5]], np.int64)}, + ) + prefill_positions, body_positions = outputs[1], outputs[3] + assert prefill_positions.shape == (3, 1, 3) + np.testing.assert_array_equal(prefill_positions, np.broadcast_to([[0, 1, 2]], (3, 1, 3))) + assert body_positions.shape == (3, 1, 1) + np.testing.assert_array_equal(body_positions, np.full((3, 1, 1), 3)) + + updated = _run( + build_decoder_step_update( + attention_dtype=ir.DataType.INT64, + position_dtype=ir.DataType.INT64, + position_sections=3, + ), + tmp_path, + { + "attention_mask": outputs[2], + "position_ids": body_positions, + }, + ) + np.testing.assert_array_equal(updated[0], [[1, 1, 1, 1, 1]]) + assert updated[1].shape == (3, 1, 1) + np.testing.assert_array_equal(updated[1], np.full((3, 1, 1), 4)) diff --git a/src/mobius/integrations/diffusers/_builder.py b/src/mobius/integrations/diffusers/_builder.py index 703636131..9fd45dde6 100644 --- a/src/mobius/integrations/diffusers/_builder.py +++ b/src/mobius/integrations/diffusers/_builder.py @@ -74,6 +74,10 @@ def _init_diffusers_class_map() -> None: from mobius.models.cogvideox import ( CogVideoXTransformer3DModel, ) + from mobius.models.cogvideox_vae import ( + AutoencoderKLCogVideoXModel, + CogVideoXVAEConfig, + ) from mobius.models.dit import DiTConfig, DiTTransformer2DModel from mobius.models.flux_sd3 import ( FluxConfig, @@ -94,7 +98,6 @@ def _init_diffusers_class_map() -> None: from mobius.models.qwen_vl import Qwen25VLCausalLMModel from mobius.models.unet import UNet2DConditionModel from mobius.models.vae import AutoencoderKLModel - from mobius.models.video_vae import VideoAutoencoderModel, VideoVAEConfig _DIFFUSERS_CLASS_MAP.update( { @@ -128,9 +131,9 @@ def _init_diffusers_class_map() -> None: "qwen-image-vae", ), "AutoencoderKLCogVideoX": ( - VideoAutoencoderModel, - VideoVAEConfig, - "vae", + AutoencoderKLCogVideoXModel, + CogVideoXVAEConfig, + "video-vae", ), "CogVideoXTransformer3DModel": ( CogVideoXTransformer3DModel, diff --git a/src/mobius/integrations/diffusers/_builder_test.py b/src/mobius/integrations/diffusers/_builder_test.py index 14b37392a..8e07bf3bb 100644 --- a/src/mobius/integrations/diffusers/_builder_test.py +++ b/src/mobius/integrations/diffusers/_builder_test.py @@ -95,6 +95,7 @@ def test_task_names_are_valid(self): "qwen-image-denoising", "qwen-image-text-encoding", "video-denoising", + "video-vae", "feature-extraction", "minimax-music3-condition", "minimax-music3-denoising", diff --git a/src/mobius/integrations/diffusers/_configs.py b/src/mobius/integrations/diffusers/_configs.py index 7fc73b0a9..f3fa64d9f 100644 --- a/src/mobius/integrations/diffusers/_configs.py +++ b/src/mobius/integrations/diffusers/_configs.py @@ -334,6 +334,7 @@ class CogVideoXConfig: max_text_seq_length: int = 226 spatial_interpolation_scale: float = 1.875 temporal_interpolation_scale: float = 1.0 + use_learned_positional_embeddings: bool = False norm_eps: float = 1e-5 # cross_attention_dim used by VideoDenoisingTask for text conditioning cross_attention_dim: int = 4096 @@ -363,6 +364,9 @@ def from_diffusers(cls, config: dict) -> CogVideoXConfig: max_text_seq_length=config.get("max_text_seq_length", 226), spatial_interpolation_scale=config.get("spatial_interpolation_scale", 1.875), temporal_interpolation_scale=config.get("temporal_interpolation_scale", 1.0), + use_learned_positional_embeddings=bool( + config.get("use_learned_positional_embeddings", False) + ), norm_eps=config.get("norm_eps", 1e-5), cross_attention_dim=text_dim, ) diff --git a/src/mobius/integrations/onnx_genai/__init__.py b/src/mobius/integrations/onnx_genai/__init__.py index ad19d0bd7..589429600 100644 --- a/src/mobius/integrations/onnx_genai/__init__.py +++ b/src/mobius/integrations/onnx_genai/__init__.py @@ -15,9 +15,13 @@ composite encoder/fusion/autoregressive-decoder pipeline. * **Speech-to-text (ASR)** — :func:`write_speech_to_text_pipeline_metadata` emits a Whisper-style cross-attention encode→decode pipeline. -* **Audio codec / multi-decoder TTS** — :func:`write_audio_codec_pipeline_metadata` - and :func:`write_tts_pipeline_metadata` emit audio-to-audio and - ``pre_embedder``-driven ``nested_autoregressive`` (Qwen3-TTS) pipelines. +* **Full-duplex speech-to-speech** — :func:`write_full_duplex_workflow_metadata` + emits one-event-per-invocation SSA with session-scoped conversational state + (Moshi / PersonaPlex). +* **Audio codec / multi-decoder TTS** — + :func:`write_audio_codec_workflow_metadata` emits typed codec SSA, while + :func:`write_tts_workflow_metadata` reports the current nested-loop induction + contract blocker precisely. * **Diffusion pipelines** — :func:`write_diffusion_pipeline_metadata` emits an iterative pipeline for a denoiser plus optional VAE / text encoder. @@ -50,45 +54,79 @@ ) from mobius.integrations.onnx_genai.inference_metadata import ( SchedulerConfig, - build_audio_codec_pipeline_metadata, + add_policy_components_to_workflow, build_diffusion_pipeline_metadata, - build_language_diffusion_pipeline_metadata, build_multimodal_pipeline_metadata, build_speech_to_text_pipeline_metadata, - build_tts_pipeline_metadata, load_diffusers_scheduler_config, - write_audio_codec_pipeline_metadata, write_diffusion_pipeline_metadata, write_multimodal_pipeline_metadata, write_speech_to_text_pipeline_metadata, - write_tts_pipeline_metadata, +) +from mobius.integrations.onnx_genai.workflow_metadata import ( + build_audio_codec_workflow_metadata, + build_decoder_workflow_metadata, + build_diffusion_workflow_metadata, + build_full_duplex_workflow_metadata, + build_image_edit_workflow_metadata, + build_language_diffusion_pipeline_metadata, + build_speculative_workflow_metadata, + build_tts_workflow_metadata, + build_video_diffusion_workflow_metadata, + build_vlm_workflow_metadata, + write_audio_codec_workflow_metadata, + write_decoder_workflow_metadata, + write_diffusion_workflow_metadata, + write_full_duplex_workflow_metadata, + write_image_edit_workflow_metadata, + write_language_diffusion_workflow_metadata, + write_speculative_workflow_metadata, + write_tts_workflow_metadata, + write_video_diffusion_workflow_metadata, + write_vlm_workflow_metadata, ) __all__ = [ "ComfyUIWorkflow", "ConversionResult", "SchedulerConfig", + "add_policy_components_to_workflow", + "build_audio_codec_workflow_metadata", + "build_full_duplex_workflow_metadata", "build_decoder_metadata", + "build_decoder_workflow_metadata", "build_diffusion_pipeline_metadata", + "build_diffusion_workflow_metadata", + "build_image_edit_workflow_metadata", "build_language_diffusion_pipeline_metadata", - "build_audio_codec_pipeline_metadata", "build_multimodal_pipeline_metadata", "build_pipeline_metadata_for_workflow", + "build_speculative_workflow_metadata", "build_speech_to_text_pipeline_metadata", - "build_tts_pipeline_metadata", + "build_tts_workflow_metadata", + "build_video_diffusion_workflow_metadata", + "build_vlm_workflow_metadata", "convert_comfyui_workflow", "decoder_metadata_from_config", - "moe_metadata_from_config", "load_diffusers_scheduler_config", + "moe_metadata_from_config", "parse_comfyui_workflow", "parse_comfyui_workflow_file", "translate_comfyui_workflow", "translate_comfyui_workflow_file", + "write_audio_codec_workflow_metadata", + "write_full_duplex_workflow_metadata", "write_decoder_metadata", + "write_decoder_workflow_metadata", "write_diffusion_pipeline_metadata", - "write_audio_codec_pipeline_metadata", + "write_diffusion_workflow_metadata", + "write_image_edit_workflow_metadata", + "write_language_diffusion_workflow_metadata", "write_multimodal_pipeline_metadata", - "write_speech_to_text_pipeline_metadata", - "write_tts_pipeline_metadata", "write_onnx_genai_config", + "write_speculative_workflow_metadata", + "write_speech_to_text_pipeline_metadata", + "write_tts_workflow_metadata", + "write_video_diffusion_workflow_metadata", + "write_vlm_workflow_metadata", ] diff --git a/src/mobius/integrations/onnx_genai/auto_export.py b/src/mobius/integrations/onnx_genai/auto_export.py index 0d9f9f793..2ada88da5 100644 --- a/src/mobius/integrations/onnx_genai/auto_export.py +++ b/src/mobius/integrations/onnx_genai/auto_export.py @@ -16,41 +16,184 @@ import os from typing import Any -import yaml +import numpy as np +import onnx_ir as ir -from mobius.integrations.onnx_genai.decoder_metadata import ( - decoder_metadata_from_config, - write_decoder_metadata, -) from mobius.integrations.onnx_genai.inference_metadata import ( + _TEXT_RUNTIME_ASSET_NAMES, SchedulerConfig, - add_explicit_package_io, + _copy_runtime_assets, load_diffusers_scheduler_config, - write_audio_codec_pipeline_metadata, - write_diffusion_pipeline_metadata, - write_multimodal_pipeline_metadata, - write_speech_to_text_pipeline_metadata, - write_tts_pipeline_metadata, + load_diffusers_vae_scaling_factor, +) +from mobius.integrations.onnx_genai.workflow_metadata import ( + write_audio_codec_workflow_metadata, + write_ctc_asr_workflow_metadata, + write_decoder_workflow_metadata, + write_diffusion_workflow_metadata, + write_encoder_embedding_workflow_metadata, + write_image_edit_workflow_metadata, + write_language_diffusion_workflow_metadata, + write_speculative_workflow_metadata, + write_speech_to_text_workflow_metadata, + write_tts_workflow_metadata, + write_vlm_workflow_metadata, ) _LOGGER = logging.getLogger(__name__) + +#: Scheduler kind -> (workflow solver component, whether the sampler rescales +#: the denoiser's input). A sampler that keeps its state variance-preserving +#: feeds the raw state to the denoiser and starts from unit-variance noise; a +#: sampler that carries state in sigma space divides by ``sqrt(sigma**2 + 1)`` +#: and starts from ``sigma_max`` scaled noise. +_DIFFUSION_SOLVERS: dict[str, tuple[str, bool]] = { + "euler": ("euler", True), + "dpmpp_2m": ("multistep", False), +} + + +def _diffusion_schedule( + scheduler: SchedulerConfig, num_inference_steps: int +) -> tuple[list[float], list[float]]: + """Materialize diffusers-compatible timesteps and sigma values.""" + if scheduler.kind not in _DIFFUSION_SOLVERS or scheduler.prediction_type != "epsilon": + raise ValueError( + "workflow diffusion currently supports deterministic epsilon schedulers " + f"{sorted(_DIFFUSION_SOLVERS)}, got kind={scheduler.kind!r}, " + f"prediction_type={scheduler.prediction_type!r}" + ) + if scheduler.kind == "dpmpp_2m" and ( + scheduler.algorithm_type != "dpmsolver++" + or scheduler.solver_order != 2 + or scheduler.solver_type != "midpoint" + or not scheduler.lower_order_final + or scheduler.final_sigmas_type != "zero" + ): + raise ValueError( + "the workflow multistep solver implements second-order dpmsolver++ with " + "midpoint updates, a lower-order final step, and a zero terminal sigma; " + f"got algorithm_type={scheduler.algorithm_type!r}, " + f"solver_order={scheduler.solver_order}, " + f"solver_type={scheduler.solver_type!r}, " + f"lower_order_final={scheduler.lower_order_final}, " + f"final_sigmas_type={scheduler.final_sigmas_type!r}" + ) + if scheduler.use_karras_sigmas or scheduler.use_exponential_sigmas: + raise ValueError( + "workflow diffusion does not yet materialize Karras or exponential sigmas" + ) + if scheduler.beta_schedule == "scaled_linear": + betas = ( + np.linspace( + np.sqrt(scheduler.beta_start), + np.sqrt(scheduler.beta_end), + scheduler.num_train_timesteps, + dtype=np.float64, + ) + ** 2 + ) + elif scheduler.beta_schedule == "linear": + betas = np.linspace( + scheduler.beta_start, + scheduler.beta_end, + scheduler.num_train_timesteps, + dtype=np.float64, + ) + else: + raise ValueError( + f"workflow diffusion does not support beta schedule {scheduler.beta_schedule!r}" + ) + training_sigmas = np.sqrt((1.0 - np.cumprod(1.0 - betas)) / np.cumprod(1.0 - betas)) + if scheduler.kind == "dpmpp_2m": + # Multistep solvers place the boundary at the terminal sigma, so the + # linspace spans one extra point and drops the trailing zero timestep. + timesteps = ( + np.linspace(0, scheduler.num_train_timesteps - 1, num_inference_steps + 1) + .round()[::-1][:-1] + .astype(np.float64) + ) + else: + timesteps = np.linspace( + scheduler.num_train_timesteps - 1, + 0, + num_inference_steps, + dtype=np.float64, + ) + sigmas = np.interp( + timesteps, + np.arange(scheduler.num_train_timesteps, dtype=np.float64), + training_sigmas, + ) + return timesteps.tolist(), [*sigmas.tolist(), 0.0] + + _DENOISER_KEYS = ("denoiser", "transformer", "unet") -def _add_explicit_io_to_file(path: str, pkg: Any, config: Any) -> None: - """Augment an emitted sidecar with roles derived from the actual ONNX ports.""" +def _flow_match_euler_schedule( + scheduler: SchedulerConfig, num_inference_steps: int, image_seq_len: int +) -> tuple[list[float], list[float]]: + """Materialize diffusers ``FlowMatchEulerDiscreteScheduler`` timesteps/sigmas. + + Reproduces ``set_timesteps(sigmas=linspace(1, 1/n, n), mu=calculate_shift(...))`` + including resolution-dependent dynamic shifting and terminal stretching, so + the baked schedule matches the pipeline that produced the reference image. + Timesteps are emitted as sigmas (``t / num_train_timesteps``) because the + Qwen Image denoiser consumes the normalized timestep directly. + """ + if scheduler.kind != "flow_match_euler": + raise ValueError( + f"image-edit workflow requires a flow-match Euler scheduler, got {scheduler.kind!r}" + ) + sigmas = np.linspace(1.0, 1.0 / num_inference_steps, num_inference_steps, dtype=np.float64) + if scheduler.use_dynamic_shifting: + base_seq_len = scheduler.base_image_seq_len or 256 + max_seq_len = scheduler.max_image_seq_len or 4096 + base_shift = scheduler.base_shift if scheduler.base_shift is not None else 0.5 + max_shift = scheduler.max_shift if scheduler.max_shift is not None else 1.15 + slope = (max_shift - base_shift) / (max_seq_len - base_seq_len) + mu = image_seq_len * slope + (base_shift - slope * base_seq_len) + if (scheduler.time_shift_type or "exponential") != "exponential": + raise ValueError( + f"unsupported flow-match time shift {scheduler.time_shift_type!r}" + ) + sigmas = np.exp(mu) / (np.exp(mu) + (1.0 / sigmas - 1.0)) + elif scheduler.shift is not None: + sigmas = scheduler.shift * sigmas / (1.0 + (scheduler.shift - 1.0) * sigmas) + if scheduler.shift_terminal is not None: + # stretch_shift_to_terminal: map the last sigma onto shift_terminal. + one_minus = 1.0 - sigmas + sigmas = 1.0 - one_minus / (one_minus[-1] / (1.0 - scheduler.shift_terminal)) + return sigmas.tolist(), [*sigmas.tolist(), 0.0] + + +def _looks_like_image_edit(pkg: Any) -> bool: + """Detect a source-image-conditioned flow-matching editing pipeline. + + Structural signals: a VAE encoder and decoder pair, plus a denoiser that + takes rank-3 packed latents and exposes a ``target_sequence_length`` port — + i.e. it consumes concatenated target+source tokens and slices its estimate + back to the target block. + """ try: - models = list(pkg.values()) + names = set(pkg.keys()) except AttributeError: - return - if not models or any(not hasattr(model, "graph") for model in models): - return - with open(path, encoding="utf-8") as handle: - metadata = yaml.safe_load(handle) - add_explicit_package_io(metadata, pkg, config) - with open(path, "w", encoding="utf-8") as handle: - yaml.safe_dump(metadata, handle, sort_keys=False) + return False + if not {"vae_encoder", "vae_decoder"} <= names: + return False + denoiser_name = next((key for key in _DENOISER_KEYS if key in names), None) + if denoiser_name is None: + return False + inputs = {value.name: value for value in pkg[denoiser_name].graph.inputs} + sample = inputs.get("sample") + return ( + "target_sequence_length" in inputs + and sample is not None + and sample.shape is not None + and len(sample.shape) == 3 + ) def _write_clip_tokenizer( @@ -116,6 +259,37 @@ def _write_clip_tokenizer( return path +def _write_text_runtime_assets( + output_dir: str, + source: str | None, + *, + revision: str | None = None, +) -> dict[str, str]: + """Emit the tokenizer *and* chat-template assets a text package needs. + + ``tokenizer.json`` alone is not enough for an instruction-tuned decoder: the + runtime applies the package's chat template to build the prompt, and without + it the raw user text (no leading BOS, no turn markers) reaches the model. + Gemma 4 answers such a prompt with unbounded repetition, so shipping the + template is a correctness requirement rather than a convenience. + + Args: + output_dir: Package directory to write the assets into. + source: Hugging Face model id or local directory holding them. + + Returns: + A mapping of asset stem to written path for every asset materialized. + """ + artifacts = _copy_runtime_assets( + output_dir, source, _TEXT_RUNTIME_ASSET_NAMES, revision=revision + ) + if "tokenizer" not in artifacts: + fallback = _write_hf_tokenizer(output_dir, source, revision=revision) + if fallback is not None: + artifacts["tokenizer"] = fallback + return artifacts + + def _write_hf_tokenizer( output_dir: str, source: str | None, @@ -204,6 +378,88 @@ def _write_hf_audio_processor( return path +def _audio_preprocessing_program( + processor_path: str | None, encoder: Any +) -> dict[str, Any] | None: + """Derive a declarative log-mel program from a HF feature-extractor config. + + The program is the executable contract the runtime audio adapter follows: + decode -> resample -> pad/trim to the fixed window -> log-mel -> normalize. + Its single output binds to the encoder's rank-3 feature input. + """ + if processor_path is None: + return None + import json + + with open(processor_path, encoding="utf-8") as handle: + config = json.load(handle) + if config.get("feature_extractor_type") != "WhisperFeatureExtractor": + _LOGGER.warning( + "Audio feature extractor %r is not a log-mel window extractor; " + "skipping declarative audio preprocessing.", + config.get("feature_extractor_type"), + ) + return None + feature_inputs = [ + value + for value in encoder.graph.inputs + if value.shape is not None and len(value.shape) == 3 + ] + if len(feature_inputs) != 1: + raise ValueError( + "audio preprocessing requires exactly one rank-3 encoder feature input, " + f"got {[value.name for value in feature_inputs]}" + ) + sampling_rate = int(config["sampling_rate"]) + num_mel_bins = int(config["feature_size"]) + n_fft = int(config["n_fft"]) + hop_length = int(config["hop_length"]) + n_samples = int(config.get("n_samples", config["chunk_length"] * sampling_rate)) + return { + "transforms": [ + {"op": "decode", "outputs": ["samples"]}, + { + "op": "resample", + "inputs": ["samples"], + "outputs": ["resampled"], + "sampling_rate": sampling_rate, + }, + { + "op": "pad", + "inputs": ["resampled"], + "outputs": ["windowed"], + "mode": "fixed_window", + "target_samples": n_samples, + "pad_value": float(config.get("padding_value", 0.0)), + }, + { + "op": "log_mel", + "inputs": ["windowed"], + "outputs": ["mel"], + "num_mel_bins": num_mel_bins, + "n_fft": n_fft, + "hop_length": hop_length, + "window": "hann", + "mel_scale": "slaney", + "sampling_rate": sampling_rate, + }, + { + "op": "normalize", + "inputs": ["mel"], + "outputs": ["features"], + "mode": "whisper_log_mel", + }, + ], + "outputs": [ + { + "source": "features", + "name": feature_inputs[0].name, + "content": "audio_features", + } + ], + } + + def _looks_like_diffusion(pkg: Any) -> bool: try: names = set(pkg.keys()) @@ -214,6 +470,25 @@ def _looks_like_diffusion(pkg: Any) -> bool: ) +def _looks_like_language_diffusion(pkg: Any) -> bool: + """Detect a full-sequence token denoiser with an executable proposal output.""" + try: + if len(pkg) != 1: + return False + model = next(iter(pkg.values())) + inputs = list(model.graph.inputs) + outputs = list(model.graph.outputs) + except (AttributeError, TypeError): + return False + return ( + len(inputs) == 1 + and inputs[0].dtype in {ir.DataType.INT32, ir.DataType.INT64} + and inputs[0].shape is not None + and len(inputs[0].shape) == 2 + and {"logits", "proposed_tokens"} <= {value.name for value in outputs} + ) + + def _looks_like_multimodal(pkg: Any) -> bool: try: names = set(pkg.keys()) @@ -229,6 +504,14 @@ def _looks_like_multimodal(pkg: Any) -> bool: ) +def _has_audio_encoder(pkg: Any) -> bool: + """Whether a package fuses audio, and so needs a feature extractor.""" + try: + return "audio_encoder" in set(pkg.keys()) + except AttributeError: + return False + + def _looks_like_speech_to_text(pkg: Any) -> bool: """Detect a cross-attention encoder-decoder ASR package (e.g. Whisper). @@ -251,6 +534,61 @@ def _looks_like_speech_to_text(pkg: Any) -> bool: return "encoder_hidden_states" in decoder_inputs +def _looks_like_ctc_asr(pkg: Any) -> bool: + """Detect a non-generative CTC ASR package. + + The signal is structural: a single ``model`` component that consumes a raw + waveform plus a sample-level mask and emits per-frame ``logits`` with no KV + cache. The absent cache is what separates CTC from a Whisper-style + autoregressive decoder that also emits ``logits``. + """ + try: + names = set(pkg.keys()) + except AttributeError: + return False + if names != {"model"}: + return False + try: + model = pkg["model"] + inputs = {value.name for value in model.graph.inputs} + outputs = {value.name for value in model.graph.outputs} + except (AttributeError, KeyError): + return False + if not {"input_values", "attention_mask"} <= inputs: + return False + if "logits" not in outputs: + return False + return not any(name.startswith("past_key_values") for name in inputs) + + +def _looks_like_encoder_embedding(pkg: Any) -> bool: + """Detect a bidirectional encoder that returns embeddings, not tokens. + + The signal is structural: a single ``model`` component that consumes + ``input_ids`` and emits ``last_hidden_state`` with no ``logits`` port and + no KV cache. The absent ``logits`` is what separates an embedding encoder + from every generative package -- there is nothing to sample, so there is no + decode step to describe. + """ + try: + names = set(pkg.keys()) + except AttributeError: + return False + if names != {"model"}: + return False + try: + model = pkg["model"] + inputs = {value.name for value in model.graph.inputs} + outputs = {value.name for value in model.graph.outputs} + except (AttributeError, KeyError): + return False + if "input_ids" not in inputs or "last_hidden_state" not in outputs: + return False + if any(name == "logits" or str(name).startswith("present") for name in outputs): + return False + return not any(str(name).startswith("past_key_values") for name in inputs) + + def _looks_like_audio_codec(pkg: Any) -> bool: """Detect an audio-to-audio neural codec package. @@ -278,33 +616,8 @@ def _looks_like_audio_codec(pkg: Any) -> bool: ) -def _audio_codec_codes_dtype(pkg: Any) -> str: - """Return the metadata dtype of the codec ``codes`` tensor (default int64).""" - # ONNX elem-type names -> onnx-genai metadata dtype tags. Float codes keep - # their precision (fp16/bf16/fp32) so the runtime binds the right buffer type. - float_dtypes = {"FLOAT": "fp32", "FLOAT16": "fp16", "BFLOAT16": "bf16"} - try: - for value in pkg["decoder"].graph.inputs: - if value.name == "codes" and value.dtype is not None: - return float_dtypes.get(value.dtype.name, "int64") - except (AttributeError, KeyError): - # Missing/partial codec structure: fall back to the documented default. - return "int64" - return "int64" - - def _looks_like_multi_decoder_tts(pkg: Any) -> bool: - """Detect a nested multi-decoder TTS package (e.g. Qwen3-TTS). - - The defining signal is a ``talker`` plus a ``code_predictor`` decoder — a - dual, nested autoregressive shape (the code_predictor expands each talker - frame's residual codebooks). When the package also carries the - ``talker_step_embedder`` pre-embedder (see :func:`_has_tts_pre_embedder`), - the dispatcher emits a runnable ``pre_embedder``-driven - ``nested_autoregressive`` contract; without it the component graph is not yet - mappable, so detection triggers a precise, actionable error rather than - mis-emitting (see DESIGN.md §20.3). - """ + """Detect a nested multi-decoder TTS package (e.g. Qwen3-TTS).""" try: names = set(pkg.keys()) except AttributeError: @@ -312,13 +625,19 @@ def _looks_like_multi_decoder_tts(pkg: Any) -> bool: return {"talker", "code_predictor"} <= names +def _looks_like_speculative(pkg: Any) -> bool: + try: + return {"proposer", "verifier"} <= set(pkg.keys()) + except AttributeError: + return False + + def _has_tts_pre_embedder(pkg: Any) -> bool: """True when a multi-decoder TTS package carries the pre-embedder component. The ``talker_step_embedder`` materializes the talker's per-step - ``inputs_embeds`` (``frame_codes [+ text_embed] -> inputs_embeds``); its - presence is what makes the package emittable to the ``pre_embedder``-driven - ``nested_autoregressive`` contract. + ``inputs_embeds`` (``frame_codes [+ text_embed] -> inputs_embeds``). It is + necessary, but not sufficient until generic loops expose induction SSA. """ try: names = set(pkg.keys()) @@ -327,44 +646,8 @@ def _has_tts_pre_embedder(pkg: Any) -> bool: return "talker_step_embedder" in names -def _tts_component_kwargs(pkg: Any, config: Any) -> dict[str, Any]: - """Derive pre-embedder-driven TTS metadata kwargs from a package + config. - - Mobius saves each component into ``/model.onnx``. ``num_code_groups`` - comes from the TTS config (the RVQ residual count per frame). - """ - tts = getattr(config, "tts", None) - num_code_groups = getattr(tts, "num_code_groups", None) if tts is not None else None - if not num_code_groups: - raise ValueError( - "TTS metadata requires config.tts.num_code_groups (RVQ codes per frame)" - ) - kwargs: dict[str, Any] = { - "num_code_groups": num_code_groups, - "talker_filename": "talker/model.onnx", - "code_predictor_filename": "code_predictor/model.onnx", - "pre_embedder_filename": "talker_step_embedder/model.onnx", - } - # Emit the prefill/trailing-text component only when the package carries it; - # otherwise the prefill-less shape (talker frame 0 + zero text_embed) is used. - try: - names = set(pkg.keys()) - except (AttributeError, TypeError): - names = set() - kwargs["prefill_embedder_filename"] = ( - "talker_prefill_embedder/model.onnx" if "talker_prefill_embedder" in names else None - ) - kwargs["activation_dtype"] = _activation_dtype_tag(config) - return kwargs - - def _activation_dtype_tag(config: Any) -> str: - """Map a model config's activation dtype to the metadata dtype tag. - - The composite dataflow edges (inputs_embeds, encoder_hidden_states, …) carry - the model's activation dtype, so metadata must reflect it (fp16/bf16 builds - would otherwise be mislabeled fp32). - """ + """Map a model config's activation dtype to the metadata dtype tag.""" dtype = getattr(config, "dtype", None) name = getattr(dtype, "name", "") or "" return {"FLOAT16": "fp16", "BFLOAT16": "bf16"}.get(name.upper(), "fp32") @@ -435,6 +718,8 @@ def write_onnx_genai_config( guidance_scale: float | None = None, source: str | None = None, revision: str | None = None, + grammar_guidance: bool = False, + adaptive_k_max: int | None = None, **kwargs: Any, ) -> dict[str, str]: """Write ``inference_metadata.yaml`` into ``output_dir`` and return its path. @@ -446,8 +731,9 @@ def write_onnx_genai_config( ===================== ============================================ ================================= Pipeline shape Structural signal (detector) Emitted ``strategy`` ===================== ============================================ ================================= + Image edit VAE pair + denoiser w/ target_sequence_len typed SSA workflow (edit loop) Diffusion denoiser / VAE present ``iterative`` - Audio codec encoder→``codes``→decoder, no cross-attn ``composite`` (two single_pass) + Audio codec encoder→``codes``→decoder, no cross-attn typed SSA workflow Multimodal VLM decoder + vision/audio encoder + fusion ``composite`` (encoders→fuse→AR) Speech-to-text (ASR) decoder consumes ``encoder_hidden_states`` ``composite`` (encode→AR) Decoder LM fallback (a config is required) bare decoder (``kv_cache`` + attn) @@ -466,18 +752,48 @@ def write_onnx_genai_config( ``scheduler`` / ``guidance_scale`` set the loop. """ os.makedirs(output_dir, exist_ok=True) - if _looks_like_diffusion(pkg): - is_qwen_image_edit = getattr(getattr(pkg, "config", None), "model_type", None) == ( - "qwen_image_edit" + if _looks_like_language_diffusion(pkg): + path = write_language_diffusion_workflow_metadata( + pkg, + output_dir, + num_inference_steps=num_inference_steps, ) - if is_qwen_image_edit: - raise ValueError( - "onnx-genai cannot execute Qwen Image Edit packages: the runtime " - "does not support source-latent packing, target/source token " - "concatenation, target-only denoiser outputs, or the required " - "Qwen true-CFG path. Export the ONNX components without " - "--runtime onnx-genai and orchestrate the pipeline directly." + artifacts = {"inference_metadata": path} + artifacts.update(_write_text_runtime_assets(output_dir, source, revision=revision)) + return artifacts + + if _looks_like_diffusion(pkg): + is_image_edit = _looks_like_image_edit(pkg) + if is_image_edit: + if scheduler is None: + scheduler = load_diffusers_scheduler_config(source) + if scheduler is None: + raise ValueError( + "image-edit workflow requires the diffusers scheduler config; " + "pass scheduler=SchedulerConfig(...) or a resolvable source" + ) + image_seq_len = kwargs.pop("image_seq_len", None) + if image_seq_len is None: + raise ValueError( + "image-edit workflow requires image_seq_len (the packed target " + "token count) to materialize the resolution-dependent schedule" + ) + timesteps, sigma_schedule = _flow_match_euler_schedule( + scheduler, num_inference_steps, int(image_seq_len) + ) + path = write_image_edit_workflow_metadata( + pkg, + output_dir, + num_inference_steps=num_inference_steps, + schedule=sigma_schedule, + timesteps=timesteps, + guidance_scale=1.0 if guidance_scale is None else guidance_scale, ) + artifacts = {"inference_metadata": path} + tokenizer_path = _write_hf_tokenizer(output_dir, source, revision=revision) + if tokenizer_path is not None: + artifacts["tokenizer"] = tokenizer_path + return artifacts if scheduler is None: scheduler = load_diffusers_scheduler_config(source, revision=revision) # Fill in component filenames from the package layout, letting any @@ -485,16 +801,45 @@ def write_onnx_genai_config( derived = _diffusion_component_kwargs(pkg) for name, value in derived.items(): kwargs.setdefault(name, value) - # Classic text-conditioned diffusion (a text encoder is present) uses - # classifier-free guidance by default; SD's canonical scale is 7.5. - if guidance_scale is None and "text_encoder_filename" in kwargs: - guidance_scale = 7.5 - path = write_diffusion_pipeline_metadata( + resolved_scheduler = scheduler or SchedulerConfig(kind="euler") + timesteps, sigma_schedule = _diffusion_schedule( + resolved_scheduler, num_inference_steps + ) + solver, scale_model_input = _DIFFUSION_SOLVERS[resolved_scheduler.kind] + # A sigma-space sampler starts from noise scaled by the largest sigma; a + # variance-preserving one starts from the unit-variance draw itself. + initial_state_scale = sigma_schedule[0] if scale_model_input else 1.0 + conditioned = "text_encoder_filename" in kwargs + if conditioned and guidance_scale is None: + raise ValueError( + "a text-conditioned diffusion package must declare its guidance: " + "pass guidance_scale=1.0 for unguided generation, or the pipeline's " + "classifier-free guidance scale to run the guided denoiser path" + ) + if guidance_scale is not None and not conditioned: + raise ValueError( + "classifier-free guidance requires a text-conditioned diffusion package" + ) + guidance = ( + None + if guidance_scale is None or np.isclose(guidance_scale, 1.0) + else float(guidance_scale) + ) + decoder_input_scale = 1.0 + scaling_factor = load_diffusers_vae_scaling_factor(source) + if scaling_factor: + decoder_input_scale = 1.0 / scaling_factor + path = write_diffusion_workflow_metadata( + pkg, output_dir, num_inference_steps=num_inference_steps, - scheduler=scheduler, - guidance_scale=guidance_scale, - **kwargs, + schedule=sigma_schedule, + timesteps=timesteps, + solver=solver, + scale_model_input=scale_model_input, + initial_state_scale=initial_state_scale, + decoder_input_scale=decoder_input_scale, + guidance_scale=guidance, ) artifacts = {"inference_metadata": path} # Emit the CLIP tokenizer.json for text-conditioned pipelines so the @@ -508,8 +853,48 @@ def write_onnx_genai_config( if _looks_like_audio_codec(pkg): # A neural codec produces tensors (waveform), not tokens, so it needs no # decoder config — emit before the config requirement below. - path = write_audio_codec_pipeline_metadata( - output_dir, codes_dtype=_audio_codec_codes_dtype(pkg) + path = write_audio_codec_workflow_metadata(pkg, output_dir) + return {"inference_metadata": path} + + if _looks_like_ctc_asr(pkg): + # CTC ASR is frame-synchronous: the encoder runs once and the transcript + # comes from the profile's decoding contract, so no decoder/KV metadata + # is produced. + ctc_config = config if config is not None else getattr(pkg, "config", None) + if ctc_config is None: + raise ValueError( + "CTC ASR metadata requires a model config (pass config=... or a " + "package carrying `.config`)" + ) + path = write_ctc_asr_workflow_metadata(pkg, output_dir, ctc_config, source=source) + artifacts = {"inference_metadata": path} + tokenizer_path = _write_hf_tokenizer(output_dir, source, revision=revision) + if tokenizer_path is not None: + artifacts["tokenizer"] = tokenizer_path + return artifacts + + if _looks_like_encoder_embedding(pkg): + # A bidirectional encoder has no logits and no cache: it runs once and + # returns one hidden vector per position. Emit before the config + # requirement below because nothing here needs a decoder config. + path = write_encoder_embedding_workflow_metadata(pkg, output_dir, config) + artifacts = {"inference_metadata": path} + tokenizer_path = _write_hf_tokenizer(output_dir, source, revision=revision) + if tokenizer_path is not None: + artifacts["tokenizer"] = tokenizer_path + return artifacts + + if _looks_like_speculative(pkg): + if kv_native_dtype is not None: + raise ValueError( + "workflow speculative export derives KV state dtype from ONNX ports; " + "kv_native_dtype overrides are unsupported" + ) + path = write_speculative_workflow_metadata( + pkg, + output_dir, + grammar_guidance=grammar_guidance, + adaptive_k_max=adaptive_k_max, ) return {"inference_metadata": path} @@ -520,92 +905,80 @@ def write_onnx_genai_config( "or a package carrying `.config`)" ) if _looks_like_multimodal(pkg): - derived = _multimodal_component_kwargs(pkg) - for name, value in derived.items(): - kwargs.setdefault(name, value) - decoder_metadata = decoder_metadata_from_config( - resolved_config, kv_native_dtype=kv_native_dtype - ) - path = write_multimodal_pipeline_metadata( + if kv_native_dtype is not None: + raise ValueError( + "workflow VLM export derives KV state dtype from ONNX ports; " + "kv_native_dtype overrides are unsupported" + ) + path = write_vlm_workflow_metadata( + pkg, output_dir, - decoder_metadata=decoder_metadata, - activation_dtype=_activation_dtype_tag(resolved_config), - **kwargs, + resolved_config, + source=source, ) - _add_explicit_io_to_file(path, pkg, resolved_config) artifacts = {"inference_metadata": path} - tokenizer_path = _write_hf_tokenizer(output_dir, source, revision=revision) - if tokenizer_path is not None: - artifacts["tokenizer"] = tokenizer_path - if "audio_encoder" in pkg: + # A multimodal package needs the processor assets as well as the + # tokenizer, because the runtime resolves image/audio preprocessing + # parameters from them. + artifacts.update(_copy_runtime_assets(output_dir, source, revision=revision)) + if "tokenizer" not in artifacts: + tokenizer_path = _write_hf_tokenizer(output_dir, source, revision=revision) + if tokenizer_path is not None: + artifacts["tokenizer"] = tokenizer_path + # A speech-language package fuses audio embeddings, so it needs the + # feature extractor too: the runtime cannot turn a waveform into the + # encoder's input without it, and no other asset carries those + # parameters. + if _has_audio_encoder(pkg): audio_processor_path = _write_hf_audio_processor( - output_dir, - source, - revision=revision, + output_dir, source, revision=revision ) if audio_processor_path is not None: artifacts["audio_processor"] = audio_processor_path return artifacts if _looks_like_speech_to_text(pkg): - encoder_outputs = {value.name for value in pkg["encoder"].graph.outputs} - decoder_inputs = {value.name for value in pkg["decoder"].graph.inputs} - kwargs.setdefault( - "encoder_attention_mask", - "encoder_attention_mask" in encoder_outputs - and "encoder_attention_mask" in decoder_inputs, - ) - decoder_metadata = decoder_metadata_from_config( - resolved_config, kv_native_dtype=kv_native_dtype - ) - path = write_speech_to_text_pipeline_metadata( + if kv_native_dtype is not None: + raise ValueError( + "workflow speech-to-text export derives KV state dtype from ONNX ports; " + "kv_native_dtype overrides are unsupported" + ) + audio_processor_path = _write_hf_audio_processor(output_dir, source, revision=revision) + path = write_speech_to_text_workflow_metadata( + pkg, output_dir, - decoder_metadata=decoder_metadata, - activation_dtype=_activation_dtype_tag(resolved_config), - **kwargs, + resolved_config, + audio_preprocessing=_audio_preprocessing_program( + audio_processor_path, pkg["encoder"] + ), ) - _add_explicit_io_to_file(path, pkg, resolved_config) artifacts = {"inference_metadata": path} - tokenizer_path = _write_hf_tokenizer(output_dir, source, revision=revision) - if tokenizer_path is not None: - artifacts["tokenizer"] = tokenizer_path - audio_processor_path = _write_hf_audio_processor( - output_dir, - source, - revision=revision, - ) + # An ASR decoder is still a text producer: ship its tokenizer and chat + # template alongside the audio processor. + artifacts.update(_write_text_runtime_assets(output_dir, source, revision=revision)) if audio_processor_path is not None: artifacts["audio_processor"] = audio_processor_path return artifacts - # A nested multi-decoder TTS stack (talker + code_predictor) uses the - # nested_autoregressive strategy. When the package also carries the - # `talker_step_embedder` pre-embedder (the real Qwen3-TTS shape), emit the - # pre-embedder-driven contract the onnx-genai runtime executes; otherwise the - # component graph is not yet mappable, so fail with a precise, actionable error. + # A nested multi-decoder TTS stack requires the generic workflow loop to expose + # its induction value. The current producer contract cannot wire step_index or + # per-group embedding selection without host preprocessing, so the workflow + # writer reports that exact contract defect. if _looks_like_multi_decoder_tts(pkg): + if kv_native_dtype is not None: + raise ValueError( + "workflow TTS export derives KV state dtype from ONNX ports; " + "kv_native_dtype overrides are unsupported" + ) if not _has_tts_pre_embedder(pkg): raise NotImplementedError( "Multi-decoder TTS packages (talker + code_predictor, e.g. Qwen3-TTS) " - "use the nested_autoregressive strategy. This package lacks the " + "require nested generic workflow loops. This package lacks the " "`talker_step_embedder` pre-embedder that materializes the talker " - "inputs_embeds, so it cannot yet be mapped to the runtime contract — " - "see onnx-genai docs/DESIGN.md §20.3 'Multi-decoder TTS'." + "inputs_embeds, so it cannot be mapped to the workflow contract." ) - decoder_metadata = decoder_metadata_from_config( - resolved_config, kv_native_dtype=kv_native_dtype - ) - path = write_tts_pipeline_metadata( - output_dir, - decoder_metadata=decoder_metadata, - **_tts_component_kwargs(pkg, resolved_config), - ) - _add_explicit_io_to_file(path, pkg, resolved_config) - artifacts = {"inference_metadata": path} - tokenizer_path = _write_hf_tokenizer(output_dir, source, revision=revision) - if tokenizer_path is not None: - artifacts["tokenizer"] = tokenizer_path - return artifacts + path = write_tts_workflow_metadata(pkg, output_dir, resolved_config) + return {"inference_metadata": path} # Fallback: a single-component decoder language model. A multi-component # package that matched none of the composite shapes above would be silently @@ -623,12 +996,17 @@ def write_onnx_genai_config( "Multi-decoder pipelines such as TTS require a dedicated emitter." ) - path = write_decoder_metadata( - output_dir, config=resolved_config, kv_native_dtype=kv_native_dtype + if kv_native_dtype is not None: + raise ValueError( + "workflow decoder export derives KV state dtype from ONNX ports; " + "kv_native_dtype overrides are unsupported" + ) + path = write_decoder_workflow_metadata( + pkg, + output_dir, + resolved_config, + sampler=str(getattr(resolved_config, "workflow_sampler", "greedy")), ) - _add_explicit_io_to_file(path, pkg, resolved_config) artifacts = {"inference_metadata": path} - tokenizer_path = _write_hf_tokenizer(output_dir, source, revision=revision) - if tokenizer_path is not None: - artifacts["tokenizer"] = tokenizer_path + artifacts.update(_write_text_runtime_assets(output_dir, source, revision=revision)) return artifacts diff --git a/src/mobius/integrations/onnx_genai/auto_export_test.py b/src/mobius/integrations/onnx_genai/auto_export_test.py index c3e9f075d..c246ad6c2 100644 --- a/src/mobius/integrations/onnx_genai/auto_export_test.py +++ b/src/mobius/integrations/onnx_genai/auto_export_test.py @@ -14,7 +14,21 @@ import yaml from mobius._configs import QuantizationConfig +from mobius._model_package import ModelPackage from mobius.integrations.onnx_genai import write_onnx_genai_config +from mobius.integrations.onnx_genai.auto_export import ( + _flow_match_euler_schedule, + _looks_like_image_edit, +) +from mobius.integrations.onnx_genai.inference_metadata import SchedulerConfig +from mobius.integrations.onnx_genai.inference_metadata_test import ( + _decoder_model, + _model, + _value, +) +from mobius.integrations.onnx_genai.workflow_metadata import ( + build_decoder_workflow_metadata, +) @dataclasses.dataclass @@ -40,20 +54,315 @@ class _DiffusionPkg(dict): pass +def _diffusion_package(*, text: bool = False): + latent = ["batch", 4, "height", "width"] + denoiser_inputs = [ + _value("sample", ir.DataType.FLOAT, latent), + _value("timestep", ir.DataType.FLOAT, ["batch"]), + ] + components = {} + if text: + denoiser_inputs.append( + _value( + "encoder_hidden_states", + ir.DataType.FLOAT, + ["batch", "prompt_sequence", 32], + ) + ) + components["text_encoder"] = _model( + "text_encoder", + [_value("input_ids", ir.DataType.INT64, ["batch", "prompt_sequence"])], + [ + ( + "encoder_hidden_states", + ir.DataType.FLOAT, + ["batch", "prompt_sequence", 32], + ) + ], + ) + denoiser = _model( + "denoiser", + denoiser_inputs, + [("noise_pred", ir.DataType.FLOAT, latent)], + ) + vae = _model( + "vae_decoder", + [_value("latent", ir.DataType.FLOAT, latent)], + [("image", ir.DataType.FLOAT, ["batch", 3, "image_height", "image_width"])], + ) + components.update({"denoiser": denoiser, "vae_decoder": vae}) + return ModelPackage(components) + + class _MultimodalPkg(dict): config = _Cfg() +@dataclasses.dataclass +class _VisionCfg: + patch_size: int = 14 + temporal_patch_size: int = 2 + merge_size: int = 1 + spatial_merge_size: int = 1 + size: dict[str, int] = dataclasses.field( + default_factory=lambda: {"shortest_edge": 224, "longest_edge": 224} + ) + + +@dataclasses.dataclass +class _VlmCfg(_Cfg): + vision: _VisionCfg = dataclasses.field(default_factory=_VisionCfg) + image_token_id: int = 32000 + eos_token_id: int = 2 + + +def _vlm_package(*, audio: bool = False): + vision = _model( + "vision_encoder", + [ + _value("pixel_values", ir.DataType.FLOAT, ["patches", 1176]), + _value("grid_thw", ir.DataType.INT64, ["images", 3]), + ], + [("image_features", ir.DataType.FLOAT, ["batch", 256, 32])], + ) + embedding_inputs = [ + _value("input_ids", ir.DataType.INT64, ["batch", "sequence"]), + _value("image_features", ir.DataType.FLOAT, ["batch", 256, 32]), + ] + components = {"vision_encoder": vision} + if audio: + components["audio_encoder"] = _model( + "audio_encoder", + [_value("input_features", ir.DataType.FLOAT, ["batch", 80, "frames"])], + [("audio_features", ir.DataType.FLOAT, ["batch", 64, 32])], + ) + embedding_inputs.append(_value("audio_features", ir.DataType.FLOAT, ["batch", 64, 32])) + embedding = _model( + "embedding", + embedding_inputs, + [("inputs_embeds", ir.DataType.FLOAT, ["batch", "sequence", 32])], + ) + decoder = _decoder_model( + [("inputs_embeds", ir.DataType.FLOAT, ["batch", "sequence", 32])], + position_shape=["batch", "sequence"], + ) + components.update({"embedding": embedding, "decoder": decoder}) + return ModelPackage(components, config=_VlmCfg()) + + +def _decoder_package(config=None): + model = _decoder_model( + [], + position_shape=["batch", "sequence"], + raw_token_input=True, + ) + return ModelPackage({"model": model}, config=config or _Cfg()) + + def test_dispatch_decoder(tmp_path): - arts = write_onnx_genai_config(object(), str(tmp_path), config=_Int4Cfg()) + package = _decoder_package(_Int4Cfg()) + arts = write_onnx_genai_config(package, str(tmp_path), config=_Int4Cfg()) with open(arts["inference_metadata"]) as handle: meta = yaml.safe_load(handle) - assert meta["model"]["attention"]["type"] == "grouped_query_attention" - assert meta["kv_cache"]["native_dtype"] == "float16" + workflow = meta["pipeline"]["workflow"] + assert "ir_version" not in workflow["manifest"] + assert "onnx_opsets" not in workflow["manifest"] + assert workflow["components"]["token_sampler"]["contract"]["id"] == ( + "onnx-genai.token-sampler" + ) + assert workflow["components"]["termination"]["contract"]["id"] == ( + "onnx-genai.termination-predicate" + ) + assert workflow["steps"][0]["kind"] == "loop" + application_inputs = { + name + for name, value in workflow["inputs"].items() + if value["source"]["kind"] == "application" + } + assert application_inputs == { + "request.prompt_lengths", + "request.eos_ids", + "request.eos_lengths", + "request.row_max_iterations", + "request.rng_counter", + } + assert workflow["inputs"]["request.prompt_lengths"]["default"] == -1 + assert [node["component"] for node in workflow["steps"][0]["setup"]] == [ + "decoder_state_initializer", + "model", + "termination_batch_initializer", + "last_token_logits", + ] + body = workflow["steps"][0]["steps"] + assert [node["kind"] for node in body].count("emit") == 1 + assert next(node for node in body if node["kind"] == "emit")["value"] == "token.body" + assert workflow["outputs"]["tokens"]["contract"]["shape"] == [ + "batch", + "generated_sequence", + ] + assert workflow["steps"][0]["iteration"] == { + "value": "loop.iteration", + "contract": {"dtype": "int64", "rank": 1, "shape": [1]}, + } + assert "iteration" not in workflow["state"] + serialized = yaml.safe_dump(workflow) + assert "initial_effects" not in serialized + assert "read_effect" not in serialized + assert "write_effect" not in serialized + assert ".read" not in serialized + assert "iteration_increment" not in workflow["components"] + assert workflow["state"]["token"]["initializer"] == "initializer.token_slot" + assert workflow["state"]["logits"] == { + "contract": { + "dtype": "float32", + "rank": 2, + "shape": ["batch", 128], + "batch_layout": {"kind": "request_aligned", "axis": 0}, + }, + "scope": "invocation", + "initializer": "decoder.setup.last_logits", + "recurrence": {"kind": "invariant"}, + } + assert (tmp_path / "policies" / "token_sampler.onnx").is_file() + + +def test_seeded_decoder_sampler_uses_request_controls_and_direct_kv_carry(): + workflow = build_decoder_workflow_metadata( + _decoder_package(), _Cfg(), sampler="seeded_categorical" + )["pipeline"]["workflow"] + sampler = workflow["components"]["token_sampler"] + assert sampler["application_overridable"] is True + assert sampler["contract"]["version"] == "2" + assert sampler["contract"]["parameters"] == { + "mode": "seeded_stochastic", + "batching": "per_row", + "inactive_rows": "preserve", + } + assert sampler["contract"]["bindings"] == { + "logits": "logits", + "token": "token", + "temperature": "temperature", + "top_k": "top_k", + "top_p": "top_p", + "min_p": "min_p", + "seed": "seed", + "counter": "counter", + "next_counter": "next_counter", + "active": "active", + "done": "done", + } + assert sampler["ports"]["inputs"]["logits"] == { + "dtype": "float32", + "rank": 2, + "shape": ["batch", "vocabulary"], + "batch_layout": {"kind": "request_aligned", "axis": 0}, + } + assert sampler["ports"]["outputs"]["token"] == { + "dtype": "int64", + "rank": 1, + "shape": ["batch"], + "batch_layout": {"kind": "request_aligned", "axis": 0}, + } + sampler_step = next( + step + for step in workflow["steps"][0]["steps"] + if step.get("component") == "token_sampler" + ) + assert sampler_step["inputs"] == { + "logits": "logits", + "temperature": "request.temperature", + "top_k": "request.top_k", + "top_p": "request.top_p", + "min_p": "request.min_p", + "seed": "request.seed", + "counter": "rng_counter", + "active": "active", + "done": "done", + } + for name in ("temperature", "top_k", "top_p", "min_p"): + assert workflow["inputs"][f"request.{name}"]["contract"]["shape"] == ["batch"] + assert workflow["inputs"]["request.prompt_lengths"]["contract"]["shape"] == ["batch"] + assert workflow["inputs"]["request.max_iterations"]["contract"]["shape"] == [1] + assert workflow["inputs"]["request.eos_ids"]["contract"]["shape"] == [ + "batch", + "num_eos", + ] + assert workflow["state"]["rng_counter"]["class"] == "semantic" + assert workflow["state"]["rng_counter"]["initializer"] == "request.rng_counter" + emit = next(node for node in workflow["steps"][0]["steps"] if node["kind"] == "emit") + assert "row_ids" not in emit + assert "emit_row_identity" not in workflow["manifest"]["capabilities"] + assert set( + workflow["serving"]["state_service"]["groups"]["decoder_cache"]["ports"]["model"] + ) == {"cache_0", "cache_1"} + assert workflow["components"]["termination"]["contract"]["version"] == "2" + assert workflow["components"]["termination"]["contract"]["parameters"] == { + "batching": "per_row", + "inactive_rows": "preserve", + } + assert set(workflow["components"]["termination"]["contract"]["bindings"]) == { + "tokens", + "active", + "eos_ids", + "eos_lengths", + "iteration", + "max_iterations", + "done", + "next_active", + "continue", + } + assert workflow["components"]["termination"]["ports"]["inputs"]["iteration"] == { + "dtype": "int64", + "rank": 1, + "shape": [1], + } + assert workflow["components"]["token_state_update"]["contract"] == { + "id": "onnx-genai.state-update", + "version": "2", + "bindings": { + "current": "current", + "update": "update", + "active": "active", + "done": "done", + "next": "next", + }, + "parameters": { + "batching": "per_row", + "inactive_rows": "preserve", + }, + } + assert not any("kv_update" in name for name in workflow["components"]) + + +def test_dispatch_language_diffusion(tmp_path): + package = ModelPackage( + { + "model": _model( + "masked_denoiser", + [_value("input_ids", ir.DataType.INT64, ["batch", "sequence"])], + [ + ("logits", ir.DataType.FLOAT, ["batch", "sequence", 128]), + ("proposed_tokens", ir.DataType.INT64, ["batch", "sequence"]), + ], + ) + }, + config=_Cfg(model_type="llada"), + ) + artifacts = write_onnx_genai_config( + package, + str(tmp_path), + num_inference_steps=12, + ) + with open(artifacts["inference_metadata"]) as handle: + metadata = yaml.safe_load(handle) + pipeline = metadata["pipeline"] + assert set(pipeline) == {"workflow"} + assert pipeline["workflow"]["inputs"]["request.max_iterations"]["default"] == 12 + assert (tmp_path / "policies" / "masked_update.onnx").is_file() def test_dispatch_diffusion(tmp_path): - pkg = _DiffusionPkg({"denoiser": object(), "vae": object()}) + pkg = _diffusion_package() arts = write_onnx_genai_config( pkg, str(tmp_path), @@ -62,65 +371,223 @@ def test_dispatch_diffusion(tmp_path): ) with open(arts["inference_metadata"]) as handle: meta = yaml.safe_load(handle) - assert meta["pipeline"]["strategy"]["kind"] == "iterative" - assert "vae" in meta["pipeline"]["models"] + workflow = meta["pipeline"]["workflow"] + assert workflow["steps"][0]["iteration"]["value"] == "loop.iteration" + assert workflow["steps"][1]["component"] == "vae_decoder" + assert "strategy" not in meta["pipeline"] + assert (tmp_path / "policies" / "solver_step.onnx").is_file() + assert (tmp_path / "policies" / "schedule_lookup.onnx").is_file() -def test_single_diffusion_component_uses_flat_model_path(tmp_path): +def test_single_diffusion_component_requires_explicit_vae(tmp_path): pkg = _DiffusionPkg({"transformer": object()}) - artifacts = write_onnx_genai_config(pkg, str(tmp_path), num_inference_steps=2) - with open(artifacts["inference_metadata"], encoding="utf-8") as handle: - metadata = yaml.safe_load(handle) - assert metadata["pipeline"]["models"]["denoiser"]["filename"] == "model.onnx" + with pytest.raises( + ValueError, + match="diffusion workflow requires distinct denoiser and VAE decoder", + ): + write_onnx_genai_config(pkg, str(tmp_path), num_inference_steps=2) + + +def _image_edit_package(): + """Build a Qwen-Image-Edit-shaped package with the real port contract. + + Mirrors ``QwenImageTask``: rank-3 packed latents, a ``target_sequence_length`` + slice port, separate image/text rotary tables, and a VAE encoder/decoder pair. + """ + tokens = ["batch", "image_sequence_length", 64] + transformer = _model( + "transformer", + [ + _value("sample", ir.DataType.FLOAT, tokens), + _value("timestep", ir.DataType.FLOAT, ["batch"]), + _value( + "encoder_hidden_states", + ir.DataType.FLOAT, + ["batch", "text_sequence_length", 32], + ), + _value( + "encoder_hidden_states_mask", + ir.DataType.BOOL, + ["batch", "text_sequence_length"], + ), + _value("image_rotary_cos", ir.DataType.FLOAT, ["image_sequence_length", 8]), + _value("image_rotary_sin", ir.DataType.FLOAT, ["image_sequence_length", 8]), + _value("text_rotary_cos", ir.DataType.FLOAT, ["text_sequence_length", 8]), + _value("text_rotary_sin", ir.DataType.FLOAT, ["text_sequence_length", 8]), + _value("target_sequence_length", ir.DataType.INT64, [1]), + ], + [("noise_pred", ir.DataType.FLOAT, tokens)], + ) + vae_encoder = _model( + "vae_encoder", + [_value("pixel_values", ir.DataType.FLOAT, ["batch", 3, 1, "height", "width"])], + [("latent_sample", ir.DataType.FLOAT, ["batch", 16, 1, "lheight", "lwidth"])], + ) + vae_decoder = _model( + "vae_decoder", + [_value("latent_sample", ir.DataType.FLOAT, ["batch", 16, 1, "lheight", "lwidth"])], + [("image", ir.DataType.FLOAT, ["batch", 3, 1, "height", "width"])], + ) + return ModelPackage( + { + "transformer": transformer, + "vae_encoder": vae_encoder, + "vae_decoder": vae_decoder, + } + ) + + +_FLOW_MATCH_SCHEDULER = { + "_class_name": "FlowMatchEulerDiscreteScheduler", + "base_image_seq_len": 256, + "max_image_seq_len": 8192, + "base_shift": 0.5, + "max_shift": 0.9, + "shift_terminal": 0.02, + "time_shift_type": "exponential", + "use_dynamic_shifting": True, +} + + +def _write_scheduler(source, config=None): + (source / "scheduler").mkdir(parents=True, exist_ok=True) + (source / "scheduler" / "scheduler_config.json").write_text( + json.dumps(_FLOW_MATCH_SCHEDULER if config is None else config), encoding="utf-8" + ) + + +def test_detects_image_edit_package_structurally(): + """Structural detection must not depend on model_type strings.""" + assert _looks_like_image_edit(_image_edit_package()) + # A plain latent-diffusion package has no VAE encoder and rank-4 samples. + assert not _looks_like_image_edit(_diffusion_package(text=True)) + # A VAE pair alone is not enough without the target-slice denoiser port. + pkg = _image_edit_package() + del pkg["transformer"] + assert not _looks_like_image_edit(pkg) + +def test_flow_match_schedule_matches_diffusers(): + """Pin the schedule against the captured Qwen-Image-Edit-2509 reference. + + These are the timesteps ``QwenImageEditPlusPipeline`` produced for the + 1216x864 reference edit (``image_seq_len=4104``, 8 steps), divided by + ``num_train_timesteps`` because the denoiser consumes normalized sigmas. + """ + scheduler = SchedulerConfig( + kind="flow_match_euler", + use_dynamic_shifting=True, + base_image_seq_len=256, + max_image_seq_len=8192, + base_shift=0.5, + max_shift=0.9, + shift_terminal=0.02, + time_shift_type="exponential", + ) + timesteps, sigmas = _flow_match_euler_schedule(scheduler, 8, 4104) + expected = [ + 1.0, + 0.9160475, + 0.8200923, + 0.7093592, + 0.5801504, + 0.4274216, + 0.2441077, + 0.02, + ] + assert timesteps == pytest.approx(expected, abs=1e-6) + assert sigmas == pytest.approx([*expected, 0.0], abs=1e-6) -def test_rejects_unsupported_qwen_image_edit_runtime_export(tmp_path): + +def test_flow_match_schedule_rejects_wrong_scheduler(): + with pytest.raises(ValueError, match="flow-match Euler scheduler"): + _flow_match_euler_schedule(SchedulerConfig(kind="euler"), 4, 4104) + + +def test_dispatch_image_edit_emits_workflow(tmp_path): + """A Qwen-Image-Edit-shaped package must dispatch to the image-edit workflow. + + Asserts the emitted pipeline actually performs the edit: encode the source + image, run two guided denoiser passes per step, combine them with true CFG, + and decode the target tokens back to pixels. + """ source = tmp_path / "source" output = tmp_path / "output" - (source / "scheduler").mkdir(parents=True) - (source / "processor").mkdir() - (source / "scheduler" / "scheduler_config.json").write_text( - json.dumps( - { - "_class_name": "FlowMatchEulerDiscreteScheduler", - "base_image_seq_len": 256, - "max_image_seq_len": 8192, - "base_shift": 0.5, - "max_shift": 0.9, - "use_dynamic_shifting": True, - } - ), - encoding="utf-8", + _write_scheduler(source) + arts = write_onnx_genai_config( + _image_edit_package(), + str(output), + source=str(source), + num_inference_steps=8, + image_seq_len=4104, + guidance_scale=4.0, ) - for filename in ( - "preprocessor_config.json", - "video_preprocessor_config.json", - "tokenizer_config.json", - "tokenizer.json", - "chat_template.jinja", - ): - (source / "processor" / filename).write_text("{}", encoding="utf-8") + with open(arts["inference_metadata"]) as handle: + meta = yaml.safe_load(handle) + workflow = meta["pipeline"]["workflow"] + + loop = next(step for step in workflow["steps"] if step["kind"] == "loop") + + # The source image is encoded and packed once, in the loop's setup block. + setup = [node["component"] for node in loop["setup"]] + assert setup.index("vae_encoder") < setup.index("pack_latents") + + body = [node["component"] for node in loop["steps"]] + # True CFG needs both a positive and a negative denoiser pass per step. + assert body.count("transformer") == 2 + assert "true_cfg" in body + assert "sequence_concat" in body + assert body.index("true_cfg") < body.index("solver_step") + assert loop["max_iterations"] == "request.max_iterations" + + # Source and target tokens are separate: only the target block is carried, + # so the loop state's sequence axis is the target length, not the denoiser's + # concatenated target+source length. + assert [cell["cell"] for cell in loop["carried"]] == ["latent", "loop_0_active"] + assert workflow["state"]["latent"]["contract"]["shape"] == [ + "batch", + "target_sequence_length", + 64, + ] - pkg = _DiffusionPkg( - { - "transformer": object(), - "text_encoder": object(), - "text_encoder_vision_encoder": object(), - "text_encoder_embedding": object(), - "vae_encoder": object(), - "vae_decoder": object(), - } + tail = [step["component"] for step in workflow["steps"] if step["kind"] == "invoke"] + assert tail == ["unpack_latents", "vae_decoder"] + + # Positive and negative conditioning cannot share a sequence contract. + inputs = workflow["inputs"] + assert ( + inputs["request.positive_encoder_hidden_states"]["contract"]["shape"][1] + == "positive_text_sequence_length" ) - pkg.config = SimpleNamespace( - model_type="qwen_image_edit", - processor_config={"patch_size": 14, "merge_size": 2}, + assert ( + inputs["request.negative_encoder_hidden_states"]["contract"]["shape"][1] + == "negative_text_sequence_length" ) - with pytest.raises(ValueError, match="cannot execute Qwen Image Edit"): + + for policy in ("solver_step", "true_cfg", "pack_latents", "unpack_latents"): + assert (output / "policies" / f"{policy}.onnx").is_file() + + +def test_image_edit_requires_image_seq_len(tmp_path): + """The schedule is resolution dependent, so it cannot be guessed.""" + source = tmp_path / "source" + _write_scheduler(source) + with pytest.raises(ValueError, match="requires image_seq_len"): write_onnx_genai_config( - pkg, - str(output), + _image_edit_package(), + str(tmp_path / "output"), source=str(source), - num_inference_steps=3, + num_inference_steps=8, + ) + + +def test_image_edit_requires_scheduler_config(tmp_path): + with pytest.raises(ValueError, match="requires the diffusers scheduler config"): + write_onnx_genai_config( + _image_edit_package(), + str(tmp_path / "output"), + num_inference_steps=8, + image_seq_len=4104, ) @@ -142,7 +609,7 @@ class _Tokenizer: monkeypatch.setattr( "transformers.AutoTokenizer.from_pretrained", lambda *args, **kwargs: _Tokenizer() ) - pkg = _DiffusionPkg({"denoiser": object(), "text_encoder": object(), "vae": object()}) + pkg = _diffusion_package(text=True) arts = write_onnx_genai_config( pkg, str(tmp_path), @@ -150,6 +617,7 @@ class _Tokenizer: vae_filename="vae.onnx", text_encoder_filename="text_encoder.onnx", source="fake/model", + guidance_scale=1.0, ) assert "tokenizer" in arts assert os.path.basename(arts["tokenizer"]) == "tokenizer.json" @@ -167,13 +635,14 @@ def _boom(*args, **kwargs): raise OSError("no tokenizer here") monkeypatch.setattr("transformers.AutoTokenizer.from_pretrained", _boom) - pkg = _DiffusionPkg({"denoiser": object(), "text_encoder": object()}) + pkg = _diffusion_package(text=True) arts = write_onnx_genai_config( pkg, str(tmp_path), num_inference_steps=20, text_encoder_filename="text_encoder.onnx", source="fake/model", + guidance_scale=1.0, ) assert "inference_metadata" in arts assert "tokenizer" not in arts @@ -188,7 +657,7 @@ def test_dispatch_diffusion_auto_reads_scheduler_from_source(tmp_path): json.dumps({"_class_name": "EulerDiscreteScheduler", "beta_schedule": "scaled_linear"}) ) out = tmp_path / "out" - pkg = _DiffusionPkg({"denoiser": object()}) + pkg = _diffusion_package() arts = write_onnx_genai_config( pkg, str(out), @@ -197,56 +666,175 @@ def test_dispatch_diffusion_auto_reads_scheduler_from_source(tmp_path): ) with open(arts["inference_metadata"]) as handle: meta = yaml.safe_load(handle) - assert meta["pipeline"]["strategy"]["scheduler_config"]["kind"] == "euler" + components = meta["pipeline"]["workflow"]["components"] + assert components["diffusion_schedule"]["ports"]["outputs"]["schedule"] == { + "dtype": "float32", + "rank": 1, + "shape": [16], + } + schedule = ir.load(out / "policies" / "diffusion_schedule.onnx") + assert list(schedule.graph.outputs[0].shape) == [16] def test_dispatch_vision_multimodal_pipeline(tmp_path): - pkg = _MultimodalPkg( - { - "decoder": object(), - "vision_encoder": object(), - "embedding": object(), - } - ) - artifacts = write_onnx_genai_config(pkg, str(tmp_path), kv_native_dtype="bf16") + pkg = _vlm_package() + artifacts = write_onnx_genai_config(pkg, str(tmp_path)) with open(artifacts["inference_metadata"]) as handle: metadata = yaml.safe_load(handle) - assert metadata["required_capabilities"] == [ - "kv_cache", - "grouped_query_attention", - ] - assert metadata["kv_cache"] == {"native_dtype": "bfloat16"} pipeline = metadata["pipeline"] - assert pipeline["models"] == { - "vision_encoder": { - "filename": "vision_encoder/model.onnx", - "type": "vision_encoder", - }, - "embedding": { - "filename": "embedding/model.onnx", - "type": "encoder", - }, - "decoder": { - "filename": "decoder/model.onnx", - "type": "decoder", - "tokenizer": "tokenizer.json", - }, + assert set(pipeline) == {"workflow"} + workflow = pipeline["workflow"] + assert workflow["manifest"]["adapter_abis"] == {"onnx-genai.image-preprocess": "1"} + assert workflow["steps"][0]["setup"][0]["component"] == "image_preprocess" + assert workflow["steps"][0]["setup"][1]["component"] == "vision_encoder" + assert workflow["steps"][0]["setup"][4]["component"] == "embedding" + assert workflow["steps"][0]["iteration"]["value"] == "loop.iteration" + assert workflow["state"]["logits"]["contract"] == { + "dtype": "float32", + "rank": 2, + "shape": ["batch", 128], + "batch_layout": {"kind": "request_aligned", "axis": 0}, } - assert pipeline["strategy"]["kind"] == "composite" + assert workflow["state"]["logits"]["initializer"] == "decoder.setup.last_logits" + assert (tmp_path / "policies" / "token_sampler.onnx").is_file() -def test_dispatch_audio_only_multimodal_pipeline(tmp_path, monkeypatch): +def test_workflow_vlm_rejects_kv_dtype_override(tmp_path): + with pytest.raises(ValueError, match="kv_native_dtype overrides are unsupported"): + write_onnx_genai_config( + _vlm_package(), + str(tmp_path), + kv_native_dtype="bf16", + ) + + +def test_text_diffusion_requires_explicit_guidance_scale(tmp_path): + with pytest.raises(ValueError, match="must declare its guidance"): + write_onnx_genai_config( + _diffusion_package(text=True), + str(tmp_path), + ) + + +def test_guidance_scale_requires_a_text_conditioned_package(tmp_path): + with pytest.raises(ValueError, match="requires a text-conditioned diffusion package"): + write_onnx_genai_config( + _diffusion_package(), + str(tmp_path), + guidance_scale=7.5, + ) + + +def test_dispatch_guided_diffusion_runs_the_denoiser_twice(tmp_path): + arts = write_onnx_genai_config( + _diffusion_package(text=True), + str(tmp_path), + num_inference_steps=4, + guidance_scale=7.5, + ) + with open(arts["inference_metadata"]) as handle: + meta = yaml.safe_load(handle) + workflow = meta["pipeline"]["workflow"] + setup = workflow["steps"][0]["setup"] + # The prompt and the negative prompt each get their own conditioning pass. + assert [node.get("component") for node in setup].count("text_encoder") == 2 + negative = next( + node + for node in setup + if node["inputs"].get("input_ids") == "request.negative_input_ids" + ) + assert negative["outputs"]["encoder_hidden_states"] == "conditioning.unconditional" + body = workflow["steps"][0]["steps"] + denoiser_calls = [node for node in body if node.get("component") == "denoiser"] + assert [node["outputs"]["noise_pred"] for node in denoiser_calls] == [ + "denoiser.unconditional", + "denoiser.conditional", + ] + combine = next(node for node in body if node.get("component") == "guidance_combine") + assert combine["inputs"] == { + "unconditional": "denoiser.unconditional", + "conditional": "denoiser.conditional", + "scale": "request.guidance_scale", + } + assert combine["outputs"] == {"estimate": "denoiser.estimate"} + assert (tmp_path / "policies" / "guidance_combine.onnx").is_file() + + +def test_dispatch_multistep_diffusion_carries_solver_history(tmp_path): + import json + + source = tmp_path / "ckpt" + (source / "scheduler").mkdir(parents=True) + (source / "scheduler" / "scheduler_config.json").write_text( + json.dumps( + { + "_class_name": "DPMSolverMultistepScheduler", + "beta_schedule": "scaled_linear", + "algorithm_type": "dpmsolver++", + "solver_order": 2, + "solver_type": "midpoint", + "lower_order_final": True, + "final_sigmas_type": "zero", + } + ) + ) + arts = write_onnx_genai_config( + _diffusion_package(text=True), + str(tmp_path / "out"), + num_inference_steps=4, + guidance_scale=1.0, + source=str(source), + ) + with open(arts["inference_metadata"]) as handle: + meta = yaml.safe_load(handle) + workflow = meta["pipeline"]["workflow"] + assert "history" in workflow["state"] + carried = {carry["cell"]: carry["next"] for carry in workflow["steps"][0]["carried"]} + assert "history.body" in carried.values() + body = workflow["steps"][0]["steps"] + solver = next(node for node in body if node.get("component") == "solver_step") + history_cell = next(cell for cell, value in carried.items() if value == "history.body") + assert solver["inputs"]["history"] == history_cell + assert solver["outputs"]["next_history"] == "history.body" + # A variance-preserving sampler feeds its state to the denoiser untouched, so + # there is no model-input rescaling component at all. + assert not any(node.get("component") == "model_input_scale" for node in body) + latent_cell = next(cell for cell, value in carried.items() if value == "latent.body") + denoiser = next(node for node in body if node.get("component") == "denoiser") + assert denoiser["inputs"]["sample"] == latent_cell + + +def test_dispatch_audio_only_multimodal_pipeline(tmp_path): # The audio-only fusion shape used by speech-language ASR models such as # qwen3_asr and fun_asr: audio_encoder -> embedding fusion -> AR decoder. - pkg = _MultimodalPkg( - { - "decoder": object(), - "audio_encoder": object(), - "embedding": object(), - } - ) + pkg = _vlm_package(audio=True) + artifacts = write_onnx_genai_config(pkg, str(tmp_path)) + + with open(artifacts["inference_metadata"]) as handle: + metadata = yaml.safe_load(handle) + + setup = metadata["pipeline"]["workflow"]["steps"][0]["setup"] + assert [node["component"] for node in setup[:3]] == [ + "image_preprocess", + "vision_encoder", + "audio_encoder", + ] + embedding = next(node for node in setup if node.get("component") == "embedding") + assert embedding["inputs"]["audio_features"] == "audio.audio_features" + + +def test_audio_package_forwards_the_pinned_revision(tmp_path, monkeypatch): + """A pinned revision must reach the asset writers, not just the build. + + Every runtime asset beside the graph — tokenizer, audio processor — is + fetched from the Hub separately from the weights. If the revision stops + being threaded, the package still builds and still validates, but its + processor silently comes from whatever the branch tip happens to be, which + is the failure a pin exists to prevent. + """ + pkg = _vlm_package(audio=True) audio_processor = tmp_path / "audio_processor.json" audio_processor.write_text("{}") calls: list[tuple[str | None, str | None]] = [] @@ -268,84 +856,21 @@ def fake_audio_processor(output_dir, source, *, revision=None): assert artifacts["audio_processor"] == str(audio_processor) assert calls == [("zai-org/GLM-ASR-Nano-2512", "pinned-revision")] - with open(artifacts["inference_metadata"]) as handle: - metadata = yaml.safe_load(handle) - - pipeline = metadata["pipeline"] - assert pipeline["models"] == { - "audio_encoder": { - "filename": "audio_encoder/model.onnx", - "type": "audio_encoder", - }, - "embedding": {"filename": "embedding/model.onnx", "type": "encoder"}, - "decoder": { - "filename": "decoder/model.onnx", - "type": "decoder", - "tokenizer": "tokenizer.json", - }, - } - assert pipeline["dataflow"] == [ - { - "from": "audio_encoder.audio_features", - "to": "embedding.audio_features", - "dtype": "fp32", - "device_transfer": False, - }, - { - "from": "embedding.inputs_embeds", - "to": "decoder.inputs_embeds", - "dtype": "fp32", - "device_transfer": False, - }, - ] - assert [stage["name"] for stage in pipeline["strategy"]["stages"]] == [ - "encode_audio", - "fuse_embeddings", - "decode", - ] - def test_dispatch_vision_and_audio_multimodal_pipeline(tmp_path): - pkg = _MultimodalPkg( - { - "decoder": object(), - "vision_encoder": object(), - "audio_encoder": object(), - "embedding": object(), - } - ) + pkg = _vlm_package(audio=True) artifacts = write_onnx_genai_config(pkg, str(tmp_path)) with open(artifacts["inference_metadata"]) as handle: metadata = yaml.safe_load(handle) - pipeline = metadata["pipeline"] - assert pipeline["dataflow"] == [ - { - "from": "vision_encoder.image_features", - "to": "embedding.image_features", - "dtype": "fp32", - "device_transfer": False, - }, - { - "from": "audio_encoder.audio_features", - "to": "embedding.audio_features", - "dtype": "fp32", - "device_transfer": False, - }, - { - "from": "embedding.inputs_embeds", - "to": "decoder.inputs_embeds", - "dtype": "fp32", - "device_transfer": False, - }, - ] - assert [stage["name"] for stage in pipeline["strategy"]["stages"]] == [ - "encode_vision", - "encode_audio", - "fuse_embeddings", - "decode", - ] + workflow = metadata["pipeline"]["workflow"] + assert set(workflow["components"]) >= { + "vision_encoder", + "audio_encoder", + "embedding", + "decoder", + } class _FakeValue: @@ -371,110 +896,134 @@ class _EncoderDecoderPkg(dict): config = _Cfg() -def test_dispatch_speech_to_text_pipeline(tmp_path): - # Whisper-style ASR: the decoder consumes encoder_hidden_states (cross-attn). - pkg = _EncoderDecoderPkg( +def _speech_package(*, encoder_mask: bool = False): + """Whisper-shaped encoder/decoder package with a real cross-attention edge.""" + encoder_inputs = [ + _value("input_features", ir.DataType.FLOAT, ["batch", 80, "audio_seq_len"]) + ] + encoder_outputs = [("encoder_hidden_states", ir.DataType.FLOAT, ["batch", 1500, 384])] + decoder_inputs = [ + _value("decoder_input_ids", ir.DataType.INT64, ["batch", "sequence_len"]), + _value("encoder_hidden_states", ir.DataType.FLOAT, ["batch", 1500, 384]), + _value("position_ids", ir.DataType.INT64, ["batch", "sequence_len"]), + _value( + "past_key_values.0.key", + ir.DataType.FLOAT, + ["batch", 6, "past_sequence_len", 64], + ), + _value( + "past_key_values.0.value", + ir.DataType.FLOAT, + ["batch", 6, "past_sequence_len", 64], + ), + ] + decoder_outputs = [ + ("logits", ir.DataType.FLOAT, ["batch", "sequence_len", 51865]), + ("present.0.key", ir.DataType.FLOAT, ["batch", 6, "total_sequence_len", 64]), + ("present.0.value", ir.DataType.FLOAT, ["batch", 6, "total_sequence_len", 64]), + ] + if encoder_mask: + encoder_inputs.append( + _value("attention_mask", ir.DataType.INT64, ["batch", "audio_seq_len"]) + ) + encoder_outputs.append(("encoder_attention_mask", ir.DataType.INT64, ["batch", 1500])) + decoder_inputs.insert( + 2, _value("encoder_attention_mask", ir.DataType.INT64, ["batch", 1500]) + ) + pkg = ModelPackage( { - "encoder": _FakeModel(["input_features"], ["encoder_hidden_states"]), - "decoder": _FakeModel(["decoder_input_ids", "encoder_hidden_states"], ["logits"]), + "encoder": _model("encoder", encoder_inputs, encoder_outputs), + "decoder": _model("decoder", decoder_inputs, decoder_outputs), } ) - artifacts = write_onnx_genai_config(pkg, str(tmp_path), kv_native_dtype="bf16") + pkg.config = SimpleNamespace(eos_token_id=50257, max_position_embeddings=448) + return pkg + + +def test_dispatch_speech_to_text_workflow(tmp_path): + # Whisper-style ASR: the decoder consumes encoder_hidden_states (cross-attn). + artifacts = write_onnx_genai_config(_speech_package(), str(tmp_path)) with open(artifacts["inference_metadata"]) as handle: metadata = yaml.safe_load(handle) - assert metadata["kv_cache"] == {"native_dtype": "bfloat16"} - pipeline = metadata["pipeline"] - assert pipeline["models"]["encoder"]["filename"] == "encoder/model.onnx" - assert pipeline["models"]["encoder"]["type"] == "encoder" - decoder_model = pipeline["models"]["decoder"] - assert decoder_model["filename"] == "decoder/model.onnx" - assert decoder_model["type"] == "decoder" - assert decoder_model["tokenizer"] == "tokenizer.json" - assert decoder_model["io"]["logits_output"] == "logits" - assert decoder_model["io"]["kv_ownership"] == "owned" - assert pipeline["dataflow"] == [ - { - "from": "encoder.encoder_hidden_states", - "to": "decoder.encoder_hidden_states", - "dtype": "fp32", - "device_transfer": False, - } - ] - stages = pipeline["strategy"]["stages"] - assert pipeline["strategy"]["kind"] == "composite" - assert [stage["name"] for stage in stages] == ["encode_audio", "decode_transcript"] - assert [stage["strategy"]["kind"] for stage in stages] == [ - "single_pass", - "autoregressive", - ] + # The redesigned schema has no legacy pipeline description at all. + assert "kv_cache" not in metadata + assert not {"models", "dataflow", "strategy", "phases"}.intersection(metadata["pipeline"]) + workflow = metadata["pipeline"]["workflow"] + assert {"encoder", "decoder"}.issubset(workflow["components"]) + + # The encoder runs once in the loop prologue and its output persists as a + # loop-invariant, request-aligned state cell the decoder reads every step. + loop = next(step for step in workflow["steps"] if step["kind"] == "loop") + setup_components = [node["component"] for node in loop["setup"]] + # The encoder conditions the prefill, so it must precede the decoder. + assert setup_components[0] == "encoder" + assert setup_components.index("decoder") > 0 + cross = workflow["state"]["cross.encoder_hidden_states"] + assert cross["contract"]["shape"] == ["batch", 1500, 384] + assert cross["contract"]["batch_layout"] == {"kind": "request_aligned", "axis": 0} + assert cross["initializer"] == "encoder.encoder_hidden_states" + assert cross["recurrence"] == {"kind": "invariant"} + carry = next( + carry for carry in loop["carried"] if carry["cell"] == "cross.encoder_hidden_states" + ) + assert carry["next"] == "cross.encoder_hidden_states" + + # The self-attention cache is a runtime-served state group; the invariant + # cross state is not, because nothing appends to it. + groups = workflow["serving"]["state_service"]["groups"] + assert set(groups) == {"decoder_cache"} + assert groups["decoder_cache"]["ports"]["decoder"]["cache_0"] == { + "input": "past_key_values.0.key", + "output": "present.0.key", + "role": "key", + "layer": 0, + } -def test_dispatch_speech_to_text_routes_encoder_mask(tmp_path): - pkg = _EncoderDecoderPkg( - { - "encoder": _FakeModel( - ["input_values", "attention_mask"], - ["encoder_hidden_states", "encoder_attention_mask"], - ), - "decoder": _FakeModel( - [ - "decoder_input_ids", - "encoder_hidden_states", - "encoder_attention_mask", - ], - ["logits"], - ), - } - ) - artifacts = write_onnx_genai_config(pkg, str(tmp_path)) +def test_dispatch_speech_to_text_rejects_kv_dtype_override(tmp_path): + with pytest.raises(ValueError, match="derives KV state dtype"): + write_onnx_genai_config(_speech_package(), str(tmp_path), kv_native_dtype="bf16") + + +def test_dispatch_speech_to_text_carries_encoder_mask_as_cross_state(tmp_path): + artifacts = write_onnx_genai_config(_speech_package(encoder_mask=True), str(tmp_path)) with open(artifacts["inference_metadata"]) as handle: metadata = yaml.safe_load(handle) - assert metadata["pipeline"]["dataflow"][1] == { - "from": "encoder.encoder_attention_mask", - "to": "decoder.encoder_attention_mask", - "dtype": "int64", - "device_transfer": False, - } + workflow = metadata["pipeline"]["workflow"] + # Every encoder output the decoder consumes becomes cross state, so an + # encoder-side mask needs no special case. + assert "cross.encoder_attention_mask" in workflow["state"] + assert workflow["state"]["cross.encoder_attention_mask"]["contract"]["dtype"] == "int64" + assert "encoder.input.attention_mask" in workflow["inputs"] def test_dispatch_audio_codec_pipeline(tmp_path): - # A neural codec: encoder outputs codes consumed by a single-pass decoder, - # with no cross-attention. It is a pure tensor pipeline (no decoder config). - pkg = _EncoderDecoderPkg( - { - "encoder": _FakeModel(["waveform"], ["codes"]), - "decoder": _FakeModel(["codes"], ["waveform"]), - } + encoder = _model( + "encoder", + [_value("waveform", ir.DataType.FLOAT, ["batch", 1, "audio_samples"])], + [("codes", ir.DataType.INT64, ["batch", 16, "frames"])], ) + decoder = _model( + "decoder", + [_value("codes", ir.DataType.INT64, ["batch", 16, "frames"])], + [("waveform", ir.DataType.FLOAT, ["batch", 1, "audio_samples"])], + ) + pkg = ModelPackage({"encoder": encoder, "decoder": decoder}) artifacts = write_onnx_genai_config(pkg, str(tmp_path)) with open(artifacts["inference_metadata"]) as handle: metadata = yaml.safe_load(handle) - # No decoder capabilities (produces tensors, not tokens). assert "model" not in metadata - pipeline = metadata["pipeline"] - assert pipeline["models"] == { - "encoder": {"filename": "encoder/model.onnx", "type": "audio_encoder"}, - "decoder": {"filename": "decoder/model.onnx", "type": "vocoder"}, - } - assert pipeline["dataflow"] == [ - { - "from": "encoder.codes", - "to": "decoder.codes", - "dtype": "int64", - "device_transfer": False, - } - ] - stages = pipeline["strategy"]["stages"] - assert [stage["strategy"]["kind"] for stage in stages] == [ - "single_pass", - "single_pass", - ] + assert not {"models", "dataflow", "strategy", "phases"}.intersection(metadata["pipeline"]) + workflow = metadata["pipeline"]["workflow"] + assert workflow["steps"][0]["outputs"] == {"codes": "codec.codes"} + assert workflow["steps"][1]["inputs"] == {"codes": "codec.codes"} + assert workflow["outputs"]["waveform"]["stage"] == "post_adapter" def test_multi_decoder_tts_without_pre_embedder_raises_precise_not_implemented(tmp_path): @@ -487,7 +1036,7 @@ def test_multi_decoder_tts_without_pre_embedder_raises_precise_not_implemented(t "embedding": _FakeModel(["text_ids"]), } ) - with pytest.raises(NotImplementedError, match="nested_autoregressive"): + with pytest.raises(NotImplementedError, match="nested generic workflow loops"): write_onnx_genai_config(pkg, str(tmp_path)) @@ -506,45 +1055,46 @@ class _TTSPkg(dict): def test_dispatch_multi_decoder_tts_with_pre_embedder(tmp_path): - # The real Qwen3-TTS shape: talker + code_predictor + talker_step_embedder - # (+ talker_prefill_embedder) emits the pre_embedder/prefill-driven - # nested_autoregressive contract. - pkg = _TTSPkg( + pkg = ModelPackage( { - "talker": _FakeModel(["inputs_embeds"], ["logits", "last_hidden_state"]), - "code_predictor": _FakeModel(["inputs_embeds"], ["logits", "codec_embeddings"]), - "talker_step_embedder": _FakeModel(["frame_codes"], ["inputs_embeds"]), - "talker_prefill_embedder": _FakeModel( - ["text_ids"], ["prefill_embeds", "trailing_text_embeds"] + "talker": _model( + "talker", + [_value("inputs_embeds", ir.DataType.FLOAT, ["batch", "sequence", 16])], + [("last_hidden_state", ir.DataType.FLOAT, ["batch", 16])], ), - "embedding": _FakeModel(["text_ids"]), - } + "code_predictor": _model( + "code_predictor", + [ + _value("last_hidden_state", ir.DataType.FLOAT, ["batch", 16]), + _value("step_index", ir.DataType.INT64, ["batch"]), + ], + [("logits", ir.DataType.FLOAT, ["batch", 64])], + ), + "talker_step_embedder": _model( + "talker_step_embedder", + [_value("frame_codes", ir.DataType.INT64, ["batch", 16])], + [("inputs_embeds", ir.DataType.FLOAT, ["batch", 1, 16])], + ), + "talker_prefill_embedder": _model( + "talker_prefill_embedder", + [_value("text_ids", ir.DataType.INT64, ["batch", "sequence"])], + [("prefill_embeds", ir.DataType.FLOAT, ["batch", "sequence", 16])], + ), + "codec": _model( + "codec", + [_value("codes", ir.DataType.INT64, ["batch", 16, "frames"])], + [("waveform", ir.DataType.FLOAT, ["batch", 1, "samples"])], + ), + }, + config=_TTSCfg(), ) artifacts = write_onnx_genai_config(pkg, str(tmp_path)) - with open(artifacts["inference_metadata"]) as handle: - metadata = yaml.safe_load(handle) - - pipeline = metadata["pipeline"] - assert set(pipeline["models"]) == { - "talker", - "talker_step_embedder", - "code_predictor", - "talker_prefill_embedder", - } - stage = pipeline["strategy"]["stages"][0]["strategy"] - assert stage["kind"] == "nested_autoregressive" - assert stage["inner_embedding_output"] == "codec_embeddings" - assert stage["pre_embedder"]["component"] == "talker_step_embedder" - assert stage["prefill_embedder"]["component"] == "talker_prefill_embedder" - assert stage["num_code_groups"] == 16 - assert pipeline["phases"]["talker_prefill_embedder"]["run_on"] == "prompt_only" - assert { - "from": "talker_step_embedder.inputs_embeds", - "to": "talker.inputs_embeds", - "dtype": "fp32", - "device_transfer": False, - } in pipeline["dataflow"] + workflow = yaml.safe_load(handle)["pipeline"]["workflow"] + outer = workflow["steps"][0] + assert outer["iteration"]["value"] == "talker.iteration" + assert outer["steps"][2]["iteration"]["value"] == "code.iteration" + assert (tmp_path / "policies" / "code_frame_update.onnx").is_file() def test_unrecognized_multi_component_package_fails_loudly(tmp_path): @@ -580,7 +1130,7 @@ def save(self, path): with mock.patch.dict("sys.modules", {"transformers": fake_tf}): artifacts = write_onnx_genai_config( - object(), str(tmp_path), config=_Cfg(), source="some/model-id" + _decoder_package(), str(tmp_path), config=_Cfg(), source="some/model-id" ) assert artifacts.get("tokenizer") == str(tmp_path / "tokenizer.json") @@ -589,6 +1139,42 @@ def save(self, path): def test_decoder_without_source_skips_tokenizer(tmp_path): - artifacts = write_onnx_genai_config(object(), str(tmp_path), config=_Cfg()) + artifacts = write_onnx_genai_config(_decoder_package(), str(tmp_path), config=_Cfg()) assert "tokenizer" not in artifacts assert not (tmp_path / "tokenizer.json").exists() + + +def test_decoder_package_ships_chat_template_assets(tmp_path): + """A text decoder package must carry the assets needed to build a prompt. + + ``tokenizer.json`` alone leaves the runtime with no chat template, so an + instruction-tuned decoder receives raw user text with no BOS and no turn + markers. Image/audio processor configs, by contrast, describe media a text + package cannot consume and must not be copied. + """ + source = tmp_path / "source" + source.mkdir() + for filename in ( + "tokenizer.json", + "tokenizer_config.json", + "special_tokens_map.json", + "chat_template.jinja", + "preprocessor_config.json", + "image_processor.json", + ): + (source / filename).write_text("{}", encoding="utf-8") + + output = tmp_path / "output" + write_onnx_genai_config( + _decoder_package(_Int4Cfg()), + str(output), + config=_Int4Cfg(), + source=str(source), + ) + + assert (output / "tokenizer.json").is_file() + assert (output / "tokenizer_config.json").is_file() + assert (output / "special_tokens_map.json").is_file() + assert (output / "chat_template.jinja").is_file() + assert not (output / "preprocessor_config.json").exists() + assert not (output / "image_processor.json").exists() diff --git a/src/mobius/integrations/onnx_genai/codec_workflow_metadata_test.py b/src/mobius/integrations/onnx_genai/codec_workflow_metadata_test.py new file mode 100644 index 000000000..ac52e262c --- /dev/null +++ b/src/mobius/integrations/onnx_genai/codec_workflow_metadata_test.py @@ -0,0 +1,210 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT License. + +from __future__ import annotations + +import dataclasses +import json +import os + +import jsonschema +import onnx_ir as ir +import pytest +import yaml + +from mobius._model_package import ModelPackage +from mobius.integrations.onnx_genai.inference_metadata_test import ( + _model, + _value, +) +from mobius.integrations.onnx_genai.workflow_metadata import ( + build_audio_codec_workflow_metadata, + build_tts_workflow_metadata, + write_audio_codec_workflow_metadata, +) +from mobius.models.qwen3_tts import Qwen3TTSForConditionalGeneration +from mobius.models.qwen3_tts_test import _TINY_CONFIG +from mobius.tasks import TTSTask + + +def _codec_package() -> ModelPackage: + encoder = _model( + "encoder", + [_value("waveform", ir.DataType.FLOAT, ["batch", 1, "audio_samples"])], + [("codes", ir.DataType.INT64, ["batch", 16, "frames"])], + ) + decoder = _model( + "decoder", + [_value("codes", ir.DataType.INT64, ["batch", 16, "frames"])], + [("waveform", ir.DataType.FLOAT, ["batch", 1, "audio_samples"])], + ) + return ModelPackage({"encoder": encoder, "decoder": decoder}) + + +def test_codec_workflow_has_typed_ssa_and_audio_emit(): + metadata = build_audio_codec_workflow_metadata(_codec_package()) + pipeline = metadata["pipeline"] + assert not {"models", "dataflow", "strategy", "phases"}.intersection(pipeline) + + workflow = pipeline["workflow"] + assert workflow["inputs"]["request.waveform"]["contract"] == { + "dtype": "float32", + "rank": 3, + "shape": ["batch", 1, "audio_samples"], + "batch_layout": {"kind": "request_aligned", "axis": 0}, + } + # An ONNX component does not restate what its artifact already says: the + # encoder's ports, dtypes and shapes live in `encoder.onnx`, and the + # workflow's own inputs are where the request contract is declared. + assert not (workflow["components"]["encoder"].get("ports") or {}) + assert not (workflow["components"]["decoder"].get("ports") or {}) + assert "effects" not in workflow["components"]["encoder"] + assert "effects" not in workflow["components"]["decoder"] + + encode, decode, emit = workflow["steps"] + assert encode["outputs"] == {"codes": "codec.codes"} + assert decode["inputs"] == {"codes": "codec.codes"} + assert emit == { + "kind": "emit", + "value": "codec.waveform", + "output": "waveform", + "mode": "replace", + } + assert workflow["outputs"]["waveform"]["role"] == "audio" + assert workflow["outputs"]["waveform"]["stage"] == "post_adapter" + + +def test_codec_workflow_rejects_incompatible_code_contracts(): + package = _codec_package() + package["decoder"] = _model( + "decoder", + [_value("codes", ir.DataType.FLOAT, ["batch", 16, "frames"])], + [("waveform", ir.DataType.FLOAT, ["batch", 1, "audio_samples"])], + ) + with pytest.raises(ValueError, match="dtypes must match"): + build_audio_codec_workflow_metadata(package) + + +def test_codec_workflow_roundtrips_yaml(tmp_path): + path = write_audio_codec_workflow_metadata(_codec_package(), str(tmp_path)) + with open(path, encoding="utf-8") as handle: + loaded = yaml.safe_load(handle) + assert loaded == build_audio_codec_workflow_metadata(_codec_package()) + + +def test_codec_workflow_matches_producer_schema(): + schema_path = os.environ.get("ONNX_GENAI_SCHEMA") + if not schema_path: + pytest.skip("set ONNX_GENAI_SCHEMA to the producer-contract schema") + with open(schema_path, encoding="utf-8") as handle: + schema = json.load(handle) + jsonschema.validate(build_audio_codec_workflow_metadata(_codec_package()), schema) + + +@dataclasses.dataclass +class _TtsSubConfig: + num_code_groups: int = 4 + + +@dataclasses.dataclass +class _TtsConfig: + tts: _TtsSubConfig = dataclasses.field(default_factory=_TtsSubConfig) + + +def _tts_package() -> ModelPackage: + talker = _model( + "talker", + [_value("inputs_embeds", ir.DataType.FLOAT, ["batch", "sequence", 16])], + [("last_hidden_state", ir.DataType.FLOAT, ["batch", 16])], + ) + predictor = _model( + "code_predictor", + [ + _value("last_hidden_state", ir.DataType.FLOAT, ["batch", 16]), + _value("step_index", ir.DataType.INT64, ["batch"]), + ], + [("logits", ir.DataType.FLOAT, ["batch", 64])], + ) + step = _model( + "talker_step_embedder", + [_value("frame_codes", ir.DataType.INT64, ["batch", 4])], + [("inputs_embeds", ir.DataType.FLOAT, ["batch", 1, 16])], + ) + prefill = _model( + "talker_prefill_embedder", + [_value("text_ids", ir.DataType.INT64, ["batch", "sequence"])], + [("prefill_embeds", ir.DataType.FLOAT, ["batch", "sequence", 16])], + ) + codec = _model( + "codec", + [_value("codes", ir.DataType.INT64, ["batch", 4, "frames"])], + [("waveform", ir.DataType.FLOAT, ["batch", 1, "samples"])], + ) + return ModelPackage( + { + "talker": talker, + "code_predictor": predictor, + "talker_step_embedder": step, + "talker_prefill_embedder": prefill, + "codec": codec, + } + ) + + +def test_tts_uses_nested_lexical_loop_induction_and_codec(): + workflow = build_tts_workflow_metadata(_tts_package(), _TtsConfig())["pipeline"][ + "workflow" + ] + outer = workflow["steps"][0] + inner = outer["steps"][2] + assert outer["iteration"]["value"] == "talker.iteration" + assert inner["iteration"]["value"] == "code.iteration" + assert inner["steps"][0]["inputs"]["step_index"] == "code.iteration" + assert workflow["steps"][-2]["component"] == "codec" + assert workflow["outputs"]["waveform"]["stage"] == "post_adapter" + + +def test_real_qwen3_tts_workflow_carries_trained_transitions_and_kv_state(): + package = TTSTask().build(Qwen3TTSForConditionalGeneration(_TINY_CONFIG), _TINY_CONFIG) + package["codec"] = _model( + "codec", + [_value("codes", ir.DataType.INT64, ["batch", 4, "frames"])], + [("waveform", ir.DataType.FLOAT, ["batch", 1, "samples"])], + ) + + workflow = build_tts_workflow_metadata(package, _TINY_CONFIG)["pipeline"]["workflow"] + + assert { + "code_predictor_prefill", + "code_predictor_step_embedder", + "talker_text_step", + }.issubset(workflow["components"]) + assert workflow["state"]["talker_cache_0"]["recurrence"]["kind"] == "bounded" + assert workflow["state"]["talker_cache_0"]["service_group"] == "talker_cache" + assert workflow["state"]["predictor_cache_0"]["service_group"] == "predictor_cache" + assert workflow["serving"]["state_service"]["groups"]["talker_cache"]["kind"] == ( + "full_attention" + ) + # Row identity is runtime-private: no published input carries a slot table. + assert not any("slot_ids" in name for name in workflow["inputs"]) + assert workflow["serving"]["state_service"]["groups"]["talker_cache"]["ports"]["talker"][ + "talker_cache_0" + ]["input"].startswith("past_key_values.") + assert workflow["state"]["predictor_cache_0"]["scope"] == "invocation" + assert ( + workflow["state"]["predictor_cache_0"]["recurrence"]["max"] + == "package.predictor_context_limit" + ) + outer = workflow["steps"][0] + setup_history = next( + node for node in outer["setup"] if node.get("component") == "code_history_append" + ) + assert setup_history["inputs"]["frame"].startswith("setup.predictor.remaining_") + + assert outer["kind"] == "loop" + inner = next(node for node in outer["steps"] if node["kind"] == "loop") + assert inner["iteration"]["value"] == "code.iteration" + assert all(carried["cell"] != "slot_ids" for carried in inner["carried"]) + assert any( + node.get("component") == "code_predictor_step_embedder" for node in inner["steps"] + ) diff --git a/src/mobius/integrations/onnx_genai/ctc_runtime.py b/src/mobius/integrations/onnx_genai/ctc_runtime.py new file mode 100644 index 000000000..e70db4d7f --- /dev/null +++ b/src/mobius/integrations/onnx_genai/ctc_runtime.py @@ -0,0 +1,366 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT License. + +"""Metadata-driven reference runtime for non-generative CTC ASR packages. + +This module executes an exported package using *only* the facts written into +``inference_metadata.yaml``. Nothing here inspects the model architecture, the +checkpoint, or the model id — every decision (sample rate, normalization, +tensor names, blank id, time axis, vocabulary) is read from the document. It +therefore doubles as an executable check that the emitted contract is complete: +if a fact is missing from the metadata, this runtime cannot produce a +transcript. + +The pipeline is frame-synchronous, not autoregressive: + + encoded audio → preprocessing program → encoder (one invocation) + → per-frame argmax → collapse repeats → drop blank → text + +Batched requests are segmented with the ``frame_lengths`` output bound by the +profile's ``decoding.lengths`` role, so a padded batch yields the same +per-row transcript as an unpadded single-row run. +""" + +from __future__ import annotations + +import os +from typing import Any + +import numpy as np +import yaml + + +class MetadataContractError(ValueError): + """Raised when the metadata document lacks a fact the runtime needs.""" + + +def load_metadata(package_dir: str, filename: str = "inference_metadata.yaml") -> dict: + """Load the one-file inference metadata document from *package_dir*.""" + path = os.path.join(package_dir, filename) + with open(path, encoding="utf-8") as handle: + return yaml.safe_load(handle) + + +def select_profile(metadata: dict, kind: str) -> tuple[str, dict]: + """Return the ``(name, profile)`` of the single profile with *kind*. + + A reader that does not understand a profile may skip it only when the + profile is ``ignorable``; a missing required profile is a contract error. + """ + matches = [ + (name, profile) + for name, profile in (metadata.get("profiles") or {}).items() + if profile.get("kind") == kind + ] + if not matches: + raise MetadataContractError(f"metadata declares no '{kind}' profile") + if len(matches) > 1: + raise MetadataContractError( + f"metadata declares {len(matches)} '{kind}' profiles; expected exactly one" + ) + return matches[0] + + +def _resample(samples: np.ndarray, source_rate: int, target_rate: int) -> np.ndarray: + """Linearly resample a mono waveform; a no-op when the rates already agree.""" + if source_rate == target_rate: + return samples + duration = samples.shape[-1] / float(source_rate) + target_length = round(duration * target_rate) + source_positions = np.arange(samples.shape[-1], dtype=np.float64) + target_positions = np.linspace(0.0, samples.shape[-1] - 1, target_length, dtype=np.float64) + return np.interp(target_positions, source_positions, samples).astype(np.float32) + + +def run_audio_preprocessing( + program: dict, + waveforms: list[np.ndarray], + sample_rate: int, +) -> dict[str, np.ndarray]: + """Execute the declared audio preprocessing program over a batch. + + Args: + program: The ``preprocessing.audio`` sub-document. + waveforms: One 1-D float array per request row, at *sample_rate*. + sample_rate: Sample rate of the supplied waveforms. + + Returns: + A mapping from workflow SSA value name to the tensor bound to it, using + the program's ``outputs`` bindings. + + Every transform is dispatched by its generic ``op`` name; an unknown op is + an error rather than a silent skip, because skipping a normalization step + would produce plausible-but-wrong logits. + """ + rows = [np.asarray(row, dtype=np.float32).reshape(-1) for row in waveforms] + target_rate = sample_rate + pad_value = 0.0 + pad_side = "right" + + for transform in program.get("transforms", []): + op = transform.get("op") + if op == "decode": + # The caller already decoded the container; the declared step exists + # so a runtime that receives raw bytes knows one is required. + pass + elif op == "resample": + target_rate = int(transform["sample_rate"]) + rows = [_resample(row, sample_rate, target_rate) for row in rows] + elif op == "downmix": + channels = int(transform.get("channels", 1)) + if channels != 1: + raise MetadataContractError( + f"reference runtime only downmixes to mono, got {channels}" + ) + elif op == "rescale": + scale = float(transform["scale"]) + rows = [row * scale for row in rows] + elif op == "zero_mean_unit_variance": + epsilon = float(transform.get("epsilon", 0.0)) + # Normalized per row, matching how a feature extractor treats each + # utterance independently — a batch-wide statistic would make a + # padded batch disagree with a single-row run. + rows = [(row - row.mean()) / np.sqrt(row.var() + epsilon) for row in rows] + elif op == "pad": + pad_value = float(transform.get("pad_value", 0.0)) + pad_side = transform.get("mode", "right") + elif op == "trim": + target_length = int(transform["target_length"]) + rows = [row[:target_length] for row in rows] + else: + raise MetadataContractError(f"unsupported audio transform op '{op}'") + + if pad_side != "right": + raise MetadataContractError( + f"reference runtime only pads on the right, got '{pad_side}'" + ) + + max_length = max(row.shape[0] for row in rows) + values = np.full((len(rows), max_length), pad_value, dtype=np.float32) + mask = np.zeros((len(rows), max_length), dtype=np.int64) + for index, row in enumerate(rows): + values[index, : row.shape[0]] = row + mask[index, : row.shape[0]] = 1 + + produced = {"values": values, "sample_mask": mask, "samples": values} + bound: dict[str, np.ndarray] = {} + for binding in program.get("outputs", []): + source = binding["source"] + if source not in produced: + raise MetadataContractError( + f"audio program output '{binding['name']}' reads undeclared value '{source}'" + ) + tensor = produced[source] + dtype = binding.get("dtype") + if dtype: + tensor = tensor.astype(_numpy_dtype(dtype), copy=False) + if int(binding.get("rank", tensor.ndim)) != tensor.ndim: + raise MetadataContractError( + f"audio program output '{binding['name']}' declares rank " + f"{binding['rank']} but produced rank {tensor.ndim}" + ) + bound[binding["name"]] = tensor + return bound + + +def _numpy_dtype(name: str) -> Any: + mapping = { + "float32": np.float32, + "float16": np.float16, + "int64": np.int64, + "int32": np.int32, + "uint8": np.uint8, + "bool": np.bool_, + } + if name not in mapping: + raise MetadataContractError(f"unsupported tensor dtype '{name}'") + return mapping[name] + + +def run_workflow( + metadata: dict, + package_dir: str, + preprocessed: dict[str, np.ndarray], + *, + providers: list[str] | None = None, +) -> dict[str, np.ndarray]: + """Execute the declared workflow steps and return its emitted outputs. + + Only ``invoke`` and ``emit`` steps are supported: a CTC workflow is a plain + sequence with no loop, no branch and no carried state. Encountering a loop + here means the package was mis-detected as frame-synchronous. + """ + import onnxruntime as ort + + workflow = metadata["pipeline"]["workflow"] + components = workflow["components"] + + # The preprocessing adapter's outputs are already materialized by the + # caller, so its invocation binds names rather than running a session. + ssa: dict[str, np.ndarray] = {} + sessions: dict[str, Any] = {} + emitted: dict[str, np.ndarray] = {} + + for step in workflow["steps"]: + kind = step.get("kind") + if kind == "invoke": + name = step["component"] + component = components[name] + implementation = component["implementation"] + if implementation["kind"] == "adapter": + for port, target in step["outputs"].items(): + if port not in preprocessed: + raise MetadataContractError( + f"adapter '{name}' output '{port}' was not produced by " + "the preprocessing program" + ) + ssa[target] = preprocessed[port] + continue + if implementation["kind"] != "onnx": + raise MetadataContractError( + f"component '{name}' uses unsupported implementation " + f"'{implementation['kind']}'" + ) + if name not in sessions: + sessions[name] = ort.InferenceSession( + os.path.join(package_dir, implementation["artifact"]), + providers=providers or ["CPUExecutionProvider"], + ) + session = sessions[name] + feed = {port: ssa[source] for port, source in step["inputs"].items()} + requested = list(step["outputs"].keys()) + results = session.run(requested, feed) + for port, value in zip(requested, results): + ssa[step["outputs"][port]] = value + elif kind == "emit": + emitted[step["output"]] = ssa[step["value"]] + else: + raise MetadataContractError( + f"CTC workflow contains unsupported step kind '{kind}'; a " + "frame-synchronous package must be a plain sequence" + ) + return emitted + + +def collapse_ctc(ids: list[int], *, blank_id: int, collapse_repeats: bool) -> list[int]: + """Collapse a frame-argmax id sequence into CTC output tokens. + + Repeats collapse *before* blanks are removed, which is what makes a doubled + letter representable: two identical letters separated by a blank survive as + two tokens, while a letter held over several frames becomes one. + """ + collapsed: list[int] = [] + previous: int | None = None + for token in ids: + if collapse_repeats and token == previous: + continue + previous = token + if token == blank_id: + continue + collapsed.append(token) + return collapsed + + +def decode_transcripts( + metadata: dict, + outputs: dict[str, np.ndarray], + *, + profile_kind: str = "transcription", +) -> dict[str, Any]: + """Turn emitted workflow outputs into per-row transcripts. + + Returns a dict with ``argmax_ids``, ``collapsed_ids`` and ``transcripts``, + one entry per batch row, so a caller can compare any decode stage against a + reference implementation. + """ + _, profile = select_profile(metadata, profile_kind) + decoding = profile.get("decoding") + if decoding is None: + raise MetadataContractError(f"'{profile_kind}' profile declares no decoding contract") + if decoding.get("kind") != "ctc": + raise MetadataContractError( + f"reference runtime decodes 'ctc' only, got '{decoding.get('kind')}'" + ) + + logits_output = profile["outputs"].get("logits") + if logits_output is None: + raise MetadataContractError("transcription profile binds no 'logits' output") + logits = outputs[logits_output] + + time_axis = int(decoding["time_axis"]) + class_axis = int(decoding["class_axis"]) + if time_axis == class_axis: + raise MetadataContractError("decoding time_axis and class_axis must differ") + blank_id = int(decoding["blank_id"]) + collapse_repeats = bool(decoding.get("collapse_repeats", False)) + + frame_ids = np.argmax(logits, axis=class_axis) # (batch, frames) + if time_axis != 1: + frame_ids = np.moveaxis(frame_ids, time_axis - (time_axis > class_axis), -1) + + lengths_role = decoding.get("lengths") + if lengths_role is not None: + lengths_output = profile["outputs"].get(lengths_role) + if lengths_output is None: + raise MetadataContractError( + f"decoding.lengths references role '{lengths_role}' that the " + "profile does not bind" + ) + lengths = np.asarray(outputs[lengths_output]).reshape(-1).astype(int) + else: + lengths = np.full(frame_ids.shape[0], frame_ids.shape[-1], dtype=int) + + vocabulary = decoding.get("vocabulary") or {} + tokens = list(vocabulary.get("tokens") or []) + ignored = set(vocabulary.get("ignored_tokens") or []) + delimiter = vocabulary.get("word_delimiter") + + argmax_ids: list[list[int]] = [] + collapsed_ids: list[list[int]] = [] + transcripts: list[str] = [] + for row in range(frame_ids.shape[0]): + valid = int(min(lengths[row], frame_ids.shape[-1])) + row_ids = frame_ids[row, :valid].tolist() + argmax_ids.append(row_ids) + collapsed = collapse_ctc(row_ids, blank_id=blank_id, collapse_repeats=collapse_repeats) + collapsed_ids.append(collapsed) + if tokens: + pieces = [ + tokens[token_id] for token_id in collapsed if tokens[token_id] not in ignored + ] + if delimiter: + # The delimiter separates words rather than emitting whitespace, + # so empty groups from leading/trailing/repeated delimiters are + # dropped instead of becoming stray spaces. + words = "".join(pieces).split(delimiter) + transcripts.append(" ".join(word for word in words if word)) + else: + transcripts.append("".join(pieces)) + else: + transcripts.append("") + return { + "argmax_ids": argmax_ids, + "collapsed_ids": collapsed_ids, + "transcripts": transcripts, + } + + +def transcribe( + package_dir: str, + waveforms: list[np.ndarray], + sample_rate: int, + *, + providers: list[str] | None = None, +) -> dict[str, Any]: + """Run the full metadata-driven CTC pipeline over a batch of waveforms.""" + metadata = load_metadata(package_dir) + program = (metadata.get("preprocessing") or {}).get("audio") + if program is None: + raise MetadataContractError("metadata declares no preprocessing.audio program") + preprocessed = run_audio_preprocessing(program, waveforms, sample_rate) + outputs = run_workflow(metadata, package_dir, preprocessed, providers=providers) + result = decode_transcripts(metadata, outputs) + result["logits"] = outputs[ + select_profile(metadata, "transcription")[1]["outputs"]["logits"] + ] + return result diff --git a/src/mobius/integrations/onnx_genai/ctc_runtime_test.py b/src/mobius/integrations/onnx_genai/ctc_runtime_test.py new file mode 100644 index 000000000..1611dc9e7 --- /dev/null +++ b/src/mobius/integrations/onnx_genai/ctc_runtime_test.py @@ -0,0 +1,311 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT License. + +"""Tests for the CTC ASR metadata contract and its metadata-driven runtime.""" + +from __future__ import annotations + +import numpy as np +import pytest + +from mobius._configs import MMSConfig +from mobius.integrations.onnx_genai import ctc_runtime +from mobius.integrations.onnx_genai.workflow_metadata import ( + build_ctc_asr_workflow_metadata, +) +from mobius.models.wav2vec2_ctc import Wav2Vec2ForCTCModel +from mobius.tasks._ctc_asr import BATCH_PADDING_SENSITIVE_KEY, CTCAsrTask + + +def _tiny_config(**overrides) -> MMSConfig: + base = { + "vocab_size": 32, + "hidden_size": 64, + "intermediate_size": 128, + "num_hidden_layers": 2, + "num_attention_heads": 4, + "conv_dim": (32, 32, 64), + "conv_kernel": (10, 3, 2), + "conv_stride": (5, 2, 2), + "conv_bias": False, + "feat_extract_norm": "group", + "do_stable_layer_norm": False, + "pad_token_id": 0, + } + base.update(overrides) + return MMSConfig(**base) + + +def _build(config: MMSConfig): + return CTCAsrTask().build(Wav2Vec2ForCTCModel(config), config) + + +class TestFeatureExtractOutputLength: + """The analytic frame count must mirror the convolution stack exactly.""" + + def test_matches_manual_convolution_arithmetic(self): + config = _tiny_config() + samples = 4000 + expected = samples + for kernel, stride in zip(config.conv_kernel, config.conv_stride): + expected = (expected - kernel) // stride + 1 + assert config.feature_extract_output_length(samples) == expected + + def test_wav2vec2_base_geometry_downsamples_by_320(self): + config = _tiny_config( + conv_dim=(512,) * 7, + conv_kernel=(10, 3, 3, 3, 3, 2, 2), + conv_stride=(5, 2, 2, 2, 2, 2, 2), + ) + # 16 kHz audio yields 50 frames per second. + assert config.feature_extract_output_length(16_000) == 49 + + def test_rejects_ragged_convolution_geometry(self): + with pytest.raises(ValueError): + _tiny_config(conv_kernel=(10, 3)) + + def test_rejects_unknown_feature_normalization(self): + with pytest.raises(ValueError): + _tiny_config(feat_extract_norm="batch") + + +class TestBatchPaddingSensitivity: + """Padding sensitivity is a property of the feature normalization.""" + + def test_group_normalization_is_padding_sensitive(self): + module = Wav2Vec2ForCTCModel(_tiny_config(feat_extract_norm="group")) + assert module.batch_padding_sensitive is True + + def test_layer_normalization_is_row_independent(self): + module = Wav2Vec2ForCTCModel(_tiny_config(feat_extract_norm="layer")) + assert module.batch_padding_sensitive is False + + def test_task_records_sensitivity_on_the_built_graph(self): + pkg = _build(_tiny_config(feat_extract_norm="group")) + assert pkg["model"].metadata_props[BATCH_PADDING_SENSITIVE_KEY] == "true" + + def test_config_class_survives_architecture_rerouting(self): + # Config resolution reaches this module by re-routing model_type + # "wav2vec2" to the "mms" registration, so the module must name its own + # config class or the convolution geometry silently reverts to defaults. + assert Wav2Vec2ForCTCModel.config_class is MMSConfig + + +class TestCtcAsrMetadata: + """The emitted document must fully describe a frame-synchronous package.""" + + @pytest.fixture + def metadata(self): + config = _tiny_config() + return build_ctc_asr_workflow_metadata(_build(config), config) + + def test_declares_audio_preprocessing_bound_to_graph_inputs(self, metadata): + audio = metadata["preprocessing"]["audio"] + assert [t["op"] for t in audio["transforms"]] == [ + "decode", + "resample", + "downmix", + "zero_mean_unit_variance", + "pad", + ] + assert {o["name"] for o in audio["outputs"]} == { + "input_values", + "attention_mask", + } + + def test_transcription_profile_carries_a_ctc_decoding_contract(self, metadata): + decoding = metadata["profiles"]["transcription"]["decoding"] + assert decoding["kind"] == "ctc" + assert decoding["blank_id"] == 0 + assert decoding["collapse_repeats"] is True + assert (decoding["time_axis"], decoding["class_axis"]) == (1, 2) + + def test_padding_sensitive_profile_binds_per_row_lengths(self, metadata): + profile = metadata["profiles"]["transcription"] + assert profile["batch_invariance"] == "padding_sensitive" + assert profile["decoding"]["lengths"] == "frame_lengths" + assert profile["outputs"]["frame_lengths"] == "frame_lengths" + + def test_row_independent_when_feature_norm_is_layer(self): + config = _tiny_config(feat_extract_norm="layer") + metadata = build_ctc_asr_workflow_metadata(_build(config), config) + assert metadata["profiles"]["transcription"]["batch_invariance"] == ("row_independent") + + def test_workflow_is_a_plain_sequence_with_no_generation_loop(self, metadata): + steps = metadata["pipeline"]["workflow"]["steps"] + assert {step["kind"] for step in steps} <= {"invoke", "emit"} + assert not metadata["pipeline"]["workflow"].get("state") + + def test_encoder_is_invoked_exactly_once(self, metadata): + steps = metadata["pipeline"]["workflow"]["steps"] + invocations = [s for s in steps if s["kind"] == "invoke"] + assert sum(1 for s in invocations if s["component"] == "encoder") == 1 + + +class TestCtcCollapse: + """CTC collapsing must fold repeats before removing blanks.""" + + def test_collapses_repeated_frames_into_one_token(self): + assert ctc_runtime.collapse_ctc( + [1, 1, 1, 2, 2], blank_id=0, collapse_repeats=True + ) == [1, 2] + + def test_blank_separated_repeats_survive_as_two_tokens(self): + assert ctc_runtime.collapse_ctc( + [1, 1, 0, 1, 1], blank_id=0, collapse_repeats=True + ) == [1, 1] + + def test_drops_blanks(self): + assert ctc_runtime.collapse_ctc([0, 0, 3, 0], blank_id=0, collapse_repeats=True) == [3] + + def test_without_repeat_collapsing_every_non_blank_frame_is_kept(self): + assert ctc_runtime.collapse_ctc([1, 1, 0, 2], blank_id=0, collapse_repeats=False) == [ + 1, + 1, + 2, + ] + + +def _decoding_metadata(**decoding_overrides) -> dict: + decoding = { + "kind": "ctc", + "blank_id": 0, + "collapse_repeats": True, + "time_axis": 1, + "class_axis": 2, + "lengths": "frame_lengths", + "vocabulary": { + "source": "inline", + "size": 5, + "tokens": ["", "|", "A", "B", ""], + "word_delimiter": "|", + "ignored_tokens": [""], + }, + } + decoding.update(decoding_overrides) + return { + "profiles": { + "transcription": { + "kind": "transcription", + "outputs": {"logits": "logits", "frame_lengths": "frame_lengths"}, + "decoding": decoding, + } + } + } + + +def _one_hot(rows: list[list[int]], classes: int = 5) -> np.ndarray: + logits = np.full((len(rows), max(len(r) for r in rows), classes), -10.0, np.float32) + for i, row in enumerate(rows): + for t, class_id in enumerate(row): + logits[i, t, class_id] = 10.0 + return logits + + +class TestMetadataDrivenDecoding: + """Rendering is driven entirely by the declared vocabulary.""" + + def test_word_delimiter_separates_words_without_adding_stray_spaces(self): + # Leading, trailing and doubled delimiters must not create empty words. + outputs = { + "logits": _one_hot([[1, 2, 1, 1, 3, 1]]), + "frame_lengths": np.array([6]), + } + decoded = ctc_runtime.decode_transcripts(_decoding_metadata(), outputs) + assert decoded["transcripts"] == ["A B"] + + def test_ignored_tokens_are_dropped_before_word_splitting(self): + outputs = { + "logits": _one_hot([[2, 4, 3]]), + "frame_lengths": np.array([3]), + } + decoded = ctc_runtime.decode_transcripts(_decoding_metadata(), outputs) + assert decoded["transcripts"] == ["AB"] + + def test_lengths_binding_segments_a_padded_batch(self): + outputs = { + "logits": _one_hot([[2, 1, 3], [3, 0, 0]]), + "frame_lengths": np.array([3, 1]), + } + decoded = ctc_runtime.decode_transcripts(_decoding_metadata(), outputs) + assert decoded["argmax_ids"] == [[2, 1, 3], [3]] + assert decoded["transcripts"] == ["A B", "B"] + + def test_absent_lengths_binding_decodes_every_frame(self): + metadata = _decoding_metadata() + del metadata["profiles"]["transcription"]["decoding"]["lengths"] + outputs = {"logits": _one_hot([[2, 0, 0]])} + decoded = ctc_runtime.decode_transcripts(metadata, outputs) + assert decoded["argmax_ids"] == [[2, 0, 0]] + + def test_unbound_lengths_role_is_rejected(self): + metadata = _decoding_metadata(lengths="missing_role") + outputs = {"logits": _one_hot([[2]]), "frame_lengths": np.array([1])} + with pytest.raises(ctc_runtime.MetadataContractError): + ctc_runtime.decode_transcripts(metadata, outputs) + + def test_non_ctc_decoding_kind_is_rejected(self): + metadata = _decoding_metadata(kind="beam_search") + outputs = {"logits": _one_hot([[2]]), "frame_lengths": np.array([1])} + with pytest.raises(ctc_runtime.MetadataContractError): + ctc_runtime.decode_transcripts(metadata, outputs) + + +class TestAudioPreprocessingProgram: + """Preprocessing must normalize per row and pad to the batch width.""" + + @staticmethod + def _program() -> dict: + return { + "transforms": [ + {"op": "decode"}, + {"op": "resample", "sample_rate": 16_000}, + {"op": "downmix", "channels": 1}, + {"op": "zero_mean_unit_variance", "epsilon": 1e-7}, + {"op": "pad", "mode": "right", "pad_value": 0.0}, + ], + "outputs": [ + { + "name": "input_values", + "source": "values", + "content": "waveform", + "dtype": "float32", + "rank": 2, + }, + { + "name": "attention_mask", + "source": "sample_mask", + "content": "validity_mask", + "dtype": "int64", + "rank": 2, + }, + ], + } + + def test_pads_to_the_longest_row_and_marks_validity(self): + bound = ctc_runtime.run_audio_preprocessing( + self._program(), [np.ones(10, np.float32), np.ones(4, np.float32)], 16_000 + ) + assert bound["input_values"].shape == (2, 10) + assert bound["attention_mask"].tolist() == [[1] * 10, [1] * 4 + [0] * 6] + + def test_normalizes_each_row_independently(self): + # A padded batch must not fold one row's statistics into another's. + rows = [np.arange(10, dtype=np.float32), np.arange(4, dtype=np.float32) * 100] + bound = ctc_runtime.run_audio_preprocessing(self._program(), rows, 16_000) + solo = ctc_runtime.run_audio_preprocessing(self._program(), rows[:1], 16_000) + np.testing.assert_allclose( + bound["input_values"][0], solo["input_values"][0], atol=1e-6 + ) + + def test_unknown_transform_is_rejected_rather_than_skipped(self): + program = self._program() + program["transforms"].append({"op": "spectrogram"}) + with pytest.raises(ctc_runtime.MetadataContractError): + ctc_runtime.run_audio_preprocessing(program, [np.ones(4, np.float32)], 16_000) + + def test_output_reading_an_undeclared_value_is_rejected(self): + program = self._program() + program["outputs"][0]["source"] = "mel" + with pytest.raises(ctc_runtime.MetadataContractError): + ctc_runtime.run_audio_preprocessing(program, [np.ones(4, np.float32)], 16_000) diff --git a/src/mobius/integrations/onnx_genai/decoder_metadata.py b/src/mobius/integrations/onnx_genai/decoder_metadata.py index 95bf26bca..5145134b1 100644 --- a/src/mobius/integrations/onnx_genai/decoder_metadata.py +++ b/src/mobius/integrations/onnx_genai/decoder_metadata.py @@ -54,22 +54,12 @@ def _canonical_float_dtype(value: Any) -> str | None: return _FLOAT_DTYPE_ALIASES.get(token) -def _infer_kv_native_dtype(config: Any) -> str | None: - """Infer KV storage dtype from the model's activation/compute dtype.""" - for name in ("activation_dtype", "compute_dtype", "dtype", "torch_dtype"): - dtype = _canonical_float_dtype(getattr(config, name, None)) - if dtype is not None: - return dtype - return None - - def build_decoder_metadata( *, num_attention_heads: int, head_dim: int, num_kv_heads: int | None = None, max_sequence_length: int | None = None, - kv_native_dtype: str | None = None, attention_type: str | None = None, sliding_window: int | None = None, sink_tokens: int | None = None, @@ -84,8 +74,6 @@ def build_decoder_metadata( num_kv_heads: Number of key/value heads (defaults to ``num_attention_heads`` = multi-head; a smaller value = GQA). max_sequence_length: Maximum total sequence length in tokens. - kv_native_dtype: KV-cache storage dtype (e.g. ``"float16"``, - ``"bfloat16"``). attention_type: Override the derived attention type (``multi_head`` / ``grouped_query_attention``). sliding_window: Sliding-window length in tokens (None = full context). @@ -146,12 +134,10 @@ def build_decoder_metadata( if mixture_of_experts is not None: model["mixture_of_experts"] = mixture_of_experts - metadata: dict[str, Any] = {"required_capabilities": capabilities, "model": model} - if kv_native_dtype: - metadata["kv_cache"] = { - "native_dtype": _canonical_float_dtype(kv_native_dtype) or kv_native_dtype - } - return metadata + # The cache's storage representation is a runtime-private choice derived + # from the exported graph (port dtypes, Q/DQ, scale ports), so the package + # deliberately declares no package-level KV dtype. + return {"required_capabilities": capabilities, "model": model} def moe_metadata_from_config( @@ -253,9 +239,7 @@ def moe_metadata_from_config( } -def decoder_metadata_from_config( - config: Any, *, kv_native_dtype: str | None = None -) -> dict[str, Any]: +def decoder_metadata_from_config(config: Any) -> dict[str, Any]: """Build decoder metadata from a Mobius ``BaseModelConfig``/``ArchitectureConfig``. Unset fields (Mobius' ``DEFAULT_INT`` sentinel) are dropped. ``head_dim`` @@ -275,16 +259,12 @@ def decoder_metadata_from_config( sliding = getattr(config, "sliding_window", None) sliding = _clean_int(sliding) if sliding not in (None, _UNSET) else None - if kv_native_dtype is None: - kv_native_dtype = _infer_kv_native_dtype(config) - return build_decoder_metadata( num_attention_heads=num_heads, head_dim=head_dim, num_kv_heads=num_kv, max_sequence_length=_clean_int(getattr(config, "max_position_embeddings", None)), sliding_window=sliding, - kv_native_dtype=kv_native_dtype, architecture=getattr(config, "architecture", None) or getattr(config, "model_type", None), mixture_of_experts=moe_metadata_from_config(config), diff --git a/src/mobius/integrations/onnx_genai/decoder_metadata_test.py b/src/mobius/integrations/onnx_genai/decoder_metadata_test.py index 59839e1e9..287cbc367 100644 --- a/src/mobius/integrations/onnx_genai/decoder_metadata_test.py +++ b/src/mobius/integrations/onnx_genai/decoder_metadata_test.py @@ -55,7 +55,6 @@ def test_grouped_query_attention(self): num_kv_heads=8, head_dim=128, max_sequence_length=131072, - kv_native_dtype="bf16", ) assert meta["required_capabilities"] == ["kv_cache", "grouped_query_attention"] att = meta["model"]["attention"] @@ -64,7 +63,8 @@ def test_grouped_query_attention(self): assert att["num_kv_heads"] == 8 assert att["head_dim"] == 128 assert meta["model"]["max_sequence_length"] == 131072 - assert meta["kv_cache"]["native_dtype"] == "bfloat16" + # The cache's storage representation is graph-derived, never declared. + assert "kv_cache" not in meta def test_multi_head_when_kv_equals_heads(self): meta = build_decoder_metadata(num_attention_heads=16, num_kv_heads=16, head_dim=64) @@ -94,7 +94,7 @@ def test_rejects_non_divisible_kv_heads(self): build_decoder_metadata(num_attention_heads=12, num_kv_heads=5, head_dim=64) def test_from_config_reads_mobius_fields(self): - meta = decoder_metadata_from_config(_FakeConfig(), kv_native_dtype="fp16") + meta = decoder_metadata_from_config(_FakeConfig()) att = meta["model"]["attention"] assert att["type"] == "grouped_query_attention" assert att["num_attention_heads"] == 32 @@ -102,16 +102,17 @@ def test_from_config_reads_mobius_fields(self): assert att["head_dim"] == 128 assert meta["model"]["max_sequence_length"] == 131072 assert meta["model"]["architecture"] == "llama" - assert meta["kv_cache"]["native_dtype"] == "float16" + assert "kv_cache" not in meta - def test_from_config_infers_fp16_kv_dtype_for_int4_weights(self): + def test_from_config_never_declares_package_level_kv_dtype(self): + """Cache storage is a runtime choice derived from the exported graph.""" cfg = _FakeConfig() cfg.dtype = ir.DataType.FLOAT16 cfg.quantization = QuantizationConfig(bits=4, quant_method="rtn") meta = decoder_metadata_from_config(cfg) - assert meta["kv_cache"]["native_dtype"] == "float16" + assert "kv_cache" not in meta def test_from_config_derives_head_dim_and_drops_unset(self): cfg = _FakeConfig(head_dim=-42, sliding_window=-42) # DEFAULT_INT sentinel @@ -137,7 +138,7 @@ def test_matches_onnx_genai_schema(self): cfg.topk_method = "noaux_tc" cfg.n_group = 4 cfg.topk_group = 2 - meta = decoder_metadata_from_config(cfg, kv_native_dtype="bf16") + meta = decoder_metadata_from_config(cfg) schema_path = _schema_path() if schema_path is None: diff --git a/src/mobius/integrations/onnx_genai/duplex_workflow_metadata_test.py b/src/mobius/integrations/onnx_genai/duplex_workflow_metadata_test.py new file mode 100644 index 000000000..9ca929887 --- /dev/null +++ b/src/mobius/integrations/onnx_genai/duplex_workflow_metadata_test.py @@ -0,0 +1,310 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT License. + +"""Tests for the full-duplex speech-to-speech workflow producer. + +The workflow describes one frame of a Moshi-family full-duplex model +(PersonaPlex, Moshi): packed audio in, packed audio out, with the whole +conversation carried in session-scoped state between invocations. +""" + +from __future__ import annotations + +import dataclasses +import json +import os + +import jsonschema +import onnx_ir as ir +import pytest +import yaml + +from mobius._model_package import ModelPackage +from mobius.integrations.onnx_genai.inference_metadata_test import _model, _value +from mobius.integrations.onnx_genai.workflow_metadata import ( + build_full_duplex_workflow_metadata, + write_full_duplex_workflow_metadata, +) + +_DELAYS = [0, 0, 1, 1, 0, 1, 1] +_CHANNELS = len(_DELAYS) +_STREAMS = 3 +_LAYERS = 2 +_CONTEXT = 250 +_DEP_LAYERS = 2 + + +@dataclasses.dataclass +class _DuplexConfig: + delays: list[int] = dataclasses.field(default_factory=lambda: list(_DELAYS)) + dep_q: int = _STREAMS + n_q: int = _STREAMS + frame_size: int = 1920 + context: int = _CONTEXT + text_initial_token_id: int = 32000 + initial_token_id: int = 2048 + + +def _duplex_package() -> ModelPackage: + encoder = _model( + "encoder", + [_value("waveform", ir.DataType.FLOAT, ["batch", 1, "audio_samples"])], + [("codes", ir.DataType.INT64, ["batch", _STREAMS, "frames"])], + ) + decoder = _model( + "decoder", + [_value("codes", ir.DataType.INT64, ["batch", _STREAMS, "frames"])], + [("waveform", ir.DataType.FLOAT, ["batch", 1, "audio_samples"])], + ) + temporal_inputs = [ + _value("input_frame", ir.DataType.INT64, ["batch", _CHANNELS, "sequence_len"]), + _value("attention_mask", ir.DataType.INT64, ["batch", "context"]), + _value("position_ids", ir.DataType.INT64, ["batch", "sequence_len"]), + ] + temporal_outputs: list[tuple[str, ir.DataType, list[int | str]]] = [ + ("hidden", ir.DataType.FLOAT, ["batch", "sequence_len", 8]), + ("text_logits", ir.DataType.FLOAT, ["batch", "sequence_len", 32]), + ] + for layer in range(_LAYERS): + for port in ("key", "value"): + temporal_inputs.append( + _value( + f"past_key_values.{layer}.{port}", + ir.DataType.FLOAT, + ["batch", 2, "past_sequence_len", 4], + ) + ) + temporal_outputs.append( + ( + f"present.{layer}.{port}", + ir.DataType.FLOAT, + ["batch", 2, "past_sequence_len + 1", 4], + ) + ) + temporal = _model("temporal", temporal_inputs, temporal_outputs) + + depformer_inputs = [ + _value("hidden", ir.DataType.FLOAT, ["batch", 1, 8]), + _value("prev_token", ir.DataType.INT64, ["batch", 1]), + _value("substep_index", ir.DataType.INT64, []), + ] + depformer_outputs: list[tuple[str, ir.DataType, list[int | str]]] = [ + ("logits", ir.DataType.FLOAT, ["batch", 1, 16]), + ] + for layer in range(_DEP_LAYERS): + for port in ("key", "value"): + depformer_inputs.append( + _value( + f"past_key_values.{layer}.{port}", + ir.DataType.FLOAT, + ["batch", 2, "past_substep_len", 4], + ) + ) + depformer_outputs.append( + ( + f"present.{layer}.{port}", + ir.DataType.FLOAT, + ["batch", 2, "past_substep_len + 1", 4], + ) + ) + depformer = _model("depformer", depformer_inputs, depformer_outputs) + return ModelPackage( + { + "encoder": encoder, + "decoder": decoder, + "temporal": temporal, + "depformer": depformer, + } + ) + + +def _workflow() -> dict: + metadata = build_full_duplex_workflow_metadata(_duplex_package(), _DuplexConfig()) + return metadata["pipeline"]["workflow"] + + +def _frame_loop(workflow: dict) -> dict: + """The outer loop that carries session state across duplex frames.""" + loops = [step for step in workflow["steps"] if step["kind"] == "loop"] + assert len(loops) == 1, "the frame loop is the only top-level step" + return loops[0] + + +def test_duplex_workflow_is_one_event_per_invocation() -> None: + workflow = _workflow() + assert workflow["inputs"]["request.audio_chunk"]["role"]["role"] == "media" + assert workflow["inputs"]["request.session_id"]["role"]["role"] == "session_id" + output = workflow["outputs"]["audio_chunk"] + assert output["role"] == "audio" + frame_loop = _frame_loop(workflow) + # One invocation is one duplex event by default; the loop exists because a + # session-scoped cell is only readable through a loop carry. + assert frame_loop["max_iterations"] == "package.frames_per_invocation" + assert workflow["inputs"]["package.frames_per_invocation"]["default"] == 1 + emit = next(step for step in frame_loop["steps"] if step["kind"] == "emit") + # A frame is only emitted once the delay ring has been primed, so the emit + # is guarded rather than unconditional. + assert emit["mode"] == "event" + assert emit["when"] == "duplex.emit" + + +def test_duplex_conversation_state_is_session_scoped() -> None: + workflow = _workflow() + state = workflow["state"] + conversational = [ + "token_cache", + "token_provided", + "offset", + "attention_mask", + "position_ids", + *[f"temporal_cache_{index}" for index in range(_LAYERS * 2)], + ] + for name in conversational: + cell = state[name] + assert cell["scope"] == "session", name + assert cell["release_boundary"] == "session", name + assert cell["management"] == "runtime", name + assert cell["session"]["policy"] == "exclusive", name + # The acoustic transformer restarts every frame, so its cache must not + # outlive the invocation. + for index in range(_DEP_LAYERS * 2): + cell = state[f"depformer_cache_{index}"] + assert cell["scope"] == "invocation" + assert "session" not in cell + + +def test_duplex_codec_prefix_state_releases_at_phase_boundary() -> None: + """Stateless codec graphs replay a prefix; it is released before the session.""" + state = _workflow()["state"] + for name in ("user_waveform", "agent_codes"): + cell = state[name] + assert cell["scope"] == "session" + assert cell["release_boundary"] == "invocation" + assert cell["recurrence"]["kind"] == "growing" + assert cell["recurrence"]["axis"] == 2 + + +def test_duplex_temporal_cache_is_bounded_by_the_context_window() -> None: + workflow = _workflow() + state = workflow["state"] + for index in range(_LAYERS * 2): + cell = state[f"temporal_cache_{index}"] + assert cell["recurrence"] == { + "kind": "bounded", + "axis": 2, + "max": "package.context_limit", + } + # A runtime-owned cache must be permutable for batch compaction. + assert cell["contract"]["batch_layout"] == {"kind": "request_aligned", "axis": 0} + assert workflow["inputs"]["package.context_limit"]["default"] == _CONTEXT + + +def test_duplex_acoustic_loop_iterates_once_per_stream() -> None: + workflow = _workflow() + frame_loop = _frame_loop(workflow) + loops = [step for step in frame_loop["steps"] if step["kind"] == "loop"] + assert len(loops) == 1, "the acoustic substep loop is the only loop in a frame" + loop = loops[0] + assert loop["max_iterations"] == "package.num_streams" + assert workflow["inputs"]["package.num_streams"]["default"] == _STREAMS + components = [step["component"] for step in loop["steps"] if step["kind"] == "invoke"] + assert components == [ + "stream_index", + "token_to_slot", + "depformer", + "last_acoustic_logits", + "token_sampler", + "frame_update", + "teacher_select", + ] + + +def test_duplex_frame_pipeline_order() -> None: + workflow = _workflow() + invokes = [ + step["component"] + for step in _frame_loop(workflow)["steps"] + if step["kind"] == "invoke" + ] + assert invokes == [ + "waveform_append", + "encoder", + "codes_tail", + "user_stream_merge", + "frame_assemble", + "temporal", + "last_text_logits", + "token_sampler", + "teacher_select", + "target_frame", + "text_frame_update", + "frame_commit", + "agent_frame_select", + "codes_append", + "decoder", + "chunk_tail", + "step_update", + "cache_length_update", + ] + + +def test_duplex_delay_pattern_is_published() -> None: + workflow = _workflow() + assert workflow["inputs"]["package.delays"]["default"] == _DELAYS + assert workflow["inputs"]["package.initial_tokens"]["default"] == [32000] + [2048] * ( + _CHANNELS - 1 + ) + + +def test_duplex_declares_temporal_kv_service() -> None: + workflow = _workflow() + serving = workflow["serving"] + assert serving["active"] == "active" + assert serving["done"] == "done" + assert serving["accepted_len"] == "accepted_len" + group = serving["state_service"]["groups"]["temporal_cache"] + assert group["sequence_axis"] == 2 + assert group["layout"] == "bnsh" + assert group["logical_lengths"] == "temporal_cache_lengths" + assert len(group["ports"]["temporal"]) == _LAYERS * 2 + for index in range(_LAYERS * 2): + assert workflow["state"][f"temporal_cache_{index}"]["service_group"] == ( + "temporal_cache" + ) + + +def test_duplex_requires_all_components() -> None: + package = _duplex_package() + del package["depformer"] + with pytest.raises(ValueError, match="missing"): + build_full_duplex_workflow_metadata(package, _DuplexConfig()) + + +def test_duplex_requires_delay_pattern() -> None: + @dataclasses.dataclass + class _NoDelays: + dep_q: int = _STREAMS + + with pytest.raises(ValueError, match="delay pattern"): + build_full_duplex_workflow_metadata(_duplex_package(), _NoDelays()) + + +def test_duplex_workflow_writes_policy_artifacts(tmp_path) -> None: + package = _duplex_package() + path = write_full_duplex_workflow_metadata(package, _DuplexConfig(), str(tmp_path)) + with open(path, encoding="utf-8") as handle: + metadata = yaml.safe_load(handle) + components = metadata["pipeline"]["workflow"]["components"] + for name in ("frame_assemble", "frame_commit", "teacher_select", "token_sampler"): + artifact = components[name]["implementation"]["artifact"] + assert os.path.isfile(os.path.join(str(tmp_path), artifact)), name + + +def test_duplex_workflow_matches_producer_schema() -> None: + schema_path = os.environ.get("ONNX_GENAI_SCHEMA") + if not schema_path or not os.path.isfile(schema_path): + pytest.skip("set ONNX_GENAI_SCHEMA to the producer-contract schema") + with open(schema_path, encoding="utf-8") as handle: + schema = json.load(handle) + metadata = build_full_duplex_workflow_metadata(_duplex_package(), _DuplexConfig()) + jsonschema.validate(metadata, schema) diff --git a/src/mobius/integrations/onnx_genai/encoder_embedding_metadata_test.py b/src/mobius/integrations/onnx_genai/encoder_embedding_metadata_test.py new file mode 100644 index 000000000..67f47a8e8 --- /dev/null +++ b/src/mobius/integrations/onnx_genai/encoder_embedding_metadata_test.py @@ -0,0 +1,179 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT License. + +"""Unit tests for the encoder-embedding workflow metadata producer. + +A bidirectional encoder is not generative. These tests pin the three facts that +separate its metadata from every decoder's: the workflow runs once, it declares +exactly the ports the artifact exposes, and it never mentions a generation +loop, a sampler or a KV cache. +""" + +from __future__ import annotations + +from typing import Any + +import pytest + +from mobius._passes import RemoveDeadGraphInputsPass +from mobius.integrations.onnx_genai.auto_export import ( + _looks_like_encoder_embedding, + write_onnx_genai_config, +) +from mobius.integrations.onnx_genai.workflow_metadata import ( + build_encoder_embedding_workflow_metadata, +) +from mobius.models.bert import BertModel +from mobius.models.bert_test import PROTBERT_TINY_CONFIG +from mobius.models.esm import EsmModel +from mobius.models.esm_test import TINY_CONFIG as ESM2_TINY_CONFIG +from mobius.tasks import FeatureExtractionTask + + +def _esm2_package(): + package = FeatureExtractionTask().build(EsmModel(ESM2_TINY_CONFIG), ESM2_TINY_CONFIG) + # ESM-2 never reads ``token_type_ids``; the export path drops the dead input. + RemoveDeadGraphInputsPass()(package["model"]) + return package + + +def _protbert_package(): + return FeatureExtractionTask().build(BertModel(PROTBERT_TINY_CONFIG), PROTBERT_TINY_CONFIG) + + +_PACKAGES = {"esm2": _esm2_package, "protbert": _protbert_package} + + +@pytest.fixture(scope="module") +def built() -> dict[str, Any]: + return {name: build() for name, build in _PACKAGES.items()} + + +@pytest.fixture(params=sorted(_PACKAGES), scope="module") +def package(request, built): + return built[request.param] + + +@pytest.fixture(scope="module") +def metadata(package) -> dict[str, Any]: + return build_encoder_embedding_workflow_metadata(package, package.config) + + +class TestDetection: + def test_an_embedding_encoder_is_recognized(self, package) -> None: + assert _looks_like_encoder_embedding(package) + + def test_a_generative_package_is_not_mistaken_for_one(self) -> None: + """A package with ``logits`` is generative however else it is shaped. + + ``logits`` is the whole signal: it is what a sampler consumes, and an + embedding encoder has none. Renaming the port on an otherwise identical + graph is what isolates that one fact. + """ + package = _esm2_package() + assert _looks_like_encoder_embedding(package) + package["model"].graph.outputs[0].name = "logits" + assert not _looks_like_encoder_embedding(package) + + def test_a_cached_decoder_is_not_mistaken_for_one(self) -> None: + package = _esm2_package() + package["model"].graph.inputs[0].name = "past_key_values.0.key" + assert not _looks_like_encoder_embedding(package) + + +class TestEncoderEmbeddingWorkflow: + def test_it_runs_exactly_once(self, metadata) -> None: + steps = metadata["pipeline"]["workflow"]["steps"] + assert [step["kind"] for step in steps] == ["invoke", "emit"] + + def test_it_carries_no_state(self, metadata) -> None: + """No loop means no carried cell; a state table here would be inert.""" + assert "state" not in metadata["pipeline"]["workflow"] + + def test_it_describes_no_generation(self, metadata) -> None: + text = repr(metadata) + for forbidden in ( + "max_output_tokens", + "eos_ids", + "sampler", + "past_key_values", + "state_service", + "logits", + ): + assert forbidden not in text, f"encoder metadata should not mention {forbidden}" + + def test_the_profile_is_an_embedding_profile(self, metadata) -> None: + profile = metadata["profiles"]["embedding"] + assert profile["kind"] == "embedding" + assert profile["outputs"] == {"last_hidden_state": "last_hidden_state"} + + def test_pooling_uses_the_sequence_axis(self, metadata) -> None: + profile = metadata["profiles"]["embedding"] + workflow = metadata["pipeline"]["workflow"] + assert profile["pooling"] == {"kind": "mean", "axis": 1, "normalize": False} + assert "request.attention_mask" in workflow["inputs"] + assert profile["batch_invariance"] == "row_independent" + + def test_every_bound_port_exists_in_the_artifact(self, metadata, package) -> None: + graph = package["model"].graph + inputs = {str(value.name) for value in graph.inputs} + outputs = {str(value.name) for value in graph.outputs} + invoke = metadata["pipeline"]["workflow"]["steps"][0] + assert set(invoke["inputs"]) == inputs + assert set(invoke["outputs"]) <= outputs + + def test_every_declared_input_is_bound(self, metadata) -> None: + workflow = metadata["pipeline"]["workflow"] + invoke = workflow["steps"][0] + assert set(workflow["inputs"]) == set(invoke["inputs"].values()) + + def test_anonymous_dynamic_dimensions_get_stable_names(self, metadata) -> None: + workflow = metadata["pipeline"]["workflow"] + contracts = [ + declaration["contract"] + for declarations in (workflow["inputs"], workflow["outputs"]) + for declaration in declarations.values() + ] + assert all(None not in contract["shape"] for contract in contracts) + + def test_the_emitted_value_is_the_invocations_output(self, metadata) -> None: + invoke, emit = metadata["pipeline"]["workflow"]["steps"] + assert emit["value"] in invoke["outputs"].values() + assert emit["output"] in metadata["pipeline"]["workflow"]["outputs"] + + +class TestPortsFollowTheArtifact: + """The two models disagree about ``token_type_ids``, and the metadata must. + + ESM-2 has no token-type embedding, so the feature-extraction task's third + input is dead and gets pruned; ProtBert reads it. A producer that copied + the task signature instead of the graph would declare a port ESM-2 does not + expose, and a runtime would fail to bind it. + """ + + def test_esm2_declares_no_token_type_ids(self, built) -> None: + metadata = build_encoder_embedding_workflow_metadata(built["esm2"]) + assert "request.token_type_ids" not in metadata["pipeline"]["workflow"]["inputs"] + + def test_protbert_declares_token_type_ids(self, built) -> None: + metadata = build_encoder_embedding_workflow_metadata(built["protbert"]) + assert "request.token_type_ids" in metadata["pipeline"]["workflow"]["inputs"] + + +class TestDispatch: + def test_the_export_entry_point_picks_the_encoder_builder(self, tmp_path) -> None: + """Without this, an encoder falls through to the decoder fallback.""" + package = _esm2_package() + artifacts = write_onnx_genai_config(package, str(tmp_path), config=package.config) + with open(artifacts["inference_metadata"], encoding="utf-8") as handle: + text = handle.read() + assert "kind: embedding" in text + assert "max_output_tokens" not in text + + +class TestRejections: + def test_a_package_without_last_hidden_state_is_refused(self) -> None: + package = _esm2_package() + package["model"].graph.outputs[0].name = "not_hidden_states" + with pytest.raises(ValueError, match="last_hidden_state"): + build_encoder_embedding_workflow_metadata(package) diff --git a/src/mobius/integrations/onnx_genai/genai_config_import.py b/src/mobius/integrations/onnx_genai/genai_config_import.py new file mode 100644 index 000000000..8800ca0fb --- /dev/null +++ b/src/mobius/integrations/onnx_genai/genai_config_import.py @@ -0,0 +1,201 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT License. + +"""Import an existing ``genai_config.json`` package into native workflow metadata. + +The onnx-genai runtime executes a single structural workflow IR +(``pipeline.workflow``). Packages published for onnxruntime-genai instead ship a +``genai_config.json`` that names ports through ``%d`` patterns and leaves the +control flow implicit. This module performs the one-way import: it reads the +declared port names, resolves them against the ONNX graphs actually present, and +emits an equivalent typed-SSA workflow plus the generation-policy ONNX +components the loop needs. + +The import is structural. Nothing here keys on a model family, a model name, or +``model.type``: an encoder-conditioned package is recognised because the config +declares an encoder whose outputs the decoder consumes. +""" + +from __future__ import annotations + +import dataclasses +import json +import os +import shutil +from typing import Any + +import onnx_ir as ir + +from mobius._model_package import ModelPackage +from mobius.integrations.onnx_genai.workflow_metadata import ( + _dump_yaml, + build_speech_to_text_workflow_metadata, +) + + +@dataclasses.dataclass(frozen=True) +class ImportedGenAiConfig: + """Minimal generation contract a ``genai_config.json`` package declares.""" + + eos_token_id: int + pad_token_id: int | None + bos_token_id: int | None + max_position_embeddings: int + vocab_size: int | None + + @property + def num_hidden_layers(self) -> int: # pragma: no cover - informational + return 0 + + +@dataclasses.dataclass(frozen=True) +class ImportResult: + output_dir: str + metadata_path: str + components: dict[str, str] + config: ImportedGenAiConfig + + +def _expand_pattern(pattern: str | None, count: int) -> list[str]: + if pattern is None: + return [] + return [pattern % index for index in range(count)] + + +def _declared_port_names(section: dict[str, Any], count: int) -> set[str]: + """Expand every ``%d`` pattern and literal name a config section declares.""" + names: set[str] = set() + for value in section.values(): + if not isinstance(value, str): + continue + if "%d" in value: + names.update(_expand_pattern(value, count)) + else: + names.add(value) + return names + + +def load_genai_config_package(source_dir: str) -> tuple[ModelPackage, ImportedGenAiConfig]: + """Load the ONNX components a ``genai_config.json`` package declares. + + Returns a package keyed by structural role (``encoder``/``decoder``) and the + generation constants the config carries. Weights stay on disk: only the graph + interface is needed to derive the workflow. + """ + config_path = os.path.join(source_dir, "genai_config.json") + with open(config_path, encoding="utf-8") as handle: + raw = json.load(handle) + model = raw["model"] + search = raw.get("search", {}) + + models: dict[str, ir.Model] = {} + filenames: dict[str, str] = {} + for role in ("encoder", "decoder"): + section = model.get(role) + if section is None: + continue + filename = section["filename"] + filenames[role] = filename + models[role] = ir.load(os.path.join(source_dir, filename)) + if "decoder" not in models: + raise ValueError(f"{config_path} declares no decoder component") + + eos = model.get("eos_token_id", 0) + if isinstance(eos, list): + eos = eos[0] if eos else 0 + config = ImportedGenAiConfig( + eos_token_id=int(eos), + pad_token_id=model.get("pad_token_id"), + bos_token_id=model.get("bos_token_id"), + max_position_embeddings=int( + model.get("context_length") or search.get("max_length") or 2048 + ), + vocab_size=model.get("vocab_size"), + ) + package = ModelPackage(models, config=config) + package.imported_filenames = filenames # type: ignore[attr-defined] + package.genai_config = raw # type: ignore[attr-defined] + return package, config + + +def unbound_decoder_ports(package: ModelPackage) -> dict[str, list[str]]: + """Report decoder graph ports the imported config never names. + + A port the config does not declare cannot be bound by an importer without + guessing, so this is reported rather than silently defaulted. + """ + raw = getattr(package, "genai_config", {}) + decoder_section = raw.get("model", {}).get("decoder", {}) + layers = int(decoder_section.get("num_hidden_layers", 0)) + declared_inputs = _declared_port_names(decoder_section.get("inputs", {}), layers) + declared_outputs = _declared_port_names(decoder_section.get("outputs", {}), layers) + decoder = package["decoder"] + return { + "inputs": sorted( + value.name for value in decoder.graph.inputs if value.name not in declared_inputs + ), + "outputs": sorted( + value.name for value in decoder.graph.outputs if value.name not in declared_outputs + ), + } + + +def import_genai_config_package( + source_dir: str, + output_dir: str, + *, + sampler: str = "greedy", + audio_preprocessing: dict[str, Any] | None = None, + link_artifacts: bool = True, +) -> ImportResult: + """Emit native workflow metadata for an existing genai_config package. + + The ONNX artifacts are referenced under their original filenames and are + materialised inside *output_dir* so the emitted package is self-contained. + Symlinks are convenient during development but a loader that refuses to + follow an artifact path outside the package root will reject them, so + ``link_artifacts=False`` copies instead. + """ + package, config = load_genai_config_package(source_dir) + if "encoder" not in package: + raise ValueError( + "genai_config import currently covers encoder-conditioned packages; " + "single-decoder packages already load through the bare model contract" + ) + filenames: dict[str, str] = package.imported_filenames # type: ignore[attr-defined] + + os.makedirs(output_dir, exist_ok=True) + if os.path.abspath(source_dir) != os.path.abspath(output_dir): + for filename in filenames.values(): + for candidate in (filename, f"{filename}.data"): + source = os.path.join(source_dir, candidate) + if not os.path.exists(source): + continue + target = os.path.join(output_dir, candidate) + if os.path.lexists(target): + continue + if link_artifacts: + try: + os.symlink(os.path.abspath(source), target) + continue + except OSError: # pragma: no cover - platform dependent + pass + shutil.copy2(source, target) + + metadata = build_speech_to_text_workflow_metadata( + package, + config, + sampler=sampler, + audio_preprocessing=audio_preprocessing, + artifacts=filenames, + ) + package.save_policy_components(output_dir) + metadata_path = os.path.join(output_dir, "inference_metadata.yaml") + with open(metadata_path, "w", encoding="utf-8") as handle: + _dump_yaml(metadata, handle) + return ImportResult( + output_dir=output_dir, + metadata_path=metadata_path, + components=dict(filenames), + config=config, + ) diff --git a/src/mobius/integrations/onnx_genai/inference_metadata.py b/src/mobius/integrations/onnx_genai/inference_metadata.py index a196c2d42..6634c5f32 100644 --- a/src/mobius/integrations/onnx_genai/inference_metadata.py +++ b/src/mobius/integrations/onnx_genai/inference_metadata.py @@ -35,13 +35,17 @@ import os import re import shutil -from collections.abc import Callable, Iterable +from collections.abc import Callable, Iterable, Sequence from functools import lru_cache from pathlib import Path from typing import Any import yaml +from mobius._constants import ( + STATIC_CACHE_KV_SEQUENCE_LENGTH, + STATIC_CACHE_WRITE_INDICES, +) from mobius._pipeline_contract import component_presence, optional_input_contract from mobius.upstream_patches import apply_asset_patches @@ -61,6 +65,19 @@ "image_processor.json", ) +#: Assets a text-only package needs. Excludes the image/audio processor +#: contracts, which would advertise media preprocessing a text package's +#: graphs cannot consume. ``chat_template.jinja`` is required, not optional: +#: instruction-tuned decoders (Gemma 4, Llama-3-Instruct, Qwen-Instruct) +#: depend on their turn markers and leading BOS, and degenerate into +#: repetition when a raw prompt reaches the model instead. +_TEXT_RUNTIME_ASSET_NAMES = tuple( + name + for name in _RUNTIME_ASSET_NAMES + if name + not in {"processor_config.json", "preprocessor_config.json", "image_processor.json"} +) + @dataclasses.dataclass(frozen=True) class _Port: @@ -113,18 +130,81 @@ def _port(value: Any) -> _Port: ) -def _shape_metadata(port: _Port) -> list[int | str | None]: +_BATCH_DIMENSION = "batch" +"""Symbolic leading dimension mobius uses for per-request batching.""" + + +def _shape_metadata(port: _Port) -> list[int | str]: """Return a YAML-safe graph shape without losing symbolic dimensions.""" - shape: list[int | str | None] = [] - for dim in port.dims: + shape: list[int | str] = [] + for axis, dim in enumerate(port.dims): if isinstance(dim, int): shape.append(dim) continue value = getattr(dim, "value", None) - shape.append(str(value) if value is not None else None) + # Metadata dimensions cannot be null. Preserve named graph dimensions; + # give anonymous dynamic dimensions a stable, port-local name instead + # of pretending they are static or serializing an invalid null. + shape.append(str(value) if value is not None else f"{port.name}_dim_{axis}") return shape +def _name_image_preprocessing_program(image: dict[str, Any]) -> None: + """Convert structural preprocessing transforms into explicit typed SSA values.""" + transforms = image["transforms"] + if all("source" in output for output in image["outputs"]) and all( + "outputs" in transform for transform in transforms + ): + # Already named: re-running would append duplicate derived transforms. + return + current: str | None = None + decoded: str | None = None + for index, transform in enumerate(transforms): + name = f"image.transform_{index}" + if transform["op"] in {"decode", "decode_rgb"}: + transform.pop("inputs", None) + decoded = name + else: + if current is None: + raise ValueError("image preprocessing must decode before transforming") + transform["inputs"] = [current] + transform["outputs"] = [name] + current = name + if current is None: + raise ValueError("image preprocessing must declare at least one transform") + + derived_ops = { + "original_size": ("emit_original_size", decoded), + "transformed_size": ("emit_transformed_size", current), + "validity_mask": ("emit_validity_mask", current), + "patch_coordinates": ("emit_patch_coordinates", current), + "grid_dimensions": ("emit_grid_coordinates", current), + } + for output in image["outputs"]: + content = output["content"] + if content == "pixels": + output["source"] = current + continue + if content not in derived_ops: + raise ValueError( + f"image preprocessing output content {content!r} has no typed SSA producer" + ) + operation, source = derived_ops[content] + if source is None: + raise ValueError( + f"image preprocessing output content {content!r} requires a decoded image" + ) + name = f"image.output_{content}" + transforms.append( + { + "op": operation, + "inputs": [source], + "outputs": [name], + } + ) + output["source"] = name + + def _port_metadata(port: _Port) -> dict[str, Any]: """Serialize one exact graph port for the executable component contract.""" return { @@ -413,6 +493,39 @@ def _match_area_grid(ports: list[_Port], values: dict[str, Any]) -> _ImageProgra ) +def _max_token_grid_transforms(config: Any, values: dict[str, Any]) -> list[dict[str, Any]]: + patch_size = int(values["patch_size"]) + merge_size = int(values["merge_size"]) + token_pixels = (patch_size * merge_size) ** 2 + declared = dict(values) + declared["size"] = { + "shortest_edge": token_pixels, + "longest_edge": int(values["max_image_tokens"]) * token_pixels, + } + return _area_grid_transforms(config, declared) + + +def _match_max_token_grid(ports: list[_Port], values: dict[str, Any]) -> _ImageProgram | None: + bindings = _match_packed_grid(ports) + if bindings is None or not all( + isinstance(values.get(key), int) + for key in ( + "patch_size", + "temporal_patch_size", + "merge_size", + "max_image_tokens", + ) + ): + return None + return _ImageProgram( + name="max_token_packed_grid", + bindings=bindings, + transforms=_max_token_grid_transforms, + token_count_source="from_grid", + summary_contents=("grid_dimensions",), + ) + + def _match_patch_budget(ports: list[_Port], values: dict[str, Any]) -> _ImageProgram | None: bindings = _match_packed_coordinates(ports) if ( @@ -459,6 +572,7 @@ def _match_dynamic_hd(ports: list[_Port], values: dict[str, Any]) -> _ImageProgr Callable[[list[_Port], dict[str, Any]], _ImageProgram | None], ... ] = ( _match_area_grid, + _match_max_token_grid, _match_patch_budget, _match_dynamic_hd, ) @@ -553,6 +667,33 @@ def _processor_values( image_processor = values.get("image_processor") if isinstance(image_processor, dict): values.update(image_processor) + processor = values.get("processor") + if isinstance(processor, dict): + for transform in processor.get("transforms", []): + operation = transform.get("operation", {}) if isinstance(transform, dict) else {} + operation_type = operation.get("type") + attrs = operation.get("attrs", {}) + if not isinstance(attrs, dict): + continue + for key, value in attrs.items(): + values.setdefault(key, value) + if operation_type == "Resize": + values.setdefault("do_resize", True) + if "min_pixels" in attrs and "max_pixels" in attrs: + values.setdefault( + "size", + { + "shortest_edge": attrs["min_pixels"], + "longest_edge": attrs["max_pixels"], + }, + ) + elif operation_type == "Rescale": + values.setdefault("do_rescale", True) + values.setdefault("rescale_factor", attrs.get("rescale_factor")) + elif operation_type == "Normalize": + values.setdefault("do_normalize", True) + values.setdefault("image_mean", attrs.get("mean")) + values.setdefault("image_std", attrs.get("std")) embedding = values.get("embd_layer") if isinstance(embedding, dict): @@ -566,6 +707,7 @@ def _processor_values( "temporal_patch_size", "spatial_merge_size", "image_crop_size", + "size", ): value = getattr(vision, name, None) if value is None: @@ -732,17 +874,17 @@ def _static_cache_io( if (layer, role) not in ports ] input_names = {port.name for port in decoder_inputs} - for control in ("write_indices", "nonpad_kv_seqlen"): + for control in (STATIC_CACHE_WRITE_INDICES, STATIC_CACHE_KV_SEQUENCE_LENGTH): if control not in input_names: missing.append(f"input.{control}") if missing: raise ValueError( - "Cannot emit model.io.static_cache because the exported TensorScatter " - f"ABI is incomplete: {missing}" + "Cannot describe the static-cache ABI because the exported " + f"TensorScatter ports are incomplete: {missing}" ) return { - "write_indices_input": "write_indices", - "kv_sequence_length_input": "nonpad_kv_seqlen", + "write_indices_input": STATIC_CACHE_WRITE_INDICES, + "kv_sequence_length_input": STATIC_CACHE_KV_SEQUENCE_LENGTH, "key_cache_inputs": [inputs[(layer, "key")] for layer in layers], "value_cache_inputs": [inputs[(layer, "value")] for layer in layers], "key_cache_outputs": [outputs[(layer, "key")] for layer in layers], @@ -1252,84 +1394,296 @@ def validate_executable_closure(pkg: Any, metadata: dict[str, Any]) -> None: ) -def add_explicit_package_io( +#: Symbolic leading dimension Mobius emits for every batched ONNX port. A port +#: that opens with it holds exactly one entry per in-flight request, which is +#: the structural fact a runtime needs to permute or drop rows. +REQUEST_AXIS_SYMBOL = "batch" + + +def request_batch_layout(shape: list[Any] | None) -> dict[str, Any] | None: + """Return the request-aligned batch layout implied by a port's shape.""" + if shape and shape[0] == REQUEST_AXIS_SYMBOL: + return {"kind": "request_aligned", "axis": 0} + return None + + +def add_policy_components_to_workflow( metadata: dict[str, Any], pkg: Any, - config: Any, ) -> dict[str, Any]: - """Attach explicit graph-port roles to emitted decoder and encoder models.""" - pipeline = metadata.get("pipeline") - if not isinstance(pipeline, dict): - component_names = list(pkg.keys()) - if len(component_names) != 1: - raise ValueError("bare decoder metadata requires exactly one graph component") - io, _ = _decoder_io(pkg[component_names[0]], set(), config) - metadata.setdefault("model", {})["io"] = io + """Reference attached ONNX policy artifacts from an existing workflow. + + This helper intentionally does not synthesize a workflow or guess bindings. + It only adds schema-defined component declarations when a producer has + already emitted the exact workflow contract. + """ + policy_components = getattr(pkg, "policy_components", {}) + if not policy_components: + return metadata + workflow = metadata.get("pipeline", {}).get("workflow") + if not isinstance(workflow, dict): return metadata + components = workflow.setdefault("components", {}) + + def semantic_contract(component: Any) -> dict[str, Any]: + contract = component.contract + contract_name, version = component.contract_id.rsplit("@", 1) + bindings = { + key: value + for key, value in contract.items() + if key + not in { + "role", + "mode", + "effect", + "rng", + "state_class", + "batching", + "inactive_rows", + } + and isinstance(value, str) + } + rng = contract.get("rng") + if isinstance(rng, dict): + bindings.update( + {key: value for key, value in rng.items() if isinstance(value, str)} + ) + declaration: dict[str, Any] = { + "id": contract_name, + "version": version, + "bindings": bindings, + } + parameters = { + key: contract[key] + for key in ("mode", "batching", "inactive_rows") + if key in contract + } + if parameters: + declaration["parameters"] = parameters + return declaration + + def tensor_contract(value: Any) -> dict[str, Any]: + port = _port(value) + dtype = { + "fp32": "float32", + "fp16": "float16", + "bf16": "bfloat16", + }.get(port.dtype, port.dtype) + shape = _shape_metadata(port) + contract: dict[str, Any] = { + "dtype": dtype, + "rank": port.rank, + "shape": shape, + } + layout = request_batch_layout(shape) + if layout is not None: + contract["batch_layout"] = layout + return contract + + for name, component in policy_components.items(): + # A policy graph is synthesized by this producer to realize the + # workflow's own control flow, so its port contracts are not a + # transcription of an external interface: they are the type annotations + # of the workflow's dataflow. A workflow value acquires its dtype, rank + # and request axis from the port that produces it, and the validator + # reads metadata without the artifacts, so a policy output that states + # no contract leaves every value derived from it untyped. + declaration = { + "implementation": { + "kind": "onnx", + "artifact": f"policies/{name}.onnx", + }, + "ports": { + "inputs": { + value.name: tensor_contract(value) + for value in component.model.graph.inputs + }, + "outputs": { + value.name: tensor_contract(value) + for value in component.model.graph.outputs + }, + }, + } + if component.contract: + declaration["contract"] = semantic_contract(component) + if component.contract.get("role") == "token_sampler": + declaration["application_overridable"] = True + components[name] = declaration + declare_request_alignment(workflow) + return metadata - models = pipeline.get("models", {}) - routed_inputs: dict[str, set[str]] = {} - for edge in pipeline.get("dataflow", []): - target = edge.get("to", "") - component, separator, port = target.partition(".") - if separator: - routed_inputs.setdefault(component, set()).add(port) - - component_ios: dict[str, dict[str, Any]] = {} - for name, model_spec in models.items(): - if name not in pkg: - continue - if model_spec.get("type") == "decoder": - io, _ = _decoder_io(pkg[name], routed_inputs.get(name, set()), config) - else: - inputs = [_port(value) for value in pkg[name].graph.inputs] - outputs = [_port(value) for value in pkg[name].graph.outputs] - io = { - "inputs": [_port_metadata(port) for port in inputs], - "outputs": [_port_metadata(port) for port in outputs], + +_BATCH_DIMENSION_NAMES = frozenset({"batch", "batch_size", "batch_dim", "b"}) + + +def declare_request_alignment(workflow: dict[str, Any]) -> None: + """Stamp the request-aligned row axis onto every batch-leading contract. + + The runtime compacts finished rows out of a batch by applying one row + permutation to every request-aligned tensor. A contract whose leading axis + is the batch symbol but that does not say so is unpermutable, so state, + component ports, and outputs would silently drift apart after the first + eviction. Deriving the declaration from the admitted graph's own batch + symbol keeps alignment a property of the model interface rather than an + annotation every workflow builder has to remember. + """ + + def stamp(contract: Any) -> None: + if not isinstance(contract, dict) or "batch_layout" in contract: + return + shape = contract.get("shape") or [] + if shape and str(shape[0]) in _BATCH_DIMENSION_NAMES: + contract["batch_layout"] = {"kind": "request_aligned", "axis": 0} + + for section in ("inputs", "outputs", "state"): + for declaration in (workflow.get(section) or {}).values(): + if isinstance(declaration, dict): + stamp(declaration.get("contract")) + for component in (workflow.get("components") or {}).values(): + ports = component.get("ports", {}) if isinstance(component, dict) else {} + for side in ("inputs", "outputs"): + for contract in (ports.get(side) or {}).values(): + stamp(contract) + # A cell backed by a state-service group is stored by the runtime, not by + # the workflow: the group owns the buffer and the eviction policy, so the + # cell also needs an explicit boundary at which the runtime may free it. + for declaration in (workflow.get("state") or {}).values(): + if isinstance(declaration, dict) and declaration.get("service_group"): + declaration.setdefault("management", "runtime") + declaration.setdefault("release_boundary", declaration.get("scope", "invocation")) + + +def add_adapter_service_to_metadata( + metadata: dict[str, Any], + pkg: Any, + output_dir: str, +) -> dict[str, Any]: + """Attach the exact generic adapter catalog and saved artifact references.""" + artifacts = getattr(pkg, "adapter_artifacts", {}) + if not artifacts: + return metadata + workflow_value = metadata.get("pipeline", {}).get("workflow") + workflow = workflow_value if isinstance(workflow_value, dict) else None + manifest = getattr(pkg, "adapter_target_manifest", None) + if manifest is None: + raise ValueError("parameter adapters require an authoritative adapter target manifest") + options = pkg.adapter_service_options + inputs = workflow.setdefault("inputs", {}) if workflow is not None else {} + + def compatible_input( + name: str, + *, + dtype: str, + shape: list[str | int], + role: str | None, + ) -> bool: + declaration = inputs.get(name) + if not isinstance(declaration, dict): + return False + contract = declaration.get("contract", {}) + semantic_role = declaration.get("role", {}) + return ( + contract.get("dtype") == dtype + and contract.get("rank") == len(shape) + and contract.get("shape") == shape + and declaration.get("required", True) + and declaration.get("source", {}).get("kind") in {"request", "application"} + and ( + role is None + or semantic_role == {"kind": "runtime", "version": "1.0", "role": role} + ) + ) + + def ensure_input( + name: str, + *, + dtype: str, + shape: list[str | int], + role: str | None, + source: dict[str, str] | None = None, + ) -> None: + if name not in inputs: + inputs[name] = { + "contract": {"dtype": dtype, "rank": len(shape), "shape": shape}, + "role": ( + {"kind": "runtime", "version": "1.0", "role": role} + if role is not None + else {"kind": "opaque"} + ), + "source": source or {"kind": "request"}, } - if model_spec.get("type") in {"encoder", "audio_encoder"}: - audio_prompt = _select_one( - inputs, lambda port: _is_float(port) and port.rank == 3 - ) - token_prompt = _select_one( - inputs, lambda port: _is_integer(port) and port.rank == 2 - ) - if audio_prompt is not None: - io["audio_features_input"] = audio_prompt.name - elif token_prompt is not None: - io["token_input"] = token_prompt.name - io["sequence_source"] = "token_ids" - model_spec["io"] = io - component_ios[name] = io - - def annotate_strategy(strategy: dict[str, Any]) -> None: - if strategy.get("kind") == "nested_autoregressive" and not strategy.get( - "inner_embedding_output" - ): - inner_name = strategy.get("inner") - inner_io = component_ios.get(inner_name, {}) - hidden_output = inner_io.get("hidden_output") - if not hidden_output: - raise ValueError( - "Cannot emit pipeline.strategy.inner_embedding_output: the inner " - f"decoder {inner_name!r} has no unique non-logits float output" - ) - strategy["inner_embedding_output"] = hidden_output - for stage in strategy.get("stages", []): - nested = stage.get("strategy") - if isinstance(nested, dict): - annotate_strategy(nested) - - strategy = pipeline.get("strategy") - if isinstance(strategy, dict): - annotate_strategy(strategy) - if "model" in metadata: - decoder_names = [ - name for name, model in models.items() if model.get("type") == "decoder" - ] - if decoder_names: - metadata["model"]["io"] = component_ios[decoder_names[0]] + if not compatible_input(name, dtype=dtype, shape=shape, role=role): + raise ValueError( + f"adapter {role} must reference a required " + "request/application-sourced " + f"{dtype}{shape} workflow input" + ) + + active = options.active + if workflow is not None: + ensure_input( + options.segments, + dtype="int64", + shape=["batch", options.max_adapters], + role="adapter_segments", + ) + ensure_input( + options.adapter_counts, + dtype="int64", + shape=["batch"], + role="adapter_counts", + ) + ensure_input( + options.scales, + dtype="float32", + shape=["batch", options.max_adapters], + role="adapter_scales", + ) + if active is not None: + ensure_input( + active, + dtype="bool", + shape=["batch"], + role="adapter_active", + ) + + catalog = pkg.save_adapter_artifacts(output_dir) + if workflow is not None: + workflow.pop("adapters", None) + metadata["adapters"] = { + "base_model_fingerprint": manifest.base_fingerprint, + "target_manifest": pkg.adapter_target_manifest_metadata(), + "discovery_fallback": options.discovery_fallback, + "selection": { + "segments": options.segments, + "adapter_counts": options.adapter_counts, + "scales": options.scales, + **({"active": active} if active is not None else {}), + "max_adapters": options.max_adapters, + }, + "application_capability": options.application_capability, + "portable_fallback": options.portable_fallback, + "cache": { + "max_entries": options.cache_max_entries, + "eviction": "lru", + }, + "planning": { + "bucket_by_adapter_set": options.bucket_by_adapter_set, + "stable_buffers": options.stable_buffers, + "invalidate_capture_on_eviction": (options.invalidate_capture_on_eviction), + }, + "artifacts": catalog, + } + if workflow is not None: + capabilities = workflow.setdefault("manifest", {}).setdefault("capabilities", []) + for capability in ("parameter_adapters", "heterogeneous_adapter_batching"): + if capability not in capabilities: + capabilities.append(capability) + # Selection inputs carry exactly one entry per in-flight request, so + # they are request-aligned by construction. Stamping here rather than + # inside `ensure_input` also covers the declarations a producer wrote + # by hand before attaching the adapter service. + declare_request_alignment(workflow) return metadata @@ -1551,7 +1905,6 @@ def build_native_vlm_package_metadata( { "name": f"run_{name}", "strategy": strategy, - "run_on": phases[name]["run_on"], } ) @@ -1589,18 +1942,19 @@ def build_native_vlm_package_metadata( if decoder_io.get("token_input") and decoder_io.get("inputs_embeds_input"): capabilities.append("dual_sequence_inputs") metadata["required_capabilities"] = capabilities - metadata.setdefault("model", {})["io"] = decoder_io metadata["preprocessing"] = { "image": { "transforms": image_program.transforms(config, processor_values), "outputs": preprocessing_outputs, } } + # Bind every declared output to the processor-local value that produces it, + # so the runtime never has to guess which transform an output came from. + _name_image_preprocessing_program(metadata["preprocessing"]["image"]) metadata["pipeline"] = { "models": models, "dataflow": dataflow, "strategy": {"kind": "composite", "stages": stages}, - "phases": phases, "vision": vision_config, } if positions is not None: @@ -1612,13 +1966,14 @@ def build_native_vlm_package_metadata( def _copy_runtime_assets( output_dir: str, source: str | None, + names: Sequence[str] = _RUNTIME_ASSET_NAMES, *, revision: str | None = None, ) -> dict[str, str]: if not source: return {} os.makedirs(output_dir, exist_ok=True) - for filename in _RUNTIME_ASSET_NAMES: + for filename in names: source_path = _source_asset_path(source, filename, revision=revision) if source_path is not None: shutil.copy2(source_path, os.path.join(output_dir, filename)) @@ -1659,7 +2014,7 @@ def _copy_runtime_assets( return { Path(filename).stem: os.path.join(output_dir, filename) - for filename in _RUNTIME_ASSET_NAMES + for filename in names if os.path.isfile(os.path.join(output_dir, filename)) } @@ -1671,7 +2026,6 @@ def write_native_vlm_package_metadata( config: Any, source: str | None = None, revision: str | None = None, - kv_native_dtype: str | None = None, filename: str = "inference_metadata.yaml", ) -> dict[str, str]: """Write native VLM metadata and the runtime's tokenizer/processor assets.""" @@ -1683,8 +2037,7 @@ def write_native_vlm_package_metadata( pkg, config=config, source=source, - revision=revision, - decoder_metadata=decoder_metadata_from_config(config, kv_native_dtype=kv_native_dtype), + decoder_metadata=decoder_metadata_from_config(config), ) os.makedirs(directory, exist_ok=True) path = os.path.join(directory, filename) @@ -1717,6 +2070,11 @@ class SchedulerConfig: time_shift_type: str | None = None invert_sigmas: bool = False stochastic_sampling: bool = False + algorithm_type: str = "dpmsolver++" + solver_order: int = 2 + solver_type: str = "midpoint" + lower_order_final: bool = True + final_sigmas_type: str = "zero" def to_metadata(self) -> dict[str, Any]: meta: dict[str, Any] = { @@ -1837,6 +2195,11 @@ def from_diffusers(cls, config: dict[str, Any]) -> SchedulerConfig: ), invert_sigmas=bool(config.get("invert_sigmas")), stochastic_sampling=bool(config.get("stochastic_sampling")), + algorithm_type=str(config.get("algorithm_type", cls.algorithm_type)), + solver_order=int(config.get("solver_order", cls.solver_order)), + solver_type=str(config.get("solver_type", cls.solver_type)), + lower_order_final=bool(config.get("lower_order_final", cls.lower_order_final)), + final_sigmas_type=str(config.get("final_sigmas_type", cls.final_sigmas_type)), ) @@ -1888,74 +2251,38 @@ def load_diffusers_scheduler_config( return None -def build_language_diffusion_pipeline_metadata( - *, - mask_token_id: int, - num_inference_steps: int, - model_filename: str = "model.onnx", - input_ids_port: str = "input_ids", - logits_port: str = "logits", - block_length: int | None = None, - temperature: float | None = None, - guidance_scale: float | None = None, -) -> dict[str, Any]: - """Build the onnx-genai ``inference_metadata`` for a masked language-diffusion model. - - For a masked (discrete) language-diffusion model (e.g. LLaDA / Dream). - - The model is a mask predictor: it takes an int64 token sequence on - ``input_ids_port`` (prompt tokens plus a masked generation region) and emits - ``[B, S, V]`` logits on ``logits_port``. onnx-genai's ``masked_diffusion`` - scheduler drives the reverse process — each step commits the highest-confidence - still-masked positions (LLaDA low-confidence remasking) via a loop-carried - ``logits -> input_ids`` self-edge, unmasking progressively. - - Args: - mask_token_id: The ``[MASK]`` token id (e.g. 126336 for LLaDA-8B). - num_inference_steps: Total reverse-process steps (``strategy.num_steps``). - model_filename: The mask-predictor ONNX filename. - input_ids_port / logits_port: Model I/O port names. - block_length: Semi-autoregressive block length in tokens. When set, the - generation region is decoded in contiguous left-to-right blocks and - ``num_inference_steps`` must be divisible by the block count. - temperature: Gumbel-max sampling temperature (default 0 = argmax). - guidance_scale: Unsupervised classifier-free guidance multiplier. LLaDA's - effective multiplier is ``cfg_scale + 1``, so pass ``cfg_scale + 1``. +def load_diffusers_vae_scaling_factor(source: str | None) -> float | None: + """Best-effort load of a diffusers ``vae/config.json`` ``scaling_factor``. - Returns: - A dict with a top-level ``pipeline`` key, ready to serialize to - ``inference_metadata.yaml``. + The latent a diffusion sampler carries is scaled by this factor before the + VAE decodes it, so the workflow needs the real value rather than a guess. + Returns ``None`` when the config cannot be read, letting the caller decide. """ - if num_inference_steps < 1: - raise ValueError("num_inference_steps must be >= 1") - if block_length is not None and block_length < 1: - raise ValueError("block_length must be >= 1") - - scheduler_config: dict[str, Any] = { - "kind": "masked_diffusion", - "mask_token_id": int(mask_token_id), - } - if temperature is not None: - scheduler_config["temperature"] = float(temperature) - if block_length is not None: - scheduler_config["block_length"] = int(block_length) - - strategy: dict[str, Any] = { - "kind": "iterative", - "denoiser": "denoiser", - "num_steps": num_inference_steps, - "scheduler_config": scheduler_config, - } - if guidance_scale is not None and not math.isclose(guidance_scale, 1.0): - strategy["guidance_scale"] = guidance_scale + if not source: + return None + raw: dict[str, Any] | None = None + local = os.path.join(source, "vae", "config.json") + if os.path.isfile(local): + try: + with open(local, encoding="utf-8") as handle: + raw = json.load(handle) + except (OSError, ValueError) as err: + _LOGGER.warning("could not read %s: %s", local, err) + return None + else: + try: + from huggingface_hub import hf_hub_download - pipeline: dict[str, Any] = { - "models": {"denoiser": {"filename": model_filename, "type": "denoiser"}}, - # Loop-carried self-edge: the emitted logits refine the token sequence. - "dataflow": [{"from": f"denoiser.{logits_port}", "to": f"denoiser.{input_ids_port}"}], - "strategy": strategy, - } - return {"pipeline": pipeline} + path = hf_hub_download(source, "vae/config.json") + with open(path, encoding="utf-8") as handle: + raw = json.load(handle) + except Exception as err: + _LOGGER.info("no diffusers VAE config for %r (%s)", source, err) + return None + # ``AutoencoderKL`` defaults this to 0.18215, and diffusers checkpoints that + # accept the default omit the key entirely. + factor = (raw or {}).get("scaling_factor", 0.18215) + return float(factor) if factor else None def build_diffusion_pipeline_metadata( @@ -2136,7 +2463,6 @@ def add_encoder( { "name": stage_name, "strategy": {"kind": "single_pass", "model": name}, - "run_on": "prompt_only", } ) phases[name] = {"run_on": "prompt_only"} @@ -2177,12 +2503,10 @@ def add_encoder( { "name": "fuse_embeddings", "strategy": {"kind": "single_pass", "model": "embedding"}, - "run_on": "prompt_only", }, { "name": "decode", "strategy": {"kind": "autoregressive", "decoder": "decoder"}, - "run_on": "every_step", }, ] ) @@ -2194,7 +2518,6 @@ def add_encoder( "models": models, "dataflow": dataflow, "strategy": {"kind": "composite", "stages": stages}, - "phases": phases, } return metadata @@ -2279,19 +2602,13 @@ def build_speech_to_text_pipeline_metadata( { "name": "encode_audio", "strategy": {"kind": "single_pass", "model": "encoder"}, - "run_on": "prompt_only", }, { "name": "decode_transcript", "strategy": {"kind": "autoregressive", "decoder": "decoder"}, - "run_on": "every_step", }, ], }, - "phases": { - "encoder": {"run_on": "prompt_only"}, - "decoder": {"run_on": "every_step"}, - }, } return metadata @@ -2311,241 +2628,6 @@ def write_speech_to_text_pipeline_metadata( return path -def build_audio_codec_pipeline_metadata( - *, - encoder_filename: str = "encoder/model.onnx", - decoder_filename: str = "decoder/model.onnx", - codes_dtype: str = "int64", -) -> dict[str, Any]: - """Build metadata for an audio-to-audio neural codec pipeline. - - This is the pure single-pass composite shape (DESIGN.md §20): an audio - encoder maps a waveform to ``codes``, and a decoder reconstructs a waveform - from those codes. Both stages run once over the shared tensor pool (there is - no autoregressive decode and no tokenizer), wired ``encoder.codes -> - decoder.codes``. - - Args: - encoder_filename: Waveform-to-codes encoder ONNX filename. - decoder_filename: Codes-to-waveform decoder ONNX filename. - codes_dtype: Metadata dtype of the ``codes`` tensor exchanged between the - two stages (neural codecs typically emit ``int64`` code indices). - - Returns: - A dict with a top-level ``pipeline`` key. No decoder capabilities are - emitted because the pipeline produces tensors, not tokens. - """ - return { - "pipeline": { - "models": { - "encoder": {"filename": encoder_filename, "type": "audio_encoder"}, - "decoder": {"filename": decoder_filename, "type": "vocoder"}, - }, - "dataflow": [ - { - "from": "encoder.codes", - "to": "decoder.codes", - "dtype": codes_dtype, - "device_transfer": False, - } - ], - "strategy": { - "kind": "composite", - "stages": [ - { - "name": "encode_waveform", - "strategy": {"kind": "single_pass", "model": "encoder"}, - "run_on": "prompt_only", - }, - { - "name": "decode_waveform", - "strategy": {"kind": "single_pass", "model": "decoder"}, - "run_on": "prompt_only", - }, - ], - }, - "phases": { - "encoder": {"run_on": "prompt_only"}, - "decoder": {"run_on": "prompt_only"}, - }, - } - } - - -def write_audio_codec_pipeline_metadata( - directory: str, - *, - filename: str = "inference_metadata.yaml", - **kwargs: Any, -) -> str: - """Build and write composite audio-codec metadata into ``directory``.""" - metadata = build_audio_codec_pipeline_metadata(**kwargs) - os.makedirs(directory, exist_ok=True) - path = os.path.join(directory, filename) - with open(path, "w", encoding="utf-8") as handle: - yaml.safe_dump(metadata, handle, sort_keys=False) - return path - - -def build_tts_pipeline_metadata( - *, - num_code_groups: int, - max_frames: int = 2000, - talker_filename: str = "talker/model.onnx", - code_predictor_filename: str = "code_predictor/model.onnx", - pre_embedder_filename: str = "talker_step_embedder/model.onnx", - prefill_embedder_filename: str | None = "talker_prefill_embedder/model.onnx", - tokenizer_filename: str = "tokenizer.json", - activation_dtype: str = "fp32", - inner_embedding_output: str = "codec_embeddings", - decoder_metadata: dict[str, Any] | None = None, -) -> dict[str, Any]: - """Build metadata for a pre-embedder-driven multi-decoder TTS pipeline. - - This is the real Qwen3-TTS shape (DESIGN.md §20.3, ``nested_autoregressive`` - with the optional ``pre_embedder`` extension): an OUTER ``talker`` AR loop - where each frame drives an INNER ``code_predictor`` AR loop of - ``num_code_groups`` steps (seeded by the talker's ``last_hidden_state``). - Unlike the plain nested shape, the talker is **not** driven by ``input_ids``: - each frame its ``inputs_embeds`` is materialized from the previous frame's - codes by the ``talker_step_embedder`` pre-embedder (``frame_codes - [+ text_embed] -> inputs_embeds``), keeping the engine generic. - - When ``prefill_embedder_filename`` is set (the default), a - ``talker_prefill_embedder`` prompt-phase component is also emitted. It maps - the tokenized prompt ``text_ids -> prefill_embeds + trailing_text_embeds``: - the runtime feeds ``prefill_embeds`` to the talker on frame 0 and threads - ``trailing_text_embeds[:, k-1, :]`` as the pre-embedder's ``text_embed`` on - frames k>=1 (see the ``prefill_embedder`` field). Pass ``None`` to emit the - prefill-less shape (talker frame 0 + ``text_embed`` fed zeros). - - The engine-driven components are emitted (``talker``, ``code_predictor``, - ``talker_step_embedder``, and ``talker_prefill_embedder`` when present). The - package's ``embedding`` and optional ``speaker_encoder`` models are internal - weight sources already folded into the pre-/prefill-embedders, so they are - not declared as pipeline models. There is **no in-package vocoder** — the - assembled ``talker.output_codes`` are decoded by a separate codec model. - - Args: - num_code_groups: Codes collected per outer frame (RVQ residual count). - max_frames: Maximum number of outer talker frames to generate. - talker_filename: Outer decoder (talker) ONNX filename. - code_predictor_filename: Inner decoder ONNX filename. - pre_embedder_filename: ``talker_step_embedder`` ONNX filename. - prefill_embedder_filename: ``talker_prefill_embedder`` ONNX filename, or - ``None`` to omit the prefill/trailing-text path. - tokenizer_filename: Tokenizer filename used by the talker. - decoder_metadata: Optional output from - :func:`decoder_metadata_from_config`; its decoder capabilities are - retained at the document top level. - - Returns: - A dict with a top-level ``pipeline`` key and any decoder capabilities. - """ - if num_code_groups < 1: - raise ValueError("num_code_groups must be at least 1") - if max_frames < 1: - raise ValueError("max_frames must be at least 1") - - models: dict[str, Any] = { - "talker": { - "filename": talker_filename, - "type": "decoder", - "tokenizer": tokenizer_filename, - }, - "talker_step_embedder": { - "filename": pre_embedder_filename, - "type": "embedding", - }, - "code_predictor": { - "filename": code_predictor_filename, - "type": "decoder", - }, - } - dataflow: list[dict[str, Any]] = [ - { - "from": "talker_step_embedder.inputs_embeds", - "to": "talker.inputs_embeds", - "dtype": activation_dtype, - "device_transfer": False, - }, - { - "from": "talker.last_hidden_state", - "to": "code_predictor.inputs_embeds", - "dtype": activation_dtype, - "device_transfer": False, - }, - ] - stage_strategy: dict[str, Any] = { - "kind": "nested_autoregressive", - "outer": "talker", - "inner": "code_predictor", - "inner_embedding_output": inner_embedding_output, - "pre_embedder": { - "component": "talker_step_embedder", - "frame_codes_input": "frame_codes", - "text_embed_input": "text_embed", - }, - "num_code_groups": num_code_groups, - "max_tokens": max_frames, - } - phases: dict[str, Any] = { - "talker": {"run_on": "every_step"}, - "talker_step_embedder": {"run_on": "on_demand"}, - "code_predictor": {"run_on": "every_step"}, - } - - if prefill_embedder_filename is not None: - models["talker_prefill_embedder"] = { - "filename": prefill_embedder_filename, - "type": "embedding", - } - # Runs once in the prompt phase; the runtime seeds the declared - # `prompt_input` with the tokenized prompt and reads the two named - # outputs from the pool. Every port is declared explicitly (the engine - # never guesses tensor names). - stage_strategy["prefill_embedder"] = { - "component": "talker_prefill_embedder", - "prompt_input": "text_ids", - "prefill_output": "prefill_embeds", - "trailing_output": "trailing_text_embeds", - } - phases["talker_prefill_embedder"] = {"run_on": "prompt_only"} - - metadata = dict(decoder_metadata or {}) - metadata["pipeline"] = { - "models": models, - "dataflow": dataflow, - "strategy": { - "kind": "composite", - "stages": [ - { - "name": "generate_codes", - "strategy": stage_strategy, - "run_on": "every_step", - }, - ], - }, - "phases": phases, - } - return metadata - - -def write_tts_pipeline_metadata( - directory: str, - *, - filename: str = "inference_metadata.yaml", - **kwargs: Any, -) -> str: - """Build and write pre-embedder-driven TTS metadata into ``directory``.""" - metadata = build_tts_pipeline_metadata(**kwargs) - os.makedirs(directory, exist_ok=True) - path = os.path.join(directory, filename) - with open(path, "w", encoding="utf-8") as handle: - yaml.safe_dump(metadata, handle, sort_keys=False) - return path - - def write_diffusion_pipeline_metadata( directory: str, *, diff --git a/src/mobius/integrations/onnx_genai/inference_metadata_test.py b/src/mobius/integrations/onnx_genai/inference_metadata_test.py index 06aef903f..8bd9c345d 100644 --- a/src/mobius/integrations/onnx_genai/inference_metadata_test.py +++ b/src/mobius/integrations/onnx_genai/inference_metadata_test.py @@ -1,6 +1,5 @@ # Copyright (c) Microsoft Corporation. # Licensed under the MIT License. - """Tests for onnx-genai diffusion inference_metadata generation.""" from __future__ import annotations @@ -22,26 +21,162 @@ declare_component_presence, declare_optional_input, ) +from mobius.generation import build_greedy_sampler from mobius.integrations.onnx_genai.inference_metadata import ( SchedulerConfig, _decoder_io, _input_source_map, + _match_max_token_grid, _port, - add_explicit_package_io, + _processor_values, + add_policy_components_to_workflow, build_diffusion_pipeline_metadata, - build_language_diffusion_pipeline_metadata, build_multimodal_pipeline_metadata, build_native_vlm_package_metadata, - build_tts_pipeline_metadata, is_native_vlm_package, load_diffusers_scheduler_config, validate_executable_closure, write_diffusion_pipeline_metadata, write_native_vlm_package_metadata, - write_tts_pipeline_metadata, ) +def test_ort_extensions_processor_config_supplies_structural_values(tmp_path): + (tmp_path / "processor_config.json").write_text( + json.dumps( + { + "processor": { + "transforms": [ + { + "operation": { + "type": "Resize", + "attrs": { + "min_pixels": 784, + "max_pixels": 2371600, + "patch_size": 14, + "merge_size": 2, + }, + } + }, + { + "operation": { + "type": "Rescale", + "attrs": {"rescale_factor": 1 / 255}, + } + }, + { + "operation": { + "type": "Normalize", + "attrs": { + "mean": [0.5, 0.5, 0.5], + "std": [0.5, 0.5, 0.5], + }, + } + }, + { + "operation": { + "type": "PatchImage", + "attrs": {"temporal_patch_size": 2}, + } + }, + ] + } + } + ), + encoding="utf-8", + ) + + values = _processor_values(str(tmp_path), object()) + + assert values["size"] == {"shortest_edge": 784, "longest_edge": 2371600} + assert values["patch_size"] == 14 + assert values["merge_size"] == 2 + assert values["temporal_patch_size"] == 2 + assert values["image_mean"] == [0.5, 0.5, 0.5] + assert values["image_std"] == [0.5, 0.5, 0.5] + + +def test_max_token_packed_grid_derives_pixel_area_bounds(): + program = _match_max_token_grid( + [ + _port(_value("pixel_values", ir.DataType.FLOAT, ["patches", 1176])), + _port(_value("image_grid_thw", ir.DataType.INT64, ["images", 3])), + ], + { + "patch_size": 14, + "temporal_patch_size": 2, + "merge_size": 2, + "max_image_tokens": 4096, + }, + ) + + assert program is not None + resize = next( + transform + for transform in program.transforms( + None, + { + "patch_size": 14, + "temporal_patch_size": 2, + "merge_size": 2, + "max_image_tokens": 4096, + }, + ) + if transform["op"] == "resize" + ) + assert resize["min_pixels"] == 784 + assert resize["max_pixels"] == 3_211_264 + + +def test_workflow_policy_components_reference_saved_onnx_artifacts(tmp_path): + package = ModelPackage({"model": _model("model", [], [])}) + package.add_policy_component("sample", build_greedy_sampler()) + package.save(str(tmp_path)) + metadata = { + "pipeline": { + "workflow": { + "manifest": {}, + "components": {}, + "graph": {"kind": "sequence", "nodes": []}, + } + } + } + + add_policy_components_to_workflow(metadata, package) + + component = metadata["pipeline"]["workflow"]["components"]["sample"] + assert component["implementation"] == { + "kind": "onnx", + "artifact": "policies/sample.onnx", + } + assert component["ports"] == { + "inputs": { + "logits": { + "dtype": "float32", + "rank": 2, + "shape": ["batch", "vocabulary"], + "batch_layout": {"kind": "request_aligned", "axis": 0}, + } + }, + "outputs": { + "token": { + "dtype": "int64", + "rank": 1, + "shape": ["batch"], + "batch_layout": {"kind": "request_aligned", "axis": 0}, + } + }, + } + assert component["contract"] == { + "id": "onnx-genai.token-sampler", + "version": "1", + "bindings": {"logits": "logits", "token": "token"}, + "parameters": {"mode": "greedy"}, + } + assert "effects" not in component + assert (tmp_path / component["implementation"]["artifact"]).is_file() + + def _onnx_genai_schema_path() -> str | None: """Locate onnx-genai's committed pipeline JSON schema, if available.""" candidates = [ @@ -257,112 +392,7 @@ def _static_cache_decoder_model() -> ir.Model: return _model("decoder", inputs, outputs) -class TestExplicitPackageIo: - def test_emits_explicit_dynamic_decoder_roles(self): - model = _decoder_model( - [], - position_shape=["batch", "sequence"], - raw_token_input=True, - kv_head_dims=[8, 16], - ) - metadata = add_explicit_package_io({"model": {}}, {"model": model}, _VlmConfig()) - io = metadata["model"]["io"] - assert io["token_input"] == "input_ids" - assert io["sequence_source"] == "token_ids" - assert io["logits_output"] == "logits" - assert io["kv_ownership"] == "owned" - assert io["kv_inputs"] == [ - "past_key_values.0.key", - "past_key_values.0.value", - "past_key_values.1.key", - "past_key_values.1.value", - ] - assert io["kv_outputs"] == [ - "present.0.key", - "present.0.value", - "present.1.key", - "present.1.value", - ] - - def test_emits_explicit_static_cache_roles_in_layer_order(self): - metadata = add_explicit_package_io( - {"model": {}}, {"model": _static_cache_decoder_model()}, _VlmConfig() - ) - io = metadata["model"]["io"] - assert io["static_cache"] == { - "write_indices_input": "write_indices", - "kv_sequence_length_input": "nonpad_kv_seqlen", - "key_cache_inputs": ["key_cache.1", "key_cache.3"], - "value_cache_inputs": ["value_cache.1", "value_cache.3"], - "key_cache_outputs": ["updated_key_cache.1", "updated_key_cache.3"], - "value_cache_outputs": [ - "updated_value_cache.1", - "updated_value_cache.3", - ], - } - assert "kv_inputs" not in io - - @pytest.mark.parametrize( - ("encoder_input", "role_field"), - [ - ( - _value( - "mel_prompt", - ir.DataType.FLOAT, - ["batch", 80, "audio_sequence"], - ), - "audio_features_input", - ), - ( - _value("prompt_tokens", ir.DataType.INT64, ["batch", "sequence"]), - "token_input", - ), - ], - ) - def test_emits_explicit_encoder_prompt_role(self, encoder_input, role_field): - encoder = _model( - "encoder", - [encoder_input], - [ - ( - "encoder_hidden_states", - ir.DataType.FLOAT, - ["batch", "encoder_sequence", 64], - ) - ], - ) - decoder = _model( - "decoder", - [ - _value("decoder_tokens", ir.DataType.INT64, ["batch", "sequence"]), - _value( - "encoder_hidden_states", - ir.DataType.FLOAT, - ["batch", "encoder_sequence", 64], - ), - ], - [("logits", ir.DataType.FLOAT, ["batch", "sequence", 128])], - ) - metadata = { - "pipeline": { - "models": { - "encoder": {"type": "encoder"}, - "decoder": {"type": "decoder"}, - }, - "dataflow": [], - "strategy": {"kind": "autoregressive", "decoder": "decoder"}, - } - } - add_explicit_package_io( - metadata, {"encoder": encoder, "decoder": decoder}, _VlmConfig() - ) - encoder_io = metadata["pipeline"]["models"]["encoder"]["io"] - assert encoder_io[role_field] == encoder_input.name - decoder_io = metadata["pipeline"]["models"]["decoder"]["io"] - assert decoder_io["token_input"] == "decoder_tokens" - assert decoder_io["encoder_hidden_states_input"] == "encoder_hidden_states" - assert decoder_io["logits_output"] == "logits" - +class TestCrossAttentionCacheSources: def test_cross_attention_cache_inputs_are_loop_state(self): inputs = [ _value("input_ids", ir.DataType.INT64, ["batch", "sequence"]), @@ -476,42 +506,6 @@ def _nested_metadata(inner_embedding_output=None): } } - def test_explicit_inner_embedding_output_is_preserved(self): - metadata = self._nested_metadata("declared_embedding") - package = self._nested_package( - [("logits", ir.DataType.FLOAT, ["batch", "sequence", 128])] - ) - add_explicit_package_io(metadata, package, _VlmConfig()) - assert ( - metadata["pipeline"]["strategy"]["inner_embedding_output"] == "declared_embedding" - ) - - def test_missing_inner_embedding_output_is_derived(self): - metadata = self._nested_metadata() - package = self._nested_package( - [ - ("logits", ir.DataType.FLOAT, ["batch", "sequence", 128]), - ("codec_embeddings", ir.DataType.FLOAT, [16, 64]), - ] - ) - add_explicit_package_io(metadata, package, _VlmConfig()) - assert metadata["pipeline"]["strategy"]["inner_embedding_output"] == "codec_embeddings" - - def test_ambiguous_inner_embedding_output_fails_actionably(self): - metadata = self._nested_metadata() - package = self._nested_package( - [ - ("logits", ir.DataType.FLOAT, ["batch", "sequence", 128]), - ("first_embedding", ir.DataType.FLOAT, [16, 64]), - ("second_embedding", ir.DataType.FLOAT, [16, 64]), - ] - ) - with pytest.raises( - ValueError, - match=r"pipeline\.strategy\.inner_embedding_output.*no unique", - ): - add_explicit_package_io(metadata, package, _VlmConfig()) - def _native_package( vision_encoder: ir.Model, @@ -757,33 +751,8 @@ def test_gemma4_routes_all_embedding_outputs(self, tmp_path): "absent": {"kind": "zeros", "shape": [0, 64]}, }, } - assert emitted_yaml["pipeline"]["phases"]["vision_encoder"] == { - "run_on": "prompt_only", - "when_present": "image", - } - assert emitted_yaml["pipeline"]["phases"]["audio_encoder"] == { - "run_on": "prompt_only", - "when_present": "audio", - } - vision_stage = next( - stage - for stage in metadata["pipeline"]["strategy"]["stages"] - if stage["strategy"].get("model") == "vision_encoder" - ) - assert vision_stage["run_on"] == "prompt_only" - audio_stage = next( - stage - for stage in metadata["pipeline"]["strategy"]["stages"] - if stage["strategy"].get("model") == "audio_encoder" - ) - assert audio_stage["run_on"] == "prompt_only" - assert metadata["pipeline"]["phases"]["embedding"] == {"run_on": "every_step"} - embedding_stage = next( - stage - for stage in metadata["pipeline"]["strategy"]["stages"] - if stage["strategy"].get("model") == "embedding" - ) - assert embedding_stage["run_on"] == "every_step" + assert "phases" not in emitted_yaml["pipeline"] + assert "phases" not in metadata["pipeline"] assert metadata["pipeline"]["models"]["embedding"]["io"]["token_input"] == "input_ids" assert metadata["pipeline"]["vision"]["token_count_source"] == "from_coordinates" assert metadata["pipeline"]["vision"]["token_pooling_factor"] == 9 @@ -795,6 +764,8 @@ def test_gemma4_routes_all_embedding_outputs(self, tmp_path): "max_patches": 2520, "pooling_kernel_size": 3, "interpolation": "bicubic", + "inputs": ["image.transform_0"], + "outputs": ["image.transform_1"], } assert ( next(transform for transform in transforms if transform["op"] == "pad")[ @@ -803,13 +774,15 @@ def test_gemma4_routes_all_embedding_outputs(self, tmp_path): == 2520 ) assert not any(transform["op"] == "normalize" for transform in transforms) - assert metadata["model"]["io"]["token_input"] == "input_ids" - assert metadata["model"]["io"]["kv_inputs"] == [ + assert "model" not in metadata or "io" not in metadata["model"] + decoder_io = metadata["pipeline"]["models"]["decoder"]["io"] + assert decoder_io["token_input"] == "input_ids" + assert decoder_io["kv_inputs"] == [ f"past_key_values.{layer}.{role}" for layer in range(3) for role in ("key", "value") ] - assert metadata["model"]["io"]["kv_outputs"] == [ + assert decoder_io["kv_outputs"] == [ f"present.{layer}.{role}" for layer in range(3) for role in ("key", "value") ] kv_inputs = { @@ -819,20 +792,23 @@ def test_gemma4_routes_all_embedding_outputs(self, tmp_path): } assert kv_inputs["past_key_values.1.key"]["shape"][-1] == 16 image_outputs = metadata["preprocessing"]["image"]["outputs"] + # Every output binds the processor-local value that produces it, so the + # runtime never has to guess which transform an output came from. assert image_outputs == [ { "name": "vision_encoder.pixel_values", "content": "pixels", "dtype": "fp32", + "source": "image.transform_4", }, { "name": "vision_encoder.pixel_position_ids", "content": "patch_coordinates", "dtype": "int64", "pad_value": -1, + "source": "image.output_patch_coordinates", }, ] - broken = copy.deepcopy(metadata) broken["pipeline"]["dataflow"] = [ edge @@ -948,7 +924,7 @@ def test_qwen_packed_grid_rank3_positions_sparse_and_fixed_state(self, tmp_path) "sections": [16, 24, 24], "processor_summaries": ["vision_encoder.image_grid_thw"], } - io = metadata["model"]["io"] + io = metadata["pipeline"]["models"]["decoder"]["io"] assert io["kv_inputs"] == [ "past_key_values.0.key", "past_key_values.0.value", @@ -1130,7 +1106,7 @@ def test_equal_shape_key_value_ports_remain_declared_kv(self, tmp_path): metadata = build_native_vlm_package_metadata( package, config=config, source=str(source) ) - io = metadata["model"]["io"] + io = metadata["pipeline"]["models"]["decoder"]["io"] assert io["kv_update"] == "append" assert io["kv_inputs"] == [ "past_key_values.0.key", @@ -1366,6 +1342,8 @@ def test_cached_gemma_processor_matches_emitted_patch_budget(self): "max_patches": 2520, "pooling_kernel_size": 3, "interpolation": "bicubic", + "inputs": ["image.transform_0"], + "outputs": ["image.transform_1"], } assert pad["target_length"] == reference["pixel_values"].shape[1] == 2520 assert patchify["channel_order"] == "channels_last" @@ -1729,78 +1707,6 @@ def test_matches_onnx_genai_json_schema(self): jsonschema.validate(instance=meta, schema=schema) -class TestLanguageDiffusionMetadata: - def test_minimal_masked_diffusion_pipeline(self): - meta = build_language_diffusion_pipeline_metadata( - mask_token_id=126336, num_inference_steps=128 - ) - pipeline = meta["pipeline"] - assert pipeline["models"]["denoiser"] == { - "filename": "model.onnx", - "type": "denoiser", - } - # Loop-carried self-edge: logits refine the token sequence. - assert pipeline["dataflow"] == [ - {"from": "denoiser.logits", "to": "denoiser.input_ids"} - ] - strategy = pipeline["strategy"] - assert strategy["kind"] == "iterative" - assert strategy["num_steps"] == 128 - assert strategy["scheduler_config"] == { - "kind": "masked_diffusion", - "mask_token_id": 126336, - } - assert "guidance_scale" not in strategy - - def test_semi_autoregressive_with_temperature_and_cfg(self): - meta = build_language_diffusion_pipeline_metadata( - mask_token_id=5, - num_inference_steps=64, - block_length=32, - temperature=0.2, - guidance_scale=2.5, # LLaDA cfg_scale=1.5 => cfg_scale + 1 - ) - strategy = meta["pipeline"]["strategy"] - assert strategy["guidance_scale"] == pytest.approx(2.5) - assert strategy["scheduler_config"]["block_length"] == 32 - assert strategy["scheduler_config"]["temperature"] == pytest.approx(0.2) - - def test_custom_ports(self): - meta = build_language_diffusion_pipeline_metadata( - mask_token_id=1, - num_inference_steps=8, - model_filename="llada.onnx", - input_ids_port="tokens", - logits_port="scores", - ) - pipeline = meta["pipeline"] - assert pipeline["models"]["denoiser"]["filename"] == "llada.onnx" - assert pipeline["dataflow"] == [{"from": "denoiser.scores", "to": "denoiser.tokens"}] - - def test_rejects_zero_steps(self): - with pytest.raises(ValueError): - build_language_diffusion_pipeline_metadata(mask_token_id=1, num_inference_steps=0) - - def test_matches_onnx_genai_json_schema(self): - schema_path = _onnx_genai_schema_path() - if schema_path is None: - pytest.skip("onnx-genai schema not found (set ONNX_GENAI_SCHEMA)") - import json - - import jsonschema - - with open(schema_path) as handle: - schema = json.load(handle) - meta = build_language_diffusion_pipeline_metadata( - mask_token_id=126336, - num_inference_steps=64, - block_length=32, - temperature=0.0, - guidance_scale=2.5, - ) - jsonschema.validate(instance=meta, schema=schema) - - class TestBuildMultimodalPipelineMetadata: def test_vision_only_pipeline(self): metadata = build_multimodal_pipeline_metadata( @@ -1844,7 +1750,6 @@ def test_vision_only_pipeline(self): "kind": "single_pass", "model": "vision_encoder", }, - "run_on": "prompt_only", }, { "name": "fuse_embeddings", @@ -1852,7 +1757,6 @@ def test_vision_only_pipeline(self): "kind": "single_pass", "model": "embedding", }, - "run_on": "prompt_only", }, { "name": "decode", @@ -1860,15 +1764,9 @@ def test_vision_only_pipeline(self): "kind": "autoregressive", "decoder": "decoder", }, - "run_on": "every_step", }, ], }, - "phases": { - "vision_encoder": {"run_on": "prompt_only"}, - "embedding": {"run_on": "prompt_only"}, - "decoder": {"run_on": "every_step"}, - }, } } @@ -1919,104 +1817,17 @@ def test_vision_and_audio_pipeline(self): { "name": "encode_vision", "strategy": {"kind": "single_pass", "model": "vision_encoder"}, - "run_on": "prompt_only", }, { "name": "encode_audio", "strategy": {"kind": "single_pass", "model": "audio_encoder"}, - "run_on": "prompt_only", }, { "name": "fuse_embeddings", "strategy": {"kind": "single_pass", "model": "embedding"}, - "run_on": "prompt_only", }, { "name": "decode", "strategy": {"kind": "autoregressive", "decoder": "decoder"}, - "run_on": "every_step", }, ] - - -class TestBuildTTSPipelineMetadata: - """Pre-embedder-driven multi-decoder TTS (Qwen3-TTS) metadata.""" - - def test_minimal_nested_autoregressive_with_pre_embedder(self): - meta = build_tts_pipeline_metadata( - num_code_groups=16, max_frames=1000, prefill_embedder_filename=None - ) - pipe = meta["pipeline"] - assert set(pipe["models"]) == {"talker", "talker_step_embedder", "code_predictor"} - assert pipe["models"]["talker"]["type"] == "decoder" - assert pipe["models"]["talker"]["tokenizer"] == "tokenizer.json" - assert pipe["models"]["talker_step_embedder"]["type"] == "embedding" - - stage = pipe["strategy"]["stages"][0]["strategy"] - assert stage["kind"] == "nested_autoregressive" - assert stage["outer"] == "talker" - assert stage["inner"] == "code_predictor" - assert stage["inner_embedding_output"] == "codec_embeddings" - assert stage["pre_embedder"]["component"] == "talker_step_embedder" - assert stage["pre_embedder"]["frame_codes_input"] == "frame_codes" - assert "prefill_embedder" not in stage - assert stage["num_code_groups"] == 16 - assert stage["max_tokens"] == 1000 - - # Required pre-embedder feed edge + inner seed edge. - assert { - "from": "talker_step_embedder.inputs_embeds", - "to": "talker.inputs_embeds", - "dtype": "fp32", - "device_transfer": False, - } in pipe["dataflow"] - assert { - "from": "talker.last_hidden_state", - "to": "code_predictor.inputs_embeds", - "dtype": "fp32", - "device_transfer": False, - } in pipe["dataflow"] - # No in-package vocoder. - assert "vocoder" not in pipe["models"] - # Pre-embedder is a loop-internal on_demand component. - assert pipe["phases"]["talker_step_embedder"]["run_on"] == "on_demand" - - def test_with_prefill_embedder(self): - # Default emits the prefill/trailing-text component (prompt phase). - meta = build_tts_pipeline_metadata(num_code_groups=16) - pipe = meta["pipeline"] - assert "talker_prefill_embedder" in pipe["models"] - assert pipe["models"]["talker_prefill_embedder"]["type"] == "embedding" - stage = pipe["strategy"]["stages"][0]["strategy"] - assert stage["prefill_embedder"]["component"] == "talker_prefill_embedder" - assert stage["prefill_embedder"]["prompt_input"] == "text_ids" - assert stage["prefill_embedder"]["prefill_output"] == "prefill_embeds" - assert stage["prefill_embedder"]["trailing_output"] == "trailing_text_embeds" - assert pipe["phases"]["talker_prefill_embedder"]["run_on"] == "prompt_only" - - def test_rejects_invalid_code_groups(self): - with pytest.raises(ValueError, match="num_code_groups"): - build_tts_pipeline_metadata(num_code_groups=0) - - def test_write_roundtrip(self, tmp_path): - path = write_tts_pipeline_metadata(str(tmp_path), num_code_groups=8) - with open(path) as handle: - loaded = yaml.safe_load(handle) - stage = loaded["pipeline"]["strategy"]["stages"][0]["strategy"] - assert stage["pre_embedder"]["component"] == "talker_step_embedder" - assert stage["pre_embedder"]["frame_codes_input"] == "frame_codes" - assert stage["num_code_groups"] == 8 - - def test_matches_onnx_genai_json_schema(self): - """Emitted TTS metadata validates against onnx-genai's published schema.""" - schema_path = _onnx_genai_schema_path() - if schema_path is None: - pytest.skip("onnx-genai schema not found (set ONNX_GENAI_SCHEMA)") - import json - - import jsonschema - - with open(schema_path) as handle: - schema = json.load(handle) - meta = build_tts_pipeline_metadata(num_code_groups=16, max_frames=2000) - jsonschema.validate(instance=meta, schema=schema) diff --git a/src/mobius/integrations/onnx_genai/performance.py b/src/mobius/integrations/onnx_genai/performance.py new file mode 100644 index 000000000..c5edc2371 --- /dev/null +++ b/src/mobius/integrations/onnx_genai/performance.py @@ -0,0 +1,221 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT License. + +"""Performance-conformance records for metadata-driven ONNX GenAI workflows.""" + +from __future__ import annotations + +from dataclasses import dataclass +from typing import Any + + +@dataclass(frozen=True) +class PerformanceComparison: + """Result of comparing one workflow run with its native control.""" + + failures: tuple[str, ...] + observations: tuple[str, ...] + + @property + def passed(self) -> bool: + return not self.failures + + +_IDENTITY_FIELDS = ( + "scenario", + "workload", + "model", + "model_hash", + "runtime_commit", + "runtime_build", + "execution_provider", + "provider_options", + "device", + "precision", + "batch_size", + "input_shapes", + "work_units", + "sampling_algorithm", + "sampling_parameters", + "rng_seed", + "rng_offset", + "kv_mode", + "kv_capacity", + "graph_capture", + "graph_capture_shape", + "warmup_count", + "measured_iterations", + "synchronization_timing", + "required_islands", +) + + +def _number(record: dict[str, Any], name: str) -> float: + value = record.get(name) + if not isinstance(value, int | float) or isinstance(value, bool): + raise TypeError(f"metric {name!r} must be numeric") + return float(value) + + +def _validate_identity(workflow: dict[str, Any], native: dict[str, Any]) -> None: + workflow_identity = workflow.get("identity") + native_identity = native.get("identity") + if not isinstance(workflow_identity, dict) or not isinstance(native_identity, dict): + raise TypeError("both records require an identity object") + missing = [ + name + for name in _IDENTITY_FIELDS + if name not in workflow_identity or name not in native_identity + ] + if missing: + raise ValueError( + "performance records omit required identity fields: " + ", ".join(missing) + ) + mismatches = [ + name for name in _IDENTITY_FIELDS if workflow_identity[name] != native_identity[name] + ] + extra_mismatches = [ + name + for name in workflow_identity.keys() | native_identity.keys() + if workflow_identity.get(name) != native_identity.get(name) + ] + mismatches.extend(name for name in extra_mismatches if name not in mismatches) + if mismatches: + raise ValueError( + "performance records are not comparable; identity differs for " + + ", ".join(mismatches) + ) + + +def _required_island_failures(record: dict[str, Any]) -> list[str]: + identity = record["identity"] + if not identity.get("graph_capture"): + return [] + required = identity.get("required_islands", []) + diagnostics = record.get("islands", []) + failures: list[str] = [] + for expected in required: + components = expected["components"] + match = next( + (island for island in diagnostics if island.get("components") == components), + None, + ) + label = " -> ".join(components) + if match is None: + failures.append(f"required execution island was not formed: {label}") + elif match.get("device") != expected["device"]: + failures.append( + f"required execution island used {match.get('device')!r}, " + f"expected {expected['device']!r}: {label}" + ) + elif _number(match, "captures") < 1 or _number(match, "replays") < 1: + reason = match.get("fallback_reason") or "capture/replay was not observed" + failures.append(f"required execution island was not replayed: {label}: {reason}") + return failures + + +def compare_performance( + workflow: dict[str, Any], + native: dict[str, Any], + *, + max_regression_fraction: float = 0.05, +) -> PerformanceComparison: + """Compare equivalent workflow/native measurements and enforce the release bar.""" + if not 0 <= max_regression_fraction < 1: + raise ValueError("max_regression_fraction must be in [0, 1)") + _validate_identity(workflow, native) + workflow_metrics = workflow.get("metrics") + native_metrics = native.get("metrics") + if not isinstance(workflow_metrics, dict) or not isinstance(native_metrics, dict): + raise TypeError("both records require a metrics object") + required_metrics = { + "throughput_unit", + "throughput", + "peak_memory_bytes", + "device_sync_count", + "host_to_device_copy_count", + "host_to_device_bytes", + "device_to_host_copy_count", + "device_to_host_bytes", + "session_boundary_count", + "kernel_launch_count", + "device_resident_intermediate_count", + "intermediate_value_count", + } + if workflow["identity"]["workload"] == "generation": + required_metrics.add("ttft_ms") + missing_metrics = [ + name + for name in sorted(required_metrics) + if name not in workflow_metrics or name not in native_metrics + ] + if missing_metrics: + raise ValueError( + "performance records omit required metrics: " + ", ".join(missing_metrics) + ) + + failures = _required_island_failures(workflow) + observations: list[str] = [] + lower_is_better = ("ttft_ms", "peak_memory_bytes") + exact_or_better = ( + "device_sync_count", + "host_to_device_copy_count", + "host_to_device_bytes", + "device_to_host_copy_count", + "device_to_host_bytes", + "session_boundary_count", + "kernel_launch_count", + ) + + throughput_unit = workflow_metrics.get("throughput_unit") + if throughput_unit != native_metrics.get("throughput_unit"): + raise ValueError("throughput units differ") + workflow_throughput = _number(workflow_metrics, "throughput") + native_throughput = _number(native_metrics, "throughput") + throughput_ratio = workflow_throughput / native_throughput + observations.append(f"throughput ratio={throughput_ratio:.4f} {throughput_unit}") + if throughput_ratio < 1 - max_regression_fraction: + failures.append( + f"throughput regressed by {(1 - throughput_ratio) * 100:.2f}% " + f"({workflow_throughput:g} vs {native_throughput:g} {throughput_unit})" + ) + + for name in lower_is_better: + if name not in workflow_metrics or name not in native_metrics: + continue + workflow_value = _number(workflow_metrics, name) + native_value = _number(native_metrics, name) + ratio = workflow_value / native_value if native_value else 1.0 + observations.append(f"{name} ratio={ratio:.4f}") + if workflow_value > native_value * (1 + max_regression_fraction): + failures.append( + f"{name} regressed by {(ratio - 1) * 100:.2f}% " + f"({workflow_value:g} vs {native_value:g})" + ) + + for name in exact_or_better: + workflow_value = _number(workflow_metrics, name) + native_value = _number(native_metrics, name) + observations.append(f"{name}={workflow_value:g} vs {native_value:g}") + if workflow_value > native_value: + failures.append(f"{name} increased ({workflow_value:g} vs {native_value:g})") + + workflow_intermediates = _number(workflow_metrics, "intermediate_value_count") + native_intermediates = _number(native_metrics, "intermediate_value_count") + if workflow_intermediates <= 0 or native_intermediates <= 0: + raise ValueError("intermediate_value_count must be positive") + workflow_residency = ( + _number(workflow_metrics, "device_resident_intermediate_count") + / workflow_intermediates + ) + native_residency = ( + _number(native_metrics, "device_resident_intermediate_count") / native_intermediates + ) + observations.append(f"device residency={workflow_residency:.4f} vs {native_residency:.4f}") + if workflow_residency < native_residency: + failures.append( + "device-resident intermediate ratio decreased " + f"({workflow_residency:.4f} vs {native_residency:.4f})" + ) + + return PerformanceComparison(tuple(failures), tuple(observations)) diff --git a/src/mobius/integrations/onnx_genai/performance_test.py b/src/mobius/integrations/onnx_genai/performance_test.py new file mode 100644 index 000000000..e1ac6171b --- /dev/null +++ b/src/mobius/integrations/onnx_genai/performance_test.py @@ -0,0 +1,141 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT License. + +from __future__ import annotations + +import copy + +import pytest + +from mobius.integrations.onnx_genai.performance import compare_performance + + +def _record(*, workflow: bool) -> dict: + return { + "identity": { + "scenario": "decoder-min-p", + "workload": "generation", + "model": "tiny-decoder", + "model_hash": "sha256:fixture", + "runtime_commit": "b2157a2", + "runtime_build": "release+cuda", + "execution_provider": "cuda", + "provider_options": {"enable_cuda_graph": True}, + "device": "cuda:0", + "precision": "float16", + "batch_size": 1, + "input_shapes": {"prompt_tokens": [1, 8]}, + "work_units": {"prompt_tokens": 8, "generated_tokens": 32}, + "sampling_algorithm": "seeded_min_p", + "sampling_parameters": {"temperature": 0.8, "min_p": 0.1}, + "rng_seed": [7], + "rng_offset": [0], + "kv_mode": "paged", + "kv_capacity": 128, + "graph_capture": True, + "graph_capture_shape": {"batch": 1, "decode_sequence": 1}, + "warmup_count": 3, + "measured_iterations": 20, + "synchronization_timing": "before_and_after_sample", + "required_islands": [ + { + "device": "cuda", + "components": [ + "decoder", + "last_token_logits", + "token_sampler", + "termination", + ], + } + ], + }, + "metrics": { + "throughput_unit": "tokens_per_second", + "throughput": 102.0 if workflow else 100.0, + "ttft_ms": 9.8 if workflow else 10.0, + "peak_memory_bytes": 1000, + "device_sync_count": 0, + "host_to_device_copy_count": 0, + "host_to_device_bytes": 0, + "device_to_host_copy_count": 0, + "device_to_host_bytes": 0, + "session_boundary_count": 1, + "kernel_launch_count": 12, + "device_resident_intermediate_count": 8, + "intermediate_value_count": 8, + }, + "islands": [ + { + "components": [ + "decoder", + "last_token_logits", + "token_sampler", + "termination", + ], + "device": "cuda", + "capture_eligible": True, + "captures": 1, + "replays": 10, + "fallback_reason": None, + } + ], + } + + +def test_equivalent_captured_workflow_meets_performance_bar(): + result = compare_performance(_record(workflow=True), _record(workflow=False)) + assert result.passed + assert "throughput ratio=1.0200 tokens_per_second" in result.observations + + +def test_reports_throughput_sync_memory_and_capture_gaps(): + workflow = _record(workflow=True) + native = _record(workflow=False) + workflow["metrics"].update( + { + "throughput": 80.0, + "ttft_ms": 12.0, + "peak_memory_bytes": 1200, + "device_sync_count": 2, + "session_boundary_count": 3, + "device_resident_intermediate_count": 6, + } + ) + workflow["islands"][0].update( + {"captures": 0, "replays": 0, "fallback_reason": "dynamic allocation"} + ) + + result = compare_performance(workflow, native) + + assert not result.passed + assert any("throughput regressed by 20.00%" in failure for failure in result.failures) + assert any("ttft_ms regressed by 20.00%" in failure for failure in result.failures) + assert any( + "peak_memory_bytes regressed by 20.00%" in failure for failure in result.failures + ) + assert any("device_sync_count increased" in failure for failure in result.failures) + assert any("session_boundary_count increased" in failure for failure in result.failures) + assert any("device-resident intermediate ratio" in failure for failure in result.failures) + assert any("dynamic allocation" in failure for failure in result.failures) + + +def test_rejects_nonidentical_benchmark_conditions(): + workflow = _record(workflow=True) + native = copy.deepcopy(_record(workflow=False)) + native["identity"]["kv_mode"] = "dense" + + with pytest.raises(ValueError, match="kv_mode"): + compare_performance(workflow, native) + + +def test_requires_memory_metrics_and_island_device(): + workflow = _record(workflow=True) + native = _record(workflow=False) + del workflow["metrics"]["peak_memory_bytes"] + with pytest.raises(ValueError, match="peak_memory_bytes"): + compare_performance(workflow, native) + + workflow = _record(workflow=True) + workflow["islands"][0]["device"] = "cpu" + result = compare_performance(workflow, native) + assert any("expected 'cuda'" in failure for failure in result.failures) diff --git a/src/mobius/integrations/onnx_genai/speech_to_text_workflow_metadata_test.py b/src/mobius/integrations/onnx_genai/speech_to_text_workflow_metadata_test.py new file mode 100644 index 000000000..76fd6d13a --- /dev/null +++ b/src/mobius/integrations/onnx_genai/speech_to_text_workflow_metadata_test.py @@ -0,0 +1,251 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT License. + +"""Tests for encoder-conditioned (Whisper-style) speech-to-text workflow metadata.""" + +from __future__ import annotations + +import json +import os +from types import SimpleNamespace + +import jsonschema +import onnx_ir as ir +import pytest +import yaml + +from mobius._model_package import ModelPackage +from mobius.integrations.onnx_genai.auto_export import _audio_preprocessing_program +from mobius.integrations.onnx_genai.inference_metadata_test import ( + _model, + _value, +) +from mobius.integrations.onnx_genai.workflow_metadata import ( + build_speech_to_text_workflow_metadata, + write_speech_to_text_workflow_metadata, +) + +_WHISPER_EXTRACTOR = { + "feature_extractor_type": "WhisperFeatureExtractor", + "sampling_rate": 16000, + "feature_size": 80, + "n_fft": 400, + "hop_length": 160, + "chunk_length": 30, + "n_samples": 480000, + "padding_value": 0.0, +} + + +def _config() -> SimpleNamespace: + return SimpleNamespace(eos_token_id=50257, max_position_embeddings=448) + + +def _speech_package(*, layers: int = 2) -> ModelPackage: + """A whisper-tiny-shaped encoder/decoder pair with self-attention cache.""" + encoder = _model( + "encoder", + [_value("input_features", ir.DataType.FLOAT, ["batch", 80, "audio_seq_len"])], + [("encoder_hidden_states", ir.DataType.FLOAT, ["batch", 1500, 384])], + ) + decoder_inputs = [ + _value("decoder_input_ids", ir.DataType.INT64, ["batch", "sequence_len"]), + _value("encoder_hidden_states", ir.DataType.FLOAT, ["batch", 1500, 384]), + _value("position_ids", ir.DataType.INT64, ["batch", "sequence_len"]), + ] + decoder_outputs = [("logits", ir.DataType.FLOAT, ["batch", "sequence_len", 51865])] + for layer in range(layers): + for kind in ("key", "value"): + decoder_inputs.append( + _value( + f"past_key_values.{layer}.{kind}", + ir.DataType.FLOAT, + ["batch", 6, "past_sequence_len", 64], + ) + ) + decoder_outputs.append( + ( + f"present.{layer}.{kind}", + ir.DataType.FLOAT, + ["batch", 6, "total_sequence_len", 64], + ) + ) + return ModelPackage( + { + "encoder": encoder, + "decoder": _model("decoder", decoder_inputs, decoder_outputs), + } + ) + + +def _workflow(metadata: dict) -> dict: + return metadata["pipeline"]["workflow"] + + +def _loop(metadata: dict) -> dict: + return next(step for step in _workflow(metadata)["steps"] if step["kind"] == "loop") + + +def test_encoder_runs_once_in_loop_setup(): + metadata = build_speech_to_text_workflow_metadata(_speech_package(), _config()) + loop = _loop(metadata) + setup = [node["component"] for node in loop["setup"]] + + # The encoder conditions the prefill, so it precedes the decoder and never + # appears in the loop body. + assert setup[0] == "encoder" + assert setup.index("decoder") > 0 + assert setup.count("encoder") == 1 + body = [node.get("component") for node in loop["steps"]] + assert "encoder" not in body + + +def test_cross_state_is_request_aligned_and_invariant(): + metadata = build_speech_to_text_workflow_metadata(_speech_package(), _config()) + workflow = _workflow(metadata) + cross = workflow["state"]["cross.encoder_hidden_states"] + + assert cross["initializer"] == "encoder.encoder_hidden_states" + assert cross["recurrence"] == {"kind": "invariant"} + # Request alignment on axis 0 is what keeps encoder states and decoder rows + # in step under batching and compaction. + assert cross["contract"]["batch_layout"] == {"kind": "request_aligned", "axis": 0} + assert cross["contract"]["shape"] == ["batch", 1500, 384] + + carry = next( + item + for item in _loop(metadata)["carried"] + if item["cell"] == "cross.encoder_hidden_states" + ) + assert carry["next"] == "cross.encoder_hidden_states" + + # Prefill reads the encoder value directly (same scope); every loop + # iteration reads the carried cell instead of rerunning the encoder. + prefill = next( + node for node in _loop(metadata)["setup"] if node.get("component") == "decoder" + ) + assert prefill["inputs"]["encoder_hidden_states"] == "encoder.encoder_hidden_states" + body_decoder = next( + node for node in _loop(metadata)["steps"] if node.get("component") == "decoder" + ) + assert body_decoder["inputs"]["encoder_hidden_states"] == "cross.encoder_hidden_states" + + +def test_self_attention_cache_is_the_only_served_group(): + metadata = build_speech_to_text_workflow_metadata(_speech_package(), _config()) + groups = _workflow(metadata)["serving"]["state_service"]["groups"] + + # The cross state is loop-invariant, so nothing appends to it and it is not + # served; only the growing self-attention cache is. + assert set(groups) == {"decoder_cache"} + ports = groups["decoder_cache"]["ports"]["decoder"] + assert ports["cache_0"] == { + "input": "past_key_values.0.key", + "output": "present.0.key", + "role": "key", + "layer": 0, + } + assert len(ports) == 4 + + +def test_audio_program_is_declared_and_bound_to_the_encoder_input(tmp_path): + processor = tmp_path / "audio_processor.json" + processor.write_text(json.dumps(_WHISPER_EXTRACTOR), encoding="utf-8") + pkg = _speech_package() + program = _audio_preprocessing_program(str(processor), pkg["encoder"]) + metadata = build_speech_to_text_workflow_metadata( + pkg, _config(), audio_preprocessing=program + ) + + # The program is document-level data, a sibling of `pipeline`. + audio = metadata["preprocessing"]["audio"] + assert [transform["op"] for transform in audio["transforms"]] == [ + "decode", + "resample", + "pad", + "log_mel", + "normalize", + ] + assert audio["transforms"][2]["target_samples"] == 480000 + assert audio["transforms"][3]["num_mel_bins"] == 80 + output = audio["outputs"][0] + assert output["name"] == "audio.input_features" + assert output["content"] == "audio_features" + assert output["contract"]["shape"] == ["batch", 80, "audio_seq_len"] + + workflow = _workflow(metadata) + assert workflow["manifest"]["adapter_abis"] == {"onnx-genai.audio-preprocess": "1"} + assert "audio_preprocessing_program" in workflow["manifest"]["capabilities"] + adapter = workflow["components"]["audio_preprocess"] + assert adapter["implementation"] == { + "kind": "adapter", + "abi": "onnx-genai.audio-preprocess", + "version": "1", + } + assert adapter["ports"]["inputs"]["encoded"]["dtype"] == "uint8" + + # Encoded bytes enter as a workflow input; the adapter invoke precedes the encoder. + assert workflow["inputs"]["request.audio"]["contract"]["dtype"] == "uint8" + setup = _loop(metadata)["setup"] + assert setup[0] == { + "kind": "invoke", + "component": "audio_preprocess", + "inputs": {"encoded": "request.audio"}, + "outputs": {"input_features": "audio.input_features"}, + } + assert setup[1]["component"] == "encoder" + assert setup[1]["inputs"]["input_features"] == "audio.input_features" + + +def test_without_a_program_features_are_a_request_input(): + metadata = build_speech_to_text_workflow_metadata(_speech_package(), _config()) + workflow = _workflow(metadata) + + assert "preprocessing" not in metadata + assert "audio_preprocess" not in workflow["components"] + assert "adapter_abis" not in workflow["manifest"] + features = workflow["inputs"]["encoder.input.input_features"] + assert features["contract"]["shape"] == ["batch", 80, "audio_seq_len"] + assert features["contract"]["batch_layout"] == {"kind": "request_aligned", "axis": 0} + + +def test_non_log_mel_extractor_declares_no_program(tmp_path): + processor = tmp_path / "audio_processor.json" + processor.write_text( + json.dumps({"feature_extractor_type": "Wav2Vec2FeatureExtractor"}), encoding="utf-8" + ) + assert _audio_preprocessing_program(str(processor), _speech_package()["encoder"]) is None + + +def test_speech_workflow_requires_encoder_and_decoder(): + pkg = ModelPackage( + { + "decoder": _speech_package()["decoder"], + } + ) + with pytest.raises(ValueError, match="encoder and decoder"): + build_speech_to_text_workflow_metadata(pkg, _config()) + + +def test_write_round_trips_the_built_metadata(tmp_path): + path = write_speech_to_text_workflow_metadata(_speech_package(), str(tmp_path), _config()) + with open(path, encoding="utf-8") as handle: + loaded = yaml.safe_load(handle) + assert loaded == build_speech_to_text_workflow_metadata(_speech_package(), _config()) + + +def test_speech_workflow_matches_producer_schema(tmp_path): + schema_path = os.environ.get("ONNX_GENAI_SCHEMA") + if not schema_path: + pytest.skip("set ONNX_GENAI_SCHEMA to the producer-contract schema") + with open(schema_path, encoding="utf-8") as handle: + schema = json.load(handle) + processor = tmp_path / "audio_processor.json" + processor.write_text(json.dumps(_WHISPER_EXTRACTOR), encoding="utf-8") + pkg = _speech_package() + metadata = build_speech_to_text_workflow_metadata( + pkg, + _config(), + audio_preprocessing=_audio_preprocessing_program(str(processor), pkg["encoder"]), + ) + jsonschema.validate(metadata, schema) diff --git a/src/mobius/integrations/onnx_genai/workflow_metadata.py b/src/mobius/integrations/onnx_genai/workflow_metadata.py new file mode 100644 index 000000000..3e3928713 --- /dev/null +++ b/src/mobius/integrations/onnx_genai/workflow_metadata.py @@ -0,0 +1,9154 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT License. + +"""ONNX GenAI workflow-IR metadata production.""" + +from __future__ import annotations + +import copy +import json +import math +import os +import re +from typing import Any + +import onnx_ir as ir +import yaml + +from mobius._constants import ( + STATIC_CACHE_KV_SEQUENCE_LENGTH, + STATIC_CACHE_LAYOUT, + STATIC_CACHE_SEQUENCE_AXIS, + STATIC_CACHE_WRITE_INDICES, +) +from mobius.generation import ( + SOLVER_BUILDERS, + PolicyCapabilities, + attach_policy_components, + build_boolean_not, + build_code_frame_update, + build_code_history_append, + build_codec_layout_transpose, + build_counter_rng_normal, + build_ddim_solver_step, + build_decoder_state_initializer, + build_decoder_step_update, + build_duplex_agent_frame_select, + build_duplex_cell_to_frame, + build_duplex_frame_assemble, + build_duplex_frame_commit, + build_duplex_stream_append, + build_duplex_stream_tail, + build_duplex_teacher_select, + build_duplex_user_stream_merge, + build_duplex_waveform_append, + build_empty_features, + build_eos_termination, + build_euler_model_input, + build_euler_solver_step, + build_flow_match_solver_step, + build_greedy_sampler, + build_guidance_combine, + build_identity_model_input, + build_integer_add, + build_integer_minimum, + build_last_token_logits, + build_model_token_cast, + build_pack_latents_2x2, + build_proposal_metrics, + build_scalar_constant, + build_scalar_integer_add, + build_schedule_constant, + build_schedule_history_append, + build_schedule_lookup, + build_seeded_categorical_sampler, + build_selective_integer_add, + build_sequence_concat, + build_sequence_length, + build_shape_constant, + build_tensor_scale, + build_termination_batch_initializer, + build_token_state_update, + build_token_to_slot, + build_true_cfg, + build_tts_decoder_state_initializer, + build_tts_decoder_step_update, + build_tts_state_initializer, + build_unpack_latents_2x2, + build_video_conv_cache_initializer, + build_video_decode_chunk, + build_video_decode_chunk_count, + build_video_latent_initializer, + build_video_latent_permute, + build_video_latent_unscale, + build_zeros_like, + rotary_axis_count, +) +from mobius.integrations.onnx_genai.inference_metadata import ( + _name_image_preprocessing_program, + _port, + _shape_metadata, + _source_asset_path, + add_adapter_service_to_metadata, + add_policy_components_to_workflow, + build_native_vlm_package_metadata, + declare_request_alignment, + request_batch_layout, +) +from mobius.tasks._ctc_asr import BATCH_PADDING_SENSITIVE_KEY + + +class _NoAliasSafeDumper(yaml.SafeDumper): + def ignore_aliases(self, data: Any) -> bool: + return True + + +def _dump_yaml(metadata: dict[str, Any], handle: Any) -> None: + yaml.dump(metadata, handle, Dumper=_NoAliasSafeDumper, sort_keys=False) + + +def _source_model_value(source: str | None, name: str, fallback: Any) -> Any: + """Resolve a value from packaged runtime metadata when available.""" + candidates: list[tuple[str, tuple[str, ...]]] = [] + if source and os.path.isdir(source): + candidates = [ + (os.path.join(source, "genai_config.json"), ("model", name)), + (os.path.join(source, "tokenizer_config.json"), (name,)), + ] + for path, keys in candidates: + if not os.path.isfile(path): + continue + try: + with open(path, encoding="utf-8") as handle: + value: Any = yaml.safe_load(handle) + for key in keys: + value = value[key] + except (OSError, TypeError, KeyError): + continue + fallback = value + break + return fallback + + +def _source_token_id(source: str | None, name: str, fallback: Any) -> int: + """Resolve a generation token ID from packaged runtime metadata when available.""" + fallback = _source_model_value(source, name, fallback) + if isinstance(fallback, list): + fallback = fallback[0] if fallback else 0 + return int(fallback or 0) + + +def _contract(value: ir.Value) -> dict[str, Any]: + port = _port(value) + dtype = {"fp16": "float16", "bf16": "bfloat16", "fp32": "float32"}.get( + port.dtype, port.dtype + ) + shape = _shape_metadata(port) + contract: dict[str, Any] = { + "dtype": dtype, + "rank": port.rank, + "shape": shape, + } + layout = request_batch_layout(shape) + if layout is not None: + contract["batch_layout"] = layout + return contract + + +def _request_aligned(contract: dict[str, Any], axis: int = 0) -> dict[str, Any]: + """Mark a contract as carrying exactly one entry per in-flight request. + + This is a structural batching fact, not a row identity: it tells the runtime + which axis to permute when it compacts the batch, while scheduler slots and + sequence handles stay runtime-private. + """ + return {**contract, "batch_layout": {"kind": "request_aligned", "axis": axis}} + + +# Translation between the port vocabulary this producer *mints* when it builds +# a graph and the runtime's architecture-neutral role vocabulary. Both sides are +# fixed vocabularies and Mobius owns one of them: the task builders in +# ``mobius.tasks`` choose these exact names, so reading them back here is a +# lookup, not an inference about a graph of unknown provenance. A port outside +# this vocabulary carries no role, because a workflow that guesses is worse than +# one that stays silent. +_PORT_ROLES: dict[str, str] = { + "input_ids": "token_ids", + "inputs_embeds": "inputs_embeds", + "attention_mask": "attention_mask", + "position_ids": "position_ids", + "logits": "logits", + "last_hidden_state": "hidden_states", + "encoder_hidden_states": "encoder_hidden_states", + "audio_features": "audio_features", +} + + +def _component( + model: ir.Model, + artifact: str, + *, + effects: tuple[str, ...] = (), +) -> dict[str, Any]: + """Declare one ONNX-backed workflow component: its artifact and port roles. + + A component declares only what its artifact cannot say about itself. The + ``.onnx`` file is shipped inside the package and is authoritative for which + ports exist and what dtype, rank and shape each one has, so transcribing + that into YAML would create a second copy of a fact the package already + carries — one that can drift from the graph and that nothing cross-checks + at rest. The runtime resolves ports against the live session instead, which + catches a name the graph does not expose rather than agreeing with a stale + echo of it. + + What no graph carries is what a port *means*. ``input_ids`` and + ``position_ids`` are both rank-2 ``int64``; nothing in the file says which + one is the autoregressive sequence. An invocation binds an SSA value to a + port, which records which value arrives but not whether it is tokens, a mask + or logits — and that second fact is what a runtime needs before it can + specialize a decode step. So ``roles`` is the whole declaration here. + + Only ports in this producer's own vocabulary get a role, and state ports + never need one: the group that carries them already names its pairs, which + is also where the fixed-capacity scatter ABI is stated. + """ + del effects + named = [str(value.name) for value in (*model.graph.inputs, *model.graph.outputs)] + roles = {name: _PORT_ROLES[name] for name in named if name in _PORT_ROLES} + declaration: dict[str, Any] = {"implementation": {"kind": "onnx", "artifact": artifact}} + if roles: + declaration["ports"] = {"roles": roles} + return declaration + + +def _grammar_adapter_component(action: str) -> dict[str, Any]: + """Declare one action of the versioned grammar-guidance adapter ABI.""" + + def port(dtype: str, shape: list[int | str]) -> dict[str, Any]: + return {"dtype": dtype, "rank": len(shape), "shape": shape} + + return { + "implementation": { + "kind": "adapter", + "abi": "onnx-genai.grammar-guidance", + "version": "1", + }, + "ports": { + "inputs": { + "state": port("int64", ["batch"]), + "tokens": port("int64", ["batch", "proposal"]), + "valid_length": port("int64", ["batch"]), + "transition_table": port("int64", ["grammar_states", "vocabulary"]), + }, + "outputs": { + "next_state": port("int64", ["batch"]), + "consumed_length": port("int64", ["batch"]), + "logits_mask": port("bool", ["batch", "vocabulary"]), + "forced_tokens": port("int64", ["batch", 1]), + "forced_length": port("int64", ["batch"]), + }, + }, + "contract": { + "id": "onnx-genai.grammar-guidance", + "version": "1", + "bindings": { + "state": "state", + "tokens": "tokens", + "valid_length": "valid_length", + "transition_table": "transition_table", + "next_state": "next_state", + "consumed_length": "consumed_length", + "logits_mask": "logits_mask", + "forced_tokens": "forced_tokens", + "forced_length": "forced_length", + }, + "parameters": {"action": action}, + }, + "effects": ["grammar"], + # The adapter keeps one grammar FSM per in-flight request, and every + # port is request-aligned on axis 0. Declaring the row scope is what + # lets the runtime drive the mandatory compact(selection)/release(row) + # ABI when the batch changes; without it the FSM rows would drift out + # of correspondence with the sequences they guide. + "row_scope": {"axis": 0, "stateful": True}, + } + + +def _effect(consumes: str, produces: str) -> dict[str, str]: + return {"consumes": consumes, "produces": produces} + + +def _publish_workflow_v1(workflow: dict[str, Any]) -> dict[str, Any]: + """Publish structured steps and logical carries without compiler bookkeeping.""" + graph = workflow.pop("graph") + workflow.pop("initial_effects", None) + for declaration in workflow.get("inputs", {}).values(): + source = declaration.get("source") + if isinstance(source, dict) and source.get("kind") == "request": + source.pop("field", None) + + # Every workflow value whose leading dimension is the batch symbol holds one + # entry per in-flight request, so declare that structurally instead of leaving + # a runtime to infer it. Graph-derived contracts already carry the layout; + # this covers the hand-written declarations the runtime compares them against + # when it validates a carry, a binding, or an emit. + def _declare_row_alignment(contract: Any) -> Any: + if ( + isinstance(contract, dict) + and "batch_layout" not in contract + and request_batch_layout(contract.get("shape")) is not None + ): + return _request_aligned(contract) + return contract + + for section in ("inputs", "outputs", "state"): + for declaration in workflow.get(section, {}).values(): + declaration["contract"] = _declare_row_alignment(declaration.get("contract")) + for component in workflow.get("components", {}).values(): + for side in ("inputs", "outputs"): + ports = component.get("ports", {}).get(side) + if not ports: + continue + for port, contract in ports.items(): + ports[port] = _declare_row_alignment(contract) + substitutions: dict[str, str] = {} + loop_index = 0 + cell_aliases = { + cell: f"{cell}_state" if cell in workflow.get("outputs", {}) else cell + for cell in workflow.get("state", {}) + } + if any(cell != alias for cell, alias in cell_aliases.items()): + workflow["state"] = { + cell_aliases[cell]: declaration for cell, declaration in workflow["state"].items() + } + + def collect_carried(node: dict[str, Any]) -> None: + if node["kind"] == "loop": + for carry in node.get("carried", []): + alias = cell_aliases.get(carry["cell"], carry["cell"]) + substitutions[carry["body_input"]] = alias + substitutions[carry["next"]] = alias + collect_carried(node["setup"]) + collect_carried(node["body"]) + elif node["kind"] == "sequence": + for child in node["nodes"]: + collect_carried(child) + elif node["kind"] == "branch": + for case in node["cases"].values(): + collect_carried(case) + if "default" in node: + collect_carried(node["default"]) + + def rewrite(value: Any) -> Any: + if isinstance(value, str): + return substitutions.get(value, value) + if isinstance(value, list): + return [rewrite(item) for item in value] + if isinstance(value, dict): + return {key: rewrite(item) for key, item in value.items()} + return value + + def convert(node: dict[str, Any]) -> dict[str, Any]: + nonlocal loop_index + kind = node["kind"] + if kind == "sequence": + return { + "kind": "sequence", + "steps": [convert(child) for child in node["nodes"]], + } + if kind == "invoke": + return { + "kind": "invoke", + "component": node["component"], + "inputs": rewrite(node.get("inputs", {})), + "outputs": rewrite(node.get("outputs", {})), + } + if kind == "emit": + result = { + "kind": "emit", + "value": rewrite(node["value"]), + "output": node["output"], + "mode": node["mode"], + } + if "axis" in node: + result["axis"] = node["axis"] + if "valid_length" in node: + result["valid_length"] = rewrite(node["valid_length"]) + if "when" in node: + result["when"] = rewrite(node["when"]) + return result + if kind == "branch": + result = { + "kind": "branch", + "predicate": rewrite(node["predicate"]), + "cases": {name: convert(case) for name, case in node["cases"].items()}, + "outputs": rewrite(node.get("outputs", {})), + } + if "default" in node: + result["default"] = convert(node["default"]) + return result + if kind == "loop": + current_loop = loop_index + loop_index += 1 + setup = node["setup"] + body = node["body"] + setup_steps = ( + [convert(child) for child in setup["nodes"]] + if setup["kind"] == "sequence" + else [convert(setup)] + ) + body_steps = ( + [convert(child) for child in body["nodes"]] + if body["kind"] == "sequence" + else [convert(body)] + ) + carried = [] + for carry in node.get("carried", []): + cell = cell_aliases.get(carry["cell"], carry["cell"]) + published_carry = { + "cell": cell, + "next": rewrite(carry["body_output"]), + } + initial = rewrite(carry["current"]) + if workflow["state"][cell]["initializer"] != initial: + published_carry["initial"] = initial + carried.append(published_carry) + active_cell = node.get("active_cell") + if active_cell is None: + active_cell = f"loop_{current_loop}_active" + active_initializer = f"package.{active_cell}" + workflow["inputs"][active_initializer] = { + "contract": {"dtype": "bool", "rank": 1, "shape": [1]}, + "role": {"kind": "opaque"}, + "source": {"kind": "literal"}, + "required": False, + "default": True, + } + workflow["state"][active_cell] = { + "contract": {"dtype": "bool", "rank": 1, "shape": [1]}, + "scope": "invocation", + "initializer": active_initializer, + "recurrence": {"kind": "invariant"}, + } + carried.append( + { + "cell": active_cell, + "next": rewrite(node["condition"]), + } + ) + result = { + "kind": "loop", + "setup": setup_steps, + "steps": body_steps, + "continue_when": active_cell, + "max_iterations": rewrite(node["max_iterations"]), + "carried": carried, + } + if "termination" in node: + result["termination"] = node["termination"] + if "iteration" in node: + result["iteration"] = node["iteration"] + return result + raise ValueError(f"unsupported workflow node kind {kind!r}") + + collect_carried(graph) + published = convert(graph) + workflow["steps"] = published["steps"] if published["kind"] == "sequence" else [published] + declare_request_alignment(workflow) + return workflow + + +def _invoke( + component: str, + inputs: dict[str, str], + outputs: dict[str, str], + _effects: dict[str, dict[str, str]] | None = None, +) -> dict[str, Any]: + return { + "kind": "invoke", + "component": component, + "inputs": inputs, + "outputs": outputs, + } + + +def build_audio_codec_workflow_metadata(pkg: Any) -> dict[str, Any]: + """Build typed SSA metadata for a waveform-to-codes-to-waveform codec.""" + names = set(pkg.keys()) + if names != {"encoder", "decoder"}: + raise ValueError( + "audio codec workflow requires exactly encoder and decoder components" + ) + encoder = pkg["encoder"] + decoder = pkg["decoder"] + if len(encoder.graph.inputs) != 1 or len(encoder.graph.outputs) != 1: + raise ValueError("codec encoder requires exactly one input and one output") + if len(decoder.graph.inputs) != 1 or len(decoder.graph.outputs) != 1: + raise ValueError("codec decoder requires exactly one input and one output") + + waveform_input = encoder.graph.inputs[0] + codes_output = encoder.graph.outputs[0] + codes_input = decoder.graph.inputs[0] + waveform_output = decoder.graph.outputs[0] + if codes_output.dtype != codes_input.dtype: + raise ValueError("codec encoder output and decoder input dtypes must match") + if _contract(codes_output) != _contract(codes_input): + raise ValueError("codec encoder output and decoder input contracts must match") + + encode_effect = "codec_encode" + decode_effect = "codec_decode" + emit_effect = "audio_emit" + workflow = { + "manifest": { + "capabilities": [ + "workflow_ssa", + "linear_effects", + "typed_emit", + ], + }, + "inputs": { + "request.waveform": { + "contract": _contract(waveform_input), + "role": {"kind": "runtime", "version": "1.0", "role": "media"}, + "source": {"kind": "request", "field": "media"}, + "required": True, + } + }, + "outputs": { + "waveform": { + "contract": _contract(waveform_output), + "role": "audio", + "stage": "post_adapter", + } + }, + "components": { + "encoder": _component( + encoder, + "encoder/model.onnx", + effects=(encode_effect,), + ), + "decoder": _component( + decoder, + "decoder/model.onnx", + effects=(decode_effect,), + ), + }, + "initial_effects": { + encode_effect: f"{encode_effect}.0", + decode_effect: f"{decode_effect}.0", + emit_effect: f"{emit_effect}.0", + }, + "graph": { + "kind": "sequence", + "nodes": [ + _invoke( + "encoder", + {waveform_input.name: "request.waveform"}, + {codes_output.name: "codec.codes"}, + {encode_effect: _effect(f"{encode_effect}.0", f"{encode_effect}.1")}, + ), + _invoke( + "decoder", + {codes_input.name: "codec.codes"}, + {waveform_output.name: "codec.waveform"}, + {decode_effect: _effect(f"{decode_effect}.0", f"{decode_effect}.1")}, + ), + { + "kind": "emit", + "value": "codec.waveform", + "output": "waveform", + "mode": "replace", + "effect_name": emit_effect, + "effect": _effect(f"{emit_effect}.0", f"{emit_effect}.1"), + }, + ], + }, + } + return {"schema_version": "v1", "pipeline": {"workflow": _publish_workflow_v1(workflow)}} + + +def write_audio_codec_workflow_metadata(pkg: Any, output_dir: str) -> str: + """Write typed SSA metadata for an audio codec package.""" + os.makedirs(output_dir, exist_ok=True) + metadata = build_audio_codec_workflow_metadata(pkg) + add_adapter_service_to_metadata(metadata, pkg, output_dir) + path = os.path.join(output_dir, "inference_metadata.yaml") + with open(path, "w", encoding="utf-8") as handle: + _dump_yaml(metadata, handle) + return path + + +def _model_cache_pairs(model: ir.Model) -> list[tuple[ir.Value, ir.Value]]: + outputs = {value.name: value for value in model.graph.outputs} + pairs = [] + for past in model.graph.inputs: + present = next( + ( + outputs.get(name) + for name in _cache_output_candidates(past.name or "") + if name in outputs + ), + None, + ) + if present is not None: + pairs.append((past, present)) + return pairs + + +def _cache_output_candidates(past_name: str) -> tuple[str, ...]: + """Names an exporter may give the output that continues a cache input. + + An appending cache renames ``past`` to ``present``; a static, indexed cache + keeps the buffer's name and prefixes the written result instead, because the + output is the same buffer rather than a longer one. + """ + return ( + past_name.replace("past_key_values", "present"), + past_name.replace("past.", "present."), + past_name.replace("past_", "present_"), + f"updated_{past_name}", + ) + + +def _constant_extent(dimension: Any) -> int | None: + """Return *dimension* as an ``int``, or ``None`` when it is symbolic. + + An :class:`ir.Shape` entry is a plain ``int`` exactly when the extent is + known; otherwise it is a ``SymbolicDim`` whose value is its name. + """ + return dimension if isinstance(dimension, int) else None + + +def _static_cache_ports(model: ir.Model) -> dict[str, Any] | None: + """Return the declared static-cache ABI of *model*, or ``None``. + + The two control ports are per-row integer vectors and so are shape-indistin- + guishable from each other; they are matched against the ABI mobius mints in + :mod:`mobius._constants`, never guessed from the graph. The buffer ports are + then whichever cache inputs the scatter addresses. + """ + inputs = {value.name: value for value in model.graph.inputs} + write_indices = inputs.get(STATIC_CACHE_WRITE_INDICES) + kv_lengths = inputs.get(STATIC_CACHE_KV_SEQUENCE_LENGTH) + if write_indices is None or kv_lengths is None: + return None + buffers = { + past.name: past + for past, present in _model_cache_pairs(model) + if present.name == f"updated_{past.name}" + } + if not buffers: + raise ValueError( + "decoder declares the static-cache control ports " + f"{STATIC_CACHE_WRITE_INDICES!r}/{STATIC_CACHE_KV_SEQUENCE_LENGTH!r} but exposes " + "no paired cache buffer to scatter into; regenerate the package with " + "updated_ outputs for every static cache input" + ) + axes = { + node.attributes.get_int("axis", 0) + for node in ir.traversal.RecursiveGraphIterator(model.graph) + if node.op_type == "TensorScatter" + } + if axes - {STATIC_CACHE_SEQUENCE_AXIS}: + raise ValueError( + f"static cache buffers are addressed on axes {sorted(axes)}, but the mobius " + f"static-cache ABI scatters along axis {STATIC_CACHE_SEQUENCE_AXIS}; the " + "declared capacity axis and the graph disagree" + ) + capacities = set() + for buffer in buffers.values(): + shape = list(buffer.shape or []) + if len(shape) <= STATIC_CACHE_SEQUENCE_AXIS: + raise ValueError( + f"static cache buffer {buffer.name!r} has rank {len(shape)}, which cannot " + f"carry a capacity on axis {STATIC_CACHE_SEQUENCE_AXIS}" + ) + capacity = _constant_extent(shape[STATIC_CACHE_SEQUENCE_AXIS]) + if capacity is None: + raise ValueError( + f"static cache buffer {buffer.name!r} declares a symbolic extent " + f"{shape[STATIC_CACHE_SEQUENCE_AXIS]!r} on its capacity axis; an " + "indexed scatter is only meaningful against one constant capacity" + ) + capacities.add(capacity) + if len(capacities) != 1: + raise ValueError( + f"static cache buffers declare conflicting capacities {sorted(capacities)}; " + "one write cursor cannot address buffers of different lengths" + ) + return { + "write_indices": STATIC_CACHE_WRITE_INDICES, + "kv_sequence_length": STATIC_CACHE_KV_SEQUENCE_LENGTH, + "buffers": buffers, + "capacity": capacities.pop(), + } + + +def _kv_storage_contract(model: ir.Model) -> dict[str, Any]: + """Derive physical KV storage from the admitted model interface. + + Shared KV is a runtime I/O-binding contract: past and present ports bind the + same full-capacity OrtValue. That is only sound when the graph's attention + operator takes the *logical* cache length as a separate input, so it can + ignore the unwritten tail of a capacity-sized buffer. ``GroupQueryAttention`` + does (``seqlens_k`` / ``total_sequence_length``), and a paged layout carries + its lengths in the block tables. + + The standard ONNX ``Attention`` operator does not: it concatenates ``past_key`` + with the current key and *derives* ``total_sequence_length`` from the past + tensor's own second-to-last dimension. Binding a capacity-sized buffer there + would both attend over unwritten slots and make the attention mask (sized to + the real length) disagree with the derived total length, which ORT rejects. + Such a graph therefore grows its cache by concatenation and must be declared + ``dynamic``. + """ + input_names = {value.name.lower() for value in model.graph.inputs} + paged = any( + marker in name + for name in input_names + for marker in ("block_table", "block_tables", "page_table", "page_tables") + ) + has_cache = bool(_model_cache_pairs(model)) + if paged: + storage = "paged" + elif has_cache and _consumes_explicit_cache_length(model): + storage = "shared_buffer" + else: + storage = "dynamic" + return { + "paging": "paged" if paged else "none", + # Row compaction is semantic for every batched KV layout: the runtime + # applies one row permutation to slot identity, KV, RNG, and loop state. + "compaction": has_cache, + "storage": storage, + } + + +#: Attention operators that accept a capacity-sized KV buffer plus an explicit +#: logical cache length, and so can be bound to a preallocated shared buffer. +_CAPACITY_ADDRESSABLE_ATTENTION = frozenset( + { + ("com.microsoft", "GroupQueryAttention"), + ("com.microsoft", "PagedAttention"), + ("com.microsoft", "SparseAttention"), + } +) + + +def _consumes_explicit_cache_length(model: ir.Model) -> bool: + """Report whether every cache consumer takes the logical cache length as input. + + Returns ``False`` when the model has no cache consumers at all, because a + graph that never reads ``past_key_values.*`` cannot promise capacity-safe + behaviour it does not exercise. + """ + cache_values = { + past.name for past, _ in _model_cache_pairs(model) if past.name is not None + } + # A static buffer's capacity safety comes from its declared write cursor and + # logical lengths, not from the attention operator's signature, so its + # scatter consumer must not veto the appending caches' storage class. + static = _static_cache_ports(model) + if static is not None: + cache_values -= set(static["buffers"]) + if not cache_values: + return False + consumers = { + (node.domain, node.op_type) + for node in ir.traversal.RecursiveGraphIterator(model.graph) + for value in node.inputs + if value is not None and value.name in cache_values + } + return bool(consumers) and consumers <= _CAPACITY_ADDRESSABLE_ATTENTION + + +def _aliasing_for_storage(storage: str) -> str: + """Translate a physical storage class into the semantic aliasing contract. + + Shared-buffer and paged layouts let the runtime bind ``present`` onto the + same allocation as ``past``; the metadata only declares that doing so is + *legal*, never that the runtime must. A growable cache returns a fresh, + longer tensor each step and must never be aliased onto its own input. + """ + return "permitted" if storage in {"shared_buffer", "paged"} else "forbidden" + + +def _annotated_alias(alias: dict[str, Any]) -> dict[str, Any]: + """Add the half and layer a state port pair carries, when the name states them. + + A layer's key buffer and its value buffer are the same shape and the same + dtype, and a cell's label is producer-chosen so its lexicographic order is + not the layer order (``cache_10`` sorts before ``cache_2``). A consumer that + paired these positionally would silently transpose two layers' caches. Both + facts are recoverable only here, from the port names this producer minted, + so both are declared on the alias. + + They are declared together or not at all: a port outside the two attention + cache ABIs — recurrent state, a convolution cache — has no half and no layer + to state, and inventing an index for it would corrupt the very ordering the + index exists to fix. + """ + name = str(alias.get("input", "")) + half = _cache_half(name) + if half is None: + return alias + return {**alias, "role": half, "layer": _cache_layer_index(name, 0)} + + +def _state_group( + *, + ports: dict[str, Any], + sequence_axis: int, + logical_lengths: str | None = None, + storage: str = "shared_buffer", + kind: str = "full_attention", + aliasing: str | None = None, + layout: str = "bnsh", +) -> dict[str, Any]: + """Describe one semantic state group of the serving state service. + + ``aliasing`` is a property of the admitted graph, not a runtime preference: + a shared full-capacity buffer lets the component write ``present`` straight + into the ``past`` binding, while a growable cache returns a fresh, longer + tensor each step and must never be aliased onto its own input. + """ + group: dict[str, Any] = { + "kind": kind, + "sequence_axis": sequence_axis, + "layout": layout, + "aliasing": (aliasing if aliasing is not None else _aliasing_for_storage(storage)), + "reuse": {"prefix_reusable": True, "evictable_prefix": False}, + "capabilities": {"snapshot": True, "fork": True}, + "ports": { + component: {cell: _annotated_alias(alias) for cell, alias in aliases.items()} + for component, aliases in ports.items() + }, + } + if logical_lengths is not None: + group["logical_lengths"] = logical_lengths + return group + + +_LAYER_STATE_KINDS = { + "sliding_attention": "sliding_attention", + "full_attention": "full_attention", + "chunked_attention": "sliding_attention", +} + + +def _state_aliasing(kv_contract: dict[str, Any]) -> str: + """Translate a physical KV storage contract into the aliasing contract.""" + return _aliasing_for_storage(str(kv_contract["storage"])) + + +def _cache_layer_index(port_name: str, fallback: int) -> int: + """Recover the decoder layer index from a cache port name. + + Both cache ABIs encode the layer in the port name — ``past_key_values.N.key`` + for an appending cache, ``key_cache.N`` for a static one — because a hybrid + decoder's cache-owning layers are a subset of its layers, so a port's + position in the port list is not its layer. + """ + match = re.search(r"\.(\d+)\.(?:key|value)$|(?:key|value)_cache\.(\d+)$", port_name) + if match is None: + return fallback + return int(match.group(1) or match.group(2)) + + +def _cache_half(port_name: str) -> str | None: + """Recover which half of a split attention cache a port carries. + + A layer's key buffer and its value buffer are the same shape and the same + dtype, so nothing downstream can tell them apart once they are in a list. + This producer minted both names, in the same two ABIs ``_cache_layer_index`` + reads, so it can say which is which; a port outside those two spellings gets + no half, which is the right answer for recurrent and latent state that has + no halves to distinguish. + """ + match = re.search(r"\.\d+\.(key|value)$|(key|value)_cache\.\d+$", port_name) + if match is None: + return None + return match.group(1) or match.group(2) + + +def _state_group_kinds(config: Any, cache_pairs: list[tuple[ir.Value, ir.Value]]) -> list[str]: + """Return the semantic ``StateKind`` of every KV cache cell. + + Hybrid models interleave sliding-window and full-attention layers, and the + two are not interchangeable: only a sliding layer's oldest positions may be + evicted. Cache-owning layers form a contiguous prefix of the decoder (a + KV-sharing suffix owns no cache at all), so the port's own layer index — + not its position in the port list — selects the layer type. + """ + layer_types = list(getattr(config, "layer_types", None) or []) + default_kind = ( + "sliding_attention" + if not layer_types and getattr(config, "sliding_window", None) + else "full_attention" + ) + kinds = [] + for index, (past, _) in enumerate(cache_pairs): + layer = _cache_layer_index(past.name or "", index // 2) + layer_type = layer_types[layer] if layer < len(layer_types) else None + kinds.append(_LAYER_STATE_KINDS.get(str(layer_type), default_kind)) + return kinds + + +def _state_service_groups( + *, + config: Any, + cache_pairs: list[tuple[ir.Value, ir.Value]], + ports: dict[str, dict[str, dict[str, str]]], + sequence_axis: int, + logical_lengths: str | None, + aliasing: str, + base_name: str, + indexed_scatter: dict[str, Any] | None = None, +) -> tuple[dict[str, Any], dict[str, str]]: + """Build ``serving.state_service.groups`` plus each cell's owning group. + + One group per semantic kind *and update discipline*: a hybrid decoder + therefore publishes distinct ``sliding_attention`` and ``full_attention`` + groups whose per-cell contracts carry their own geometry (Gemma 4's global + layers are double-wide), and a decoder that appends some caches while + scattering others into fixed buffers keeps those apart too, because the + valid region of an appended buffer is its shape while the valid region of a + scattered one is a declared prefix. The group declares *semantics* only — + eviction legality, aliasing legality, layout — never a storage class, + allocator, or compaction algorithm, which are the runtime's to choose. + + ``indexed_scatter`` describes the static, fixed-capacity buffers: which + cache inputs they are, the constant capacity they were built against, the + cell carrying each row's write cursor, and the port that receives it. + """ + indexed_scatter = indexed_scatter or {} + indexed_inputs = set(indexed_scatter.get("buffers", ())) + kinds = _state_group_kinds(config, cache_pairs) + scattered = [(past.name or "") in indexed_inputs for past, _ in cache_pairs] + # Group identity is (semantic kind, update discipline); the suffix only + # appears when more than one identity is present, so a homogeneous decoder + # keeps publishing exactly one group under its base name. + identities = [ + (kind, "indexed_scatter" if is_scattered else "append") + for kind, is_scattered in zip(kinds, scattered) + ] + distinct = sorted(set(identities)) + names = { + identity: (base_name if len(distinct) == 1 else f"{base_name}_{identity[0]}") + for identity in distinct + } + if len({*names.values()}) != len(distinct): + names = {identity: f"{base_name}_{identity[0]}_{identity[1]}" for identity in distinct} + cell_group = {} + grouped_ports: dict[str, dict[str, dict[str, dict[str, Any]]]] = { + name: {} for name in names.values() + } + for index, identity in enumerate(identities): + cell = f"cache_{index}" + cell_group[cell] = names[identity] + # A cell's label is producer-chosen and its lexicographic order is not the + # layer order, and a layer's key and value buffers are shape-identical, so + # each alias carries the half and layer its port name states. + for component, aliases in ports.items(): + for cell, alias in aliases.items(): + grouped_ports[cell_group[cell]].setdefault(component, {})[cell] = _annotated_alias( + alias + ) + groups = {} + for kind, update in distinct: + is_scattered = update == "indexed_scatter" + name = names[(kind, update)] + group: dict[str, Any] = { + "kind": kind, + "sequence_axis": (STATIC_CACHE_SEQUENCE_AXIS if is_scattered else sequence_axis), + "layout": STATIC_CACHE_LAYOUT if is_scattered else "bnsh", + } + group_lengths = indexed_scatter["logical_lengths"] if is_scattered else logical_lengths + if group_lengths: + group["logical_lengths"] = group_lengths + if is_scattered: + group["update"] = { + "kind": "indexed_scatter", + "write_indices": indexed_scatter["write_indices"], + "capacity": indexed_scatter["capacity"], + "write_indices_ports": dict.fromkeys( + grouped_ports[name], indexed_scatter["port"] + ), + # The graph-visible valid length is a second rank-1 integer + # vector sitting beside the destinations, so it is equally + # shape-indistinguishable and equally must be named rather than + # inferred. Together these two entries and ``ports`` below are + # the whole scatter ABI, which is why the package needs no + # second copy of it outside the workflow. + "kv_length_ports": dict.fromkeys( + grouped_ports[name], indexed_scatter["kv_length_port"] + ), + } + # A scatter writes through its buffer by construction: the written result + # *is* the input allocation, so aliasing is legal for every static group + # regardless of what the appending caches in the same graph can do. + group["aliasing"] = "permitted" if is_scattered else aliasing + group["reuse"] = { + "prefix_reusable": True, + # Dropping the oldest positions is only semantics-preserving + # for a windowed layer; a full-attention layer that loses its + # prefix silently answers a different question. + "evictable_prefix": kind == "sliding_attention", + } + group["ports"] = grouped_ports[name] + groups[name] = group + return groups, cell_group + + +def _build_real_tts_workflow_metadata(pkg: Any, config: Any) -> dict[str, Any]: + """Build the weight-bearing Qwen3-TTS talker/predictor/codec workflow.""" + talker = pkg["talker"] + predictor = pkg["code_predictor"] + prefill_embedder = pkg["talker_prefill_embedder"] + predictor_indices = pkg["code_predictor_indices"] + codec_name = next( + (name for name in ("codec", "vocoder", "decoder") if name in pkg), + None, + ) + if codec_name is None: + raise ValueError("real TTS workflow requires a merged codec/vocoder decoder component") + codec = pkg[codec_name] + num_groups = int(getattr(getattr(config, "tts", config), "num_code_groups", 16)) + if num_groups < 2: + raise ValueError("real TTS workflow requires at least two code groups") + + prompt = next(iter(prefill_embedder.graph.inputs)) + prefill = _find_port(prefill_embedder.graph.outputs, "prefill") + trailing = _find_port(prefill_embedder.graph.outputs, "trailing") + talker_embed = _find_port(talker.graph.inputs, "inputs_embeds") + talker_mask = _find_port(talker.graph.inputs, "attention_mask") + talker_position = _find_port(talker.graph.inputs, "position_ids") + talker_logits = _find_port(talker.graph.outputs, "logits") + talker_hidden = _find_port(talker.graph.outputs, "hidden") + predictor_embed = _find_port(predictor.graph.inputs, "inputs_embeds") + predictor_step = _find_port(predictor.graph.inputs, "step_index") + predictor_mask = _find_port(predictor.graph.inputs, "attention_mask") + predictor_position = _find_port(predictor.graph.inputs, "position_ids") + predictor_logits = _find_port(predictor.graph.outputs, "logits") + codec_embeddings = _find_port(predictor.graph.outputs, "codec_embeddings") + codec_input = next(iter(codec.graph.inputs), None) + waveform = next(iter(codec.graph.outputs), None) + required_ports = ( + prefill, + trailing, + talker_embed, + talker_mask, + talker_position, + talker_logits, + talker_hidden, + predictor_embed, + predictor_step, + predictor_mask, + predictor_position, + predictor_logits, + codec_embeddings, + codec_input, + waveform, + ) + if any(value is None for value in required_ports): + raise ValueError("real TTS package is missing a required typed transition port") + assert prefill is not None and trailing is not None + assert talker_embed is not None and talker_mask is not None + assert talker_position is not None and talker_logits is not None + assert talker_hidden is not None and predictor_embed is not None + assert predictor_step is not None and predictor_mask is not None + assert predictor_position is not None and predictor_logits is not None + assert codec_embeddings is not None and codec_input is not None and waveform is not None + + talker_caches = _model_cache_pairs(talker) + predictor_caches = _model_cache_pairs(predictor) + talker_kv = _kv_storage_contract(talker) + predictor_kv = _kv_storage_contract(predictor) + attach_policy_components(pkg, PolicyCapabilities()) + pkg.add_policy_component("last_token_logits", build_last_token_logits()) + pkg.add_policy_component( + "setup_talker_sampler", build_greedy_sampler(effect="setup_talker_sample") + ) + pkg.add_policy_component( + "setup_predictor_sampler", + build_greedy_sampler(effect="setup_predictor_sample"), + ) + pkg.add_policy_component("talker_sampler", build_greedy_sampler(effect="talker_sample")) + pkg.add_policy_component( + "predictor_prefill_sampler", + build_greedy_sampler(effect="predictor_prefill_sample"), + ) + pkg.add_policy_component( + "predictor_body_sampler", + build_greedy_sampler(effect="predictor_body_sample"), + ) + pkg.add_policy_component("continue_predicate", build_boolean_not()) + pkg.add_policy_component("tts_state_initializer", build_tts_state_initializer(num_groups)) + pkg.add_policy_component("token_to_slot", build_token_to_slot()) + pkg.add_policy_component( + "code_frame_update", build_code_frame_update(num_groups, scalar_index=True) + ) + pkg.add_policy_component("code_history_append", build_code_history_append(num_groups)) + if talker_caches or predictor_caches: + pkg.add_policy_component("cache_length_update", build_integer_add()) + pkg.add_policy_component( + "talker_state_initializer", + build_tts_decoder_state_initializer( + talker, + graph_name="talker_state_initializer", + embedding_input=talker_embed.name, + attention_mask_input=talker_mask.name, + position_ids_input=talker_position.name, + cache_inputs=[past.name for past, _ in talker_caches], + ), + ) + pkg.add_policy_component( + "predictor_state_initializer", + build_tts_decoder_state_initializer( + predictor, + graph_name="predictor_state_initializer", + embedding_input=predictor_embed.name, + attention_mask_input=predictor_mask.name, + position_ids_input=predictor_position.name, + cache_inputs=[past.name for past, _ in predictor_caches], + ), + ) + pkg.add_policy_component( + "talker_step_update", + build_tts_decoder_step_update( + graph_name="talker_step_update", + attention_dtype=talker_mask.dtype, + position_dtype=talker_position.dtype, + position_rank=len(talker_position.shape or []), + ), + ) + pkg.add_policy_component( + "predictor_step_update", + build_tts_decoder_step_update( + graph_name="predictor_step_update", + attention_dtype=predictor_mask.dtype, + position_dtype=predictor_position.dtype, + position_rank=len(predictor_position.shape or []), + ), + ) + codec_group_major = ( + codec_input.shape is not None + and len(codec_input.shape) == 3 + and str(list(codec_input.shape)[1]) == str(num_groups) + ) + if codec_group_major: + pkg.add_policy_component("codec_layout", build_codec_layout_transpose(num_groups)) + + batch = _contract(prompt)["shape"][0] + batch_int = _request_aligned({"dtype": "int64", "rank": 1, "shape": [batch]}) + batch_bool = _request_aligned({"dtype": "bool", "rank": 1, "shape": [batch]}) + control_int = {"dtype": "int64", "rank": 1, "shape": [1]} + inputs = { + "request.prompt_tokens": { + "contract": _contract(prompt), + "role": {"kind": "runtime", "version": "1.0", "role": "prompt_tokens"}, + "source": {"kind": "request", "field": "prompt_tokens"}, + "required": True, + }, + "request.max_iterations": { + "contract": control_int, + "role": {"kind": "runtime", "version": "1.0", "role": "max_output_tokens"}, + "source": {"kind": "request", "field": "max_output_tokens"}, + "required": True, + }, + "package.false": { + "contract": _request_aligned(batch_bool), + "role": {"kind": "opaque"}, + "source": {"kind": "literal"}, + "required": False, + "default": False, + }, + "package.zero_scalar": { + "contract": {"dtype": "int64", "rank": 0, "shape": []}, + "role": {"kind": "opaque"}, + "source": {"kind": "literal"}, + "required": False, + "default": 0, + }, + "package.one_scalar": { + "contract": {"dtype": "int64", "rank": 0, "shape": []}, + "role": {"kind": "opaque"}, + "source": {"kind": "literal"}, + "required": False, + "default": 1, + }, + "package.remaining_groups": { + "contract": control_int, + "role": {"kind": "opaque"}, + "source": {"kind": "literal"}, + "required": False, + "default": num_groups - 2, + }, + "package.predictor_context_limit": { + "contract": control_int, + "role": {"kind": "opaque"}, + "source": {"kind": "literal"}, + "required": False, + "default": num_groups, + }, + "package.predictor_mask_limit": { + "contract": control_int, + "role": {"kind": "opaque"}, + "source": {"kind": "literal"}, + "required": False, + "default": num_groups + 1, + }, + "package.talker_context_limit": { + "contract": control_int, + "role": {"kind": "opaque"}, + "source": {"kind": "literal"}, + "required": False, + "default": int(getattr(config, "max_position_embeddings", 4096)), + }, + "package.one_control": { + "contract": control_int, + "role": {"kind": "opaque"}, + "source": {"kind": "literal"}, + "required": False, + "default": 1, + }, + "package.zero_batch": { + "contract": batch_int, + "role": {"kind": "opaque"}, + "source": {"kind": "literal"}, + "required": False, + "default": 0, + }, + "package.one_batch": { + "contract": batch_int, + "role": {"kind": "opaque"}, + "source": {"kind": "literal"}, + "required": False, + "default": 1, + }, + "package.true": { + "contract": batch_bool, + "role": {"kind": "opaque"}, + "source": {"kind": "literal"}, + "required": False, + "default": True, + }, + } + for iteration in range(num_groups - 2): + inputs[f"package.setup_predictor_iteration_{iteration}"] = { + "contract": {"dtype": "int64", "rank": 0, "shape": []}, + "role": {"kind": "opaque"}, + "source": {"kind": "literal"}, + "required": False, + "default": iteration, + } + + talker_setup_inputs = { + talker_embed.name: "tts.prefill_embeds", + talker_mask.name: f"talker.initializer.{talker_mask.name}", + talker_position.name: f"talker.initializer.{talker_position.name}", + **{past.name: f"talker.initializer.{past.name}" for past, _ in talker_caches}, + } + talker_body_inputs = { + talker_embed.name: "talker.step_embeds", + talker_mask.name: "state.talker_mask.body", + talker_position.name: "state.talker_position.body", + **{ + past.name: f"state.talker_cache_{i}.body" + for i, (past, _) in enumerate(talker_caches) + }, + } + talker_setup_outputs = { + talker_logits.name: "talker.setup.logits", + talker_hidden.name: "talker.setup.hidden", + **{present.name: f"talker.setup.{present.name}" for _, present in talker_caches}, + } + talker_body_outputs = { + talker_logits.name: "talker.body.logits", + talker_hidden.name: "talker.body.hidden", + **{present.name: f"talker.body.{present.name}" for _, present in talker_caches}, + } + + def predictor_outputs(prefix: str) -> dict[str, str]: + return { + predictor_logits.name: f"{prefix}.logits", + codec_embeddings.name: f"{prefix}.codec_embeddings", + **{present.name: f"{prefix}.{present.name}" for _, present in predictor_caches}, + } + + predictor_body_inputs = { + predictor_embed.name: "predictor.body.inputs_embeds", + predictor_step.name: "predictor.body.step_index", + predictor_mask.name: "state.predictor_mask.inner", + predictor_position.name: "state.predictor_position.inner", + **{ + past.name: f"state.predictor_cache_{i}.inner" + for i, (past, _) in enumerate(predictor_caches) + }, + } + + def frame_generation_nodes(prefix: str, hidden: str, logits: str) -> list[dict[str, Any]]: + if prefix == "setup": + talker_sampler = "setup_talker_sampler" + talker_effect = "setup_talker_sample" + predictor_sampler = "setup_predictor_sampler" + predictor_effect = "setup_predictor_sample" + else: + talker_sampler = "talker_sampler" + talker_effect = "talker_sample" + predictor_sampler = "predictor_prefill_sampler" + predictor_effect = "predictor_prefill_sample" + initializer = f"{prefix}.predictor.initializer" + return [ + _invoke( + "last_token_logits", + {"logits": logits}, + {"last_logits": f"{prefix}.group0_logits"}, + ), + _invoke( + talker_sampler, + {"logits": f"{prefix}.group0_logits"}, + {"token": f"{prefix}.group0"}, + {talker_effect: _effect(f"{talker_effect}.0", f"{talker_effect}.1")}, + ), + _invoke( + "token_to_slot", + { + "token": f"{prefix}.group0", + }, + {"slot": f"{prefix}.group0_slot"}, + ), + _invoke( + "embedding", + { + "text_ids": "request.prompt_tokens", + "codec_ids": f"{prefix}.group0_slot", + }, + { + "text_embeds": f"{prefix}.unused_text_embeds", + "codec_embeds": f"{prefix}.group0_embed", + }, + ), + _invoke( + "code_predictor_prefill", + { + "talker_hidden": hidden, + "group_0_embed": f"{prefix}.group0_embed", + }, + {"inputs_embeds": f"{prefix}.predictor_prefill"}, + ), + _invoke( + "predictor_state_initializer", + {"prefill_embeds": f"{prefix}.predictor_prefill"}, + { + predictor_mask.name: f"{initializer}.{predictor_mask.name}", + predictor_position.name: f"{initializer}.{predictor_position.name}", + "body_attention_mask": f"{initializer}.body_attention_mask", + "body_position_ids": f"{initializer}.body_position_ids", + **{ + past.name: f"{initializer}.{past.name}" for past, _ in predictor_caches + }, + }, + ), + _invoke( + "code_predictor", + { + predictor_embed.name: f"{prefix}.predictor_prefill", + predictor_step.name: "package.zero_scalar", + predictor_mask.name: f"{initializer}.{predictor_mask.name}", + predictor_position.name: f"{initializer}.{predictor_position.name}", + **{ + past.name: f"{initializer}.{past.name}" for past, _ in predictor_caches + }, + }, + predictor_outputs(f"{prefix}.predictor"), + ), + _invoke( + "last_token_logits", + {"logits": f"{prefix}.predictor.logits"}, + {"last_logits": f"{prefix}.group1_logits"}, + ), + _invoke( + predictor_sampler, + {"logits": f"{prefix}.group1_logits"}, + {"token": f"{prefix}.group1"}, + { + predictor_effect: _effect( + f"{predictor_effect}.0", + f"{predictor_effect}.1", + ) + }, + ), + _invoke( + "code_frame_update", + { + "frame_codes": "initializer.frame_codes", + "token": f"{prefix}.group0", + "index": "package.zero_scalar", + }, + {"next_frame": f"{prefix}.frame_group0"}, + ), + _invoke( + "code_frame_update", + { + "frame_codes": f"{prefix}.frame_group0", + "token": f"{prefix}.group1", + "index": "package.one_scalar", + }, + {"next_frame": f"{prefix}.frame_prefill"}, + ), + ] + + setup_completion_nodes: list[dict[str, Any]] = [] + setup_frame = "setup.frame_prefill" + setup_token = "setup.group1" + setup_mask = "setup.predictor.initializer.body_attention_mask" + setup_position = "setup.predictor.initializer.body_position_ids" + setup_caches = [f"setup.predictor.{present.name}" for _, present in predictor_caches] + for iteration in range(num_groups - 2): + prefix = f"setup.predictor.remaining_{iteration}" + setup_completion_nodes.extend( + [ + _invoke( + "code_predictor_indices", + {"iteration": (f"package.setup_predictor_iteration_{iteration}")}, + { + "embedding_index": f"{prefix}.embedding_index", + "step_index": f"{prefix}.step_index", + "frame_index": f"{prefix}.frame_index", + }, + ), + _invoke( + "code_predictor_step_embedder", + { + "codec_embeddings": "setup.predictor.codec_embeddings", + "token": setup_token, + "embedding_index": f"{prefix}.embedding_index", + }, + {"inputs_embeds": f"{prefix}.inputs_embeds"}, + ), + _invoke( + "code_predictor", + { + predictor_embed.name: f"{prefix}.inputs_embeds", + predictor_step.name: f"{prefix}.step_index", + predictor_mask.name: setup_mask, + predictor_position.name: setup_position, + **{ + past.name: setup_caches[index] + for index, (past, _) in enumerate(predictor_caches) + }, + }, + predictor_outputs(prefix), + ), + _invoke( + "last_token_logits", + {"logits": f"{prefix}.logits"}, + {"last_logits": f"{prefix}.last_logits"}, + ), + _invoke( + "predictor_body_sampler", + {"logits": f"{prefix}.last_logits"}, + {"token": f"{prefix}.token"}, + { + "predictor_body_sample": _effect( + f"predictor_body_sample.{iteration}", + f"predictor_body_sample.{iteration + 1}", + ) + }, + ), + _invoke( + "code_frame_update", + { + "frame_codes": setup_frame, + "token": f"{prefix}.token", + "index": f"{prefix}.frame_index", + }, + {"next_frame": f"{prefix}.frame"}, + ), + _invoke( + "predictor_step_update", + { + "attention_mask": setup_mask, + "position_ids": setup_position, + }, + { + "next_attention_mask": f"{prefix}.mask", + "next_position_ids": f"{prefix}.position", + }, + ), + ] + ) + setup_frame = f"{prefix}.frame" + setup_token = f"{prefix}.token" + setup_mask = f"{prefix}.mask" + setup_position = f"{prefix}.position" + setup_caches = [f"{prefix}.{present.name}" for _, present in predictor_caches] + + inner_body = { + "kind": "sequence", + "nodes": [ + _invoke( + "code_predictor_indices", + {"iteration": "code.iteration"}, + { + "embedding_index": "predictor.body.embedding_index", + "step_index": "predictor.body.step_index", + "frame_index": "predictor.body.frame_index", + }, + ), + *( + [ + _invoke( + "cache_length_update", + { + "left": "state.predictor_cache_lengths.inner", + "right": "package.one_batch", + }, + {"total": "predictor_cache_lengths.next"}, + ) + ] + if predictor_caches + else [] + ), + _invoke( + "code_predictor_step_embedder", + { + "codec_embeddings": "frame.predictor.codec_embeddings", + "token": "state.code_token.inner", + "embedding_index": "predictor.body.embedding_index", + }, + {"inputs_embeds": "predictor.body.inputs_embeds"}, + ), + _invoke( + "code_predictor", predictor_body_inputs, predictor_outputs("predictor.body") + ), + _invoke( + "last_token_logits", + {"logits": "predictor.body.logits"}, + {"last_logits": "predictor.body.last_logits"}, + ), + _invoke( + "predictor_body_sampler", + {"logits": "predictor.body.last_logits"}, + {"token": "code.token"}, + { + "predictor_body_sample": _effect( + f"predictor_body_sample.{num_groups - 2}", + f"predictor_body_sample.{num_groups - 1}", + ) + }, + ), + _invoke( + "code_frame_update", + { + "frame_codes": "state.frame.inner", + "token": "code.token", + "index": "predictor.body.frame_index", + }, + {"next_frame": "frame.inner"}, + ), + _invoke( + "predictor_step_update", + { + "attention_mask": "state.predictor_mask.inner", + "position_ids": "state.predictor_position.inner", + }, + { + "next_attention_mask": "predictor.mask.inner", + "next_position_ids": "predictor.position.inner", + }, + ), + _invoke( + "continue_predicate", + {"done": "package.false"}, + {"continue": "code.continue"}, + ), + ], + } + inner_carried = [ + { + "cell": "frame", + "current": "frame.frame_prefill", + "body_input": "state.frame.inner", + "body_output": "frame.inner", + "next": "frame.completed", + "read_effect": _effect("state:frame.0", "state:frame.read"), + "write_effect": _effect("state:frame.read", "state:frame.1"), + }, + { + "cell": "code_token", + "current": "frame.group1", + "body_input": "state.code_token.inner", + "body_output": "code.token", + "next": "code_token.final", + "read_effect": _effect("state:code_token.0", "state:code_token.read"), + "write_effect": _effect("state:code_token.read", "state:code_token.1"), + }, + { + "cell": "predictor_mask", + "current": "frame.predictor.initializer.body_attention_mask", + "body_input": "state.predictor_mask.inner", + "body_output": "predictor.mask.inner", + "next": "predictor.mask.final", + "read_effect": _effect("state:predictor_mask.0", "state:predictor_mask.read"), + "write_effect": _effect("state:predictor_mask.read", "state:predictor_mask.1"), + }, + { + "cell": "predictor_position", + "current": "frame.predictor.initializer.body_position_ids", + "body_input": "state.predictor_position.inner", + "body_output": "predictor.position.inner", + "next": "predictor.position.final", + "read_effect": _effect( + "state:predictor_position.0", "state:predictor_position.read" + ), + "write_effect": _effect( + "state:predictor_position.read", "state:predictor_position.1" + ), + }, + ] + if predictor_caches: + inner_carried.append( + { + "cell": "predictor_cache_lengths", + "current": "package.zero_batch", + "body_input": "state.predictor_cache_lengths.inner", + "body_output": "predictor_cache_lengths.next", + "next": "predictor_cache_lengths.final", + } + ) + for index, (_, present) in enumerate(predictor_caches): + inner_carried.append( + { + "cell": f"predictor_cache_{index}", + "current": f"frame.predictor.{present.name}", + "body_input": f"state.predictor_cache_{index}.inner", + "body_output": f"predictor.body.{present.name}", + "next": f"predictor.cache_{index}.final", + "read_effect": _effect( + f"state:predictor_cache_{index}.0", f"state:predictor_cache_{index}.read" + ), + "write_effect": _effect( + f"state:predictor_cache_{index}.read", f"state:predictor_cache_{index}.1" + ), + } + ) + inner_loop = { + "kind": "loop", + "setup": { + "kind": "sequence", + "nodes": [ + _invoke( + "continue_predicate", + {"done": "package.false"}, + {"continue": "code.setup.continue"}, + ) + ], + }, + "body": inner_body, + "condition": "code.continue", + "max_iterations": "package.remaining_groups", + "iteration": { + "value": "code.iteration", + "contract": _contract(next(iter(predictor_indices.graph.inputs))), + }, + "carried": inner_carried, + } + + outer_body_nodes = [ + _invoke( + "talker_text_step", + { + "trailing_text_embeds": "tts.trailing_text_embeds", + "iteration": "talker.iteration", + }, + {"text_embed": "talker.text_embed"}, + ), + _invoke( + "talker_step_embedder", + { + "frame_codes": "state.last_frame.outer", + "text_embed": "talker.text_embed", + }, + {"inputs_embeds": "talker.step_embeds"}, + ), + _invoke("talker", talker_body_inputs, talker_body_outputs), + *frame_generation_nodes("frame", "talker.body.hidden", "talker.body.logits"), + inner_loop, + _invoke( + "code_history_append", + {"history": "state.history.outer", "frame": "frame.completed"}, + {"next_history": "history.outer"}, + ), + _invoke( + "talker_step_update", + { + "attention_mask": "state.talker_mask.body", + "position_ids": "state.talker_position.body", + }, + { + "next_attention_mask": "talker.mask.body", + "next_position_ids": "talker.position.body", + }, + ), + *( + [ + _invoke( + "cache_length_update", + { + "left": "state.talker_cache_lengths.body", + "right": "package.one_batch", + }, + {"total": "talker_cache_lengths.next"}, + ), + _invoke( + "cache_length_update", + {"left": "package.zero_batch", "right": "package.one_batch"}, + {"total": "accepted_len.next"}, + ), + ] + if talker_caches or predictor_caches + else [] + ), + _invoke( + "continue_predicate", + {"done": "package.false"}, + {"continue": "talker.continue"}, + ), + ] + + setup_nodes = [ + _invoke( + "tts_state_initializer", + {"prompt_tokens": "request.prompt_tokens"}, + { + "frame_codes": "initializer.frame_codes", + "token_slot": "initializer.token_slot", + "code_history": "initializer.code_history", + }, + ), + _invoke( + "talker_prefill_embedder", + {prompt.name: "request.prompt_tokens"}, + { + prefill.name: "tts.prefill_embeds", + trailing.name: "tts.trailing_text_embeds", + }, + ), + _invoke( + "talker_state_initializer", + {"prefill_embeds": "tts.prefill_embeds"}, + { + talker_mask.name: f"talker.initializer.{talker_mask.name}", + talker_position.name: f"talker.initializer.{talker_position.name}", + "body_attention_mask": "talker.initializer.body_attention_mask", + "body_position_ids": "talker.initializer.body_position_ids", + **{past.name: f"talker.initializer.{past.name}" for past, _ in talker_caches}, + }, + ), + _invoke("talker", talker_setup_inputs, talker_setup_outputs), + *frame_generation_nodes("setup", "talker.setup.hidden", "talker.setup.logits"), + *setup_completion_nodes, + _invoke( + "code_history_append", + {"history": "initializer.code_history", "frame": setup_frame}, + {"next_history": "history.setup"}, + ), + _invoke( + "continue_predicate", + {"done": "package.false"}, + {"continue": "talker.setup.continue"}, + ), + ] + + state = { + "last_frame": { + "contract": {"dtype": "int64", "rank": 2, "shape": [batch, num_groups]}, + "scope": "invocation", + "initializer": setup_frame, + "recurrence": {"kind": "invariant"}, + }, + "history": { + "contract": {"dtype": "int64", "rank": 3, "shape": [batch, "frames", num_groups]}, + "scope": "invocation", + "initializer": "history.setup", + "recurrence": { + "kind": "growing", + "axis": 1, + "increment": "package.one_control", + "max": "package.talker_context_limit", + }, + }, + "talker_mask": { + "contract": { + "dtype": _contract(talker_mask)["dtype"], + "rank": 2, + "shape": [batch, "talker_context"], + }, + "scope": "invocation", + "initializer": "talker.initializer.body_attention_mask", + "recurrence": { + "kind": "growing", + "axis": 1, + "increment": "package.one_control", + "max": "package.talker_context_limit", + }, + }, + "talker_position": { + "contract": _contract( + next( + value + for value in pkg.policy_components[ + "talker_state_initializer" + ].model.graph.outputs + if value.name == "body_position_ids" + ) + ), + "scope": "invocation", + "initializer": "talker.initializer.body_position_ids", + "recurrence": {"kind": "invariant"}, + }, + "frame": { + "contract": {"dtype": "int64", "rank": 2, "shape": [batch, num_groups]}, + "scope": "invocation", + "initializer": "frame.frame_prefill", + "recurrence": {"kind": "invariant"}, + }, + "code_token": { + "contract": batch_int, + "scope": "invocation", + "initializer": "frame.group1", + "recurrence": {"kind": "invariant"}, + }, + "predictor_mask": { + "contract": { + "dtype": _contract(predictor_mask)["dtype"], + "rank": 2, + "shape": [batch, "predictor_context"], + }, + "scope": "invocation", + "initializer": "frame.predictor.initializer.body_attention_mask", + "recurrence": { + "kind": "growing", + "axis": 1, + "increment": "package.one_control", + "max": "package.predictor_mask_limit", + }, + }, + "predictor_position": { + "contract": _contract( + next( + value + for value in pkg.policy_components[ + "predictor_state_initializer" + ].model.graph.outputs + if value.name == "body_position_ids" + ) + ), + "scope": "invocation", + "initializer": "frame.predictor.initializer.body_position_ids", + "recurrence": {"kind": "invariant"}, + }, + "active": { + "contract": batch_bool, + "class": "semantic", + "scope": "invocation", + "initializer": "package.true", + "recurrence": {"kind": "invariant"}, + }, + "done": { + "contract": batch_bool, + "class": "semantic", + "scope": "invocation", + "initializer": "package.false", + "recurrence": {"kind": "invariant"}, + }, + "accepted_len": { + "contract": batch_int, + "class": "semantic", + "scope": "invocation", + "initializer": "package.zero_batch", + "recurrence": {"kind": "invariant"}, + }, + "talker_cache_lengths": { + "contract": batch_int, + "class": "semantic", + "scope": "invocation", + "initializer": "package.zero_batch", + "recurrence": {"kind": "invariant"}, + }, + "predictor_cache_lengths": { + "contract": batch_int, + "class": "semantic", + "scope": "invocation", + "initializer": "package.zero_batch", + "recurrence": {"kind": "invariant"}, + }, + } + for index, (past, present) in enumerate(talker_caches): + state[f"talker_cache_{index}"] = { + "contract": _request_aligned(_contract(past)), + "class": "semantic", + "scope": "invocation", + "initializer": f"talker.setup.{present.name}", + "recurrence": { + "kind": "bounded", + "axis": 2, + "max": "package.talker_context_limit", + }, + "service_group": "talker_cache", + "management": "runtime", + "release_boundary": "invocation", + } + for index, (past, present) in enumerate(predictor_caches): + state[f"predictor_cache_{index}"] = { + "contract": _request_aligned(_contract(past)), + "class": "semantic", + "scope": "invocation", + "initializer": f"frame.predictor.{present.name}", + "recurrence": { + "kind": "bounded", + "axis": 2, + "max": "package.predictor_context_limit", + }, + "service_group": "predictor_cache", + "management": "runtime", + "release_boundary": "invocation", + } + + outer_carried = [ + { + "cell": "active", + "current": "package.true", + "body_input": "state.active.body", + "body_output": "state.active.body", + "next": "state.active.final", + }, + { + "cell": "done", + "current": "package.false", + "body_input": "state.done.body", + "body_output": "state.done.body", + "next": "state.done.final", + }, + { + "cell": "accepted_len", + "current": "package.zero_batch", + "body_input": "state.accepted_len.body", + "body_output": ( + "accepted_len.next" + if talker_caches or predictor_caches + else "state.accepted_len.body" + ), + "next": "state.accepted_len.final", + }, + { + "cell": "talker_cache_lengths", + "current": "package.zero_batch", + "body_input": "state.talker_cache_lengths.body", + "body_output": ( + "talker_cache_lengths.next" + if talker_caches or predictor_caches + else "state.talker_cache_lengths.body" + ), + "next": "state.talker_cache_lengths.final", + }, + { + "cell": "last_frame", + "current": "setup.frame_prefill", + "body_input": "state.last_frame.outer", + "body_output": "frame.completed", + "next": "last_frame.final", + "read_effect": _effect("state:last_frame.0", "state:last_frame.read"), + "write_effect": _effect("state:last_frame.read", "state:last_frame.1"), + }, + { + "cell": "history", + "current": "history.setup", + "body_input": "state.history.outer", + "body_output": "history.outer", + "next": "history.final", + "read_effect": _effect("state:history.0", "state:history.read"), + "write_effect": _effect("state:history.read", "state:history.1"), + }, + { + "cell": "talker_mask", + "current": "talker.initializer.body_attention_mask", + "body_input": "state.talker_mask.body", + "body_output": "talker.mask.body", + "next": "talker.mask.final", + "read_effect": _effect("state:talker_mask.0", "state:talker_mask.read"), + "write_effect": _effect("state:talker_mask.read", "state:talker_mask.1"), + }, + { + "cell": "talker_position", + "current": "talker.initializer.body_position_ids", + "body_input": "state.talker_position.body", + "body_output": "talker.position.body", + "next": "talker.position.final", + "read_effect": _effect("state:talker_position.0", "state:talker_position.read"), + "write_effect": _effect("state:talker_position.read", "state:talker_position.1"), + }, + ] + for index, (_, present) in enumerate(talker_caches): + outer_carried.append( + { + "cell": f"talker_cache_{index}", + "current": f"talker.setup.{present.name}", + "body_input": f"state.talker_cache_{index}.body", + "body_output": f"talker.body.{present.name}", + "next": f"talker.cache_{index}.final", + "read_effect": _effect( + f"state:talker_cache_{index}.0", f"state:talker_cache_{index}.read" + ), + "write_effect": _effect( + f"state:talker_cache_{index}.read", f"state:talker_cache_{index}.1" + ), + } + ) + initial_effects = { + "setup_talker_sample": "setup_talker_sample.0", + "setup_predictor_sample": "setup_predictor_sample.0", + "talker_sample": "talker_sample.0", + "predictor_prefill_sample": "predictor_prefill_sample.0", + "predictor_body_sample": "predictor_body_sample.0", + "emit": "emit.0", + } + for cell in state: + initial_effects[f"state:{cell}"] = f"state:{cell}.0" + + codec_value = "history.final" + final_nodes: list[dict[str, Any]] = [] + if codec_group_major: + final_nodes.append( + _invoke("codec_layout", {"history": "history.final"}, {"codes": "codec.codes"}) + ) + codec_value = "codec.codes" + final_nodes.extend( + [ + _invoke( + codec_name, {codec_input.name: codec_value}, {waveform.name: "tts.waveform"} + ), + { + "kind": "emit", + "value": "tts.waveform", + "output": "waveform", + "mode": "replace", + "effect_name": "emit", + "effect": _effect("emit.0", "emit.1"), + }, + ] + ) + workflow = { + "manifest": { + "capabilities": [ + "workflow_ssa", + "linear_effects", + "nested_control_flow", + "loop_induction_values", + "typed_emit", + *( + ["serving_service_contract", "bounded_state_recurrence"] + if talker_caches or predictor_caches + else [] + ), + ], + }, + "inputs": inputs, + "outputs": { + "waveform": { + "contract": _contract(waveform), + "role": "audio", + "stage": "post_adapter", + } + }, + "components": { + name: _component(model, _artifact(name, len(pkg))) for name, model in pkg.items() + }, + "state": state, + **( + { + "serving": { + "active": "active", + "done": "done", + "accepted_len": "accepted_len", + "state_service": { + "groups": { + **( + { + "talker_cache": _state_group( + sequence_axis=2, + logical_lengths="talker_cache_lengths", + storage=talker_kv["storage"], + ports={ + "talker": { + f"talker_cache_{index}": { + "input": past.name, + "output": present.name, + } + for index, (past, present) in enumerate( + talker_caches + ) + } + }, + ) + } + if talker_caches + else {} + ), + **( + { + "predictor_cache": _state_group( + sequence_axis=2, + logical_lengths="predictor_cache_lengths", + storage=predictor_kv["storage"], + ports={ + "code_predictor": { + f"predictor_cache_{index}": { + "input": past.name, + "output": present.name, + } + for index, (past, present) in enumerate( + predictor_caches + ) + } + }, + ) + } + if predictor_caches + else {} + ), + }, + }, + } + } + if talker_caches or predictor_caches + else {} + ), + "initial_effects": initial_effects, + "graph": { + "kind": "sequence", + "nodes": [ + { + "kind": "loop", + "setup": {"kind": "sequence", "nodes": setup_nodes}, + "body": {"kind": "sequence", "nodes": outer_body_nodes}, + "condition": "talker.continue", + "max_iterations": "request.max_iterations", + "iteration": {"value": "talker.iteration", "contract": batch_int}, + "carried": outer_carried, + }, + *final_nodes, + ], + }, + } + metadata = { + "schema_version": "v1", + "pipeline": {"workflow": _publish_workflow_v1(workflow)}, + } + add_policy_components_to_workflow(metadata, pkg) + return metadata + + +def build_tts_workflow_metadata(pkg: Any, config: Any) -> dict[str, Any]: + """Build nested talker/code-predictor loops with lexical induction SSA.""" + real_transition_components = { + "embedding", + "code_predictor_prefill", + "code_predictor_step_embedder", + "code_predictor_indices", + "talker_text_step", + } + if real_transition_components <= set(pkg.keys()): + return _build_real_tts_workflow_metadata(pkg, config) + required = { + "talker", + "code_predictor", + "talker_step_embedder", + "talker_prefill_embedder", + } + missing = sorted(required.difference(pkg.keys())) + if missing: + raise ValueError(f"TTS workflow is missing required components: {missing}") + talker = pkg["talker"] + predictor = pkg["code_predictor"] + step_embedder = pkg["talker_step_embedder"] + prefill_embedder = pkg["talker_prefill_embedder"] + num_groups = int(getattr(getattr(config, "tts", config), "num_code_groups", 16)) + prompt_input = next( + ( + value + for value in prefill_embedder.graph.inputs + if value.dtype in {ir.DataType.INT32, ir.DataType.INT64} + and value.shape is not None + and len(value.shape) == 2 + ), + None, + ) + prefill_output = _find_port(prefill_embedder.graph.outputs, "prefill", "embeds") + step_frame_input = _find_port(step_embedder.graph.inputs, "frame", "codes") + step_output = _find_port(step_embedder.graph.outputs, "embeds") + talker_embed_input = _find_port(talker.graph.inputs, "embeds") + talker_hidden = _find_port(talker.graph.outputs, "hidden") + predictor_hidden_input = _find_port(predictor.graph.inputs, "hidden", "embeds") + predictor_step_input = _find_port(predictor.graph.inputs, "step", "index") + predictor_logits = _find_port(predictor.graph.outputs, "logits") + if None in ( + prompt_input, + prefill_output, + step_frame_input, + step_output, + talker_embed_input, + talker_hidden, + predictor_hidden_input, + predictor_step_input, + predictor_logits, + ): + raise ValueError("TTS components do not expose the required typed ports") + assert prompt_input is not None + assert prefill_output is not None + assert step_frame_input is not None + assert step_output is not None + assert talker_embed_input is not None + assert talker_hidden is not None + assert predictor_hidden_input is not None + assert predictor_step_input is not None + assert predictor_logits is not None + if predictor_logits.shape is None or len(predictor_logits.shape) != 2: + raise ValueError("TTS code predictor logits must be rank 2") + predictor_step_contract = _contract(predictor_step_input) + scalar_code_index = predictor_step_contract["rank"] == 0 + if predictor_step_contract["dtype"] != "int64" or predictor_step_contract["rank"] not in { + 0, + 1, + }: + raise ValueError("TTS code predictor step index must be scalar or batch int64") + + codec_name = next( + (name for name in ("codec", "vocoder", "decoder") if name in pkg), + None, + ) + codec = pkg[codec_name] if codec_name is not None else None + codec_input = next(iter(codec.graph.inputs), None) if codec is not None else None + waveform_output = next(iter(codec.graph.outputs), None) if codec is not None else None + if codec is None or codec_input is None or waveform_output is None: + raise ValueError("TTS workflow requires a codec or vocoder component") + + attach_policy_components(pkg, PolicyCapabilities(sampler="greedy")) + pkg.add_policy_component("continue_predicate", build_boolean_not()) + pkg.add_policy_component("tts_state_initializer", build_tts_state_initializer(num_groups)) + pkg.add_policy_component( + "code_frame_update", + build_code_frame_update(num_groups, scalar_index=scalar_code_index), + ) + pkg.add_policy_component("code_history_append", build_code_history_append(num_groups)) + codec_group_major = ( + codec_input.shape is not None + and len(codec_input.shape) == 3 + and str(getattr(list(codec_input.shape)[1], "value", list(codec_input.shape)[1])) + == str(num_groups) + ) + if codec_group_major: + pkg.add_policy_component("codec_layout", build_codec_layout_transpose(num_groups)) + + batch = _contract(prompt_input)["shape"][0] + batch_int = _request_aligned({"dtype": "int64", "rank": 1, "shape": [batch]}) + batch_bool = _request_aligned({"dtype": "bool", "rank": 1, "shape": [batch]}) + control_int = {"dtype": "int64", "rank": 1, "shape": [1]} + inputs: dict[str, Any] = { + "request.prompt_tokens": { + "contract": _contract(prompt_input), + "role": {"kind": "runtime", "version": "1.0", "role": "prompt_tokens"}, + "source": {"kind": "request", "field": "prompt_tokens"}, + "required": True, + }, + "request.max_iterations": { + "contract": control_int, + "role": { + "kind": "runtime", + "version": "1.0", + "role": "max_output_tokens", + }, + "source": {"kind": "request", "field": "max_output_tokens"}, + "required": True, + }, + "package.code_groups": { + "contract": control_int, + "role": {"kind": "opaque"}, + "source": {"kind": "literal"}, + "required": False, + "default": num_groups, + }, + "package.false": { + "contract": batch_bool, + "role": {"kind": "opaque"}, + "source": {"kind": "literal"}, + "required": False, + "default": False, + }, + "package.one": { + "contract": control_int, + "role": {"kind": "opaque"}, + "source": {"kind": "literal"}, + "required": False, + "default": 1, + }, + } + + def bind_remaining( + component: str, + values: Any, + known: dict[str, str], + ) -> dict[str, str]: + result = dict(known) + for value in values: + if value.name in result: + continue + name = f"request.{component}.{value.name}" + inputs[name] = { + "contract": _contract(value), + "role": {"kind": "opaque"}, + "source": {"kind": "application", "name": f"{component}.{value.name}"}, + "required": True, + } + result[value.name] = name + return result + + def bind_outputs(values: Any, bound: dict[str, str], prefix: str) -> dict[str, str]: + result = dict(bound) + for value in values: + result.setdefault(value.name, f"{prefix}.{value.name}") + return result + + prefill_inputs = bind_remaining( + "talker_prefill_embedder", + prefill_embedder.graph.inputs, + {prompt_input.name: "request.prompt_tokens"}, + ) + talker_setup_inputs = bind_remaining( + "talker", + talker.graph.inputs, + {talker_embed_input.name: "talker.prefill_embeds"}, + ) + step_inputs = bind_remaining( + "talker_step_embedder", + step_embedder.graph.inputs, + {step_frame_input.name: "state.last_frame.outer"}, + ) + talker_body_inputs = bind_remaining( + "talker", + talker.graph.inputs, + {talker_embed_input.name: "talker.step_embeds"}, + ) + predictor_inputs = bind_remaining( + "code_predictor", + predictor.graph.inputs, + { + predictor_hidden_input.name: "talker.body.hidden", + predictor_step_input.name: "code.iteration", + }, + ) + + frame_contract = { + "dtype": "int64", + "rank": 2, + "shape": [batch, num_groups], + } + history_contract = { + "dtype": "int64", + "rank": 3, + "shape": [batch, "frames", num_groups], + } + initial_effects = { + "sample": "sample.0", + "emit": "emit.0", + "state:last_frame": "state:last_frame.0", + "state:frame": "state:frame.0", + "state:history": "state:history.0", + } + inner_loop = { + "kind": "loop", + "setup": { + "kind": "sequence", + "nodes": [ + _invoke( + "continue_predicate", + {"done": "package.false"}, + {"continue": "code.setup.continue"}, + ) + ], + }, + "body": { + "kind": "sequence", + "nodes": [ + _invoke( + "code_predictor", + predictor_inputs, + bind_outputs( + predictor.graph.outputs, + {predictor_logits.name: "code.logits"}, + "code_predictor.body", + ), + ), + _invoke( + "token_sampler", + {"logits": "code.logits"}, + {"token": "code.token"}, + {"sample": _effect("sample.0", "sample.1")}, + ), + _invoke( + "code_frame_update", + { + "frame_codes": "state.frame.inner", + "token": "code.token", + "index": "code.iteration", + }, + {"next_frame": "frame.inner"}, + ), + _invoke( + "continue_predicate", + {"done": "package.false"}, + {"continue": "code.continue"}, + ), + ], + }, + "condition": "code.continue", + "max_iterations": "package.code_groups", + "iteration": { + "value": "code.iteration", + "contract": predictor_step_contract, + }, + "carried": [ + { + "cell": "frame", + "current": "initializer.frame_codes", + "body_input": "state.frame.inner", + "body_output": "frame.inner", + "next": "frame.completed", + "read_effect": _effect("state:frame.0", "state:frame.read"), + "write_effect": _effect("state:frame.read", "state:frame.1"), + } + ], + } + outer_body = { + "kind": "sequence", + "nodes": [ + _invoke( + "talker_step_embedder", + step_inputs, + bind_outputs( + step_embedder.graph.outputs, + {step_output.name: "talker.step_embeds"}, + "talker_step_embedder.body", + ), + ), + _invoke( + "talker", + talker_body_inputs, + bind_outputs( + talker.graph.outputs, + {talker_hidden.name: "talker.body.hidden"}, + "talker.body", + ), + ), + inner_loop, + _invoke( + "code_history_append", + { + "history": "state.history.outer", + "frame": "frame.completed", + }, + {"next_history": "history.outer"}, + ), + _invoke( + "continue_predicate", + {"done": "package.false"}, + {"continue": "talker.continue"}, + ), + ], + } + setup_outputs = bind_outputs( + prefill_embedder.graph.outputs, + {prefill_output.name: "talker.prefill_embeds"}, + "talker_prefill_embedder.setup", + ) + if talker_hidden.name in {value.name for value in talker.graph.outputs}: + talker_setup_outputs = bind_outputs( + talker.graph.outputs, + {talker_hidden.name: "talker.prefill.hidden"}, + "talker.setup", + ) + else: + talker_setup_outputs = {} + outer_loop = { + "kind": "loop", + "setup": { + "kind": "sequence", + "nodes": [ + _invoke( + "tts_state_initializer", + {"prompt_tokens": "request.prompt_tokens"}, + { + "frame_codes": "initializer.frame_codes", + "code_history": "initializer.code_history", + }, + ), + _invoke("talker_prefill_embedder", prefill_inputs, setup_outputs), + _invoke("talker", talker_setup_inputs, talker_setup_outputs), + ], + }, + "body": outer_body, + "condition": "talker.continue", + "max_iterations": "request.max_iterations", + "iteration": {"value": "talker.iteration", "contract": batch_int}, + "carried": [ + { + "cell": "last_frame", + "current": "initializer.frame_codes", + "body_input": "state.last_frame.outer", + "body_output": "frame.completed", + "next": "state.last_frame.final", + "read_effect": _effect("state:last_frame.0", "state:last_frame.read"), + "write_effect": _effect("state:last_frame.read", "state:last_frame.1"), + }, + { + "cell": "history", + "current": "initializer.code_history", + "body_input": "state.history.outer", + "body_output": "history.outer", + "next": "history.final", + "read_effect": _effect("state:history.0", "state:history.read"), + "write_effect": _effect("state:history.read", "state:history.1"), + }, + ], + } + codec_value = "codec.codes" + post_nodes: list[dict[str, Any]] = [outer_loop] + if codec_group_major: + post_nodes.append( + _invoke( + "codec_layout", + {"history": "history.final"}, + {"codes": codec_value}, + ) + ) + else: + codec_value = "history.final" + post_nodes.extend( + [ + _invoke( + codec_name, + {codec_input.name: codec_value}, + bind_outputs( + codec.graph.outputs, + {waveform_output.name: "tts.waveform"}, + "codec.final", + ), + ), + { + "kind": "emit", + "value": "tts.waveform", + "output": "waveform", + "mode": "replace", + "effect_name": "emit", + "effect": _effect("emit.0", "emit.1"), + }, + ] + ) + workflow = { + "manifest": { + "capabilities": [ + "workflow_ssa", + "linear_effects", + "nested_control_flow", + "loop_induction_values", + "typed_emit", + ], + }, + "inputs": inputs, + "outputs": { + "waveform": { + "contract": _contract(waveform_output), + "role": "audio", + "stage": "post_adapter", + } + }, + "components": { + name: _component(model, _artifact(name, len(pkg))) for name, model in pkg.items() + }, + "state": { + "last_frame": { + "contract": frame_contract, + "scope": "invocation", + "initializer": "initializer.frame_codes", + "recurrence": {"kind": "invariant"}, + }, + "frame": { + "contract": frame_contract, + "scope": "invocation", + "initializer": "initializer.frame_codes", + "recurrence": {"kind": "invariant"}, + }, + "history": { + "contract": history_contract, + "scope": "invocation", + "initializer": "initializer.code_history", + "recurrence": { + "kind": "growing", + "axis": 1, + "increment": "package.one", + "max": "request.max_iterations", + }, + }, + }, + "initial_effects": initial_effects, + "graph": {"kind": "sequence", "nodes": post_nodes}, + } + metadata = { + "schema_version": "v1", + "pipeline": {"workflow": _publish_workflow_v1(workflow)}, + } + add_policy_components_to_workflow(metadata, pkg) + return metadata + + +def write_tts_workflow_metadata(pkg: Any, output_dir: str, config: Any) -> str: + """Build and save an executable TTS workflow package.""" + metadata = build_tts_workflow_metadata(pkg, config) + os.makedirs(output_dir, exist_ok=True) + pkg.save_policy_components(output_dir) + add_adapter_service_to_metadata(metadata, pkg, output_dir) + path = os.path.join(output_dir, "inference_metadata.yaml") + with open(path, "w", encoding="utf-8") as handle: + _dump_yaml(metadata, handle) + return path + + +def _artifact(name: str, package_size: int) -> str: + return f"{name}/model.onnx" if package_size > 1 else "model.onnx" + + +def _find_port(values: Any, *fragments: str) -> ir.Value | None: + return next( + (value for value in values if any(part in value.name.lower() for part in fragments)), + None, + ) + + +def _contracts_compatible(left: ir.Value, right: ir.Value) -> bool: + """Return whether symbolic tensor contracts can unify by dtype and fixed dims.""" + if left.dtype != right.dtype or left.shape is None or right.shape is None: + return False + left_dims = list(left.shape) + right_dims = list(right.shape) + if len(left_dims) != len(right_dims): + return False + return all( + not isinstance(left_dim, int) + or not isinstance(right_dim, int) + or left_dim == right_dim + for left_dim, right_dim in zip(left_dims, right_dims) + ) + + +def _accumulated_contract(contract: dict[str, Any], symbol: str) -> dict[str, Any]: + """Rename the trailing axis of an append-mode output to a private symbol. + + ``mode: append`` emits concatenate chunks along the last axis, so the final + workflow value has ``steps * chunk`` entries there. Reusing the chunk's + symbol (for example ``width``) would rebind it to the accumulated extent and + contradict every other value that shares it, so the accumulating axis gets a + symbol of its own. + """ + shape = list(contract["shape"]) + shape[-1] = symbol + return {**contract, "shape": shape} + + +def _cache_cell(port: str) -> str: + """State-cell name for a ``conv_cache.`` decoder port.""" + return "conv_cache_" + port[len("conv_cache.") :].replace(".", "_") + + +CONV_CACHE_SCALE_METADATA = "mobius.conv_cache.spatial_scale." + + +def _cache_spatial_scale(model: Any, value: ir.Value) -> int: + """Spatial upsampling a conv-cache port has undergone relative to the latent. + + A causal video decoder caches activations at several resolutions, and the + workflow has to allocate the empty first-chunk caches at exactly those + resolutions or the decoder's concatenation fails. The producing task records + the ratio on the model, because symbolic dimension names are not a reliable + channel: shape inference is free to replace a declared ``8*latent_height`` + with an anonymous symbol when it unifies the port with an internal value. + """ + recorded = model.metadata_props.get(f"{CONV_CACHE_SCALE_METADATA}{value.name}") + if recorded is not None: + return int(recorded) + dimension = str(list(value.shape)[3]) + return int(dimension.split("*")[0]) if "*" in dimension else 1 + + +def build_diffusion_workflow_metadata( + pkg: Any, + *, + num_inference_steps: int, + schedule: list[float] | None = None, + timesteps: list[float] | None = None, + solver: str = "euler", + scale_model_input: bool = True, + initial_state_scale: float = 1.0, + decoder_input_scale: float = 1.0, + guidance_scale: float | None = None, + latent_source: str = "application", + latent_row_shape: list[int] | None = None, +) -> dict[str, Any]: + """Build a fixed-schedule diffusion workflow with explicit latent state. + + Everything the reference sampler hides in Python attributes becomes an + explicit part of the workflow: the sigma schedule and timestep table are + constant components, the step index is the loop induction value, a + multistep solver's previous data estimate is a declared state cell, and the + RNG counter is an ordinary integer tensor. Classifier-free guidance is two + denoiser invocations plus a combine component rather than a hidden batch + doubling, so every value stays request-aligned on axis 0. + """ + if num_inference_steps < 1: + raise ValueError("num_inference_steps must be >= 1") + if solver not in SOLVER_BUILDERS: + raise ValueError(f"unsupported diffusion solver {solver!r}") + if latent_source not in {"application", "seed"}: + raise ValueError(f"unsupported diffusion latent source {latent_source!r}") + names = set(pkg.keys()) + denoiser_name = next( + (name for name in ("denoiser", "transformer", "unet") if name in names), + None, + ) + vae_name = next( + (name for name in ("vae_decoder", "decoder", "vae") if name in names), + None, + ) + if denoiser_name is None or vae_name is None or denoiser_name == vae_name: + raise ValueError("diffusion workflow requires distinct denoiser and VAE decoder") + denoiser = pkg[denoiser_name] + vae = pkg[vae_name] + sample_input = _find_port(denoiser.graph.inputs, "sample", "latent", "hidden_states") + timestep_input = _find_port(denoiser.graph.inputs, "timestep", "time") + estimate_output = next(iter(denoiser.graph.outputs), None) + vae_input = _find_port(vae.graph.inputs, "latent", "sample") + vae_output = next(iter(vae.graph.outputs), None) + if None in (sample_input, timestep_input, estimate_output, vae_input, vae_output): + raise ValueError( + "diffusion components do not expose sample/timestep/estimate/VAE ports" + ) + assert sample_input is not None + assert timestep_input is not None + assert estimate_output is not None + assert vae_input is not None + assert vae_output is not None + if len(sample_input.shape or []) != 4 or _contract(sample_input) != _contract( + estimate_output + ): + raise ValueError("diffusion workflow requires matching rank-4 latent/estimate") + if _contract(vae_input) != _contract(sample_input): + raise ValueError("VAE latent input must match the solver latent contract") + + text_name = next( + (name for name in ("text_encoder", "text_encoder_2") if name in names), + None, + ) + text_encoder = pkg[text_name] if text_name is not None else None + conditioning_input = next( + ( + value + for value in denoiser.graph.inputs + if value is not sample_input + and value is not timestep_input + and ("encoder" in value.name or "context" in value.name) + ), + None, + ) + conditioning_output = None + if text_encoder is not None and conditioning_input is not None: + conditioning_output = next( + ( + value + for value in text_encoder.graph.outputs + if _contract(value) == _contract(conditioning_input) + ), + next(iter(text_encoder.graph.outputs), None), + ) + conditioned = text_encoder is not None and conditioning_output is not None + if guidance_scale is not None and not conditioned: + raise ValueError("classifier-free guidance requires a conditioned denoiser") + if latent_source == "seed" and not latent_row_shape: + raise ValueError("a seeded latent initializer requires an explicit row shape") + + solver_component = SOLVER_BUILDERS[solver](sample_input.dtype) + solver_ports = {value.name for value in solver_component.model.graph.inputs} + carries_history = "history" in solver_ports + + attach_policy_components(pkg, PolicyCapabilities()) + pkg.add_policy_component("solver_step", solver_component) + pkg.add_policy_component("continue_predicate", build_boolean_not()) + if scale_model_input: + pkg.add_policy_component( + "model_input_scale", build_euler_model_input(sample_input.dtype) + ) + schedule_values = schedule or [ + 1.0 - index / num_inference_steps for index in range(num_inference_steps + 1) + ] + timestep_values = timesteps or schedule_values[:-1] + if len(schedule_values) != num_inference_steps + 1: + raise ValueError( + "diffusion solver schedule must contain num_inference_steps + 1 values" + ) + if len(timestep_values) != num_inference_steps: + raise ValueError("diffusion timesteps must contain num_inference_steps values") + pkg.add_policy_component("diffusion_schedule", build_schedule_constant(schedule_values)) + pkg.add_policy_component("diffusion_timesteps", build_schedule_constant(timestep_values)) + pkg.add_policy_component("schedule_lookup", build_schedule_lookup(timestep_input.dtype)) + # A sampler whose state already lives in the denoiser's space, and a VAE whose + # latents are unnormalized, need no rescaling step at all; only emit the + # constant and the multiply that the pipeline actually performs. + scales_initial_state = not math.isclose(initial_state_scale, 1.0) + scales_decoder_input = not math.isclose(decoder_input_scale, 1.0) + if scales_initial_state or scales_decoder_input: + pkg.add_policy_component("tensor_scale", build_tensor_scale(sample_input.dtype)) + if scales_initial_state: + pkg.add_policy_component( + "initial_state_scale", build_scalar_constant(initial_state_scale) + ) + if scales_decoder_input: + pkg.add_policy_component( + "decoder_input_scale", build_scalar_constant(decoder_input_scale) + ) + if carries_history: + pkg.add_policy_component("history_initializer", build_zeros_like(sample_input.dtype)) + if guidance_scale is not None: + pkg.add_policy_component( + "guidance_combine", build_guidance_combine(sample_input.dtype) + ) + if latent_source == "seed": + pkg.add_policy_component( + "latent_row_shape", build_shape_constant(list(latent_row_shape or ())) + ) + pkg.add_policy_component("latent_noise", build_counter_rng_normal(sample_input.dtype)) + + batch = _contract(sample_input)["shape"][0] + latent_contract = _contract(sample_input) + row_int = _request_aligned({"dtype": "int64", "rank": 1, "shape": [batch]}) + row_float = _request_aligned({"dtype": "float32", "rank": 1, "shape": [batch]}) + batch_int = _request_aligned({"dtype": "int64", "rank": 1, "shape": [batch]}) + batch_bool = _request_aligned({"dtype": "bool", "rank": 1, "shape": [batch]}) + control_int = {"dtype": "int64", "rank": 1, "shape": [1]} + inputs: dict[str, Any] = { + "request.max_iterations": { + "contract": control_int, + "role": {"kind": "runtime", "version": "1.0", "role": "max_iterations"}, + "source": {"kind": "request", "field": "max_iterations"}, + "required": False, + "default": num_inference_steps, + }, + "package.false": { + "contract": batch_bool, + "role": {"kind": "opaque"}, + "source": {"kind": "literal"}, + "required": False, + "default": False, + }, + } + setup_nodes: list[dict[str, Any]] = [ + _invoke("diffusion_schedule", {}, {"schedule": "diffusion.schedule"}), + _invoke("diffusion_timesteps", {}, {"schedule": "diffusion.timesteps"}), + ] + if scales_initial_state: + setup_nodes.append( + _invoke("initial_state_scale", {}, {"value": "diffusion.initial_scale"}) + ) + if scales_decoder_input: + setup_nodes.append( + _invoke("decoder_input_scale", {}, {"value": "diffusion.decoder_scale"}) + ) + outputs: dict[str, Any] = { + "image": { + "contract": _contract(vae_output), + "role": "image", + "stage": "pre_adapter", + }, + "latent": { + "contract": latent_contract, + "role": "tensor", + "stage": "pre_adapter", + }, + "noise_estimate": { + "contract": _accumulated_contract(latent_contract, "noise_estimate_width"), + "role": "tensor", + "stage": "pre_adapter", + }, + "latent_trajectory": { + "contract": _accumulated_contract(latent_contract, "trajectory_width"), + "role": "tensor", + "stage": "pre_adapter", + }, + } + + if latent_source == "application": + inputs["request.noise"] = { + "contract": latent_contract, + "role": {"kind": "opaque"}, + "source": {"kind": "application", "name": "noise"}, + "required": True, + "externally_suppliable": True, + } + noise_value = "request.noise" + else: + inputs["request.seed"] = { + "contract": row_int, + "role": {"kind": "runtime", "version": "1.0", "role": "seed"}, + "source": {"kind": "request", "field": "seed"}, + "required": False, + "default": 0, + "externally_suppliable": True, + } + inputs["package.rng_offset"] = { + "contract": row_int, + "role": {"kind": "opaque"}, + "source": {"kind": "literal"}, + "required": False, + "default": 0, + } + noise_value = "diffusion.noise" + setup_nodes.append( + _invoke("latent_row_shape", {}, {"shape": "diffusion.latent_row_shape"}) + ) + setup_nodes.append( + _invoke( + "latent_noise", + { + "seed": "request.seed", + "offset": "package.rng_offset", + "row_shape": "diffusion.latent_row_shape", + }, + {"noise": noise_value, "next_offset": "diffusion.rng_offset"}, + ) + ) + outputs["rng_offset"] = { + "contract": row_int, + "role": "tensor", + "stage": "pre_adapter", + } + + initial_state_value = noise_value + if scales_initial_state: + initial_state_value = "diffusion.initial_state" + setup_nodes.append( + _invoke( + "tensor_scale", + {"tensor": noise_value, "scale": "diffusion.initial_scale"}, + {"scaled": initial_state_value}, + ) + ) + + conditioning_value = None + unconditional_value = None + if conditioned: + assert text_encoder is not None + assert conditioning_output is not None + conditioning_value = "conditioning.hidden_states" + text_inputs = {} + for index, value in enumerate(text_encoder.graph.inputs): + name = f"request.{value.name}" + inputs[name] = { + "contract": _contract(value), + "role": ( + {"kind": "runtime", "version": "1.0", "role": "prompt_tokens"} + if index == 0 + else {"kind": "opaque"} + ), + "source": ( + {"kind": "request", "field": "prompt_tokens"} + if index == 0 + else {"kind": "application", "name": value.name} + ), + "required": True, + "externally_suppliable": True, + } + text_inputs[value.name] = name + setup_nodes.append( + _invoke( + text_name, + text_inputs, + {conditioning_output.name: conditioning_value}, + ) + ) + if guidance_scale is not None: + unconditional_value = "conditioning.unconditional" + negative_inputs = {} + for value in text_encoder.graph.inputs: + name = f"request.negative_{value.name}" + inputs[name] = { + "contract": _contract(value), + "role": {"kind": "opaque"}, + "source": {"kind": "application", "name": f"negative_{value.name}"}, + "required": True, + "externally_suppliable": True, + } + negative_inputs[value.name] = name + inputs["request.guidance_scale"] = { + "contract": row_float, + "role": {"kind": "opaque"}, + "source": {"kind": "application", "name": "guidance_scale"}, + "required": False, + "default": float(guidance_scale), + } + setup_nodes.append( + _invoke( + text_name, + negative_inputs, + {conditioning_output.name: unconditional_value}, + ) + ) + + state: dict[str, Any] = { + "latent": { + "contract": latent_contract, + "scope": "invocation", + "initializer": initial_state_value, + "recurrence": {"kind": "invariant"}, + } + } + carried: list[dict[str, Any]] = [ + { + "cell": "latent", + "current": initial_state_value, + "body_input": "state.latent.body", + "body_output": "latent.body", + "next": "latent.final", + "read_effect": _effect("state:latent.0", "state:latent.read"), + "write_effect": _effect("state:latent.read", "state:latent.1"), + } + ] + if carries_history: + setup_nodes.append( + _invoke( + "history_initializer", + {"reference": initial_state_value}, + {"zeros": "diffusion.initial_history"}, + ) + ) + state["history"] = { + "contract": latent_contract, + "scope": "invocation", + "initializer": "diffusion.initial_history", + "recurrence": {"kind": "invariant"}, + } + carried.append( + { + "cell": "history", + "current": "diffusion.initial_history", + "body_input": "state.history.body", + "body_output": "history.body", + "next": "history.final", + "read_effect": _effect("state:history.0", "state:history.read"), + "write_effect": _effect("state:history.read", "state:history.1"), + } + ) + setup_nodes.append( + _invoke( + "continue_predicate", + {"done": "package.false"}, + {"continue": "setup.continue"}, + ) + ) + + body_nodes: list[dict[str, Any]] = [ + _invoke( + "schedule_lookup", + {"schedule": "diffusion.timesteps", "step": "loop.iteration"}, + {"timestep": "diffusion.timestep"}, + ) + ] + if scale_model_input: + model_input_value = "diffusion.model_input" + body_nodes.append( + _invoke( + "model_input_scale", + { + "sample": "state.latent.body", + "step": "loop.iteration", + "schedule": "diffusion.schedule", + }, + {"model_input": model_input_value}, + ) + ) + else: + model_input_value = "state.latent.body" + + def denoiser_call(conditioning: str | None, estimate: str) -> dict[str, Any]: + call_inputs = { + sample_input.name: model_input_value, + timestep_input.name: "diffusion.timestep", + } + if conditioning_input is not None and conditioning is not None: + call_inputs[conditioning_input.name] = conditioning + return _invoke(denoiser_name, call_inputs, {estimate_output.name: estimate}) + + if guidance_scale is None: + body_nodes.append(denoiser_call(conditioning_value, "denoiser.estimate")) + else: + body_nodes.append(denoiser_call(unconditional_value, "denoiser.unconditional")) + body_nodes.append(denoiser_call(conditioning_value, "denoiser.conditional")) + body_nodes.append( + _invoke( + "guidance_combine", + { + "unconditional": "denoiser.unconditional", + "conditional": "denoiser.conditional", + "scale": "request.guidance_scale", + }, + {"estimate": "denoiser.estimate"}, + ) + ) + + solver_inputs = { + "sample": "state.latent.body", + "step": "loop.iteration", + "schedule": "diffusion.schedule", + } + solver_inputs["estimate" if carries_history else "derivative"] = "denoiser.estimate" + solver_outputs = {"next_state": "latent.body"} + if carries_history: + solver_inputs["history"] = "state.history.body" + solver_outputs["next_history"] = "history.body" + body_nodes.append(_invoke("solver_step", solver_inputs, solver_outputs)) + body_nodes.extend( + [ + { + "kind": "emit", + "value": "denoiser.estimate", + "output": "noise_estimate", + "mode": "append", + "effect_name": "emit", + "effect": _effect("emit.0", "emit.1"), + }, + { + "kind": "emit", + "value": "latent.body", + "output": "latent_trajectory", + "mode": "append", + "effect_name": "emit", + "effect": _effect("emit.1", "emit.2"), + }, + _invoke( + "continue_predicate", + {"done": "package.false"}, + {"continue": "loop.continue"}, + ), + ] + ) + + decoder_input_value = "latent.final" + tail_nodes: list[dict[str, Any]] = [] + if scales_decoder_input: + decoder_input_value = "diffusion.decoder_input" + tail_nodes.append( + _invoke( + "tensor_scale", + {"tensor": "latent.final", "scale": "diffusion.decoder_scale"}, + {"scaled": decoder_input_value}, + ) + ) + tail_nodes.extend( + [ + _invoke( + vae_name, + {vae_input.name: decoder_input_value}, + {vae_output.name: "vae.image"}, + ), + { + "kind": "emit", + "value": "latent.final", + "output": "latent", + "mode": "replace", + "effect_name": "emit", + "effect": _effect("emit.2", "emit.3"), + }, + { + "kind": "emit", + "value": "vae.image", + "output": "image", + "mode": "replace", + "effect_name": "emit", + "effect": _effect("emit.3", "emit.4"), + }, + ] + ) + if latent_source == "seed": + tail_nodes.append( + { + "kind": "emit", + "value": "diffusion.rng_offset", + "output": "rng_offset", + "mode": "replace", + "effect_name": "emit", + "effect": _effect("emit.4", "emit.5"), + } + ) + + workflow = { + "manifest": { + "capabilities": [ + "workflow_ssa", + "linear_effects", + "nested_control_flow", + "loop_induction_values", + "typed_emit", + ], + }, + "inputs": inputs, + "outputs": outputs, + "components": { + name: _component(model, _artifact(name, len(pkg))) for name, model in pkg.items() + }, + "state": state, + "initial_effects": { + "emit": "emit.0", + **{f"state:{cell}": f"state:{cell}.0" for cell in state}, + }, + "graph": { + "kind": "sequence", + "nodes": [ + { + "kind": "loop", + "setup": {"kind": "sequence", "nodes": setup_nodes}, + "body": {"kind": "sequence", "nodes": body_nodes}, + "condition": "loop.continue", + "max_iterations": "request.max_iterations", + "iteration": {"value": "loop.iteration", "contract": batch_int}, + "carried": carried, + }, + *tail_nodes, + ], + }, + } + metadata = { + "schema_version": "v1", + "pipeline": {"workflow": _publish_workflow_v1(workflow)}, + } + add_policy_components_to_workflow(metadata, pkg) + return metadata + + +def write_diffusion_workflow_metadata( + pkg: Any, + output_dir: str, + *, + num_inference_steps: int, + schedule: list[float] | None = None, + timesteps: list[float] | None = None, + solver: str = "euler", + scale_model_input: bool = True, + initial_state_scale: float = 1.0, + decoder_input_scale: float = 1.0, + guidance_scale: float | None = None, + latent_source: str = "application", + latent_row_shape: list[int] | None = None, +) -> str: + os.makedirs(output_dir, exist_ok=True) + metadata = build_diffusion_workflow_metadata( + pkg, + num_inference_steps=num_inference_steps, + schedule=schedule, + timesteps=timesteps, + solver=solver, + scale_model_input=scale_model_input, + initial_state_scale=initial_state_scale, + decoder_input_scale=decoder_input_scale, + guidance_scale=guidance_scale, + latent_source=latent_source, + latent_row_shape=latent_row_shape, + ) + pkg.save_policy_components(output_dir) + add_adapter_service_to_metadata(metadata, pkg, output_dir) + path = os.path.join(output_dir, "inference_metadata.yaml") + with open(path, "w", encoding="utf-8") as handle: + _dump_yaml(metadata, handle) + return path + + +def _application_input(value: ir.Value, name: str | None = None) -> dict[str, Any]: + """Declare a workflow input supplied by the host application.""" + return { + "contract": _contract(value), + "role": {"kind": "opaque"}, + "source": {"kind": "application", "name": name or value.name}, + "required": True, + } + + +def build_image_edit_workflow_metadata( + pkg: Any, + *, + num_inference_steps: int, + schedule: list[float], + timesteps: list[float], + guidance_scale: float, +) -> dict[str, Any]: + """Build a flow-matching image-edit workflow with true classifier-free guidance. + + Emitted pipeline (Qwen Image Edit and any package with the same component + shape):: + + vae_encoder(source pixels) -> pack -> source tokens [setup] + loop: + timestep = timesteps[i] + model_input = concat([target tokens, source tokens], 1) + cond = denoiser(model_input, positive prompt) + uncond = denoiser(model_input, negative prompt) + estimate = true_cfg(cond, uncond) + target tokens = target + (sigma[i+1] - sigma[i]) * estimate + unpack(target tokens) -> vae_decoder -> image + + The denoiser slices its own output back to the target token count using the + ``target_sequence_length`` port, so the loop state stays rank-3 and the + source tokens stay loop-invariant. + """ + if num_inference_steps < 1: + raise ValueError("num_inference_steps must be >= 1") + if len(schedule) != num_inference_steps + 1: + raise ValueError("image-edit schedule must contain num_inference_steps + 1 values") + if len(timesteps) != num_inference_steps: + raise ValueError("image-edit timesteps must contain num_inference_steps values") + + denoiser_name = "transformer" + denoiser = pkg[denoiser_name] + encoder = pkg["vae_encoder"] + decoder = pkg["vae_decoder"] + + ports = {value.name: value for value in denoiser.graph.inputs} + sample_input = ports["sample"] + timestep_input = ports["timestep"] + estimate_output = denoiser.graph.outputs[0] + encoder_input = encoder.graph.inputs[0] + encoder_output = encoder.graph.outputs[0] + decoder_input = decoder.graph.inputs[0] + decoder_output = decoder.graph.outputs[0] + dtype = sample_input.dtype + + # Loop state is the target token block only; the denoiser input additionally + # carries the source tokens, so the two contracts differ in sequence length. + latent_contract = { + "dtype": _contract(sample_input)["dtype"], + "rank": 3, + "shape": ["batch", "target_sequence_length", _contract(sample_input)["shape"][2]], + } + + attach_policy_components(pkg, PolicyCapabilities()) + pkg.add_policy_component("solver_step", build_flow_match_solver_step(dtype)) + pkg.add_policy_component("continue_predicate", build_boolean_not()) + pkg.add_policy_component("diffusion_schedule", build_schedule_constant(schedule)) + pkg.add_policy_component("diffusion_timesteps", build_schedule_constant(timesteps)) + pkg.add_policy_component("schedule_lookup", build_schedule_lookup(timestep_input.dtype)) + pkg.add_policy_component("pack_latents", build_pack_latents_2x2(dtype)) + pkg.add_policy_component("unpack_latents", build_unpack_latents_2x2(dtype)) + pkg.add_policy_component("sequence_concat", build_sequence_concat(dtype)) + pkg.add_policy_component("true_cfg", build_true_cfg(dtype, guidance_scale=guidance_scale)) + + batch = latent_contract["shape"][0] + batch_int = {"dtype": "int64", "rank": 1, "shape": [batch]} + batch_bool = {"dtype": "bool", "rank": 1, "shape": [batch]} + control_int = {"dtype": "int64", "rank": 1, "shape": [1]} + + conditioning_ports = ("encoder_hidden_states", "encoder_hidden_states_mask") + rotary_ports = ("image_rotary_cos", "image_rotary_sin") + text_rotary_ports = ("text_rotary_cos", "text_rotary_sin") + + inputs: dict[str, Any] = { + "request.latent": { + "contract": latent_contract, + "role": {"kind": "opaque"}, + "source": {"kind": "application", "name": "latent"}, + "required": True, + }, + "request.source_pixels": _application_input(encoder_input, "source_pixels"), + "request.target_sequence_length": _application_input(ports["target_sequence_length"]), + "request.latent_height": { + "contract": control_int, + "role": {"kind": "opaque"}, + "source": {"kind": "application", "name": "latent_height"}, + "required": True, + }, + "request.latent_width": { + "contract": control_int, + "role": {"kind": "opaque"}, + "source": {"kind": "application", "name": "latent_width"}, + "required": True, + }, + "request.max_iterations": { + "contract": control_int, + "role": {"kind": "runtime", "version": "1.0", "role": "max_iterations"}, + "source": {"kind": "request", "field": "max_iterations"}, + "required": False, + "default": num_inference_steps, + }, + "package.false": { + "contract": batch_bool, + "role": {"kind": "opaque"}, + "source": {"kind": "literal"}, + "required": False, + "default": False, + }, + } + for port_name in rotary_ports: + inputs[f"request.{port_name}"] = _application_input(ports[port_name]) + # Positive and negative conditioning are separate application inputs: the two + # prompts tokenize to different lengths, so they cannot share a contract dim. + for prefix in ("positive", "negative"): + for port_name in conditioning_ports + text_rotary_ports: + contract = _contract(ports[port_name]) + contract["shape"] = [ + f"{prefix}_text_sequence_length" + if isinstance(dim, str) and "text" in dim + else dim + for dim in contract["shape"] + ] + inputs[f"request.{prefix}_{port_name}"] = { + "contract": contract, + "role": {"kind": "opaque"}, + "source": {"kind": "application", "name": f"{prefix}_{port_name}"}, + "required": True, + } + + setup_nodes: list[dict[str, Any]] = [ + _invoke("diffusion_schedule", {}, {"schedule": "diffusion.schedule"}), + _invoke("diffusion_timesteps", {}, {"schedule": "diffusion.timesteps"}), + _invoke( + "vae_encoder", + {encoder_input.name: "request.source_pixels"}, + {encoder_output.name: "source.latent"}, + ), + _invoke( + "pack_latents", + {"latent_sample": "source.latent"}, + {"packed_latent": "source.tokens"}, + ), + _invoke( + "continue_predicate", + {"done": "package.false"}, + {"continue": "setup.continue"}, + ), + ] + + def denoise(prefix: str, output: str) -> dict[str, Any]: + feeds = { + sample_input.name: "diffusion.model_input", + timestep_input.name: "diffusion.timestep", + "target_sequence_length": "request.target_sequence_length", + } + for port_name in rotary_ports: + feeds[port_name] = f"request.{port_name}" + for port_name in conditioning_ports + text_rotary_ports: + feeds[port_name] = f"request.{prefix}_{port_name}" + return _invoke(denoiser_name, feeds, {estimate_output.name: output}) + + body_nodes: list[dict[str, Any]] = [ + _invoke( + "schedule_lookup", + {"schedule": "diffusion.timesteps", "step": "loop.iteration"}, + {"timestep": "diffusion.timestep"}, + ), + _invoke( + "sequence_concat", + {"target": "state.latent.body", "source": "source.tokens"}, + {"sequence": "diffusion.model_input"}, + ), + denoise("positive", "denoiser.conditional"), + denoise("negative", "denoiser.unconditional"), + _invoke( + "true_cfg", + { + "conditional": "denoiser.conditional", + "unconditional": "denoiser.unconditional", + }, + {"estimate": "denoiser.estimate"}, + ), + _invoke( + "solver_step", + { + "sample": "state.latent.body", + "derivative": "denoiser.estimate", + "step": "loop.iteration", + "schedule": "diffusion.schedule", + }, + {"next_state": "latent.body"}, + {"solver": _effect("solver.0", "solver.1")}, + ), + _invoke( + "continue_predicate", + {"done": "package.false"}, + {"continue": "loop.continue"}, + ), + ] + + latent_effect = "state:latent" + workflow = { + "manifest": { + "capabilities": [ + "workflow_ssa", + "linear_effects", + "nested_control_flow", + "loop_induction_values", + "typed_emit", + ], + }, + "inputs": inputs, + "outputs": { + "image": { + "contract": _contract(decoder_output), + "role": "image", + "stage": "pre_adapter", + } + }, + "components": { + name: _component(model, _artifact(name, len(pkg))) for name, model in pkg.items() + }, + "state": { + "latent": { + "contract": latent_contract, + "scope": "invocation", + "initializer": "request.latent", + "recurrence": {"kind": "invariant"}, + } + }, + "initial_effects": { + "solver": "solver.0", + latent_effect: f"{latent_effect}.0", + "emit": "emit.0", + }, + "graph": { + "kind": "sequence", + "nodes": [ + { + "kind": "loop", + "setup": {"kind": "sequence", "nodes": setup_nodes}, + "body": {"kind": "sequence", "nodes": body_nodes}, + "condition": "loop.continue", + "max_iterations": "request.max_iterations", + "iteration": {"value": "loop.iteration", "contract": batch_int}, + "carried": [ + { + "cell": "latent", + "current": "request.latent", + "body_input": "state.latent.body", + "body_output": "latent.body", + "next": "latent.final", + "read_effect": _effect( + f"{latent_effect}.0", f"{latent_effect}.read" + ), + "write_effect": _effect( + f"{latent_effect}.read", f"{latent_effect}.1" + ), + } + ], + }, + _invoke( + "unpack_latents", + { + "packed_latent": "latent.final", + "height": "request.latent_height", + "width": "request.latent_width", + }, + {"latent_sample": "vae.latent"}, + ), + _invoke( + "vae_decoder", + {decoder_input.name: "vae.latent"}, + {decoder_output.name: "vae.image"}, + ), + { + "kind": "emit", + "value": "vae.image", + "output": "image", + "mode": "replace", + "effect_name": "emit", + "effect": _effect("emit.0", "emit.1"), + }, + ], + }, + } + metadata = { + "schema_version": "v1", + "pipeline": {"workflow": _publish_workflow_v1(workflow)}, + } + add_policy_components_to_workflow(metadata, pkg) + return metadata + + +def write_image_edit_workflow_metadata( + pkg: Any, + output_dir: str, + *, + num_inference_steps: int, + schedule: list[float], + timesteps: list[float], + guidance_scale: float, +) -> str: + os.makedirs(output_dir, exist_ok=True) + metadata = build_image_edit_workflow_metadata( + pkg, + num_inference_steps=num_inference_steps, + schedule=schedule, + timesteps=timesteps, + guidance_scale=guidance_scale, + ) + pkg.save_policy_components(output_dir) + add_adapter_service_to_metadata(metadata, pkg, output_dir) + path = os.path.join(output_dir, "inference_metadata.yaml") + with open(path, "w", encoding="utf-8") as handle: + _dump_yaml(metadata, handle) + return path + + +def build_video_diffusion_workflow_metadata( + pkg: Any, + *, + num_inference_steps: int, + schedule: list[float] | None = None, + timesteps: list[float] | None = None, + init_noise_sigma: float = 1.0, + scaling_factor: float = 1.0, + latent_permutation: list[int] | None = None, + solver: str = "euler", + clip_sample_range: float | None = None, +) -> dict[str, Any]: + """Build a text-to-video diffusion workflow over rank-5 temporal latents. + + The shape of the workflow differs from the image path in ways that are + intrinsic to video rather than cosmetic: + + - the latent and the published frames carry a temporal axis, so every + contract is rank 5 and no stage may assume a single frame; + - the scheduler's timestep history is carried state, not telemetry; + - the decoder is causal and is invoked once per latent-frame chunk, with the + convolution caches as runtime-owned state released at the end of the + invocation; + - frames are published incrementally, appending on the temporal axis as each + chunk is decoded, so a consumer sees frames before the clip is finished. + """ + if num_inference_steps < 1: + raise ValueError("num_inference_steps must be >= 1") + names = set(pkg.keys()) + denoiser_name = next( + (name for name in ("denoiser", "transformer", "unet") if name in names), None + ) + vae_name = next( + (name for name in ("vae_decoder", "decoder", "vae") if name in names), None + ) + if denoiser_name is None or vae_name is None or denoiser_name == vae_name: + raise ValueError("video workflow requires distinct denoiser and VAE decoder") + denoiser = pkg[denoiser_name] + vae = pkg[vae_name] + sample_input = _find_port(denoiser.graph.inputs, "sample", "latent", "hidden_states") + timestep_input = _find_port(denoiser.graph.inputs, "timestep", "time") + estimate_output = next(iter(denoiser.graph.outputs), None) + vae_input = _find_port(vae.graph.inputs, "latent_sample", "latent", "sample") + vae_output = next( + (value for value in vae.graph.outputs if not value.name.startswith("conv_cache")), + None, + ) + if None in (sample_input, timestep_input, estimate_output, vae_input, vae_output): + raise ValueError("video components do not expose sample/timestep/estimate/VAE ports") + assert sample_input is not None + assert timestep_input is not None + assert estimate_output is not None + assert vae_input is not None + assert vae_output is not None + if len(sample_input.shape or []) != 5: + raise ValueError( + "video diffusion workflow requires a rank-5 [batch, frames, channels, " + "height, width] latent; use build_diffusion_workflow_metadata for images" + ) + if _contract(sample_input) != _contract(estimate_output): + raise ValueError("video workflow requires matching latent/estimate contracts") + if len(vae_input.shape or []) != 5 or len(vae_output.shape or []) != 5: + raise ValueError("video VAE decode must be rank 5 on both the latent and the frames") + + cache_ports = [ + value.name for value in vae.graph.inputs if value.name.startswith("conv_cache.") + ] + cache_outputs = { + value.name[len("conv_cache_out.") :]: value.name + for value in vae.graph.outputs + if value.name.startswith("conv_cache_out.") + } + if not cache_ports or set(cache_ports) != {f"conv_cache.{name}" for name in cache_outputs}: + raise ValueError("causal video decode requires paired conv_cache/conv_cache_out ports") + cache_entries = [ + ( + port, + int(list(next(v for v in vae.graph.inputs if v.name == port).shape)[1]), + _cache_spatial_scale(vae, next(v for v in vae.graph.inputs if v.name == port)), + ) + for port in cache_ports + ] + + text_name = next( + (name for name in ("text_encoder", "text_encoder_2") if name in names), None + ) + text_encoder = pkg[text_name] if text_name is not None else None + conditioning_input = next( + ( + value + for value in denoiser.graph.inputs + if value is not sample_input + and value is not timestep_input + and ("encoder" in value.name or "context" in value.name) + ), + None, + ) + conditioning_output = None + if text_encoder is not None and conditioning_input is not None: + conditioning_output = next( + ( + value + for value in text_encoder.graph.outputs + if _contract(value) == _contract(conditioning_input) + ), + next(iter(text_encoder.graph.outputs), None), + ) + + if solver not in ("euler", "ddim"): + raise ValueError("video solver must be 'euler' or 'ddim'") + latent_dims = ["batch", "frames", "channels", "height", "width"] + attach_policy_components(pkg, PolicyCapabilities()) + if solver == "ddim": + # DDIM defines scale_model_input as the identity and consumes cumulative + # alphas rather than sigmas. + pkg.add_policy_component( + "model_input", build_identity_model_input(sample_input.dtype, latent_dims) + ) + pkg.add_policy_component( + "solver_step", + build_ddim_solver_step( + sample_input.dtype, latent_dims, clip_sample_range=clip_sample_range + ), + ) + else: + pkg.add_policy_component( + "model_input", build_euler_model_input(sample_input.dtype, latent_dims) + ) + pkg.add_policy_component( + "solver_step", build_euler_solver_step(sample_input.dtype, latent_dims) + ) + pkg.add_policy_component("continue_predicate", build_boolean_not()) + pkg.add_policy_component( + "video_latent_init", + build_video_latent_initializer(sample_input.dtype, init_noise_sigma), + ) + pkg.add_policy_component( + "schedule_history_append", build_schedule_history_append(timestep_input.dtype) + ) + pkg.add_policy_component( + "video_latent_permute", + build_video_latent_permute(latent_permutation or [0, 2, 1, 3, 4]), + ) + pkg.add_policy_component( + "video_latent_unscale", build_video_latent_unscale(scaling_factor) + ) + pkg.add_policy_component("video_decode_chunks", build_video_decode_chunk_count()) + pkg.add_policy_component("video_decode_chunk", build_video_decode_chunk()) + pkg.add_policy_component( + "video_conv_cache_init", build_video_conv_cache_initializer(cache_entries) + ) + + schedule_values = schedule or [ + 1.0 - index / num_inference_steps for index in range(num_inference_steps + 1) + ] + timestep_values = timesteps or schedule_values[:-1] + if len(schedule_values) != num_inference_steps + 1: + raise ValueError("video solver schedule must contain num_inference_steps + 1 values") + if len(timestep_values) != num_inference_steps: + raise ValueError("video timesteps must contain num_inference_steps values") + pkg.add_policy_component("diffusion_schedule", build_schedule_constant(schedule_values)) + pkg.add_policy_component("diffusion_timesteps", build_schedule_constant(timestep_values)) + pkg.add_policy_component("schedule_lookup", build_schedule_lookup(timestep_input.dtype)) + + latent_contract = _request_aligned(_contract(sample_input)) + batch = latent_contract["shape"][0] + batch_int = {"dtype": "int64", "rank": 1, "shape": [batch]} + batch_bool = {"dtype": "bool", "rank": 1, "shape": [batch]} + control_int = {"dtype": "int64", "rank": 1, "shape": [1]} + history_contract = _request_aligned( + { + "dtype": _contract(timestep_input)["dtype"], + "rank": 2, + "shape": [batch, "scheduler_history"], + } + ) + + inputs: dict[str, Any] = { + "request.noise": { + "contract": latent_contract, + "role": {"kind": "opaque"}, + "source": {"kind": "application", "name": "noise"}, + "required": True, + }, + "request.max_iterations": { + "contract": control_int, + "role": {"kind": "runtime", "version": "1.0", "role": "max_iterations"}, + "source": {"kind": "request", "field": "max_iterations"}, + "required": False, + "default": num_inference_steps, + }, + "package.false": { + "contract": batch_bool, + "role": {"kind": "opaque"}, + "source": {"kind": "literal"}, + "required": False, + "default": False, + }, + "package.one_control": { + "contract": control_int, + "role": {"kind": "opaque"}, + "source": {"kind": "literal"}, + "required": False, + "default": 1, + }, + "package.history_limit": { + "contract": control_int, + "role": {"kind": "opaque"}, + "source": {"kind": "literal"}, + "required": False, + "default": num_inference_steps, + }, + "package.cache_frames": { + "contract": control_int, + "role": {"kind": "opaque"}, + "source": {"kind": "literal"}, + "required": False, + "default": 2, + }, + } + + setup_nodes: list[dict[str, Any]] = [ + _invoke("diffusion_schedule", {}, {"schedule": "diffusion.schedule"}), + _invoke("diffusion_timesteps", {}, {"schedule": "diffusion.timesteps"}), + ] + conditioning_value = None + if text_encoder is not None and conditioning_output is not None: + text_inputs = {} + for index, value in enumerate(text_encoder.graph.inputs): + name = f"request.{value.name}" + inputs[name] = { + "contract": _request_aligned(_contract(value)), + "role": ( + {"kind": "runtime", "version": "1.0", "role": "prompt_tokens"} + if index == 0 + else {"kind": "opaque"} + ), + "source": ( + {"kind": "request", "field": "prompt_tokens"} + if index == 0 + else {"kind": "application", "name": value.name} + ), + "required": True, + } + text_inputs[value.name] = name + conditioning_value = "conditioning.hidden_states" + setup_nodes.append( + _invoke(text_name, text_inputs, {conditioning_output.name: conditioning_value}) + ) + if conditioning_input is not None and conditioning_value is None: + # No text encoder ships with the package, so the prompt embedding is + # supplied by the application. Conditioning stays a declared input + # rather than an implicit constant: an unconditioned video model would + # simply have no such port on its denoiser. + conditioning_value = f"request.{conditioning_input.name}" + inputs[conditioning_value] = { + "contract": _request_aligned(_contract(conditioning_input)), + "role": {"kind": "opaque"}, + "source": {"kind": "application", "name": conditioning_input.name}, + "required": True, + } + setup_nodes.append( + _invoke( + "continue_predicate", + {"done": "package.false"}, + {"continue": "setup.continue"}, + ) + ) + + denoiser_inputs = { + sample_input.name: "diffusion.model_input", + timestep_input.name: "diffusion.timestep", + } + if conditioning_input is not None and conditioning_value is not None: + denoiser_inputs[conditioning_input.name] = conditioning_value + + body_nodes: list[dict[str, Any]] = [ + _invoke( + "schedule_lookup", + {"schedule": "diffusion.timesteps", "step": "loop.iteration"}, + {"timestep": "diffusion.timestep"}, + ), + _invoke( + "model_input", + { + "sample": "state.latent.body", + "step": "loop.iteration", + "schedule": "diffusion.schedule", + }, + {"model_input": "diffusion.model_input"}, + ), + _invoke(denoiser_name, denoiser_inputs, {estimate_output.name: "denoiser.estimate"}), + _invoke( + "solver_step", + { + "sample": "state.latent.body", + "derivative": "denoiser.estimate", + "step": "loop.iteration", + "schedule": "diffusion.schedule", + }, + {"next_state": "latent.body"}, + {"solver": _effect("solver.0", "solver.1")}, + ), + _invoke( + "schedule_history_append", + {"history": "state.scheduler_history.body", "timestep": "diffusion.timestep"}, + {"next": "scheduler_history.body"}, + ), + _invoke( + "continue_predicate", + {"done": "package.false"}, + {"continue": "loop.continue"}, + ), + ] + + decode_body: list[dict[str, Any]] = [ + _invoke( + "video_decode_chunk", + {"latent": "decode.latent", "step": "decode.iteration"}, + {"chunk": "decode.chunk"}, + ), + _invoke( + vae_name, + { + vae_input.name: "decode.chunk", + **{port: f"state.{_cache_cell(port)}.body" for port in cache_ports}, + }, + { + vae_output.name: "decode.frames", + **{ + cache_outputs[name]: f"{_cache_cell(f'conv_cache.{name}')}.body" + for name in cache_outputs + }, + }, + ), + { + "kind": "emit", + "value": "decode.frames", + "output": "video", + "mode": "append", + "axis": 2, + "effect_name": "emit", + "effect": _effect("emit.0", "emit.1"), + }, + _invoke( + "continue_predicate", + {"done": "package.false"}, + {"continue": "decode.loop.continue"}, + ), + ] + + frames_contract = _request_aligned(_contract(vae_output)) + workflow = { + "manifest": { + "capabilities": [ + "workflow_ssa", + "linear_effects", + "nested_control_flow", + "loop_induction_values", + "typed_emit", + "bounded_state_recurrence", + ], + }, + "inputs": inputs, + "outputs": { + "video": { + "contract": frames_contract, + "role": "video", + "stage": "pre_adapter", + } + }, + "components": { + name: _component(model, _artifact(name, len(pkg))) for name, model in pkg.items() + }, + "state": { + "latent": { + "contract": latent_contract, + "scope": "invocation", + "initializer": "latent.initial", + "recurrence": {"kind": "invariant"}, + }, + "scheduler_history": { + "contract": history_contract, + "scope": "invocation", + "initializer": "scheduler.history.initial", + "recurrence": { + "kind": "growing", + "axis": 1, + "increment": "package.one_control", + "max": "package.history_limit", + }, + }, + **{ + _cache_cell(port): { + "contract": _request_aligned( + _contract(next(v for v in vae.graph.inputs if v.name == port)) + ), + "scope": "invocation", + "initializer": f"{_cache_cell(port)}.initial", + "recurrence": { + "kind": "bounded", + "axis": 2, + "max": "package.cache_frames", + }, + "management": "runtime", + "release_boundary": "invocation", + } + for port in cache_ports + }, + }, + "initial_effects": { + "solver": "solver.0", + "state:latent": "state:latent.0", + "state:scheduler_history": "state:scheduler_history.0", + **{ + f"state:{_cache_cell(port)}": f"state:{_cache_cell(port)}.0" + for port in cache_ports + }, + "emit": "emit.0", + }, + "graph": { + "kind": "sequence", + "nodes": [ + _invoke( + "video_latent_init", + {"noise": "request.noise"}, + {"latent": "latent.initial", "history": "scheduler.history.initial"}, + ), + { + "kind": "loop", + "setup": {"kind": "sequence", "nodes": setup_nodes}, + "body": {"kind": "sequence", "nodes": body_nodes}, + "condition": "loop.continue", + "max_iterations": "request.max_iterations", + "iteration": {"value": "loop.iteration", "contract": batch_int}, + "carried": [ + { + "cell": "latent", + "current": "latent.initial", + "body_input": "state.latent.body", + "body_output": "latent.body", + "next": "latent.final", + "read_effect": _effect("state:latent.0", "state:latent.read"), + "write_effect": _effect("state:latent.read", "state:latent.1"), + }, + { + "cell": "scheduler_history", + "current": "scheduler.history.initial", + "body_input": "state.scheduler_history.body", + "body_output": "scheduler_history.body", + "next": "scheduler_history.final", + "read_effect": _effect( + "state:scheduler_history.0", "state:scheduler_history.read" + ), + "write_effect": _effect( + "state:scheduler_history.read", "state:scheduler_history.1" + ), + }, + ], + }, + _invoke( + "video_latent_permute", + {"latent": "latent.final"}, + {"permuted": "decode.latent_permuted"}, + ), + _invoke( + "video_latent_unscale", + {"latent": "decode.latent_permuted"}, + {"unscaled": "decode.latent"}, + ), + _invoke( + "video_decode_chunks", + {"latent": "decode.latent"}, + {"count": "decode.chunks"}, + ), + _invoke( + "video_conv_cache_init", + {"latent": "decode.latent"}, + {port: f"{_cache_cell(port)}.initial" for port in cache_ports}, + ), + { + "kind": "loop", + "setup": { + "kind": "sequence", + "nodes": [ + _invoke( + "continue_predicate", + {"done": "package.false"}, + {"continue": "decode.setup.continue"}, + ) + ], + }, + "body": {"kind": "sequence", "nodes": decode_body}, + "condition": "decode.loop.continue", + "max_iterations": "decode.chunks", + "iteration": {"value": "decode.iteration", "contract": batch_int}, + "carried": [ + { + "cell": _cache_cell(port), + "current": f"{_cache_cell(port)}.initial", + "body_input": f"state.{_cache_cell(port)}.body", + "body_output": f"{_cache_cell(port)}.body", + "next": f"{_cache_cell(port)}.final", + "read_effect": _effect( + f"state:{_cache_cell(port)}.0", + f"state:{_cache_cell(port)}.read", + ), + "write_effect": _effect( + f"state:{_cache_cell(port)}.read", + f"state:{_cache_cell(port)}.1", + ), + } + for port in cache_ports + ], + }, + ], + }, + } + metadata = { + "schema_version": "v1", + "pipeline": {"workflow": _publish_workflow_v1(workflow)}, + } + add_policy_components_to_workflow(metadata, pkg) + return metadata + + +def write_video_diffusion_workflow_metadata( + pkg: Any, + output_dir: str, + **kwargs: Any, +) -> str: + os.makedirs(output_dir, exist_ok=True) + metadata = build_video_diffusion_workflow_metadata(pkg, **kwargs) + pkg.save_policy_components(output_dir) + add_adapter_service_to_metadata(metadata, pkg, output_dir) + path = os.path.join(output_dir, "inference_metadata.yaml") + with open(path, "w", encoding="utf-8") as handle: + _dump_yaml(metadata, handle) + return path + + +def build_vlm_workflow_metadata( + pkg: Any, + config: Any, + *, + source: str | None = None, +) -> dict[str, Any]: + """Build a vision/audio-encoder to embedding to decoder SSA workflow.""" + required = {"vision_encoder", "embedding", "decoder"} + missing = sorted(required.difference(pkg.keys())) + if missing: + raise ValueError(f"VLM workflow is missing required components: {missing}") + vision = pkg["vision_encoder"] + embedding = pkg["embedding"] + decoder = pkg["decoder"] + token_input = next( + ( + value + for value in embedding.graph.inputs + if value.dtype in {ir.DataType.INT32, ir.DataType.INT64} + and value.shape is not None + and len(value.shape) == 2 + ), + None, + ) + embedding_output = next( + ( + value + for value in embedding.graph.outputs + if value.shape is not None and len(value.shape) == 3 + ), + None, + ) + decoder_embed_input = ( + next( + ( + value + for value in decoder.graph.inputs + if embedding_output is not None + and _contracts_compatible(value, embedding_output) + ), + None, + ) + if embedding_output is not None + else None + ) + logits_output = _find_port(decoder.graph.outputs, "logits") + if None in (token_input, embedding_output, decoder_embed_input, logits_output): + raise ValueError("VLM workflow requires token->embedding->decoder logits ports") + assert token_input is not None + assert embedding_output is not None + assert decoder_embed_input is not None + assert logits_output is not None + embedding_output.shape = decoder_embed_input.shape + embedding_inputs_by_name = {value.name: value for value in embedding.graph.inputs} + for value in vision.graph.outputs: + target = embedding_inputs_by_name.get(value.name) + if target is not None and _contracts_compatible(value, target): + value.shape = target.shape + + decoder_outputs = {value.name: value for value in decoder.graph.outputs} + cache_pairs: list[tuple[ir.Value, ir.Value]] = [] + for value in decoder.graph.inputs: + present = next( + ( + decoder_outputs.get(name) + for name in _cache_output_candidates(value.name or "") + if name in decoder_outputs + ), + None, + ) + if present is not None: + if present.shape is None: + present.shape = value.shape + cache_pairs.append((value, present)) + cache_names = {value.name for value, _ in cache_pairs} + # A multimodal decoder can be hybrid: sliding layers keep a growing cache + # while full-attention layers scatter into fixed buffers. Both disciplines + # are described side by side rather than one being folded into the other. + static_cache = _static_cache_ports(decoder) + decoder_kv = _kv_storage_contract(decoder) + rank2_integer = [ + value + for value in decoder.graph.inputs + if value.name not in cache_names + and value.dtype == ir.DataType.INT64 + and value.shape is not None + and len(value.shape) == 2 + ] + attention_input = _find_port(rank2_integer, "mask") + position_input = _find_port(rank2_integer, "position") + if attention_input is None: + raise ValueError("VLM decoder requires an attention-mask input") + fixed_capacity = bool(cache_pairs) and decoder_kv["storage"] == "shared_buffer" + # A static cache carries its own per-row length on a graph port, so the + # prompt length must be materialized for it whether or not the shared-buffer + # attention-mask discipline also applies. + tracks_cache_lengths = fixed_capacity or static_cache is not None + + legacy = build_native_vlm_package_metadata(pkg, config=config, source=source) + preprocessing = legacy.get("preprocessing") + if not preprocessing or "image" not in preprocessing: + raise ValueError("VLM workflow requires declared image preprocessing") + _name_image_preprocessing_program(preprocessing["image"]) + image_outputs = preprocessing["image"]["outputs"] + vision_inputs = {value.name: value for value in vision.graph.inputs} + adapter_outputs: dict[str, Any] = {} + preprocessing_values: dict[str, str] = {} + for output in image_outputs: + endpoint = output["name"] + port_name = endpoint.split(".", 1)[-1] + if port_name not in vision_inputs: + raise ValueError(f"preprocessing output {endpoint!r} has no vision input") + output["contract"] = _contract(vision_inputs[port_name]) + output["dtype"] = output["contract"]["dtype"] + output["name"] = f"image.{port_name}" + adapter_outputs[port_name] = output["contract"] + preprocessing_values[port_name] = output["name"] + vision_feature_outputs = [ + value + for value in vision.graph.outputs + if value.name in embedding_inputs_by_name + and value.shape is not None + and len(value.shape) == 2 + and isinstance(list(value.shape)[-1], int) + ] + text_only_vision = vision_feature_outputs[0] if len(vision_feature_outputs) == 1 else None + + attach_policy_components( + pkg, + PolicyCapabilities( + sampler="greedy", + eos_termination=True, + token_state_update=True, + ), + ) + pkg.add_policy_component("last_token_logits", build_last_token_logits(logits_output.dtype)) + if text_only_vision is not None: + pkg.add_policy_component( + "empty_image_features", + build_empty_features( + text_only_vision.dtype, + int(list(text_only_vision.shape)[-1]), + ), + ) + pkg.add_policy_component( + "decoder_state_initializer", + build_decoder_state_initializer( + decoder, + token_input=None, + prompt_dtype=token_input.dtype, + attention_mask_input=attention_input.name, + position_ids_input=position_input.name if position_input is not None else None, + cache_inputs=sorted(cache_names), + fixed_capacity=fixed_capacity, + ragged=True, + write_indices_output=( + static_cache["write_indices"] if static_cache is not None else None + ), + ), + ) + pkg.add_policy_component( + "decoder_step_update", + build_decoder_step_update( + attention_dtype=attention_input.dtype, + position_dtype=position_input.dtype if position_input is not None else None, + fixed_capacity=fixed_capacity, + position_sections=( + rotary_axis_count(position_input) if position_input is not None else None + ), + ), + ) + pkg.add_policy_component("token_sampler", build_seeded_categorical_sampler()) + pkg.add_policy_component("termination", build_eos_termination(row_selective=True)) + pkg.add_policy_component( + "termination_batch_initializer", + build_termination_batch_initializer(), + ) + pkg.add_policy_component( + "token_state_update", + build_token_state_update(row_selective=True), + ) + pkg.add_policy_component("token_to_slot", build_token_to_slot()) + pkg.add_policy_component("generated_length_update", build_selective_integer_add()) + if cache_pairs: + pkg.add_policy_component("cache_length_update", build_selective_integer_add()) + + batch = _contract(token_input)["shape"][0] + batch_int = _request_aligned({"dtype": "int64", "rank": 1, "shape": [batch]}) + batch_bool = _request_aligned({"dtype": "bool", "rank": 1, "shape": [batch]}) + control_int = {"dtype": "int64", "rank": 1, "shape": [1]} + eos = _source_token_id(source, "eos_token_id", getattr(config, "eos_token_id", 0)) + inputs: dict[str, Any] = { + "request.prompt_tokens": { + "contract": _contract(token_input), + "role": {"kind": "runtime", "version": "1.0", "role": "prompt_tokens"}, + "source": {"kind": "request", "field": "prompt_tokens"}, + "required": True, + }, + "request.image": { + "contract": {"dtype": "uint8", "rank": 1, "shape": ["encoded_bytes"]}, + "role": {"kind": "runtime", "version": "1.0", "role": "media"}, + "source": {"kind": "request", "field": "media"}, + "required": text_only_vision is None, + **( + {"present_as": "request.image_present"} if text_only_vision is not None else {} + ), + }, + "request.max_iterations": { + "contract": control_int, + "role": { + "kind": "runtime", + "version": "1.0", + "role": "max_output_tokens", + }, + "source": {"kind": "request", "field": "max_output_tokens"}, + "required": True, + }, + "request.prompt_lengths": { + "contract": batch_int, + "role": {"kind": "opaque"}, + "source": {"kind": "application", "name": "prompt_lengths"}, + "required": False, + "default": -1, + }, + "request.eos_ids": { + "contract": {"dtype": "int64", "rank": 2, "shape": [batch, "num_eos"]}, + "role": {"kind": "opaque"}, + "source": {"kind": "application", "name": "eos_ids"}, + "required": False, + "default": eos, + }, + "request.eos_lengths": { + "contract": batch_int, + "role": {"kind": "opaque"}, + "source": {"kind": "application", "name": "eos_lengths"}, + "required": False, + "default": 1, + }, + "request.row_max_iterations": { + "contract": batch_int, + "role": {"kind": "opaque"}, + "source": {"kind": "application", "name": "row_max_iterations"}, + "required": False, + "default": -1, + }, + "package.eos_ids": { + "contract": {"dtype": "int64", "rank": 1, "shape": [1]}, + "role": {"kind": "opaque"}, + "source": {"kind": "literal"}, + "required": False, + "default": eos, + }, + "package.max_context": { + "contract": control_int, + "role": {"kind": "opaque"}, + "source": {"kind": "literal"}, + "required": False, + "default": int( + _source_model_value( + source, + "context_length", + getattr(config, "max_position_embeddings", 4096), + ) + ), + }, + **( + { + # The capacity a static graph was built against is a graph fact, + # not a deployment budget: it bounds legal write destinations. + "package.cache_capacity": { + "contract": control_int, + "role": {"kind": "opaque"}, + "source": {"kind": "literal"}, + "required": False, + "default": int(static_cache["capacity"]), + } + } + if static_cache is not None + else {} + ), + "package.one": { + "contract": batch_int, + "role": {"kind": "opaque"}, + "source": {"kind": "literal"}, + "required": False, + "default": 1, + }, + "package.one_step": { + # A growing-recurrence increment advances the whole invocation's + # state axis by one, so it is an invocation-scoped control scalar, + # not a per-row value like ``package.one``/``package.one_token``. + "contract": control_int, + "role": {"kind": "opaque"}, + "source": {"kind": "literal"}, + "required": False, + "default": 1, + }, + "package.active": { + "contract": batch_bool, + "role": {"kind": "opaque"}, + "source": {"kind": "literal"}, + "required": False, + "default": True, + }, + "package.false": { + "contract": batch_bool, + "role": {"kind": "opaque"}, + "source": {"kind": "literal"}, + "required": False, + "default": False, + }, + "package.zero_batch": { + "contract": batch_int, + "role": {"kind": "opaque"}, + "source": {"kind": "literal"}, + "required": False, + "default": 0, + }, + } + inputs.update( + { + "request.temperature": { + "contract": {"dtype": "float32", "rank": 1, "shape": [batch]}, + "role": { + "kind": "runtime", + "version": "1.0", + "role": "sampling_temperature", + }, + "source": {"kind": "request", "field": "sampling_temperature"}, + "required": False, + "default": 1.0, + }, + "request.top_k": { + "contract": batch_int, + "role": { + "kind": "runtime", + "version": "1.0", + "role": "sampling_top_k", + }, + "source": {"kind": "request", "field": "sampling_top_k"}, + "required": False, + "default": 1, + }, + "request.top_p": { + "contract": {"dtype": "float32", "rank": 1, "shape": [batch]}, + "role": { + "kind": "runtime", + "version": "1.0", + "role": "sampling_top_p", + }, + "source": {"kind": "request", "field": "sampling_top_p"}, + "required": False, + "default": 1.0, + }, + "request.min_p": { + "contract": {"dtype": "float32", "rank": 1, "shape": [batch]}, + "role": { + "kind": "runtime", + "version": "1.0", + "role": "sampling_min_p", + }, + "source": {"kind": "request", "field": "sampling_min_p"}, + "required": False, + "default": 0.0, + }, + "request.seed": { + "contract": batch_int, + "role": {"kind": "runtime", "version": "1.0", "role": "seed"}, + "source": {"kind": "request", "field": "seed"}, + "required": False, + "default": 0, + }, + "request.rng_counter": { + "contract": batch_int, + "role": {"kind": "opaque"}, + "source": {"kind": "application", "name": "rng_counter"}, + "required": False, + "default": 0, + }, + } + ) + vision_invoke_inputs = { + name: preprocessing_values[name] + for name in vision_inputs + if name in preprocessing_values + } + vision_outputs = {value.name: f"vision.{value.name}" for value in vision.graph.outputs} + audio_setup_nodes: list[dict[str, Any]] = [] + audio_outputs: dict[str, str] = {} + if "audio_encoder" in pkg: + audio = pkg["audio_encoder"] + audio_inputs = {} + for value in audio.graph.inputs: + name = f"request.audio.{value.name}" + inputs[name] = { + "contract": _contract(value), + "role": {"kind": "opaque"}, + "source": {"kind": "application", "name": f"audio.{value.name}"}, + "required": True, + } + audio_inputs[value.name] = name + audio_outputs = {value.name: f"audio.{value.name}" for value in audio.graph.outputs} + audio_setup_nodes.append(_invoke("audio_encoder", audio_inputs, audio_outputs)) + embedding_setup_inputs: dict[str, str] = {token_input.name: "request.prompt_tokens"} + embedding_body_inputs: dict[str, str] = {token_input.name: "token.body"} + produced_features = {value.name: f"vision.{value.name}" for value in vision.graph.outputs} + produced_features.update(audio_outputs) + for value in embedding.graph.inputs: + if value is token_input: + continue + if value.name in produced_features: + embedding_setup_inputs[value.name] = produced_features[value.name] + embedding_body_inputs[value.name] = produced_features[value.name] + else: + name = f"request.{value.name}" + inputs[name] = { + "contract": _contract(value), + "role": {"kind": "opaque"}, + "source": {"kind": "application", "name": value.name}, + "required": True, + } + embedding_setup_inputs[value.name] = name + embedding_body_inputs[value.name] = name + + setup_decoder_inputs = { + decoder_embed_input.name: "embedding.setup.embeds", + attention_input.name: f"initializer.{attention_input.name}", + } + body_decoder_inputs = { + decoder_embed_input.name: "embedding.body.embeds", + attention_input.name: ( + "decoder_step.body_attention_mask" + if fixed_capacity + else "state.attention_mask.body" + ), + } + if position_input is not None: + setup_decoder_inputs[position_input.name] = f"initializer.{position_input.name}" + body_decoder_inputs[position_input.name] = "state.position_ids.body" + if static_cache is not None: + # Prefill scatters from slot zero; each decode step writes at the row's + # current logical length, which is the same cursor the group declares. + setup_decoder_inputs[static_cache["write_indices"]] = ( + f"initializer.{static_cache['write_indices']}" + ) + setup_decoder_inputs[static_cache["kv_sequence_length"]] = "initializer.cache_lengths" + body_decoder_inputs[static_cache["write_indices"]] = "state.cache_lengths.body" + body_decoder_inputs[static_cache["kv_sequence_length"]] = "cache_lengths.next" + for past, _ in cache_pairs: + setup_decoder_inputs[past.name] = f"initializer.{past.name}" + body_decoder_inputs[past.name] = f"state.{past.name}.body" + + setup_decoder_outputs = {logits_output.name: "decoder.setup.logits"} + body_decoder_outputs = {logits_output.name: "decoder.body.logits"} + logits_contract = _contract(logits_output) + last_logits_contract = { + "dtype": "float32", + "rank": 2, + "shape": [logits_contract["shape"][0], logits_contract["shape"][-1]], + } + state: dict[str, Any] = { + "token": { + "contract": {"dtype": "int64", "rank": 2, "shape": [batch, 1]}, + "scope": "invocation", + "initializer": "initializer.token_slot", + "recurrence": {"kind": "invariant"}, + }, + "logits": { + "contract": last_logits_contract, + "scope": "invocation", + "initializer": "decoder.setup.last_logits", + "recurrence": {"kind": "invariant"}, + }, + "generated_lengths": { + "contract": batch_int, + "class": "semantic", + "scope": "invocation", + "initializer": "initializer.generated_lengths", + "recurrence": {"kind": "invariant"}, + }, + "attention_mask": { + "contract": { + "dtype": _contract(attention_input)["dtype"], + "rank": 2, + "shape": [batch, "context"], + }, + "scope": "invocation", + "initializer": ( + f"initializer.{attention_input.name}" + if fixed_capacity + else "initializer.body_attention_mask" + ), + "recurrence": ( + {"kind": "invariant"} + if fixed_capacity + else { + "kind": "growing", + "axis": 1, + "increment": "package.one_step", + "max": "package.max_context", + } + ), + }, + "active": { + "contract": batch_bool, + "class": "semantic", + "scope": "invocation", + "initializer": "package.active", + "recurrence": {"kind": "invariant"}, + }, + "done": { + "contract": batch_bool, + "class": "semantic", + "scope": "invocation", + "initializer": "package.false", + "recurrence": {"kind": "invariant"}, + }, + "accepted_len": { + "contract": batch_int, + "class": "semantic", + "scope": "invocation", + "initializer": "package.zero_batch", + "recurrence": {"kind": "invariant"}, + }, + "cache_lengths": { + "contract": batch_int, + "class": "semantic", + "scope": "invocation", + "initializer": ( + "initializer.cache_lengths" if tracks_cache_lengths else "package.zero_batch" + ), + "recurrence": {"kind": "invariant"}, + }, + "rng_counter": { + "contract": batch_int, + "class": "semantic", + "scope": "invocation", + "initializer": "request.rng_counter", + "recurrence": {"kind": "invariant"}, + }, + } + state_specs = [ + ( + "token", + "initializer.token_slot", + "state.token.body", + "token.body", + "state.token.final", + ), + ( + "logits", + "decoder.setup.last_logits", + "state.logits.body", + "decoder.body.last_logits", + "state.logits.final", + ), + ( + "attention_mask", + ( + f"initializer.{attention_input.name}" + if fixed_capacity + else "initializer.body_attention_mask" + ), + "state.attention_mask.body", + "decoder_step.body_attention_mask", + "state.attention_mask.final", + ), + ( + "generated_lengths", + "initializer.generated_lengths", + "state.generated_lengths.body", + "token.next_lengths", + "state.generated_lengths.final", + ), + ( + "rng_counter", + "request.rng_counter", + "state.rng_counter.body", + "sample.next_counter", + "state.rng_counter.final", + ), + ( + "active", + "package.active", + "state.active.body", + "loop.next_active", + "state.active.final", + ), + ("done", "package.false", "state.done.body", "loop.done", "state.done.final"), + ( + "accepted_len", + "package.zero_batch", + "state.accepted_len.body", + "accepted_len.next", + "state.accepted_len.final", + ), + ( + "cache_lengths", + "initializer.cache_lengths" if tracks_cache_lengths else "package.zero_batch", + "state.cache_lengths.body", + "cache_lengths.next", + "state.cache_lengths.final", + ), + ] + if position_input is not None: + state["position_ids"] = { + "contract": { + "dtype": _contract(position_input)["dtype"], + "rank": 2, + "shape": [batch, 1], + }, + "scope": "invocation", + "initializer": "initializer.body_position_ids", + "recurrence": {"kind": "invariant"}, + } + state_specs.append( + ( + "position_ids", + "initializer.body_position_ids", + "state.position_ids.body", + "decoder_step.body_position_ids", + "state.position_ids.final", + ) + ) + for index, (past, present) in enumerate(cache_pairs): + cell = f"cache_{index}" + # A scattered buffer overwrites cells inside a capacity fixed at export, + # so its extent never changes; only an appending cache grows. + scattered = static_cache is not None and past.name in static_cache["buffers"] + state[cell] = { + "contract": _request_aligned(_contract(past)), + "scope": "invocation", + "initializer": f"decoder.setup.{present.name}", + "recurrence": ( + {"kind": "invariant"} + if scattered + else { + "kind": "bounded", + "axis": next( + ( + axis + for axis, dimension in enumerate(_contract(past)["shape"]) + if "sequence" in str(dimension) + ), + 2, + ), + "max": "package.max_context", + } + ), + "management": "runtime", + "release_boundary": "invocation", + } + setup_decoder_outputs[present.name] = f"decoder.setup.{present.name}" + body_decoder_outputs[present.name] = f"decoder.body.{present.name}" + state_specs.append( + ( + cell, + f"decoder.setup.{present.name}", + f"state.{past.name}.body", + f"decoder.body.{present.name}", + f"state.{past.name}.final", + ) + ) + # Hybrid VL decoders (Gemma 3/4) publish one group per attention kind so a + # global layer's prefix is never evicted with the sliding layers'. + vlm_state_groups, vlm_cell_groups = _state_service_groups( + config=getattr(config, "text", config), + cache_pairs=cache_pairs, + ports={ + "decoder": { + f"cache_{index}": {"input": past.name, "output": present.name} + for index, (past, present) in enumerate(cache_pairs) + } + }, + sequence_axis=next( + ( + axis + for axis, dimension in enumerate(_contract(cache_pairs[0][0])["shape"]) + if "sequence" in str(dimension) + ), + 2, + ) + if cache_pairs + else 2, + logical_lengths="cache_lengths", + aliasing=_state_aliasing(decoder_kv), + base_name="decoder_cache", + indexed_scatter=( + { + "buffers": static_cache["buffers"], + "capacity": "package.cache_capacity", + # The write cursor and the logical length are one quantity: a + # row's next write lands exactly where its valid prefix ends. + "write_indices": "cache_lengths", + "logical_lengths": "cache_lengths", + "port": static_cache["write_indices"], + "kv_length_port": static_cache["kv_sequence_length"], + } + if static_cache is not None + else None + ), + ) + for cell, group_name in vlm_cell_groups.items(): + state[cell]["service_group"] = group_name + carried = [] + initial_effects = { + "sample": "sample.0", + "termination": "termination.0", + "state": "state.0", + "emit": "emit.0", + } + for cell, current, body_input, body_output, final in state_specs: + effect = f"state:{cell}" + initial_effects[effect] = f"{effect}.0" + carried.append( + { + "cell": cell, + "current": current, + "body_input": body_input, + "body_output": body_output, + "next": final, + "read_effect": _effect(f"{effect}.0", f"{effect}.read"), + "write_effect": _effect(f"{effect}.read", f"{effect}.1"), + } + ) + + if text_only_vision is not None: + feature_name = text_only_vision.name + vision_setup_nodes: list[dict[str, Any]] = [ + { + "kind": "branch", + "predicate": "request.image_present", + "cases": { + "true": { + "kind": "sequence", + "nodes": [ + _invoke( + "image_preprocess", + {"encoded": "request.image"}, + dict(preprocessing_values), + ), + _invoke( + "vision_encoder", + vision_invoke_inputs, + { + **vision_outputs, + feature_name: f"vision.with_media.{feature_name}", + }, + ), + ], + }, + "false": _invoke( + "empty_image_features", + {}, + {"features": f"vision.empty.{feature_name}"}, + ), + }, + "outputs": { + f"vision.{feature_name}": { + "cases": { + "true": f"vision.with_media.{feature_name}", + "false": f"vision.empty.{feature_name}", + } + } + }, + } + ] + else: + vision_setup_nodes = [ + _invoke( + "image_preprocess", + {"encoded": "request.image"}, + dict(preprocessing_values), + ), + _invoke("vision_encoder", vision_invoke_inputs, vision_outputs), + ] + + setup = { + "kind": "sequence", + "nodes": [ + *vision_setup_nodes, + *audio_setup_nodes, + _invoke( + "decoder_state_initializer", + { + "prompt_tokens": "request.prompt_tokens", + "prompt_lengths": "request.prompt_lengths", + **({"max_iterations": "request.max_iterations"} if fixed_capacity else {}), + }, + { + attention_input.name: f"initializer.{attention_input.name}", + "body_attention_mask": "initializer.body_attention_mask", + "token_slot": "initializer.token_slot", + "generated_lengths": "initializer.generated_lengths", + **( + {"cache_lengths": "initializer.cache_lengths"} + if tracks_cache_lengths + else {} + ), + **( + { + static_cache["write_indices"]: ( + f"initializer.{static_cache['write_indices']}" + ) + } + if static_cache is not None + else {} + ), + **( + { + position_input.name: f"initializer.{position_input.name}", + "body_position_ids": "initializer.body_position_ids", + } + if position_input is not None + else {} + ), + **{name: f"initializer.{name}" for name in sorted(cache_names)}, + }, + ), + _invoke( + "termination_batch_initializer", + { + "input_eos_ids": "request.eos_ids", + "input_eos_lengths": "request.eos_lengths", + "input_max_iterations": "request.row_max_iterations", + "fallback_max_iterations": "request.max_iterations", + "active": "package.active", + }, + { + "row_eos_ids": "termination.eos_ids", + "eos_lengths": "termination.eos_lengths", + "max_iterations": "termination.max_iterations", + }, + ), + _invoke( + "embedding", + embedding_setup_inputs, + {embedding_output.name: "embedding.setup.embeds"}, + ), + _invoke("decoder", setup_decoder_inputs, setup_decoder_outputs), + _invoke( + "last_token_logits", + {"logits": "decoder.setup.logits"}, + {"last_logits": "decoder.setup.last_logits"}, + ), + ], + } + decoder_step_invoke = _invoke( + "decoder_step_update", + { + "attention_mask": "state.attention_mask.body", + **({"logical_length": "state.cache_lengths.body"} if fixed_capacity else {}), + **( + {"position_ids": "state.position_ids.body"} + if position_input is not None + else {} + ), + }, + { + "next_attention_mask": "decoder_step.body_attention_mask", + **( + {"next_position_ids": "decoder_step.body_position_ids"} + if position_input is not None + else {} + ), + }, + ) + body = { + "kind": "sequence", + "nodes": [ + _invoke( + "token_sampler", + { + "logits": "state.logits.body", + "temperature": "request.temperature", + "top_k": "request.top_k", + "top_p": "request.top_p", + "min_p": "request.min_p", + "seed": "request.seed", + "counter": "state.rng_counter.body", + "active": "state.active.body", + "done": "state.done.body", + }, + {"token": "sample.body", "next_counter": "sample.next_counter"}, + {"sample": _effect("sample.0", "sample.1")}, + ), + _invoke( + "termination", + { + "tokens": "sample.body", + "eos_ids": "termination.eos_ids", + "eos_lengths": "termination.eos_lengths", + "iteration": "loop.iteration", + "max_iterations": "termination.max_iterations", + "active": "state.active.body", + }, + { + "done": "loop.done", + "next_active": "loop.next_active", + "continue": "loop.continue", + }, + {"termination": _effect("termination.0", "termination.1")}, + ), + *( + [ + _invoke( + "cache_length_update", + { + "left": "state.cache_lengths.body", + "right": "package.one", + "active": "state.active.body", + "done": "state.done.body", + }, + {"total": "cache_lengths.next"}, + ), + _invoke( + "cache_length_update", + { + "left": "package.zero_batch", + "right": "package.one", + "active": "state.active.body", + "done": "state.done.body", + }, + {"total": "accepted_len.next"}, + ), + ] + if cache_pairs + else [] + ), + _invoke("token_to_slot", {"token": "sample.body"}, {"slot": "sample.slot"}), + _invoke( + "generated_length_update", + { + "left": "state.generated_lengths.body", + "right": "package.one", + "active": "state.active.body", + "done": "state.done.body", + }, + {"total": "token.next_lengths"}, + ), + _invoke( + "generated_length_update", + { + "left": "package.zero_batch", + "right": "package.one", + "active": "state.active.body", + "done": "state.done.body", + }, + {"total": "token.emitted_length"}, + ), + _invoke( + "token_state_update", + { + "current": "state.token.body", + "update": "sample.slot", + "active": "state.active.body", + "done": "state.done.body", + }, + {"next": "token.body"}, + {"state": _effect("state.0", "state.1")}, + ), + { + "kind": "emit", + "value": "token.body", + "output": "tokens", + "mode": "append", + "when": "state.active.body", + "valid_length": "token.emitted_length", + "effect_name": "emit", + "effect": _effect("emit.0", "emit.1"), + }, + *([decoder_step_invoke] if fixed_capacity else []), + _invoke( + "embedding", + embedding_body_inputs, + {embedding_output.name: "embedding.body.embeds"}, + ), + _invoke("decoder", body_decoder_inputs, body_decoder_outputs), + _invoke( + "last_token_logits", + {"logits": "decoder.body.logits"}, + {"last_logits": "decoder.body.last_logits"}, + ), + *([] if fixed_capacity else [decoder_step_invoke]), + ], + } + components = { + name: _component(model, _artifact(name, len(pkg))) for name, model in pkg.items() + } + components["image_preprocess"] = { + "implementation": { + "kind": "adapter", + "abi": "onnx-genai.image-preprocess", + "version": "1", + }, + "ports": { + "inputs": { + "encoded": { + "dtype": "uint8", + "rank": 1, + "shape": ["encoded_bytes"], + } + }, + "outputs": adapter_outputs, + }, + } + workflow = { + "manifest": { + "adapter_abis": {"onnx-genai.image-preprocess": "1"}, + "capabilities": [ + "workflow_ssa", + "linear_effects", + "nested_control_flow", + "loop_induction_values", + "typed_emit", + "emit_valid_length", + *(["input_presence"] if text_only_vision is not None else []), + *( + ["serving_service_contract", "bounded_state_recurrence"] + if cache_pairs + else [] + ), + ], + }, + "inputs": inputs, + "outputs": { + "tokens": { + "contract": _request_aligned( + { + "dtype": "int64", + "rank": 2, + "shape": [batch, "generated_sequence"], + } + ), + "role": "tokens", + "stage": "pre_adapter", + } + }, + "components": components, + "state": state, + **( + { + "serving": { + "active": "active", + "done": "done", + "accepted_len": "accepted_len", + "state_service": {"groups": vlm_state_groups}, + } + } + if cache_pairs + else {} + ), + "initial_effects": initial_effects, + "graph": { + "kind": "loop", + "setup": setup, + "body": body, + "condition": "loop.continue", + "termination": "generation_eos", + "max_iterations": "request.max_iterations", + "iteration": { + "value": "loop.iteration", + "contract": {"dtype": "int64", "rank": 1, "shape": [1]}, + }, + "carried": carried, + }, + } + metadata = { + "schema_version": "v1", + "preprocessing": preprocessing, + "pipeline": {"workflow": _publish_workflow_v1(workflow)}, + } + add_policy_components_to_workflow(metadata, pkg) + return metadata + + +def write_vlm_workflow_metadata( + pkg: Any, + output_dir: str, + config: Any, + *, + source: str | None = None, +) -> str: + os.makedirs(output_dir, exist_ok=True) + metadata = build_vlm_workflow_metadata(pkg, config, source=source) + pkg.save_policy_components(output_dir) + add_adapter_service_to_metadata(metadata, pkg, output_dir) + path = os.path.join(output_dir, "inference_metadata.yaml") + with open(path, "w", encoding="utf-8") as handle: + _dump_yaml(metadata, handle) + return path + + +def build_speculative_workflow_metadata( + pkg: Any, + config: Any | None = None, + *, + grammar_guidance: bool = False, + adaptive_k_max: int | None = None, +) -> dict[str, Any]: + """Build proposer/verifier workflow with branch phi and effect joins.""" + config = config or getattr(pkg, "config", None) + if not {"proposer", "verifier"} <= set(pkg.keys()): + raise ValueError("speculative workflow requires proposer and verifier") + proposer = pkg["proposer"] + verifier = pkg["verifier"] + verifier_kv = _kv_storage_contract(verifier) + proposer_input = next( + ( + value + for value in proposer.graph.inputs + if value.dtype in {ir.DataType.INT32, ir.DataType.INT64} + and value.shape is not None + and len(value.shape) == 2 + ), + None, + ) + proposed_tokens = _find_port(proposer.graph.outputs, "proposed", "tokens") + proposal_scores = _find_port(proposer.graph.outputs, "scores", "logits") + verifier_token_input = next( + ( + value + for value in verifier.graph.inputs + if proposed_tokens is not None and _contract(value) == _contract(proposed_tokens) + ), + None, + ) + target_scores = _find_port(verifier.graph.outputs, "scores", "logits") + if None in (proposer_input, proposed_tokens, verifier_token_input, target_scores): + raise ValueError("speculative components do not expose compatible token/score ports") + assert proposer_input is not None + assert proposed_tokens is not None + assert verifier_token_input is not None + assert target_scores is not None + proposal_budget_input = next( + ( + value + for value in proposer.graph.inputs + if value is not proposer_input + and value.dtype == ir.DataType.INT64 + and value.shape is not None + and len(value.shape) == 1 + and any(term in value.name.lower() for term in ("budget", "proposal_k", "draft_k")) + ), + None, + ) + if adaptive_k_max is not None and proposal_budget_input is None: + raise ValueError( + "adaptive speculative workflow requires a rank-1 proposer budget input" + ) + if _contract(proposer_input) != _contract(proposed_tokens): + raise ValueError( + "representative speculative workflow requires fixed token-block shape" + ) + + attach_policy_components( + pkg, + PolicyCapabilities( + speculative_acceptance=True, + grammar_guidance=grammar_guidance, + adaptive_k_max=adaptive_k_max, + ), + ) + if grammar_guidance: + pkg.add_policy_component("grammar_length", build_integer_minimum()) + pkg.add_policy_component("grammar_sampler_logits", build_last_token_logits()) + if adaptive_k_max is None: + pkg.add_policy_component("proposal_length", build_sequence_length()) + if adaptive_k_max is not None: + pkg.add_policy_component("proposal_metrics", build_proposal_metrics()) + pkg.add_policy_component("cache_length_update", build_integer_add()) + batch = _contract(proposer_input)["shape"][0] + batch_int = _request_aligned({"dtype": "int64", "rank": 1, "shape": [batch]}) + batch_bool = _request_aligned({"dtype": "bool", "rank": 1, "shape": [batch]}) + control_int = {"dtype": "int64", "rank": 1, "shape": [1]} + inputs: dict[str, Any] = { + "request.tokens": { + "contract": _contract(proposer_input), + "role": {"kind": "runtime", "version": "1.0", "role": "prompt_tokens"}, + "source": {"kind": "request", "field": "prompt_tokens"}, + "required": True, + }, + "request.seed": { + "contract": batch_int, + "role": {"kind": "runtime", "version": "1.0", "role": "seed"}, + "source": {"kind": "request", "field": "seed"}, + "required": False, + "default": 0, + }, + "request.max_iterations": { + "contract": control_int, + "role": { + "kind": "runtime", + "version": "1.0", + "role": "max_output_tokens", + }, + "source": {"kind": "request", "field": "max_output_tokens"}, + "required": True, + }, + "package.zero": { + "contract": batch_int, + "role": {"kind": "opaque"}, + "source": {"kind": "literal"}, + "required": False, + "default": 0, + }, + "package.one": { + "contract": batch_int, + "role": {"kind": "opaque"}, + "source": {"kind": "literal"}, + "required": False, + "default": 1, + }, + "package.max_context": { + "contract": control_int, + "role": {"kind": "opaque"}, + "source": {"kind": "literal"}, + "required": False, + "default": int(getattr(config, "max_position_embeddings", 4096)), + }, + "package.false": { + "contract": batch_bool, + "role": {"kind": "opaque"}, + "source": {"kind": "literal"}, + "required": False, + "default": False, + }, + "package.active": { + "contract": batch_bool, + "role": {"kind": "opaque"}, + "source": {"kind": "literal"}, + "required": False, + "default": True, + }, + "request.cache_lengths": { + "contract": batch_int, + "role": {"kind": "opaque"}, + "source": {"kind": "application", "name": "serving.cache_lengths"}, + "required": False, + "default": 0, + }, + } + if grammar_guidance: + inputs.update( + { + "request.grammar_state": { + "contract": batch_int, + "role": {"kind": "opaque"}, + "source": { + "kind": "application", + "name": "grammar.initial_state", + }, + "required": True, + }, + "request.grammar_transition_table": { + "contract": { + "dtype": "int64", + "rank": 2, + "shape": ["grammar_states", "vocabulary"], + }, + "role": {"kind": "opaque"}, + "source": { + "kind": "application", + "name": "grammar.transition_table", + }, + "required": True, + }, + } + ) + if adaptive_k_max is not None: + estimate_slots = 4 * (adaptive_k_max + 1) + 4 + inputs.update( + { + "request.adaptive_k": { + "contract": batch_int, + "role": {"kind": "opaque"}, + "source": {"kind": "application", "name": "adaptive.current_k"}, + "required": False, + "default": 1, + }, + "request.adaptive_estimates": { + "contract": { + "dtype": "float32", + "rank": 2, + "shape": [batch, estimate_slots], + }, + "role": {"kind": "opaque"}, + "source": { + "kind": "application", + "name": "adaptive.estimates", + }, + "required": False, + "default": 0.0, + }, + "request.draft_ms": { + "contract": { + "dtype": "float32", + "rank": 1, + "shape": [batch], + }, + "role": {"kind": "opaque"}, + "source": {"kind": "application", "name": "telemetry.draft_ms"}, + "required": True, + }, + "request.target_ms": { + "contract": { + "dtype": "float32", + "rank": 1, + "shape": [batch], + }, + "role": {"kind": "opaque"}, + "source": {"kind": "application", "name": "telemetry.target_ms"}, + "required": True, + }, + } + ) + + proposer_inputs = {proposer_input.name: "state.tokens.body"} + for value in proposer.graph.inputs: + if value is proposer_input: + continue + if value is proposal_budget_input: + proposer_inputs[value.name] = "state.proposal_k.body" + continue + name = f"request.proposer.{value.name}" + inputs[name] = { + "contract": _contract(value), + "role": {"kind": "opaque"}, + "source": {"kind": "application", "name": f"proposer.{value.name}"}, + "required": True, + } + proposer_inputs[value.name] = name + verifier_inputs = {verifier_token_input.name: "proposal.tokens"} + verifier_outputs = {target_scores.name: "target.scores"} + cache_pairs: list[tuple[ir.Value, ir.Value]] = [] + verifier_output_map = {value.name: value for value in verifier.graph.outputs} + for value in verifier.graph.inputs: + if value is verifier_token_input: + continue + present = next( + ( + verifier_output_map.get(name) + for name in ( + value.name.replace("past_key_values", "present"), + value.name.replace("past.", "present."), + ) + if name in verifier_output_map + ), + None, + ) + if present is not None: + cell = f"cache_{len(cache_pairs)}" + cache_pairs.append((value, present)) + name = f"request.verifier.{value.name}" + inputs[name] = { + "contract": _contract(value), + "role": {"kind": "opaque"}, + "source": {"kind": "application", "name": f"verifier.{value.name}"}, + "required": True, + } + verifier_inputs[value.name] = f"state.{cell}.body" + verifier_outputs[present.name] = f"verifier.{present.name}" + else: + name = f"request.verifier.{value.name}" + inputs[name] = { + "contract": _contract(value), + "role": {"kind": "opaque"}, + "source": {"kind": "application", "name": f"verifier.{value.name}"}, + "required": True, + } + verifier_inputs[value.name] = name + + proposal_outputs = {proposed_tokens.name: "proposal.tokens"} + acceptance_inputs = { + "target_scores": "target.scores", + "proposed_tokens": "proposal.tokens", + "seed": "request.seed", + "offset": "state.rng_offset.body", + } + if proposal_scores is not None: + proposal_outputs[proposal_scores.name] = "proposal.scores" + proposal_measure_nodes: list[dict[str, Any]] = [] + if adaptive_k_max is not None: + proposal_measure_nodes.append( + _invoke( + "proposal_metrics", + { + "proposed_tokens": "proposal.tokens", + "requested_k": "state.proposal_k.body", + }, + { + "evaluated": "proposal.evaluated", + "filled_proposal_budget": "proposal.filled_budget", + }, + ) + ) + elif grammar_guidance: + proposal_measure_nodes.append( + _invoke( + "proposal_length", + {"tokens": "proposal.tokens"}, + {"length": "proposal.evaluated"}, + ) + ) + grammar_pre_nodes: list[dict[str, Any]] = [] + if grammar_guidance: + grammar_pre_nodes.extend( + [ + _invoke( + "grammar_clone", + { + "state": "state.grammar.body", + "tokens": "proposal.tokens", + "valid_length": "package.zero", + "transition_table": "request.grammar_transition_table", + }, + { + "next_state": "grammar.clone.state", + "consumed_length": "grammar.clone.consumed", + "logits_mask": "grammar.clone.mask", + "forced_tokens": "grammar.clone.forced", + "forced_length": "grammar.clone.forced_length", + }, + {"grammar": _effect("grammar.0", "grammar.clone")}, + ), + _invoke( + "grammar_lookahead", + { + "state": "grammar.clone.state", + "tokens": "proposal.tokens", + "valid_length": "proposal.evaluated", + "transition_table": "request.grammar_transition_table", + }, + { + "next_state": "grammar.lookahead.state", + "consumed_length": "grammar.valid_length", + "logits_mask": "grammar.lookahead.mask", + "forced_tokens": "grammar.lookahead.forced", + "forced_length": "grammar.lookahead.forced_length", + }, + {"grammar": _effect("grammar.clone", "grammar.lookahead")}, + ), + ] + ) + emit_length = "acceptance.length" + grammar_post_nodes: list[dict[str, Any]] = [] + if grammar_guidance: + emit_length = "grammar.committed_length" + grammar_post_nodes.extend( + [ + _invoke( + "grammar_length", + { + "left": "acceptance.length", + "right": "grammar.valid_length", + }, + {"minimum": "grammar.committed_length"}, + ), + _invoke( + "grammar_commit", + { + "state": "state.grammar.body", + "tokens": "acceptance.tokens", + "valid_length": "grammar.committed_length", + "transition_table": "request.grammar_transition_table", + }, + { + "next_state": "grammar.next", + "consumed_length": "grammar.committed", + "logits_mask": "grammar.mask", + "forced_tokens": "grammar.forced", + "forced_length": "grammar.forced_length", + }, + {"grammar": _effect("grammar.lookahead", "grammar.commit")}, + ), + _invoke( + "grammar_sampler_logits", + {"logits": "target.scores"}, + {"last_logits": "grammar.sampler_logits"}, + ), + _invoke( + "grammar_guidance", + { + "logits": "grammar.sampler_logits", + "logits_mask": "grammar.mask", + "forced_tokens": "grammar.forced", + "forced_length": "grammar.forced_length", + }, + {"token": "grammar.token"}, + ), + ] + ) + adaptive_nodes: list[dict[str, Any]] = [] + if adaptive_k_max is not None: + committed_metric = ( + "grammar.committed_length" if grammar_guidance else "acceptance.length" + ) + adaptive_nodes.append( + _invoke( + "adaptive_k", + { + "current_k": "state.proposal_k.body", + "accepted": committed_metric, + "evaluated": "proposal.evaluated", + "committed_tokens": committed_metric, + "filled_proposal_budget": "proposal.filled_budget", + "draft_ms": "request.draft_ms", + "target_ms": "request.target_ms", + "estimates": "state.adaptive_estimates.body", + }, + { + "next_k": "adaptive.next_k", + "next_estimates": "adaptive.next_estimates", + }, + {"adaptive": _effect("adaptive.0", "adaptive.1")}, + ) + ) + cache_next_outputs = { + f"cache_{index}.next": f"verifier.{present.name}" + for index, (_past, present) in enumerate(cache_pairs) + } + body_nodes = [ + _invoke("proposer", proposer_inputs, proposal_outputs), + *proposal_measure_nodes, + *grammar_pre_nodes, + _invoke("verifier", verifier_inputs, verifier_outputs), + _invoke( + "speculative_acceptance", + acceptance_inputs, + { + "accepted_tokens": "acceptance.tokens", + "accepted_len": "acceptance.length", + "done": "acceptance.done", + "continue": "acceptance.continue", + "next_offset": "rng_offset.body", + "rollback_len": "acceptance.rollback_length", + }, + {"verify": _effect("verify.0", "verify.1")}, + ), + *grammar_post_nodes, + _invoke( + "cache_length_update", + { + "left": "state.cache_lengths.body", + "right": emit_length, + }, + {"total": "cache_lengths.next"}, + ), + *adaptive_nodes, + { + "kind": "emit", + "value": "acceptance.tokens", + "valid_length": emit_length, + "output": "tokens", + "mode": "append", + "effect_name": "emit", + "effect": _effect("emit.0", "emit.1"), + }, + ] + if grammar_guidance: + body_nodes.append( + { + "kind": "emit", + "value": "grammar.token", + "valid_length": "grammar.forced_length", + "output": "tokens", + "mode": "append", + "effect_name": "emit", + "effect": _effect("emit.1", "emit.2"), + } + ) + state = { + "tokens": { + "contract": _contract(proposer_input), + "class": "semantic", + "scope": "invocation", + "initializer": "request.tokens", + "recurrence": {"kind": "invariant"}, + }, + "rng_offset": { + "contract": batch_int, + "class": "semantic", + "scope": "invocation", + "initializer": "package.zero", + "recurrence": {"kind": "invariant"}, + }, + "active": { + "contract": batch_bool, + "class": "semantic", + "scope": "invocation", + "initializer": "package.active", + "recurrence": {"kind": "invariant"}, + }, + "done": { + "contract": batch_bool, + "class": "semantic", + "scope": "invocation", + "initializer": "package.false", + "recurrence": {"kind": "invariant"}, + }, + "accepted_len": { + "contract": batch_int, + "class": "semantic", + "scope": "invocation", + "initializer": "package.zero", + "recurrence": {"kind": "invariant"}, + }, + "cache_lengths": { + "contract": batch_int, + "class": "semantic", + "scope": "invocation", + "initializer": "request.cache_lengths", + "recurrence": {"kind": "invariant"}, + }, + } + state_specs = [ + ( + "tokens", + "request.tokens", + "state.tokens.body", + "acceptance.tokens", + "state.tokens.final", + ), + ( + "rng_offset", + "package.zero", + "state.rng_offset.body", + "rng_offset.body", + "state.rng_offset.final", + ), + ( + "active", + "package.active", + "state.active.body", + "acceptance.continue", + "state.active.final", + ), + ( + "done", + "package.false", + "state.done.body", + "acceptance.done", + "state.done.final", + ), + ( + "accepted_len", + "package.zero", + "state.accepted_len.body", + emit_length, + "state.accepted_len.final", + ), + ( + "cache_lengths", + "request.cache_lengths", + "state.cache_lengths.body", + "cache_lengths.next", + "state.cache_lengths.final", + ), + ] + if grammar_guidance: + state["grammar"] = { + "contract": batch_int, + "class": "semantic", + "scope": "invocation", + "initializer": "request.grammar_state", + "recurrence": {"kind": "invariant"}, + } + state_specs.append( + ( + "grammar", + "request.grammar_state", + "state.grammar.body", + "grammar.next", + "state.grammar.final", + ) + ) + if adaptive_k_max is not None: + state["proposal_k"] = { + "contract": batch_int, + "class": "advisory", + "scope": "invocation", + "initializer": "request.adaptive_k", + "recurrence": {"kind": "invariant"}, + } + state["adaptive_estimates"] = { + "contract": inputs["request.adaptive_estimates"]["contract"], + "class": "advisory", + "scope": "invocation", + "initializer": "request.adaptive_estimates", + "recurrence": {"kind": "invariant"}, + } + state_specs.extend( + [ + ( + "proposal_k", + "request.adaptive_k", + "state.proposal_k.body", + "adaptive.next_k", + "state.proposal_k.final", + ), + ( + "adaptive_estimates", + "request.adaptive_estimates", + "state.adaptive_estimates.body", + "adaptive.next_estimates", + "state.adaptive_estimates.final", + ), + ] + ) + kv_ports: dict[str, Any] = {} + kv_sequence_axis = 2 + for index, (past, present) in enumerate(cache_pairs): + cell = f"cache_{index}" + initializer = f"request.verifier.{past.name}" + kv_sequence_axis = next( + ( + axis + for axis, dimension in enumerate(_contract(past)["shape"]) + if "sequence" in str(dimension) + ), + 2, + ) + state[cell] = { + "contract": _request_aligned(_contract(past)), + "class": "semantic", + "scope": "invocation", + "initializer": initializer, + "recurrence": { + "kind": "bounded", + "axis": kv_sequence_axis, + "max": "package.max_context", + }, + "service_group": "verifier_cache", + "management": "runtime", + "release_boundary": "invocation", + } + kv_ports[cell] = {"input": past.name, "output": present.name} + state_specs.append( + ( + cell, + initializer, + f"state.{cell}.body", + cache_next_outputs[f"{cell}.next"], + f"state.{cell}.final", + ) + ) + initial_effects = { + "verify": "verify.0", + "emit": "emit.0", + "state": "branch.state.in", + } + if grammar_guidance: + initial_effects["grammar"] = "grammar.0" + if adaptive_k_max is not None: + initial_effects["adaptive"] = "adaptive.0" + for index in range(len(cache_pairs)): + initial_effects[f"rollback_cache_{index}"] = f"rollback.cache_{index}.0" + initial_effects[f"branch:cache_{index}"] = f"branch.cache_{index}.in" + carried = [] + for cell, current, body_input, body_output, final in state_specs: + effect = f"state:{cell}" + initial_effects[effect] = f"{effect}.0" + carried.append( + { + "cell": cell, + "current": current, + "body_input": body_input, + "body_output": body_output, + "next": final, + "read_effect": _effect(f"{effect}.0", f"{effect}.read"), + "write_effect": _effect(f"{effect}.read", f"{effect}.1"), + } + ) + workflow = { + "manifest": { + "capabilities": [ + "workflow_ssa", + "linear_effects", + "nested_control_flow", + "loop_induction_values", + "typed_emit", + "emit_valid_length", + "bounded_state_recurrence", + "serving_service_contract", + ], + }, + "inputs": inputs, + "outputs": { + "tokens": { + "contract": _request_aligned( + { + **_contract(proposed_tokens), + "shape": [ + *_contract(proposed_tokens)["shape"][:-1], + "accepted_sequence", + ], + } + ), + "role": "tokens", + "stage": "pre_adapter", + }, + }, + "components": { + name: _component(model, _artifact(name, len(pkg))) for name, model in pkg.items() + }, + "state": state, + "serving": { + "active": "active", + "done": "done", + "accepted_len": "accepted_len", + "state_service": { + "groups": { + "verifier_cache": _state_group( + sequence_axis=kv_sequence_axis, + logical_lengths="cache_lengths", + storage=verifier_kv["storage"], + ports={"verifier": kv_ports}, + ) + }, + }, + }, + "initial_effects": initial_effects, + "graph": { + "kind": "loop", + "setup": {"kind": "sequence", "nodes": []}, + "body": {"kind": "sequence", "nodes": body_nodes}, + "condition": "acceptance.continue", + "active_cell": "active", + "max_iterations": "request.max_iterations", + "iteration": {"value": "speculative.iteration", "contract": batch_int}, + "carried": carried, + }, + } + if grammar_guidance: + workflow["manifest"]["adapter_abis"] = {"onnx-genai.grammar-guidance": "1"} + workflow["manifest"]["capabilities"].append("grammar_guidance_adapter") + workflow["components"].update( + { + "grammar_clone": _grammar_adapter_component("clone"), + "grammar_lookahead": _grammar_adapter_component("lookahead"), + "grammar_commit": _grammar_adapter_component("commit"), + } + ) + # The grammar adapter's three actions are a speculation protocol: + # `clone` snapshots the FSM before a proposal, `lookahead` advances the + # snapshot, and `commit` applies only the accepted prefix. Abandoning a + # rejected proposal is therefore a transaction abort rather than an + # unrecoverable side effect, and the explicit `clone` action is exactly + # what makes the domain safe to enter speculatively. + workflow["effects"] = { + "grammar": { + "retry": "transactional", + "speculation_safety": {"kind": "clonable"}, + } + } + if adaptive_k_max is not None: + workflow["manifest"]["capabilities"].extend( + ["adaptive_proposal_budget", "advisory_state"] + ) + metadata = { + "schema_version": "v1", + "pipeline": {"workflow": _publish_workflow_v1(workflow)}, + } + add_policy_components_to_workflow(metadata, pkg) + return metadata + + +def write_speculative_workflow_metadata( + pkg: Any, + output_dir: str, + *, + grammar_guidance: bool = False, + adaptive_k_max: int | None = None, +) -> str: + os.makedirs(output_dir, exist_ok=True) + metadata = build_speculative_workflow_metadata( + pkg, + grammar_guidance=grammar_guidance, + adaptive_k_max=adaptive_k_max, + ) + pkg.save_policy_components(output_dir) + add_adapter_service_to_metadata(metadata, pkg, output_dir) + path = os.path.join(output_dir, "inference_metadata.yaml") + with open(path, "w", encoding="utf-8") as handle: + _dump_yaml(metadata, handle) + return path + + +def _cross_state_bindings( + encoder: ir.Model, decoder: ir.Model +) -> dict[str, tuple[str, ir.Value]]: + """Map decoder inputs that a conditioning encoder produces once per request. + + The mapping is purely structural. A decoder input is bound to an encoder + output when the names match exactly (``encoder_hidden_states``) or when the + decoder's ``past``/``pass``-side spelling rewrites to the encoder's + ``present``-side spelling (``past_key_cross_0`` <- ``present_key_cross_0``), + which is the same past/present rewrite already used for self-attention KV. + No model family, tensor role, or vendor name is consulted. + """ + encoder_outputs = {value.name: value for value in encoder.graph.outputs} + bindings: dict[str, tuple[str, ir.Value]] = {} + for value in decoder.graph.inputs: + if value.name is None: + continue + candidates = ( + value.name, + value.name.replace("past_key_values", "present"), + value.name.replace("past.", "present."), + value.name.replace("past_", "present_"), + ) + produced = next( + ((name, encoder_outputs[name]) for name in candidates if name in encoder_outputs), + None, + ) + if produced is not None: + bindings[value.name] = produced + return bindings + + +def build_speech_to_text_workflow_metadata( + pkg: Any, + config: Any, + *, + sampler: str = "greedy", + audio_preprocessing: dict[str, Any] | None = None, + artifacts: dict[str, str] | None = None, +) -> dict[str, Any]: + """Build the typed SSA workflow for an encoder-conditioned decoder package. + + The encoder runs once per request inside the loop ``setup``; every value it + produces that the decoder consumes becomes a shape-invariant workflow state + cell carried unchanged through the decode loop. That is what keeps encoder + states and decoder rows aligned under batching: the cross state is + request-aligned on the same axis as the tokens and is permuted by the same + compaction that permutes the self-attention cache. + """ + if set(pkg.keys()) != {"encoder", "decoder"}: + raise ValueError( + "speech-to-text workflow requires exactly encoder and decoder components, " + f"got {sorted(pkg.keys())}" + ) + return _build_autoregressive_workflow_metadata( + pkg, + config, + sampler=sampler, + encoder_name="encoder", + audio_preprocessing=audio_preprocessing, + artifacts=artifacts, + ) + + +def build_decoder_workflow_metadata( + pkg: Any, + config: Any, + *, + sampler: str = "greedy", +) -> dict[str, Any]: + """Build the exact workflow-policy contract for an autoregressive decoder.""" + if len(pkg) != 1: + raise ValueError("decoder workflow requires exactly one neural component") + return _build_autoregressive_workflow_metadata(pkg, config, sampler=sampler) + + +def _build_autoregressive_workflow_metadata( + pkg: Any, + config: Any, + *, + sampler: str = "greedy", + encoder_name: str | None = None, + audio_preprocessing: dict[str, Any] | None = None, + artifacts: dict[str, str] | None = None, +) -> dict[str, Any]: + """Build an autoregressive decode loop, optionally conditioned by an encoder. + + ``artifacts`` overrides the package-relative ONNX path of a component, which + lets an importer describe an existing on-disk layout without renaming files. + """ + encoder = pkg[encoder_name] if encoder_name is not None else None + decoder_items = [(name, model) for name, model in pkg.items() if name != encoder_name] + if len(decoder_items) != 1: + raise ValueError("workflow requires exactly one autoregressive component") + decoder_name, decoder = decoder_items[0] + cross_bindings = _cross_state_bindings(encoder, decoder) if encoder is not None else {} + if encoder is not None and not cross_bindings: + raise ValueError( + "encoder-conditioned workflow requires at least one decoder input produced " + "by the encoder" + ) + inputs = list(decoder.graph.inputs) + outputs = list(decoder.graph.outputs) + decoder_kv_contract = _kv_storage_contract(decoder) + token_input = next( + ( + value + for value in inputs + if ("input_ids" in value.name or "token" in value.name) + and value.dtype in {ir.DataType.INT32, ir.DataType.INT64} + and value.shape is not None + and len(value.shape) == 2 + ), + next( + ( + value + for value in inputs + if value.dtype in {ir.DataType.INT32, ir.DataType.INT64} + and value.shape is not None + and len(value.shape) == 2 + ), + None, + ), + ) + logits_output = next( + ( + value + for value in outputs + if value.dtype + in { + ir.DataType.FLOAT, + ir.DataType.FLOAT16, + ir.DataType.BFLOAT16, + ir.DataType.DOUBLE, + } + and value.shape is not None + and len(value.shape) == 3 + ), + None, + ) + if token_input is None or logits_output is None: + raise ValueError( + "decoder workflow requires rank-2 token input and rank-3 logits output" + ) + attach_policy_components( + pkg, + PolicyCapabilities( + sampler=sampler, + eos_termination=True, + token_state_update=True, + ), + ) + pkg.add_policy_component("last_token_logits", build_last_token_logits(logits_output.dtype)) + + output_by_suffix = {value.name: value for value in outputs} + cache_pairs: list[tuple[ir.Value, ir.Value]] = [] + for value in inputs: + if value.name in cross_bindings: + continue + present = next( + ( + output_by_suffix.get(name) + for name in _cache_output_candidates(value.name or "") + if name in output_by_suffix + ), + None, + ) + if present is not None: + cache_pairs.append((value, present)) + cache_names = {past.name for past, _ in cache_pairs} + # A static cache is a fixed-capacity buffer the graph scatters into at + # declared destinations, so it needs a write cursor and a per-row valid + # length that no shape can carry. Its two control ports are integer vectors + # and shape-indistinguishable, hence read from the exporter's declared ABI. + static_cache = _static_cache_ports(decoder) + static_control_names = ( + {static_cache["write_indices"], static_cache["kv_sequence_length"]} + if static_cache is not None + else set() + ) + integer_rank2 = [ + value + for value in inputs + if value is not token_input + and value.name not in cache_names + and value.dtype in {ir.DataType.INT32, ir.DataType.INT64} + and value.shape is not None + and len(value.shape) == 2 + ] + attention_input = next( + ( + value + for value in integer_rank2 + if value.name not in cross_bindings + and ( + "mask" in value.name + or "past" in str(getattr(list(value.shape)[1], "value", list(value.shape)[1])) + ) + ), + None, + ) + position_input = next( + ( + value + for value in integer_rank2 + if value is not attention_input + and value.name not in cross_bindings + and "position" in value.name + ), + next( + ( + value + for value in integer_rank2 + if value is not attention_input and value.name not in cross_bindings + ), + None, + ), + ) + derived_names = cache_names | set(cross_bindings) | static_control_names + if attention_input is not None: + derived_names.add(attention_input.name) + if position_input is not None: + derived_names.add(position_input.name) + unsupported = [ + value.name + for value in inputs + if value is not token_input and value.name not in derived_names + ] + if unsupported: + raise ValueError(f"decoder workflow has unsupported non-request inputs: {unsupported}") + # A shared full-capacity KV buffer needs an attention mask to convey each + # row's logical length; a mask-free decoder must grow its cache instead. + fixed_capacity = ( + bool(cache_pairs) + and decoder_kv_contract["storage"] == "shared_buffer" + and attention_input is not None + ) + # A static cache carries its own per-row length on a graph port, so the + # prompt length has to be materialized for it whether or not an attention + # mask is also present. + tracks_cache_lengths = fixed_capacity or static_cache is not None + pkg.add_policy_component( + "decoder_state_initializer", + build_decoder_state_initializer( + decoder, + token_input=token_input.name, + attention_mask_input=attention_input.name if attention_input is not None else None, + position_ids_input=position_input.name if position_input is not None else None, + cache_inputs=sorted(cache_names), + fixed_capacity=fixed_capacity, + ragged=bool(cache_pairs), + write_indices_output=( + static_cache["write_indices"] if static_cache is not None else None + ), + ), + ) + if attention_input is not None: + pkg.add_policy_component( + "decoder_step_update", + build_decoder_step_update( + attention_dtype=attention_input.dtype, + position_dtype=position_input.dtype if position_input is not None else None, + fixed_capacity=fixed_capacity, + position_sections=rotary_axis_count(position_input) + if position_input is not None + else None, + ), + ) + elif position_input is not None: + pkg.add_policy_component( + "decoder_step_update", + build_decoder_step_update( + attention_dtype=None, + position_dtype=position_input.dtype, + fixed_capacity=False, + position_sections=rotary_axis_count(position_input), + ), + ) + needs_token_cast = token_input.dtype != ir.DataType.INT64 + if needs_token_cast: + pkg.add_policy_component("model_token_cast", build_model_token_cast(token_input.dtype)) + + workflow_inputs: dict[str, Any] = {} + setup_decoder_inputs: dict[str, str] = {} + body_decoder_inputs: dict[str, str] = {} + for value in inputs: + if value.name in derived_names: + continue + name = f"request.{value.name}" + if value is token_input: + role = { + "kind": "runtime", + "version": "1.0", + "role": "prompt_tokens", + } + source = {"kind": "request", "field": "prompt_tokens"} + else: + role = {"kind": "opaque"} + source = {"kind": "application", "name": value.name} + workflow_inputs[name] = { + "contract": _contract(value), + "role": role, + "source": source, + "required": True, + } + setup_decoder_inputs[value.name] = name + body_decoder_inputs[value.name] = name + + # Encoder inputs are request inputs of the whole workflow: the encoder runs + # once in the loop setup and its results persist as invariant state. + encoder_invoke_inputs: dict[str, str] = {} + encoder_invoke_outputs: dict[str, str] = {} + audio_adapter_outputs: dict[str, Any] = {} + audio_values: dict[str, str] = {} + audio_program: dict[str, Any] | None = None + if encoder is not None: + assert encoder_name is not None + encoder_inputs_by_name = {value.name: value for value in encoder.graph.inputs} + if audio_preprocessing is not None: + audio_program = copy.deepcopy(audio_preprocessing) + for binding in audio_program["outputs"]: + port_name = binding["name"] + if port_name not in encoder_inputs_by_name: + raise ValueError( + f"audio preprocessing output {port_name!r} has no encoder input" + ) + contract = _contract(encoder_inputs_by_name[port_name]) + binding["contract"] = contract + binding["dtype"] = contract["dtype"] + binding["name"] = f"audio.{port_name}" + audio_adapter_outputs[port_name] = contract + audio_values[port_name] = binding["name"] + for value in encoder.graph.inputs: + ssa = f"encoder.input.{value.name}" + if value.name in audio_values: + # The declared preprocessing program produces this encoder input, + # so the workflow consumes decoded audio bytes instead of features. + encoder_invoke_inputs[value.name] = audio_values[value.name] + continue + workflow_inputs[ssa] = { + "contract": _contract(value), + "role": {"kind": "opaque"}, + "source": {"kind": "application", "name": value.name}, + "required": True, + # An application may reuse a previously computed encoder input + # instead of recomputing the feature extraction. + "externally_suppliable": True, + } + encoder_invoke_inputs[value.name] = ssa + for _decoder_input, (encoder_output, _) in sorted(cross_bindings.items()): + encoder_invoke_outputs[encoder_output] = f"encoder.{encoder_output}" + if audio_program is not None: + # Raw encoded audio is the request-level input; the adapter decodes and + # turns it into the encoder feature tensor declared by the program. + workflow_inputs["request.audio"] = { + "contract": {"dtype": "uint8", "rank": 1, "shape": ["encoded_bytes"]}, + "role": {"kind": "runtime", "version": "1.0", "role": "media"}, + "source": {"kind": "request", "field": "media"}, + "required": True, + } + batch_dimension = _shape_metadata(_port(token_input))[0] + batch_int = _request_aligned({"dtype": "int64", "rank": 1, "shape": [batch_dimension]}) + batch_bool = _request_aligned({"dtype": "bool", "rank": 1, "shape": [batch_dimension]}) + control_int = {"dtype": "int64", "rank": 1, "shape": [1]} + eos_token_id = getattr(config, "eos_token_id", 0) + if isinstance(eos_token_id, list): + eos_token_id = eos_token_id[0] if eos_token_id else 0 + eos_token_id = int(eos_token_id or 0) + workflow_inputs.update( + { + "request.max_iterations": { + "contract": control_int, + "role": { + "kind": "runtime", + "version": "1.0", + "role": "max_output_tokens", + }, + "source": {"kind": "request", "field": "max_output_tokens"}, + "required": True, + }, + "package.eos_ids": { + "contract": {"dtype": "int64", "rank": 1, "shape": ["E"]}, + "role": {"kind": "opaque"}, + "source": {"kind": "literal"}, + "required": True, + "default": eos_token_id, + }, + "package.one_token": { + "contract": batch_int, + "role": {"kind": "opaque"}, + "source": {"kind": "literal"}, + "required": False, + "default": 1, + }, + "package.one_step": { + # A growing-recurrence increment advances the whole invocation's + # state axis by one, so it is an invocation-scoped control scalar, + # not a per-row value like ``package.one``/``package.one_token``. + "contract": control_int, + "role": {"kind": "opaque"}, + "source": {"kind": "literal"}, + "required": False, + "default": 1, + }, + "package.max_context": { + "contract": control_int, + "role": {"kind": "opaque"}, + "source": {"kind": "literal"}, + "required": False, + "default": int(getattr(config, "max_position_embeddings", 4096)), + }, + } + ) + if static_cache is not None: + # The capacity a static graph was built against is a graph fact, not a + # deployment budget: it bounds legal write destinations and nothing else. + workflow_inputs["package.cache_capacity"] = { + "contract": control_int, + "role": {"kind": "opaque"}, + "source": {"kind": "literal"}, + "required": False, + "default": int(static_cache["capacity"]), + } + if cache_pairs: + workflow_inputs.update( + { + "request.prompt_lengths": { + "contract": batch_int, + "role": {"kind": "opaque"}, + "source": {"kind": "application", "name": "prompt_lengths"}, + "required": False, + "default": -1, + }, + "request.eos_ids": { + "contract": { + "dtype": "int64", + "rank": 2, + "shape": [batch_dimension, "num_eos"], + }, + "role": {"kind": "opaque"}, + "source": {"kind": "application", "name": "eos_ids"}, + "required": False, + "default": eos_token_id, + }, + "request.eos_lengths": { + "contract": batch_int, + "role": {"kind": "opaque"}, + "source": {"kind": "application", "name": "eos_lengths"}, + "required": False, + "default": 1, + }, + "request.row_max_iterations": { + "contract": batch_int, + "role": {"kind": "opaque"}, + "source": { + "kind": "application", + "name": "row_max_iterations", + }, + "required": False, + "default": -1, + }, + } + ) + stochastic_sampler = sampler != "greedy" + sampler_with_rng = stochastic_sampler or bool(cache_pairs) + if sampler_with_rng: + workflow_inputs.update( + { + "request.temperature": { + "contract": { + "dtype": "float32", + "rank": 1, + "shape": [batch_dimension], + }, + "role": { + "kind": "runtime", + "version": "1.0", + "role": "sampling_temperature", + }, + "source": { + "kind": "request", + "field": "sampling_temperature", + }, + "required": False, + "default": 1.0, + }, + "request.top_k": { + "contract": batch_int, + "role": { + "kind": "runtime", + "version": "1.0", + "role": "sampling_top_k", + }, + "source": {"kind": "request", "field": "sampling_top_k"}, + "required": False, + "default": 0 if stochastic_sampler else 1, + }, + "request.top_p": { + "contract": { + "dtype": "float32", + "rank": 1, + "shape": [batch_dimension], + }, + "role": { + "kind": "runtime", + "version": "1.0", + "role": "sampling_top_p", + }, + "source": {"kind": "request", "field": "sampling_top_p"}, + "required": False, + "default": 1.0, + }, + "request.min_p": { + "contract": { + "dtype": "float32", + "rank": 1, + "shape": [batch_dimension], + }, + "role": { + "kind": "runtime", + "version": "1.0", + "role": "sampling_min_p", + }, + "source": { + "kind": "request", + "field": "sampling_min_p", + }, + "required": False, + "default": 0.0, + }, + "request.seed": { + "contract": batch_int, + "role": { + "kind": "runtime", + "version": "1.0", + "role": "seed", + }, + "source": {"kind": "request", "field": "seed"}, + "required": False, + "default": 0, + }, + "request.rng_counter": { + "contract": batch_int, + "role": {"kind": "opaque"}, + "source": {"kind": "application", "name": "rng_counter"}, + "required": False, + "default": 0, + }, + } + ) + if cache_pairs: + pkg.add_policy_component("cache_length_update", build_selective_integer_add()) + pkg.add_policy_component( + "token_sampler", + build_seeded_categorical_sampler(), + ) + pkg.add_policy_component("termination", build_eos_termination(row_selective=True)) + pkg.add_policy_component( + "termination_batch_initializer", + build_termination_batch_initializer(), + ) + pkg.add_policy_component( + "token_state_update", + build_token_state_update(row_selective=True), + ) + pkg.add_policy_component("token_to_slot", build_token_to_slot()) + pkg.add_policy_component("generated_length_update", build_selective_integer_add()) + workflow_inputs.update( + { + "package.active": { + "contract": batch_bool, + "role": {"kind": "opaque"}, + "source": {"kind": "literal"}, + "required": False, + "default": True, + }, + "package.not_done": { + "contract": batch_bool, + "role": {"kind": "opaque"}, + "source": {"kind": "literal"}, + "required": False, + "default": False, + }, + "package.cache_lengths": { + "contract": batch_int, + "role": {"kind": "opaque"}, + "source": {"kind": "literal"}, + "required": False, + "default": 0, + }, + "package.zero_batch": { + "contract": batch_int, + "role": {"kind": "opaque"}, + "source": {"kind": "literal"}, + "required": False, + "default": 0, + }, + } + ) + + for value in inputs: + if value is token_input: + continue + if value.name in cache_names: + body_decoder_inputs[value.name] = f"state.{value.name}.body" + setup_decoder_inputs[value.name] = f"initializer.{value.name}" + if attention_input is not None: + setup_decoder_inputs[attention_input.name] = f"initializer.{attention_input.name}" + body_decoder_inputs[attention_input.name] = ( + "decoder_step.body_attention_mask" + if fixed_capacity + else "state.attention_mask.body" + ) + if position_input is not None: + setup_decoder_inputs[position_input.name] = f"initializer.{position_input.name}" + body_decoder_inputs[position_input.name] = "state.position_ids.body" + if static_cache is not None: + # Prefill scatters the whole prompt from slot zero and ends with + # ``prompt_length`` valid entries. Each decode step then writes at the + # row's current logical length and ends one entry longer — except for a + # finished row, whose length does not advance, so the slot it just wrote + # stays outside its valid prefix and is reclaimed by its next write. + setup_decoder_inputs[static_cache["write_indices"]] = ( + f"initializer.{static_cache['write_indices']}" + ) + setup_decoder_inputs[static_cache["kv_sequence_length"]] = "initializer.cache_lengths" + body_decoder_inputs[static_cache["write_indices"]] = "state.cache_lengths.body" + body_decoder_inputs[static_cache["kv_sequence_length"]] = "cache_lengths.next" + # Cross state is produced once by the encoder and read unchanged by every + # decode step, so setup binds the encoder result and the body binds the + # invariant carried cell that holds it. + for decoder_input, (encoder_output, _) in sorted(cross_bindings.items()): + setup_decoder_inputs[decoder_input] = f"encoder.{encoder_output}" + body_decoder_inputs[decoder_input] = f"state.cross.{decoder_input}.body" + body_decoder_inputs[token_input.name] = ( + "model_token.body" if needs_token_cast else "token.body" + ) + + setup_decoder_outputs = {logits_output.name: "decoder.setup.logits"} + body_decoder_outputs = {logits_output.name: "decoder.body.logits"} + logits_contract = _contract(logits_output) + last_logits_contract = { + "dtype": "float32", + "rank": 2, + "shape": [logits_contract["shape"][0], logits_contract["shape"][-1]], + } + state: dict[str, Any] = { + "token": { + "contract": { + "dtype": "int64", + "rank": 2, + "shape": [batch_dimension, 1], + }, + "scope": "invocation", + "initializer": "initializer.token_slot", + "recurrence": {"kind": "invariant"}, + }, + "logits": { + "contract": last_logits_contract, + "scope": "invocation", + "initializer": "decoder.setup.last_logits", + "recurrence": {"kind": "invariant"}, + }, + } + if cache_pairs: + state.update( + { + "generated_lengths": { + "contract": batch_int, + "class": "semantic", + "scope": "invocation", + "initializer": "initializer.generated_lengths", + "recurrence": {"kind": "invariant"}, + }, + "active": { + "contract": batch_bool, + "class": "semantic", + "scope": "invocation", + "initializer": "package.active", + "recurrence": {"kind": "invariant"}, + }, + "done": { + "contract": batch_bool, + "class": "semantic", + "scope": "invocation", + "initializer": "package.not_done", + "recurrence": {"kind": "invariant"}, + }, + "accepted_len": { + "contract": batch_int, + "class": "semantic", + "scope": "invocation", + "initializer": "package.zero_batch", + "recurrence": {"kind": "invariant"}, + }, + "cache_lengths": { + "contract": batch_int, + "class": "semantic", + "scope": "invocation", + "initializer": ( + "initializer.cache_lengths" + if tracks_cache_lengths + else "package.cache_lengths" + ), + "recurrence": {"kind": "invariant"}, + }, + } + ) + initial_effects = { + "sample": "sample.0", + "termination": "termination.0", + "state": "state.0", + "emit": "emit.0", + "state:token": "state:token.0", + "state:logits": "state:logits.0", + } + carried = [ + { + "cell": "token", + "current": "initializer.token_slot", + "body_input": "state.token.body", + "body_output": "token.body", + "next": "token.final", + "read_effect": _effect("state:token.0", "state:token.read"), + "write_effect": _effect("state:token.read", "state:token.1"), + } + ] + carried.extend( + [ + { + "cell": "logits", + "current": "decoder.setup.last_logits", + "body_input": "state.logits.body", + "body_output": "decoder.body.last_logits", + "next": "state.logits.final", + "read_effect": _effect("state:logits.0", "state:logits.read"), + "write_effect": _effect("state:logits.read", "state:logits.1"), + }, + ] + ) + if cache_pairs: + carried.extend( + [ + { + "cell": "generated_lengths", + "current": "initializer.generated_lengths", + "body_input": "state.generated_lengths.body", + "body_output": "token.next_lengths", + "next": "state.generated_lengths.final", + }, + { + "cell": "active", + "current": "package.active", + "body_input": "state.active.body", + "body_output": "loop.next_active", + "next": "state.active.final", + }, + { + "cell": "done", + "current": "package.not_done", + "body_input": "state.done.body", + "body_output": "loop.done", + "next": "state.done.final", + }, + { + "cell": "cache_lengths", + "current": ( + "initializer.cache_lengths" + if tracks_cache_lengths + else "package.cache_lengths" + ), + "body_input": "state.cache_lengths.body", + "body_output": "cache_lengths.next", + "next": "state.cache_lengths.final", + }, + { + "cell": "accepted_len", + "current": "package.zero_batch", + "body_input": "state.accepted_len.body", + "body_output": "accepted_len.next", + "next": "state.accepted_len.final", + }, + ] + ) + if sampler_with_rng: + state["rng_counter"] = { + "contract": batch_int, + "scope": "invocation", + "class": "semantic", + "initializer": "request.rng_counter", + "recurrence": {"kind": "invariant"}, + } + initial_effects["state:rng_counter"] = "state:rng_counter.0" + carried.append( + { + "cell": "rng_counter", + "current": "request.rng_counter", + "body_input": "state.rng_counter.body", + "body_output": "sample.next_counter", + "next": "state.rng_counter.final", + "read_effect": _effect("state:rng_counter.0", "state:rng_counter.read"), + "write_effect": _effect("state:rng_counter.read", "state:rng_counter.1"), + } + ) + decoder_state_specs: dict[str, tuple[dict[str, Any], str, str, dict[str, Any]]] = {} + if attention_input is not None: + decoder_state_specs["attention_mask"] = ( + { + "dtype": _contract(attention_input)["dtype"], + "rank": 2, + "shape": [batch_dimension, "context"], + }, + ( + f"initializer.{attention_input.name}" + if fixed_capacity + else "initializer.body_attention_mask" + ), + "decoder_step.body_attention_mask", + ( + {"kind": "invariant"} + if fixed_capacity + else { + "kind": "growing", + "axis": 1, + "increment": "package.one_step", + "max": "package.max_context", + } + ), + ) + if position_input is not None: + decoder_state_specs["position_ids"] = ( + { + "dtype": _contract(position_input)["dtype"], + "rank": 2, + "shape": [batch_dimension, 1], + }, + "initializer.body_position_ids", + "decoder_step.body_position_ids", + {"kind": "invariant"}, + ) + for cell, (contract, current, body_output, recurrence) in decoder_state_specs.items(): + effect_name = f"state:{cell}" + initial_effects[effect_name] = f"{effect_name}.0" + state[cell] = { + "contract": contract, + "scope": "invocation", + "initializer": current, + "recurrence": recurrence, + } + carried.append( + { + "cell": cell, + "current": current, + "body_input": f"state.{cell}.body", + "body_output": body_output, + "next": f"state.{cell}.final", + "read_effect": _effect(f"{effect_name}.0", f"{effect_name}.read"), + "write_effect": _effect(f"{effect_name}.read", f"{effect_name}.1"), + } + ) + decoder_kv_ports: dict[str, Any] = {} + decoder_cache_cells: list[str] = [] + decoder_kv_axis = 2 + for cache_index, (past, present) in enumerate(cache_pairs): + # Generated-length state is orthogonal to the admitted cache ABI and + # must not renumber stable cache service cells. + cell = f"cache_{cache_index}" + setup_value = f"decoder.setup.{present.name}" + body_value = f"decoder.body.{present.name}" + setup_decoder_outputs[present.name] = setup_value + body_decoder_outputs[present.name] = body_value + decoder_kv_axis = next( + ( + index + for index, dimension in enumerate(_contract(past)["shape"]) + if "sequence" in str(dimension) + ), + 2, + ) + # A scattered buffer never changes shape: every step overwrites cells + # inside a capacity fixed at export, so its extent is invariant and the + # logical prefix is carried separately by ``cache_lengths``. + scattered = static_cache is not None and past.name in static_cache["buffers"] + state[cell] = { + "contract": _request_aligned(_contract(past)), + "scope": "invocation", + "initializer": setup_value, + "recurrence": ( + {"kind": "invariant"} + if scattered + else { + "kind": "bounded", + "axis": decoder_kv_axis, + "max": "package.max_context", + } + ), + # Binding a cell to a state service group hands its storage to the + # runtime, which then owns allocation, compaction, and release. + "management": "runtime", + "release_boundary": "invocation", + } + decoder_kv_ports[cell] = {"input": past.name, "output": present.name} + decoder_cache_cells.append(cell) + effect_name = f"state:{cell}" + initial_effects[effect_name] = f"{effect_name}.0" + carried.append( + { + "cell": cell, + "current": setup_value, + "body_input": f"state.{past.name}.body", + "body_output": body_value, + "next": f"state.{past.name}.final", + "read_effect": _effect(f"{effect_name}.0", f"{effect_name}.read"), + "write_effect": _effect(f"{effect_name}.read", f"{effect_name}.1"), + } + ) + + # Cross state: produced once by the encoder, shape-invariant, carried + # unchanged. The identity carry keeps the cell request-aligned on the same + # axis as the tokens so row compaction permutes encoder state and decoder + # rows together. + for decoder_input, (encoder_output, encoder_value) in sorted(cross_bindings.items()): + cell = f"cross.{decoder_input}" + setup_value = f"encoder.{encoder_output}" + body_input = f"state.cross.{decoder_input}.body" + contract = _contract(encoder_value) + contract["batch_layout"] = {"kind": "request_aligned", "axis": 0} + state[cell] = { + "contract": contract, + "class": "semantic", + "scope": "invocation", + "initializer": setup_value, + "recurrence": {"kind": "invariant"}, + } + effect_name = f"state:{cell}" + initial_effects[effect_name] = f"{effect_name}.0" + carried.append( + { + "cell": cell, + "current": setup_value, + "body_input": body_input, + "body_output": body_input, + "next": f"state.cross.{decoder_input}.final", + "read_effect": _effect(f"{effect_name}.0", f"{effect_name}.read"), + "write_effect": _effect(f"{effect_name}.read", f"{effect_name}.1"), + } + ) + # One semantic state group per attention kind. A hybrid decoder (Gemma 3/4, + # Gemma 3n, ...) publishes distinct sliding and full-attention groups so the + # runtime never evicts a global layer's prefix, and so each group's cells + # carry their own geometry (Gemma 4's global layers are double-wide). + decoder_state_groups, decoder_cell_groups = _state_service_groups( + config=config, + cache_pairs=cache_pairs, + ports={decoder_name: decoder_kv_ports}, + sequence_axis=decoder_kv_axis, + # A mask-free decoder grows its cache instead of writing into a shared + # full-capacity buffer: it publishes no logical lengths and its present + # ports must never be aliased onto the past bindings. + logical_lengths="cache_lengths" if fixed_capacity else None, + aliasing=_aliasing_for_storage( + decoder_kv_contract["storage"] if fixed_capacity else "growable" + ), + base_name="decoder_cache", + indexed_scatter=( + { + "buffers": static_cache["buffers"], + "capacity": "package.cache_capacity", + # The write cursor and the logical length are the same quantity: + # a row's next write lands exactly where its valid prefix ends. + "write_indices": "cache_lengths", + "logical_lengths": "cache_lengths", + "port": static_cache["write_indices"], + "kv_length_port": static_cache["kv_sequence_length"], + } + if static_cache is not None + else None + ), + ) + for cell in decoder_cache_cells: + state[cell]["service_group"] = decoder_cell_groups[cell] + + setup = { + "kind": "sequence", + "nodes": [ + *( + [ + _invoke( + "audio_preprocess", + {"encoded": "request.audio"}, + dict(audio_values), + ) + ] + if audio_program is not None + else [] + ), + *( + [_invoke(encoder_name, encoder_invoke_inputs, encoder_invoke_outputs)] + if encoder is not None + else [] + ), + _invoke( + "decoder_state_initializer", + { + "prompt_tokens": f"request.{token_input.name}", + **({"prompt_lengths": "request.prompt_lengths"} if cache_pairs else {}), + **({"max_iterations": "request.max_iterations"} if fixed_capacity else {}), + }, + { + **( + { + attention_input.name: f"initializer.{attention_input.name}", + "body_attention_mask": "initializer.body_attention_mask", + } + if attention_input is not None + else {} + ), + "token_slot": "initializer.token_slot", + **( + {"generated_lengths": "initializer.generated_lengths"} + if cache_pairs + else {} + ), + **( + {"cache_lengths": "initializer.cache_lengths"} + if tracks_cache_lengths + else {} + ), + **( + { + static_cache["write_indices"]: ( + f"initializer.{static_cache['write_indices']}" + ) + } + if static_cache is not None + else {} + ), + **( + { + position_input.name: f"initializer.{position_input.name}", + "body_position_ids": "initializer.body_position_ids", + } + if position_input is not None + else {} + ), + **{name: f"initializer.{name}" for name in sorted(cache_names)}, + }, + ), + _invoke(decoder_name, setup_decoder_inputs, setup_decoder_outputs), + *( + [ + _invoke( + "termination_batch_initializer", + { + "input_eos_ids": "request.eos_ids", + "input_eos_lengths": "request.eos_lengths", + "input_max_iterations": "request.row_max_iterations", + "fallback_max_iterations": "request.max_iterations", + "active": "package.active", + }, + { + "row_eos_ids": "termination.eos_ids", + "eos_lengths": "termination.eos_lengths", + "max_iterations": "termination.max_iterations", + }, + ) + ] + if cache_pairs + else [] + ), + _invoke( + "last_token_logits", + {"logits": "decoder.setup.logits"}, + {"last_logits": "decoder.setup.last_logits"}, + ), + ], + } + has_step_update = attention_input is not None or position_input is not None + decoder_step_invoke = _invoke( + "decoder_step_update", + { + **( + {"attention_mask": "state.attention_mask.body"} + if attention_input is not None + else {} + ), + **({"logical_length": "state.cache_lengths.body"} if fixed_capacity else {}), + **( + {"position_ids": "state.position_ids.body"} + if position_input is not None + else {} + ), + }, + { + **( + {"next_attention_mask": "decoder_step.body_attention_mask"} + if attention_input is not None + else {} + ), + **( + {"next_position_ids": "decoder_step.body_position_ids"} + if position_input is not None + else {} + ), + }, + ) + body = { + "kind": "sequence", + "nodes": [ + _invoke( + "token_sampler", + { + "logits": "state.logits.body", + **( + { + "temperature": "request.temperature", + "top_k": "request.top_k", + "top_p": "request.top_p", + "min_p": "request.min_p", + "seed": "request.seed", + "counter": "state.rng_counter.body", + } + if sampler_with_rng + else {} + ), + **( + { + "active": "state.active.body", + "done": "state.done.body", + } + if cache_pairs + else {} + ), + }, + { + "token": "sample.body", + **({"next_counter": "sample.next_counter"} if sampler_with_rng else {}), + }, + {"sample": _effect("sample.0", "sample.1")}, + ), + *( + [ + _invoke( + "token_to_slot", + {"token": "sample.body"}, + {"slot": "sample.slot"}, + ), + _invoke( + "generated_length_update", + { + "left": "state.generated_lengths.body", + "right": "package.one_token", + "active": "state.active.body", + "done": "state.done.body", + }, + {"total": "token.next_lengths"}, + ), + _invoke( + "generated_length_update", + { + "left": "package.zero_batch", + "right": "package.one_token", + "active": "state.active.body", + "done": "state.done.body", + }, + {"total": "token.emitted_length"}, + ), + ] + if cache_pairs + else [] + ), + _invoke( + "token_state_update", + { + "current": "state.token.body", + "update": "sample.slot" if cache_pairs else "sample.body", + **( + { + "active": "state.active.body", + "done": "state.done.body", + } + if cache_pairs + else {} + ), + }, + {"next": "token.body"}, + {"state": _effect("state.0", "state.1")}, + ), + *( + [ + _invoke( + "model_token_cast", + {"token": "token.body"}, + {"model_token": "model_token.body"}, + ) + ] + if needs_token_cast + else [] + ), + _invoke( + "termination", + { + "tokens": "sample.body", + "eos_ids": ("termination.eos_ids" if cache_pairs else "package.eos_ids"), + **({"eos_lengths": "termination.eos_lengths"} if cache_pairs else {}), + "iteration": "loop.iteration", + "max_iterations": ( + "termination.max_iterations" + if cache_pairs + else "request.max_iterations" + ), + **({"active": "state.active.body"} if cache_pairs else {}), + }, + { + "done": "loop.done", + "continue": "loop.continue", + **({"next_active": "loop.next_active"} if cache_pairs else {}), + }, + {"termination": _effect("termination.0", "termination.1")}, + ), + *( + [ + _invoke( + "cache_length_update", + { + "left": "state.cache_lengths.body", + "right": "package.one_token", + "active": "state.active.body", + "done": "state.done.body", + }, + {"total": "cache_lengths.next"}, + ), + _invoke( + "cache_length_update", + { + "left": "package.zero_batch", + "right": "package.one_token", + "active": "state.active.body", + "done": "state.done.body", + }, + {"total": "accepted_len.next"}, + ), + ] + if cache_pairs + else [] + ), + { + "kind": "emit", + "value": "token.body", + "output": "tokens", + "mode": "append", + **( + { + "when": "state.active.body", + "valid_length": "token.emitted_length", + } + if cache_pairs + else {} + ), + "effect_name": "emit", + "effect": _effect("emit.0", "emit.1"), + }, + *([decoder_step_invoke] if has_step_update and fixed_capacity else []), + _invoke(decoder_name, body_decoder_inputs, body_decoder_outputs), + _invoke( + "last_token_logits", + {"logits": "decoder.body.logits"}, + {"last_logits": "decoder.body.last_logits"}, + ), + *([decoder_step_invoke] if has_step_update and not fixed_capacity else []), + ], + } + + artifacts = artifacts or {} + use_subfolders = len(pkg) > 1 + artifact = artifacts.get( + decoder_name, f"{decoder_name}/model.onnx" if use_subfolders else "model.onnx" + ) + workflow = { + "manifest": { + **( + {"adapter_abis": {"onnx-genai.audio-preprocess": "1"}} + if audio_program is not None + else {} + ), + "capabilities": [ + "workflow_ssa", + "linear_effects", + "nested_control_flow", + "typed_emit", + "emit_valid_length", + "loop_induction_values", + *(["audio_preprocessing_program"] if audio_program is not None else []), + *(["serving_service_contract"] if cache_pairs else []), + *(["bounded_state_recurrence"] if cache_pairs else []), + ], + }, + "inputs": workflow_inputs, + "outputs": { + "tokens": { + "contract": _request_aligned( + { + "dtype": "int64", + "rank": 2, + "shape": [batch_dimension, "generated_sequence"], + } + ), + "role": "tokens", + "stage": "pre_adapter", + } + }, + "components": { + decoder_name: _component(decoder, artifact), + **( + { + encoder_name: _component( + encoder, + artifacts.get(encoder_name, f"{encoder_name}/model.onnx"), + ) + } + if encoder is not None + else {} + ), + **( + { + "audio_preprocess": { + "implementation": { + "kind": "adapter", + "abi": "onnx-genai.audio-preprocess", + "version": "1", + }, + "ports": { + "inputs": { + "encoded": { + "dtype": "uint8", + "rank": 1, + "shape": ["encoded_bytes"], + } + }, + "outputs": audio_adapter_outputs, + }, + } + } + if audio_program is not None + else {} + ), + }, + "state": state, + **( + { + "serving": { + "active": "active", + "done": "done", + "accepted_len": "accepted_len", + "state_service": {"groups": decoder_state_groups}, + } + } + if cache_pairs + else {} + ), + "initial_effects": initial_effects, + "graph": { + "kind": "loop", + "setup": setup, + "body": body, + "condition": "loop.continue", + "termination": "generation_eos", + **({"active_cell": "active"} if cache_pairs else {}), + "max_iterations": "request.max_iterations", + "iteration": { + "value": "loop.iteration", + "contract": {"dtype": "int64", "rank": 1, "shape": [1]}, + }, + "carried": carried, + }, + } + metadata = { + "schema_version": "1.0", + **({"preprocessing": {"audio": audio_program}} if audio_program is not None else {}), + "pipeline": {"workflow": _publish_workflow_v1(workflow)}, + } + add_policy_components_to_workflow(metadata, pkg) + return metadata + + +def build_language_diffusion_pipeline_metadata( + pkg: Any, + *, + num_inference_steps: int, +) -> dict[str, Any]: + """Build a generic SSA workflow for a masked language-diffusion model.""" + if num_inference_steps < 1: + raise ValueError("num_inference_steps must be >= 1") + if len(pkg) != 1: + raise ValueError("language-diffusion workflow requires exactly one neural component") + denoiser_name, denoiser = next(iter(pkg.items())) + if len(denoiser.graph.inputs) != 1: + raise ValueError("language-diffusion denoiser requires exactly one token input") + + token_input = denoiser.graph.inputs[0] + logits_output = next( + (value for value in denoiser.graph.outputs if value.name == "logits"), + None, + ) + proposal_output = next( + (value for value in denoiser.graph.outputs if value.name == "proposed_tokens"), + None, + ) + if ( + token_input.dtype not in {ir.DataType.INT32, ir.DataType.INT64} + or token_input.shape is None + or len(token_input.shape) != 2 + or logits_output is None + or logits_output.shape is None + or len(logits_output.shape) != 3 + or proposal_output is None + or proposal_output.shape is None + or len(proposal_output.shape) != 2 + ): + raise ValueError( + "language-diffusion workflow requires token [B,T], logits [B,T,V], " + "and proposed_tokens [B,T] ports" + ) + + attach_policy_components(pkg, PolicyCapabilities(masked_update=True)) + + token_contract = _contract(token_input) + mask_contract = { + "dtype": "bool", + "rank": 2, + "shape": token_contract["shape"], + } + batch_dimension = token_contract["shape"][0] + batch_int = _request_aligned({"dtype": "int64", "rank": 1, "shape": [batch_dimension]}) + control_int = {"dtype": "int64", "rank": 1, "shape": [1]} + inputs = { + "request.input_ids": { + "contract": token_contract, + "role": { + "kind": "runtime", + "version": "1.0", + "role": "prompt_tokens", + }, + "source": {"kind": "request", "field": "prompt_tokens"}, + "required": True, + }, + "request.mask": { + "contract": mask_contract, + "role": {"kind": "opaque"}, + "source": {"kind": "application", "name": "masked_positions"}, + "required": True, + }, + "request.seed": { + "contract": batch_int, + "role": { + "kind": "runtime", + "version": "1.0", + "role": "seed", + }, + "source": {"kind": "request", "field": "seed"}, + "required": False, + "default": 0, + }, + "request.rng_offset": { + "contract": batch_int, + "role": {"kind": "opaque"}, + "source": {"kind": "application", "name": "rng_offset"}, + "required": False, + "default": 0, + }, + "request.max_iterations": { + "contract": control_int, + "role": { + "kind": "runtime", + "version": "1.0", + "role": "max_iterations", + }, + "source": {"kind": "request", "field": "max_iterations"}, + "required": False, + "default": num_inference_steps, + }, + "package.num_steps": { + "contract": batch_int, + "role": {"kind": "opaque"}, + "source": {"kind": "literal"}, + "required": False, + "default": num_inference_steps, + }, + } + + def denoiser_invoke(tokens: str, prefix: str) -> dict[str, Any]: + return _invoke( + denoiser_name, + {token_input.name: tokens}, + { + logits_output.name: f"{prefix}.logits", + proposal_output.name: f"{prefix}.proposal", + }, + ) + + def update_invoke( + tokens: str, + mask: str, + iteration: str, + offset: str, + logits: str, + proposal: str, + prefix: str, + effect_in: str, + effect_out: str, + ) -> dict[str, Any]: + return _invoke( + "masked_update", + { + "current_tokens": tokens, + "proposed_tokens": proposal, + "logits": logits, + "masked": mask, + "step": iteration, + "total_steps": "package.num_steps", + "seed": "request.seed", + "offset": offset, + }, + { + "next_state": f"{prefix}.tokens", + "next_mask": f"{prefix}.mask", + "next_offset": f"{prefix}.rng_offset", + "done": f"{prefix}.done", + "continue": f"{prefix}.continue", + }, + {"update": _effect(effect_in, effect_out)}, + ) + + setup = { + "kind": "sequence", + "nodes": [denoiser_invoke("request.input_ids", "denoiser.setup")], + } + body = { + "kind": "sequence", + "nodes": [ + update_invoke( + "state.tokens.body", + "state.mask.body", + "loop.iteration", + "state.rng_offset.body", + "state.logits.body", + "state.proposal.body", + "denoiser.body", + "update.0", + "update.1", + ), + { + "kind": "emit", + "value": "denoiser.body.tokens", + "output": "tokens", + "mode": "replace", + "effect_name": "emit", + "effect": _effect("emit.0", "emit.1"), + }, + denoiser_invoke("denoiser.body.tokens", "denoiser.body"), + ], + } + + state_specs = { + "tokens": (token_contract, "request.input_ids", "request.input_ids"), + "mask": (mask_contract, "request.mask", "request.mask"), + "rng_offset": (batch_int, "request.rng_offset", "request.rng_offset"), + "logits": ( + _contract(logits_output), + "denoiser.setup.logits", + "denoiser.setup.logits", + ), + "proposal": ( + _contract(proposal_output), + "denoiser.setup.proposal", + "denoiser.setup.proposal", + ), + } + state: dict[str, Any] = {} + carried: list[dict[str, Any]] = [] + initial_effects = {"update": "update.0", "emit": "emit.0"} + for name, (contract, initializer, current) in state_specs.items(): + effect_name = f"state:{name}" + initial_effects[effect_name] = f"{effect_name}.0" + state[name] = { + "contract": contract, + "scope": "invocation", + "initializer": initializer, + "recurrence": {"kind": "invariant"}, + } + carried.append( + { + "cell": name, + "current": current, + "body_input": f"state.{name}.body", + "body_output": f"denoiser.body.{name}", + "next": f"state.{name}.final", + "read_effect": _effect(f"{effect_name}.0", f"{effect_name}.read"), + "write_effect": _effect(f"{effect_name}.read", f"{effect_name}.1"), + } + ) + + workflow = { + "manifest": { + "capabilities": [ + "workflow_ssa", + "linear_effects", + "nested_control_flow", + "typed_emit", + "emit_valid_length", + "loop_induction_values", + ], + }, + "inputs": inputs, + "outputs": { + "tokens": { + "contract": token_contract, + "role": "tokens", + "stage": "pre_adapter", + } + }, + "components": {denoiser_name: _component(denoiser, "model.onnx")}, + "state": state, + "initial_effects": initial_effects, + "graph": { + "kind": "loop", + "setup": setup, + "body": body, + "condition": "denoiser.body.continue", + "max_iterations": "request.max_iterations", + "iteration": { + "value": "loop.iteration", + "contract": {"dtype": "int64", "rank": 1, "shape": [1]}, + }, + "carried": carried, + }, + } + metadata = { + "schema_version": "1.0", + "pipeline": {"workflow": _publish_workflow_v1(workflow)}, + } + add_policy_components_to_workflow(metadata, pkg) + return metadata + + +def write_decoder_workflow_metadata( + pkg: Any, + output_dir: str, + config: Any, + *, + sampler: str = "greedy", +) -> str: + """Write decoder workflow metadata and policy artifacts.""" + os.makedirs(output_dir, exist_ok=True) + metadata = build_decoder_workflow_metadata(pkg, config, sampler=sampler) + pkg.save_policy_components(output_dir) + add_adapter_service_to_metadata(metadata, pkg, output_dir) + path = os.path.join(output_dir, "inference_metadata.yaml") + with open(path, "w", encoding="utf-8") as handle: + _dump_yaml(metadata, handle) + return path + + +def write_speech_to_text_workflow_metadata( + pkg: Any, + output_dir: str, + config: Any, + *, + sampler: str = "greedy", + audio_preprocessing: dict[str, Any] | None = None, +) -> str: + """Write encoder-conditioned decode workflow metadata and policy artifacts.""" + os.makedirs(output_dir, exist_ok=True) + metadata = build_speech_to_text_workflow_metadata( + pkg, + config, + sampler=sampler, + audio_preprocessing=audio_preprocessing, + ) + pkg.save_policy_components(output_dir) + add_adapter_service_to_metadata(metadata, pkg, output_dir) + path = os.path.join(output_dir, "inference_metadata.yaml") + with open(path, "w", encoding="utf-8") as handle: + _dump_yaml(metadata, handle) + return path + + +def write_language_diffusion_workflow_metadata( + pkg: Any, + output_dir: str, + *, + num_inference_steps: int, +) -> str: + """Write masked language-diffusion workflow metadata and policy artifacts.""" + os.makedirs(output_dir, exist_ok=True) + metadata = build_language_diffusion_pipeline_metadata( + pkg, + num_inference_steps=num_inference_steps, + ) + pkg.save_policy_components(output_dir) + add_adapter_service_to_metadata(metadata, pkg, output_dir) + path = os.path.join(output_dir, "inference_metadata.yaml") + with open(path, "w", encoding="utf-8") as handle: + _dump_yaml(metadata, handle) + return path + + +_AUDIO_PREPROCESS_ABI = "onnx-genai.audio-preprocess" +_AUDIO_PREPROCESS_ABI_VERSION = "1" + + +def _audio_preprocess_component( + values_contract: dict[str, Any], + mask_contract: dict[str, Any], +) -> dict[str, Any]: + """Declare the versioned audio-preprocessing adapter component. + + The adapter turns request-supplied encoded audio bytes into the encoder's + waveform tensor and its sample-level validity mask. Its ports are declared + so the runtime can type-check the binding without knowing which model family + produced the package. + """ + return { + "implementation": { + "kind": "adapter", + "abi": _AUDIO_PREPROCESS_ABI, + "version": _AUDIO_PREPROCESS_ABI_VERSION, + }, + "ports": { + "inputs": { + "encoded": {"dtype": "uint8", "rank": 1, "shape": ["bytes"]}, + }, + "outputs": { + "input_values": values_contract, + "attention_mask": mask_contract, + }, + }, + "contract": { + "id": _AUDIO_PREPROCESS_ABI, + "version": _AUDIO_PREPROCESS_ABI_VERSION, + "bindings": { + "encoded": "encoded", + "input_values": "input_values", + "attention_mask": "attention_mask", + }, + }, + "effects": ["audio_preprocess"], + } + + +def _ctc_vocabulary(source: str | None, vocab_size: int) -> dict[str, Any]: + """Describe the class-id → string table used to render a transcript. + + The table is inlined when the source checkpoint's ``vocab.json`` is + reachable so the document is self-contained; otherwise the profile points at + the packaged tokenizer. + """ + vocabulary: dict[str, Any] = {"source": "tokenizer", "size": vocab_size} + path = _source_asset_path(source, "vocab.json") if source else None + if path is None: + return vocabulary + try: + with open(path, encoding="utf-8") as handle: + table = json.load(handle) + except (OSError, ValueError): + return vocabulary + if not isinstance(table, dict) or not table: + return vocabulary + tokens = [""] * (max(int(index) for index in table.values()) + 1) + for token, index in table.items(): + tokens[int(index)] = token + if len(tokens) != vocab_size: + return vocabulary + vocabulary = { + "source": "inline", + "size": vocab_size, + "tokens": tokens, + } + if "|" in table: + vocabulary["word_delimiter"] = "|" + ignored = [token for token in ("", "", "", "") if token in table] + if ignored: + vocabulary["ignored_tokens"] = ignored + return vocabulary + + +def build_ctc_asr_workflow_metadata( + pkg: Any, + config: Any, + *, + source: str | None = None, + artifact: str = "model.onnx", +) -> dict[str, Any]: + """Build one-file metadata for a non-generative CTC ASR package. + + A CTC acoustic model is frame-synchronous: the encoder runs exactly once and + emits one class distribution per frame. The workflow is therefore a plain + sequence with no loop and no carried state, and the transcript is recovered + by the ``transcription`` profile's decoding contract rather than by a + generation loop. + + Args: + pkg: The built :class:`ModelPackage`; must hold a single ``model``. + config: The resolved architecture config (supplies vocabulary size and + the CTC blank id). + source: HuggingFace model id or local directory used to inline the + decoding vocabulary. + artifact: Encoder artifact path relative to the package root. + + Returns: + A metadata document with ``preprocessing.audio``, ``profiles`` and a + single-step ``pipeline.workflow``. + """ + if "model" not in pkg: + raise ValueError("CTC ASR workflow requires a 'model' component") + model = pkg["model"] + + graph_inputs = {value.name: value for value in model.graph.inputs} + graph_outputs = {value.name: value for value in model.graph.outputs} + for required in ("input_values", "attention_mask"): + if required not in graph_inputs: + raise ValueError(f"CTC ASR encoder must declare input '{required}'") + if "logits" not in graph_outputs: + raise ValueError("CTC ASR encoder must declare output 'logits'") + has_frame_lengths = "frame_lengths" in graph_outputs + + values_contract = _contract(graph_inputs["input_values"]) + mask_contract = _contract(graph_inputs["attention_mask"]) + logits_contract = _contract(graph_outputs["logits"]) + + sample_rate = int(getattr(getattr(config, "audio", None), "sampling_rate", 0) or 16_000) + # Shape inference may leave the class axis unknown (or the whole shape + # absent), so fall back to the config rather than emitting a vocabulary + # whose declared size silently disagrees with the graph. + logits_shape = logits_contract.get("shape") or [] + inferred_classes = logits_shape[-1] if logits_shape else None + vocab_size = int(inferred_classes or getattr(config, "vocab_size", 0) or 0) + if vocab_size <= 0: + raise ValueError( + "CTC ASR metadata requires a vocabulary size; the 'logits' output " + "declares no static class axis and the config has no vocab_size" + ) + blank_id = int(getattr(config, "pad_token_id", 0) or 0) + + workflow_outputs = { + "logits": { + "contract": logits_contract, + "role": "tensor", + "stage": "post_adapter", + } + } + emit_nodes = [ + { + "kind": "emit", + "value": "encoder.logits", + "output": "logits", + "mode": "replace", + } + ] + profile_outputs = {"logits": "logits"} + if has_frame_lengths: + workflow_outputs["frame_lengths"] = { + "contract": _contract(graph_outputs["frame_lengths"]), + "role": "tensor", + "stage": "post_adapter", + } + emit_nodes.append( + { + "kind": "emit", + "value": "encoder.frame_lengths", + "output": "frame_lengths", + "mode": "replace", + } + ) + profile_outputs["frame_lengths"] = "frame_lengths" + + workflow = { + "manifest": { + "adapter_abis": {_AUDIO_PREPROCESS_ABI: _AUDIO_PREPROCESS_ABI_VERSION}, + "capabilities": ["workflow_ssa", "linear_effects", "typed_emit"], + }, + "effects": { + # Both steps are pure functions of their inputs: decoding audio and + # running the encoder observe nothing external, so replay is always + # safe. Speculation safety is irrelevant here because a CTC + # workflow has no speculative region, but it is declared explicitly + # rather than left to a default. + "audio_preprocess": {"retry": "pure", "speculation_safety": {"kind": "clonable"}}, + "encode": {"retry": "pure", "speculation_safety": {"kind": "clonable"}}, + }, + "inputs": { + "request.audio": { + "contract": {"dtype": "uint8", "rank": 1, "shape": ["bytes"]}, + "role": {"kind": "runtime", "version": "1.0", "role": "media"}, + "source": {"kind": "request", "field": "media"}, + "required": True, + } + }, + "outputs": workflow_outputs, + "components": { + "audio_preprocess": _audio_preprocess_component(values_contract, mask_contract), + "encoder": _component(model, artifact, effects=("encode",)), + }, + "initial_effects": { + "audio_preprocess": "audio_preprocess.0", + "encode": "encode.0", + }, + "graph": { + "kind": "sequence", + "nodes": [ + _invoke( + "audio_preprocess", + {"encoded": "request.audio"}, + { + "input_values": "audio.input_values", + "attention_mask": "audio.attention_mask", + }, + ), + _invoke( + "encoder", + { + "input_values": "audio.input_values", + "attention_mask": "audio.attention_mask", + }, + { + name: f"encoder.{name}" + for name in ("logits", "frame_lengths") + if name in graph_outputs + }, + ), + *emit_nodes, + ], + }, + } + + decoding: dict[str, Any] = { + "kind": "ctc", + "blank_id": blank_id, + "collapse_repeats": True, + "time_axis": 1, + "class_axis": 2, + "vocabulary": _ctc_vocabulary(source, vocab_size), + } + if has_frame_lengths: + decoding["lengths"] = "frame_lengths" + + # A feature extractor that reduces over the padded time axis makes a row's + # values depend on the width of the batch it was padded into. The fact is + # recorded by the task on the built graph; when nobody stated it we leave + # the field absent rather than claim rows are independent. + normalization = getattr(config, "feat_extract_norm", None) + if normalization == "group": + batch_invariance = "padding_sensitive" + elif normalization == "layer": + batch_invariance = "row_independent" + else: + recorded = model.metadata_props.get(BATCH_PADDING_SENSITIVE_KEY) + batch_invariance = ( + None + if recorded is None + else ("padding_sensitive" if recorded == "true" else "row_independent") + ) + + profile: dict[str, Any] = { + "kind": "transcription", + "version": "1.0", + "requirement": "required", + "outputs": profile_outputs, + "decoding": decoding, + } + # The claim is only checkable when the package also publishes per-row + # lengths; without them a reader cannot isolate a row's valid region. + if batch_invariance == "row_independent" or has_frame_lengths: + if batch_invariance is not None: + profile["batch_invariance"] = batch_invariance + + return { + "schema_version": "v1", + "preprocessing": { + "audio": { + "transforms": [ + {"op": "decode", "outputs": ["samples"]}, + {"op": "resample", "sample_rate": sample_rate}, + {"op": "downmix", "channels": 1}, + {"op": "zero_mean_unit_variance", "epsilon": 1e-7}, + { + "op": "pad", + "mode": "right", + "pad_value": 0.0, + "outputs": ["values", "sample_mask"], + }, + ], + "outputs": [ + { + "name": "input_values", + "source": "values", + "content": "waveform", + "dtype": values_contract["dtype"], + "rank": values_contract["rank"], + }, + { + "name": "attention_mask", + "source": "sample_mask", + "content": "validity_mask", + "dtype": mask_contract["dtype"], + "rank": mask_contract["rank"], + }, + ], + } + }, + "profiles": { + "transcription": profile, + }, + "pipeline": {"workflow": _publish_workflow_v1(workflow)}, + } + + +def write_ctc_asr_workflow_metadata( + pkg: Any, + output_dir: str, + config: Any, + *, + source: str | None = None, +) -> str: + """Write one-file CTC ASR metadata into *output_dir*.""" + os.makedirs(output_dir, exist_ok=True) + metadata = build_ctc_asr_workflow_metadata(pkg, config, source=source) + path = os.path.join(output_dir, "inference_metadata.yaml") + with open(path, "w", encoding="utf-8") as handle: + _dump_yaml(metadata, handle) + return path + + +_ENCODER_EMBEDDING_INPUT_ROLES: dict[str, str] = { + "input_ids": "prompt_tokens", + "attention_mask": "attention_mask", + "token_type_ids": "token_type_ids", + "position_ids": "position_ids", +} + + +def build_encoder_embedding_workflow_metadata( + pkg: Any, + config: Any = None, + *, + artifact: str = "model.onnx", +) -> dict[str, Any]: + """Build one-file metadata for a bidirectional encoder that emits embeddings. + + An encoder such as BERT, ESM-2 or ProtBert is not generative: it reads the + whole sequence at once and returns one hidden vector per position. There is + no next-token step, no KV cache and nothing to sample, so the workflow is a + plain sequence with a single invocation and no carried state. Describing + such a package with decoder metadata would publish a generation loop that + the artifact cannot execute -- it has no ``logits`` port to sample and + re-feeding its output would be meaningless -- so this builder exists to + keep the metadata a truthful description of the graph. + + Args: + pkg: The built :class:`ModelPackage`; must hold a single ``model``. + config: The resolved architecture config. Unused today; accepted so the + dispatch site can pass it uniformly with the other builders. + artifact: Encoder artifact path relative to the package root. + + Returns: + A metadata document with an ``embedding`` profile and a single-step + ``pipeline.workflow``. + """ + del config + if "model" not in pkg: + raise ValueError("encoder embedding workflow requires a 'model' component") + model = pkg["model"] + + graph_inputs = {str(value.name): value for value in model.graph.inputs} + graph_outputs = {str(value.name): value for value in model.graph.outputs} + if "input_ids" not in graph_inputs: + raise ValueError("encoder embedding graph must declare input 'input_ids'") + if "last_hidden_state" not in graph_outputs: + raise ValueError("encoder embedding graph must declare output 'last_hidden_state'") + + # Declare exactly the ports the artifact exposes. ESM-2 has no token type + # embedding, so its graph carries no ``token_type_ids``; BERT-family + # encoders do. Reading the graph rather than the task signature is what + # keeps the two packages describable by one builder. + workflow_inputs: dict[str, Any] = {} + invoke_inputs: dict[str, str] = {} + for name, role in _ENCODER_EMBEDDING_INPUT_ROLES.items(): + if name not in graph_inputs: + continue + input_role: dict[str, Any] + input_source: dict[str, Any] + if role == "prompt_tokens": + input_role = {"kind": "runtime", "version": "1.0", "role": role} + input_source = {"kind": "request"} + else: + # The portable runtime-role vocabulary intentionally does not + # encode architecture-specific auxiliary graph inputs. + input_role = {"kind": "opaque"} + input_source = {"kind": "application", "name": f"request.{name}"} + declaration: dict[str, Any] = { + "contract": _contract(graph_inputs[name]), + "role": input_role, + "source": input_source, + # Every port here is a graph input of a single-invocation workflow, + # so a runtime must bind all of them; none is optional. + "required": True, + } + workflow_inputs[f"request.{name}"] = declaration + invoke_inputs[name] = f"request.{name}" + + workflow_outputs: dict[str, Any] = {} + invoke_outputs: dict[str, str] = {} + emit_nodes: list[dict[str, Any]] = [] + profile_outputs: dict[str, str] = {} + for name in ("last_hidden_state", "pooler_output"): + if name not in graph_outputs: + continue + workflow_outputs[name] = { + "contract": _contract(graph_outputs[name]), + "role": "tensor", + "stage": "post_adapter", + } + invoke_outputs[name] = f"encoder.{name}" + emit_nodes.append( + { + "kind": "emit", + "value": f"encoder.{name}", + "output": name, + "mode": "replace", + } + ) + profile_outputs[name] = name + + workflow = { + "manifest": { + "capabilities": ["workflow_ssa", "linear_effects", "typed_emit"], + }, + "effects": { + # One pure call: the encoder observes nothing outside its inputs, + # so a retry replays it exactly and a speculative clone is safe. + "encode": {"retry": "pure", "speculation_safety": {"kind": "clonable"}}, + }, + "inputs": workflow_inputs, + "outputs": workflow_outputs, + "components": {"encoder": _component(model, artifact, effects=("encode",))}, + "initial_effects": {"encode": "encode.0"}, + "graph": { + "kind": "sequence", + "nodes": [ + _invoke("encoder", invoke_inputs, invoke_outputs), + *emit_nodes, + ], + }, + } + + profile: dict[str, Any] = { + "kind": "embedding", + "version": "1.0", + "requirement": "required", + "outputs": profile_outputs, + } + if "attention_mask" in graph_inputs: + # The portable pooling profile describes the supported sequence + # reduction; the application remains responsible for supplying the + # architecture-specific attention mask to the graph. + profile["pooling"] = { + "kind": "mean", + "axis": 1, + "normalize": False, + } + profile["batch_invariance"] = "row_independent" + + return { + "schema_version": "v1", + "profiles": {"embedding": profile}, + "pipeline": {"workflow": _publish_workflow_v1(workflow)}, + } + + +def write_encoder_embedding_workflow_metadata( + pkg: Any, + output_dir: str, + config: Any = None, +) -> str: + """Write one-file encoder-embedding metadata into *output_dir*.""" + os.makedirs(output_dir, exist_ok=True) + metadata = build_encoder_embedding_workflow_metadata(pkg, config) + path = os.path.join(output_dir, "inference_metadata.yaml") + with open(path, "w", encoding="utf-8") as handle: + _dump_yaml(metadata, handle) + return path + + +def _ir_dtype(value: ir.Value) -> ir.DataType: + """Return the exact ONNX element type of a graph port.""" + dtype = value.dtype + if dtype is None: + raise ValueError(f"port {value.name!r} has no element type") + return dtype + + +def _duplex_delays(config: Any) -> list[int]: + """Read the per-stream delay pattern of a Moshi-family full-duplex LM.""" + delays = getattr(config, "delays", None) + if delays is None: + raise ValueError("full-duplex workflow requires a delay pattern on the config") + return [int(value) for value in delays] + + +def build_full_duplex_workflow_metadata(pkg: Any, config: Any) -> dict[str, Any]: + """Build typed SSA metadata for one event of a full-duplex speech workflow. + + Full-duplex speech-to-speech models (Moshi, PersonaPlex) consume and produce + audio simultaneously at a fixed frame rate. One invocation of this workflow is + exactly one frame: it accepts one packed audio chunk plus a session ID, and it + emits at most one packed audio chunk. Everything that must survive between + frames -- the temporal transformer KV cache, the delay ring buffer, the frame + offset, and the codec prefixes -- is declared as ``session``-scoped state with + a ``session`` release boundary and an exclusive session lease, so a runtime can + resume the conversation on the next invocation without replaying history. + + The acoustic (depformer) loop is the only loop in the graph. Its KV cache is + ``invocation``-scoped because the upstream model resets it on every frame. + + Components: + + * ``encoder`` / ``decoder`` -- the Mimi-style neural audio codec. + * ``temporal`` -- the frame-rate transformer over the interleaved token frame. + * ``depformer`` -- the per-frame acoustic transformer over ``num_streams`` substeps. + * ``frame_assemble`` / ``frame_commit`` -- the delay ring-buffer bookkeeping. + * ``teacher_select`` -- prefers an externally supplied token over a sampled one. + * ``token_sampler`` -- greedy sampling shared by the text and acoustic heads. + """ + required = {"encoder", "decoder", "temporal", "depformer"} + missing = required - set(pkg.keys()) + if missing: + raise ValueError( + f"full-duplex workflow requires components {sorted(required)}; " + f"missing {sorted(missing)}" + ) + encoder = pkg["encoder"] + decoder = pkg["decoder"] + temporal = pkg["temporal"] + depformer = pkg["depformer"] + + delays = _duplex_delays(config) + channels = len(delays) + max_delay = max(delays) + cache_length = max_delay + 3 + num_streams = int(getattr(config, "dep_q", channels // 2)) + audio_streams = int(getattr(config, "n_q", num_streams)) + frame_size = int(getattr(config, "frame_size", 1920)) + initial_tokens = [int(getattr(config, "text_initial_token_id", 32000))] + [ + int(getattr(config, "initial_token_id", 2048)) + ] * (channels - 1) + + waveform_input = _find_port(encoder.graph.inputs, "waveform") + codes_output = _find_port(encoder.graph.outputs, "codes") + codes_input = _find_port(decoder.graph.inputs, "codes") + waveform_output = _find_port(decoder.graph.outputs, "waveform") + input_frame = _find_port(temporal.graph.inputs, "input_frame") + temporal_mask = _find_port(temporal.graph.inputs, "attention_mask") + temporal_position = _find_port(temporal.graph.inputs, "position_ids") + hidden = _find_port(temporal.graph.outputs, "hidden") + text_logits = _find_port(temporal.graph.outputs, "text_logits") + dep_hidden = _find_port(depformer.graph.inputs, "hidden") + dep_prev = _find_port(depformer.graph.inputs, "prev_token") + dep_index = _find_port(depformer.graph.inputs, "substep_index") + dep_logits = _find_port(depformer.graph.outputs, "logits") + ports = ( + waveform_input, + codes_output, + codes_input, + waveform_output, + input_frame, + temporal_mask, + temporal_position, + hidden, + text_logits, + dep_hidden, + dep_prev, + dep_index, + dep_logits, + ) + if any(port is None for port in ports): + raise ValueError("full-duplex workflow is missing a required component port") + temporal_caches = _model_cache_pairs(temporal) + depformer_caches = _model_cache_pairs(depformer) + if not temporal_caches: + raise ValueError("full-duplex workflow requires a temporal KV cache") + if not depformer_caches: + raise ValueError("full-duplex workflow requires a depformer KV cache") + + attach_policy_components(pkg, PolicyCapabilities(sampler="greedy")) + pkg.add_policy_component( + "frame_assemble", + build_duplex_frame_assemble(channels=channels, cache_length=cache_length), + ) + pkg.add_policy_component( + "frame_commit", + build_duplex_frame_commit( + channels=channels, cache_length=cache_length, max_delay=max_delay + ), + ) + pkg.add_policy_component("teacher_select", build_duplex_teacher_select(channels=channels)) + pkg.add_policy_component( + "frame_update", build_code_frame_update(channels, scalar_index=True) + ) + pkg.add_policy_component( + "waveform_append", build_duplex_waveform_append(dtype=_ir_dtype(waveform_input)) + ) + pkg.add_policy_component("codes_append", build_duplex_stream_append(streams=audio_streams)) + pkg.add_policy_component("codes_tail", build_duplex_stream_tail(streams=audio_streams)) + pkg.add_policy_component( + "chunk_tail", + build_duplex_stream_tail(streams=1, dtype=_ir_dtype(waveform_output)), + ) + + batch = _contract(input_frame)["shape"][0] + request_aligned = {"kind": "request_aligned", "axis": 0} + control_int = {"dtype": "int64", "rank": 0, "shape": []} + scalar_bool = {"dtype": "bool", "rank": 0, "shape": []} + batch_bool = { + "dtype": "bool", + "rank": 1, + "shape": [batch], + "batch_layout": request_aligned, + } + batch_int = { + "dtype": "int64", + "rank": 1, + "shape": [batch], + "batch_layout": request_aligned, + } + loop_flag = {"dtype": "bool", "rank": 1, "shape": [1]} + frame_contract = {"dtype": "int64", "rank": 2, "shape": [batch, channels]} + ring_contract = {"dtype": "int64", "rank": 3, "shape": [batch, channels, cache_length]} + ring_flags = {"dtype": "bool", "rank": 3, "shape": [batch, channels, cache_length]} + + inputs: dict[str, Any] = { + "request.audio_chunk": { + "contract": _contract(waveform_input), + "role": {"kind": "runtime", "version": "1.0", "role": "media"}, + "source": {"kind": "request", "field": "media"}, + "required": True, + }, + "request.session_id": { + "contract": {"dtype": "int64", "rank": 1, "shape": [batch]}, + "role": {"kind": "runtime", "version": "1.0", "role": "session_id"}, + "source": {"kind": "request", "field": "session_id"}, + "required": True, + }, + "package.stream_tokens": { + "contract": frame_contract, + "role": {"kind": "opaque"}, + "source": {"kind": "literal"}, + "required": False, + # -1 marks "this stream has nothing to contribute this frame", so the + # broadcast default is a frame in which the model predicts everything. + "default": -1, + }, + "package.frames_per_invocation": { + "contract": control_int, + "role": {"kind": "opaque"}, + "source": {"kind": "literal"}, + "required": False, + # One event per invocation by default; a runtime that batches several + # codec frames into a single call raises this without changing the graph. + "default": 1, + }, + "package.true_batch": { + "contract": batch_bool, + "role": {"kind": "opaque"}, + "source": {"kind": "literal"}, + "required": False, + "default": True, + }, + "package.false_batch": { + "contract": batch_bool, + "role": {"kind": "opaque"}, + "source": {"kind": "literal"}, + "required": False, + "default": False, + }, + "package.zero_batch": { + "contract": batch_int, + "role": {"kind": "opaque"}, + "source": {"kind": "literal"}, + "required": False, + "default": 0, + }, + "package.one_batch": { + "contract": batch_int, + "role": {"kind": "opaque"}, + "source": {"kind": "literal"}, + "required": False, + "default": 1, + }, + "package.delays": { + "contract": {"dtype": "int64", "rank": 1, "shape": [channels]}, + "role": {"kind": "opaque"}, + "source": {"kind": "literal"}, + "required": False, + "default": delays, + }, + "package.initial_tokens": { + "contract": {"dtype": "int64", "rank": 1, "shape": [channels]}, + "role": {"kind": "opaque"}, + "source": {"kind": "literal"}, + "required": False, + "default": initial_tokens, + }, + "package.num_streams": { + "contract": control_int, + "role": {"kind": "opaque"}, + "source": {"kind": "literal"}, + "required": False, + "default": num_streams, + }, + "package.one_frame": { + "contract": control_int, + "role": {"kind": "opaque"}, + "source": {"kind": "literal"}, + "required": False, + "default": 1, + }, + "package.frame_size": { + "contract": control_int, + "role": {"kind": "opaque"}, + "source": {"kind": "literal"}, + "required": False, + "default": frame_size, + }, + "package.false": { + "contract": scalar_bool, + "role": {"kind": "opaque"}, + "source": {"kind": "literal"}, + "required": False, + "default": False, + }, + } + + session_lease = {"policy": "exclusive", "optimistic_metadata_version": False} + + def session_cell( + contract: dict[str, Any], + initializer: str, + recurrence: dict[str, Any], + *, + release: str = "session", + ) -> dict[str, Any]: + return { + "contract": contract, + "scope": "session", + "initializer": initializer, + "recurrence": recurrence, + "management": "runtime", + "release_boundary": release, + "session": session_lease, + } + + invariant = {"kind": "invariant"} + state: dict[str, Any] = { + "token_cache": session_cell(ring_contract, "package.token_cache_init", invariant), + "token_provided": session_cell(ring_flags, "package.token_provided_init", invariant), + "offset": session_cell(control_int, "package.offset_init", invariant), + "attention_mask": session_cell( + { + "dtype": _contract(temporal_mask)["dtype"], + "rank": 2, + "shape": [batch, "context"], + }, + "package.attention_mask_init", + { + "kind": "growing", + "axis": 1, + "increment": "package.one_frame", + "max": "package.context_limit", + }, + ), + "position_ids": session_cell( + _contract(temporal_position), "package.position_ids_init", invariant + ), + "user_waveform": session_cell( + { + "dtype": _contract(waveform_input)["dtype"], + "rank": 3, + "shape": [batch, 1, "user_samples"], + }, + "package.user_waveform_init", + { + "kind": "growing", + "axis": 2, + "increment": "package.frame_size", + "max": "package.codec_prefix_limit", + }, + release="invocation", + ), + "agent_codes": session_cell( + { + "dtype": _contract(codes_input)["dtype"], + "rank": 3, + "shape": [batch, audio_streams, "agent_frames"], + }, + "package.agent_codes_init", + { + "kind": "growing", + "axis": 2, + "increment": "package.one_frame", + "max": "package.codec_frame_limit", + }, + release="invocation", + ), + } + inputs["package.context_limit"] = { + "contract": control_int, + "role": {"kind": "opaque"}, + "source": {"kind": "literal"}, + "required": False, + "default": int(getattr(config, "context", 3000)), + } + inputs["package.codec_prefix_limit"] = { + "contract": control_int, + "role": {"kind": "opaque"}, + "source": {"kind": "literal"}, + "required": False, + "default": int(getattr(config, "context", 3000)) * frame_size, + } + inputs["package.codec_frame_limit"] = { + "contract": control_int, + "role": {"kind": "opaque"}, + "source": {"kind": "literal"}, + "required": False, + "default": int(getattr(config, "context", 3000)), + } + zero_default = {"int64": 0, "int32": 0, "bool": False} + for name, initial in ( + ("package.token_cache_init", ring_contract), + ("package.token_provided_init", ring_flags), + ("package.attention_mask_init", state["attention_mask"]["contract"]), + ("package.position_ids_init", _contract(temporal_position)), + ("package.user_waveform_init", state["user_waveform"]["contract"]), + ("package.agent_codes_init", state["agent_codes"]["contract"]), + ): + inputs[name] = { + "contract": initial, + "role": {"kind": "opaque"}, + "source": {"kind": "literal"}, + "required": False, + # A scalar default broadcasts across the whole contract, so an empty + # session starts from an all-zero (or all-false) tensor. + "default": zero_default.get(initial["dtype"], 0.0), + } + inputs["package.offset_init"] = { + "contract": control_int, + "role": {"kind": "opaque"}, + "source": {"kind": "literal"}, + "required": False, + "default": 1, + } + + for index, (past, present) in enumerate(temporal_caches): + # A runtime that owns the cache must be able to permute and compact rows, + # so a service-bound cell has to declare its request axis explicitly. + cache_contract = {**_contract(past), "batch_layout": request_aligned} + state[f"temporal_cache_{index}"] = session_cell( + cache_contract, + f"package.temporal_cache_{index}_init", + { + # The temporal cache is a sliding context window: it grows one + # frame per invocation but never past ``context``, so it is + # bounded rather than unboundedly growing. + "kind": "bounded", + "axis": 2, + "max": "package.context_limit", + }, + ) + state[f"temporal_cache_{index}"]["service_group"] = "temporal_cache" + state[f"temporal_cache_{index}"]["class"] = "semantic" + inputs[f"package.temporal_cache_{index}_init"] = { + "contract": cache_contract, + "role": {"kind": "opaque"}, + "source": {"kind": "literal"}, + "required": False, + "default": 0.0, + } + del present + state["temporal_cache_lengths"] = { + "contract": batch_int, + "class": "semantic", + "scope": "session", + "initializer": "package.zero_batch", + "recurrence": invariant, + "session": session_lease, + } + state["active"] = { + "contract": batch_bool, + "class": "semantic", + "scope": "invocation", + "initializer": "package.true_batch", + "recurrence": invariant, + } + state["done"] = { + "contract": batch_bool, + "class": "semantic", + "scope": "invocation", + "initializer": "package.false_batch", + "recurrence": invariant, + } + state["accepted_len"] = { + "contract": batch_int, + "class": "semantic", + "scope": "invocation", + "initializer": "package.zero_batch", + "recurrence": invariant, + } + for index, (past, _) in enumerate(depformer_caches): + state[f"depformer_cache_{index}"] = { + "contract": _contract(past), + "scope": "invocation", + "initializer": f"package.depformer_cache_{index}_init", + "recurrence": { + "kind": "growing", + "axis": 2, + "increment": "package.one_frame", + "max": "package.num_streams", + }, + } + inputs[f"package.depformer_cache_{index}_init"] = { + "contract": _contract(past), + "role": {"kind": "opaque"}, + "source": {"kind": "literal"}, + "required": False, + "default": 0.0, + } + state["acoustic_frame"] = { + "contract": frame_contract, + "scope": "invocation", + "initializer": "duplex.target_frame", + "recurrence": invariant, + } + state["prev_token"] = { + "contract": {"dtype": "int64", "rank": 1, "shape": [batch]}, + "scope": "invocation", + "initializer": "duplex.text_token", + "recurrence": invariant, + } + + depformer_inputs = { + dep_hidden.name: "duplex.hidden", + dep_prev.name: "duplex.prev_token_slot", + dep_index.name: "duplex.substep", + **{ + past.name: f"state.depformer_cache_{index}.body" + for index, (past, _) in enumerate(depformer_caches) + }, + } + depformer_outputs = { + dep_logits.name: "duplex.acoustic_logits", + **{ + present.name: f"duplex.depformer_cache_{index}" + for index, (_, present) in enumerate(depformer_caches) + }, + } + + acoustic_loop = { + "kind": "loop", + "setup": {"kind": "sequence", "nodes": []}, + "body": { + "kind": "sequence", + "nodes": [ + _invoke( + "token_to_slot", + {"token": "state.prev_token.body"}, + {"slot": "duplex.prev_token_slot"}, + ), + _invoke("depformer", depformer_inputs, depformer_outputs), + _invoke( + "last_acoustic_logits", + {"logits": "duplex.acoustic_logits"}, + {"last_logits": "duplex.acoustic_last"}, + ), + _invoke( + "token_sampler", + {"logits": "duplex.acoustic_last"}, + {"token": "duplex.acoustic_sampled"}, + ), + _invoke( + "frame_update", + { + "frame_codes": "state.acoustic_frame.body", + "token": "duplex.acoustic_sampled", + "index": "duplex.acoustic_stream", + }, + {"next_frame": "duplex.acoustic_frame_next"}, + ), + _invoke( + "teacher_select", + { + "target": "duplex.target", + "target_provided": "duplex.target_provided", + "sampled": "duplex.acoustic_sampled", + "index": "duplex.acoustic_stream", + }, + {"token": "duplex.acoustic_prev_next"}, + ), + ], + }, + "condition": "package.loop_active", + "max_iterations": "package.num_streams", + "iteration": {"value": "duplex.substep", "contract": control_int}, + "carried": [ + { + "cell": "acoustic_frame", + "current": "duplex.target_frame", + "body_input": "state.acoustic_frame.body", + "body_output": "duplex.acoustic_frame_next", + "next": "duplex.acoustic_frame_final", + }, + { + "cell": "prev_token", + "current": "duplex.text_token", + "body_input": "state.prev_token.body", + "body_output": "duplex.acoustic_prev_next", + "next": "duplex.prev_token_final", + }, + *[ + { + "cell": f"depformer_cache_{index}", + "current": f"package.depformer_cache_{index}_init", + "body_input": f"state.depformer_cache_{index}.body", + "body_output": f"duplex.depformer_cache_{index}", + "next": f"duplex.depformer_cache_{index}_final", + } + for index in range(len(depformer_caches)) + ], + ], + } + inputs["package.loop_active"] = { + "contract": loop_flag, + "role": {"kind": "opaque"}, + "source": {"kind": "literal"}, + "required": False, + "default": True, + } + inputs["package.acoustic_stream_offset"] = { + "contract": control_int, + "role": {"kind": "opaque"}, + "source": {"kind": "literal"}, + "required": False, + "default": 1, + } + acoustic_loop["body"]["nodes"].insert( + 0, + _invoke( + "stream_index", + {"left": "duplex.substep", "right": "package.acoustic_stream_offset"}, + {"total": "duplex.acoustic_stream"}, + ), + ) + pkg.add_policy_component("stream_index", build_scalar_integer_add()) + pkg.add_policy_component("last_text_logits", build_last_token_logits()) + pkg.add_policy_component("last_acoustic_logits", build_last_token_logits()) + pkg.add_policy_component("token_to_slot", build_token_to_slot()) + + temporal_inputs = { + input_frame.name: "duplex.input_frame", + temporal_mask.name: "state.attention_mask.body", + temporal_position.name: "state.position_ids.body", + **{ + past.name: f"state.temporal_cache_{index}.body" + for index, (past, _) in enumerate(temporal_caches) + }, + } + temporal_outputs = { + hidden.name: "duplex.hidden", + text_logits.name: "duplex.text_logits", + **{ + present.name: f"duplex.temporal_cache_{index}" + for index, (_, present) in enumerate(temporal_caches) + }, + } + + graph = { + "kind": "sequence", + "nodes": [], + } + frame_body = graph["nodes"] + frame_body.extend( + [ + # 1. packed audio in: grow the codec prefix and encode the newest frame. + _invoke( + "waveform_append", + {"prefix": "state.user_waveform.body", "chunk": "request.audio_chunk"}, + {"next_prefix": "duplex.user_waveform_next"}, + ), + _invoke( + "encoder", + {waveform_input.name: "duplex.user_waveform_next"}, + {codes_output.name: "duplex.user_codes"}, + ), + _invoke( + "codes_tail", + {"prefix": "duplex.user_codes", "count": "package.one_frame"}, + {"tail": "duplex.user_frame_codes"}, + ), + # 2. delay ring buffer: write every supplied stream, read one frame. + _invoke( + "user_stream_merge", + { + "frame_codes": "package.stream_tokens", + "codes": "duplex.user_frame_codes", + }, + {"stream_tokens": "duplex.stream_tokens"}, + ), + _invoke( + "frame_assemble", + { + "token_cache": "state.token_cache.body", + "token_provided": "state.token_provided.body", + "offset": "state.offset.body", + "stream_tokens": "duplex.stream_tokens", + "delays": "package.delays", + "initial_tokens": "package.initial_tokens", + }, + { + "next_token_cache": "duplex.cache_assembled", + "next_token_provided": "duplex.provided_assembled", + "input_frame": "duplex.input_frame", + "target": "duplex.target", + "target_provided": "duplex.target_provided", + }, + ), + # 3. frame-rate temporal transformer over the interleaved frame. + _invoke("temporal", temporal_inputs, temporal_outputs), + _invoke( + "last_text_logits", + {"logits": "duplex.text_logits"}, + {"last_logits": "duplex.text_last"}, + ), + _invoke( + "token_sampler", + {"logits": "duplex.text_last"}, + {"token": "duplex.text_sampled"}, + ), + _invoke( + "teacher_select", + { + "target": "duplex.target", + "target_provided": "duplex.target_provided", + "sampled": "duplex.text_sampled", + "index": "package.text_stream", + }, + {"token": "duplex.text_token"}, + ), + _invoke( + "target_frame", + {"target": "duplex.target"}, + {"frame": "duplex.target_frame"}, + ), + # 4. acoustic loop: one substep per acoustic stream, KV reset per frame. + acoustic_loop, + _invoke( + "text_frame_update", + { + "frame_codes": "duplex.acoustic_frame_final", + "token": "duplex.text_token", + "index": "package.text_stream", + }, + {"next_frame": "duplex.completed_frame"}, + ), + # 5. commit the frame and undo the per-stream delays. + _invoke( + "frame_commit", + { + "token_cache": "duplex.cache_assembled", + "token_provided": "duplex.provided_assembled", + "offset": "state.offset.body", + "frame": "duplex.completed_frame", + "delays": "package.delays", + }, + { + "next_token_cache": "duplex.cache_committed", + "next_token_provided": "duplex.provided_committed", + "out_frame": "duplex.out_frame", + "next_offset": "duplex.next_offset", + "emit": "duplex.emit", + }, + ), + # 6. packed audio out: grow the agent code prefix and decode it. + _invoke( + "agent_frame_select", + {"frame": "duplex.out_frame"}, + {"codes": "duplex.agent_frame"}, + ), + _invoke( + "codes_append", + {"prefix": "state.agent_codes.body", "frame": "duplex.agent_frame"}, + {"next_prefix": "duplex.agent_codes_next"}, + ), + _invoke( + "decoder", + {codes_input.name: "duplex.agent_codes_next"}, + {waveform_output.name: "duplex.agent_waveform"}, + ), + _invoke( + "chunk_tail", + {"prefix": "duplex.agent_waveform", "count": "package.frame_size"}, + {"tail": "duplex.agent_chunk"}, + ), + _invoke( + "step_update", + { + "attention_mask": "state.attention_mask.body", + "position_ids": "state.position_ids.body", + }, + { + "next_attention_mask": "duplex.attention_mask_next", + "next_position_ids": "duplex.position_ids_next", + }, + ), + { + "kind": "emit", + "value": "duplex.agent_chunk", + "output": "audio_chunk", + "mode": "event", + "when": "duplex.emit", + }, + ] + ) + # Session-resident cells are only readable inside a loop carry, so the whole + # frame body is a loop. ``package.frames_per_invocation`` defaults to 1, which + # makes one invocation exactly one duplex event; a runtime that hands several + # codec frames to a single call raises it without changing the graph. + session_carries = [ + # A duplex conversation has no generation-side termination predicate: it + # runs until the session lease is released, so liveness is invariant. + # These are carried first because the temporal cache recurrence quotes + # ``accepted_len`` as its per-row growth increment. + ("active", "package.true_batch", "package.true_batch"), + ("done", "package.false_batch", "package.false_batch"), + ("accepted_len", "package.zero_batch", "package.one_batch"), + ("token_cache", "package.token_cache_init", "duplex.cache_committed"), + ("token_provided", "package.token_provided_init", "duplex.provided_committed"), + ("offset", "package.offset_init", "duplex.next_offset"), + ("attention_mask", "package.attention_mask_init", "duplex.attention_mask_next"), + ("position_ids", "package.position_ids_init", "duplex.position_ids_next"), + ("user_waveform", "package.user_waveform_init", "duplex.user_waveform_next"), + ("agent_codes", "package.agent_codes_init", "duplex.agent_codes_next"), + *[ + ( + f"temporal_cache_{index}", + f"package.temporal_cache_{index}_init", + f"duplex.temporal_cache_{index}", + ) + for index in range(len(temporal_caches)) + ], + ("temporal_cache_lengths", "package.zero_batch", "duplex.temporal_cache_lengths"), + ] + frame_body.append( + _invoke( + "cache_length_update", + { + "left": "state.temporal_cache_lengths.body", + "right": "package.one_batch", + }, + {"total": "duplex.temporal_cache_lengths"}, + ) + ) + pkg.add_policy_component("cache_length_update", build_integer_add()) + frame_loop = { + "kind": "loop", + "setup": {"kind": "sequence", "nodes": []}, + "body": {"kind": "sequence", "nodes": frame_body}, + "condition": "package.loop_active", + "max_iterations": "package.frames_per_invocation", + "iteration": {"value": "duplex.frame_index", "contract": batch_int}, + "carried": [ + { + "cell": cell, + "current": current, + "body_input": f"state.{cell}.body", + "body_output": produced, + "next": f"duplex.{cell}_final", + } + for cell, current, produced in session_carries + ], + } + graph = {"kind": "sequence", "nodes": [frame_loop]} + + pkg.add_policy_component( + "user_stream_merge", + build_duplex_user_stream_merge(channels=channels, streams=audio_streams), + ) + pkg.add_policy_component("target_frame", build_duplex_cell_to_frame(channels=channels)) + pkg.add_policy_component( + "agent_frame_select", + build_duplex_agent_frame_select(channels=channels, streams=audio_streams), + ) + pkg.add_policy_component( + "text_frame_update", build_code_frame_update(channels, scalar_index=True) + ) + pkg.add_policy_component( + "step_update", + build_decoder_step_update( + attention_dtype=_ir_dtype(temporal_mask), + position_dtype=_ir_dtype(temporal_position), + ), + ) + inputs["package.text_stream"] = { + "contract": control_int, + "role": {"kind": "opaque"}, + "source": {"kind": "literal"}, + "required": False, + "default": 0, + } + + workflow = { + "manifest": { + "capabilities": [ + "workflow_ssa", + "linear_effects", + "typed_emit", + "streaming_emit", + "nested_control_flow", + "loop_induction_values", + "bounded_state_recurrence", + "serving_service_contract", + "session_state_lease", + ], + }, + "inputs": inputs, + "outputs": { + "audio_chunk": { + "contract": { + **_contract(waveform_output), + "batch_layout": request_aligned, + }, + "role": "audio", + "stage": "post_adapter", + } + }, + "components": { + "encoder": _component(encoder, "encoder/model.onnx"), + "decoder": _component(decoder, "decoder/model.onnx"), + "temporal": _component(temporal, "temporal/model.onnx"), + "depformer": _component(depformer, "depformer/model.onnx"), + }, + "state": state, + "graph": graph, + } + metadata = { + "schema_version": "v1", + "pipeline": {"workflow": _publish_workflow_v1(workflow)}, + } + add_policy_components_to_workflow(metadata, pkg) + _annotate_duplex_state_service(metadata, config, temporal_caches) + return metadata + + +def _annotate_duplex_state_service( + metadata: dict[str, Any], + config: Any, + temporal_caches: list[tuple[ir.Value, ir.Value]], +) -> None: + """Publish the semantic contract of the session-resident temporal KV group. + + A full-duplex conversation has no termination predicate: it runs until the + session is released. The serving contract therefore points ``active``/``done`` + at the session-scoped liveness cells rather than at a generation-loop flag. + """ + workflow = metadata["pipeline"]["workflow"] + context = int(getattr(config, "context", 0)) + workflow["serving"] = { + "active": "active", + "done": "done", + "accepted_len": "accepted_len", + "state_service": { + "groups": { + "temporal_cache": { + "kind": "sliding_attention" if context else "full_attention", + "sequence_axis": 2, + "layout": "bnsh", + "logical_lengths": "temporal_cache_lengths", + "aliasing": "permitted", + "reuse": {"prefix_reusable": True, "evictable_prefix": False}, + "capabilities": {"snapshot": True, "fork": False}, + "ports": { + "temporal": { + f"temporal_cache_{index}": _annotated_alias( + {"input": past.name, "output": present.name} + ) + for index, (past, present) in enumerate(temporal_caches) + } + }, + } + } + }, + } + + +def write_full_duplex_workflow_metadata(pkg: Any, config: Any, output_dir: str) -> str: + """Write full-duplex workflow metadata and its policy artifacts.""" + os.makedirs(output_dir, exist_ok=True) + metadata = build_full_duplex_workflow_metadata(pkg, config) + pkg.save_policy_components(output_dir) + add_adapter_service_to_metadata(metadata, pkg, output_dir) + path = os.path.join(output_dir, "inference_metadata.yaml") + with open(path, "w", encoding="utf-8") as handle: + _dump_yaml(metadata, handle) + return path diff --git a/src/mobius/integrations/onnx_genai/workflow_metadata_test.py b/src/mobius/integrations/onnx_genai/workflow_metadata_test.py new file mode 100644 index 000000000..525a2e7aa --- /dev/null +++ b/src/mobius/integrations/onnx_genai/workflow_metadata_test.py @@ -0,0 +1,874 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT License. + +from __future__ import annotations + +import json +from pathlib import Path + +import onnx_ir as ir +import pytest +import yaml + +from mobius._model_package import ModelPackage +from mobius.integrations.onnx_genai.inference_metadata_test import ( + _model, + _native_package, + _VlmConfig, +) +from mobius.integrations.onnx_genai.workflow_metadata import ( + _kv_storage_contract, + _static_cache_ports, + build_language_diffusion_pipeline_metadata, + build_speculative_workflow_metadata, + build_video_diffusion_workflow_metadata, + build_vlm_workflow_metadata, + write_speculative_workflow_metadata, + write_vlm_workflow_metadata, +) + + +def test_speculative_writer_saves_policy_artifacts(tmp_path): + write_speculative_workflow_metadata(_speculative_package(), str(tmp_path)) + assert (tmp_path / "policies" / "speculative_acceptance.onnx").is_file() + assert (tmp_path / "policies" / "cache_length_update.onnx").is_file() + + +def test_speculative_writer_saves_guidance_and_adaptive_artifacts(tmp_path): + write_speculative_workflow_metadata( + _speculative_package(adaptive=True), + str(tmp_path), + grammar_guidance=True, + adaptive_k_max=4, + ) + assert (tmp_path / "policies" / "grammar_guidance.onnx").is_file() + assert (tmp_path / "policies" / "adaptive_k.onnx").is_file() + + +def test_speculative_emit_uses_accepted_prefix_length(): + workflow = build_speculative_workflow_metadata(_speculative_package())["pipeline"][ + "workflow" + ] + emit = next(node for node in workflow["steps"][0]["steps"] if node["kind"] == "emit") + assert emit["valid_length"] == "acceptance.length" + # Row identity is runtime-private: the published emit step must not name it. + assert "row_ids" not in emit + assert "emit_valid_length" in workflow["manifest"]["capabilities"] + assert "emit_row_identity" not in workflow["manifest"]["capabilities"] + assert not any("slot_ids" in name for name in workflow["inputs"]) + assert workflow["outputs"]["tokens"]["contract"]["shape"][-1] == "accepted_sequence" + assert workflow["state"]["cache_0"]["recurrence"] == { + "kind": "bounded", + "axis": 2, + "max": "package.max_context", + } + assert workflow["state"]["cache_0"]["service_group"] == "verifier_cache" + assert workflow["serving"]["accepted_len"] == "accepted_len" + assert workflow["inputs"]["package.max_context"]["default"] == 4096 + + +def _value(name: str, dtype: ir.DataType, shape: list[int | str]) -> ir.Value: + return ir.Value(name=name, type=ir.TensorType(dtype), shape=ir.Shape(shape)) + + +def test_vlm_preprocessing_is_explicit_typed_ssa(tmp_path): + source = tmp_path / "source" + source.mkdir() + (source / "config.json").write_text( + json.dumps({"use_hd_transform": True}), encoding="utf-8" + ) + (source / "preprocessor_config.json").write_text( + json.dumps( + { + "dynamic_hd": 1, + "crop_size": 16, + "include_thumbnail": False, + "thumbnail_order": "none", + "mask_patch_size": 1, + } + ), + encoding="utf-8", + ) + vision = _model( + "vision_encoder", + [ + _value("pixel_values", ir.DataType.FLOAT, [1, 3, 16, 16]), + _value("image_sizes", ir.DataType.INT64, [1, 2]), + _value("image_attention_mask", ir.DataType.FLOAT, [1, 16, 16]), + ], + [("image_features", ir.DataType.FLOAT, [1, 64])], + ) + + metadata = build_vlm_workflow_metadata( + _native_package(vision, _VlmConfig()), + _VlmConfig(), + source=str(source), + ) + image = metadata["preprocessing"]["image"] + declared = {name for transform in image["transforms"] for name in transform["outputs"]} + assert "inputs" not in image["transforms"][0] + assert all("outputs" in transform for transform in image["transforms"]) + assert all(output["source"] in declared for output in image["outputs"]) + + +def _decoder_with_capacity_addressable_attention( + inputs: list[ir.Value], + output_specs: list[tuple[str, ir.DataType, list[int | str]]], +) -> ir.Model: + """Build a decoder whose only cache consumer is ``GroupQueryAttention``. + + Capacity-preallocated KV storage is a property of the attention operator, + not of the port names: only an operator that takes the logical cache length + as a separate input can ignore the unwritten tail of a capacity-sized + buffer. A decoder wired through plain ``Identity`` nodes therefore describes + a *dynamic* cache, so an artifact meant to stand in for a shared-buffer + decoder has to name the operator that makes sharing sound. + """ + outputs = [_value(*spec) for spec in output_specs] + by_name = {output.name: output for output in outputs} + logits = outputs[0] + past_by_name = {value.name: value for value in inputs} + nodes = [ir.Node("", "Identity", [inputs[0]], outputs=[logits], name="emit_logits")] + layer = 0 + while f"past_key_values.{layer}.key" in past_by_name: + nodes.append( + ir.Node( + "com.microsoft", + "GroupQueryAttention", + [ + inputs[0], + past_by_name[f"past_key_values.{layer}.key"], + past_by_name[f"past_key_values.{layer}.value"], + ], + outputs=[ + by_name[f"present.{layer}.key"], + by_name[f"present.{layer}.value"], + ], + name=f"attention_{layer}", + ) + ) + layer += 1 + graph = ir.Graph( + inputs=inputs, + outputs=outputs, + nodes=nodes, + name="decoder", + opset_imports={"": 21, "com.microsoft": 1}, + ) + return ir.Model(graph, ir_version=10) + + +def test_vlm_writer_derives_real_decoder_contract_from_artifact(tmp_path): + source = tmp_path / "source" + source.mkdir() + (source / "config.json").write_text( + json.dumps({"use_hd_transform": True}), encoding="utf-8" + ) + (source / "genai_config.json").write_text( + json.dumps( + { + "model": {"eos_token_id": 200001, "context_length": 131072}, + "search": {"past_present_share_buffer": True}, + } + ), + encoding="utf-8", + ) + (source / "preprocessor_config.json").write_text( + json.dumps( + { + "dynamic_hd": 1, + "crop_size": 16, + "include_thumbnail": False, + "thumbnail_order": "none", + "mask_patch_size": 1, + } + ), + encoding="utf-8", + ) + + decoder_inputs = [ + _value("inputs_embeds", ir.DataType.BFLOAT16, ["batch", "sequence", 6656]), + _value( + "attention_mask", + ir.DataType.INT64, + ["batch", "past_sequence + sequence"], + ), + ] + decoder_outputs = [("logits", ir.DataType.BFLOAT16, ["batch", "sequence", 202048])] + for layer in range(52): + cache_shape = ["batch", 2, "past_sequence", 128] + present_shape = ["batch", 2, "total_sequence", 128] + for kind in ("key", "value"): + decoder_inputs.append( + _value( + f"past_key_values.{layer}.{kind}", + ir.DataType.BFLOAT16, + cache_shape, + ) + ) + decoder_outputs.append( + ( + f"present.{layer}.{kind}", + ir.DataType.BFLOAT16, + present_shape, + ) + ) + decoder = _decoder_with_capacity_addressable_attention(decoder_inputs, decoder_outputs) + vision = _model( + "vision_encoder", + [ + _value("pixel_values", ir.DataType.FLOAT, [1, 3, 16, 16]), + _value("image_sizes", ir.DataType.INT64, [1, 2]), + _value("image_attention_mask", ir.DataType.FLOAT, [1, 16, 16]), + ], + [("image_features", ir.DataType.BFLOAT16, ["image_tokens", 6656])], + ) + embedding = _model( + "embedding", + [ + _value("input_ids", ir.DataType.INT64, ["batch", "sequence"]), + _value( + "image_features", + ir.DataType.BFLOAT16, + ["image_tokens", 6656], + ), + ], + [("inputs_embeds", ir.DataType.BFLOAT16, ["batch", "sequence", 6656])], + ) + package = ModelPackage( + {"decoder": decoder, "vision_encoder": vision, "embedding": embedding} + ) + + # Deliberately tiny config values must never override admitted artifact I/O. + config = _VlmConfig() + config.eos_token_id = 2 + path = write_vlm_workflow_metadata( + package, + str(tmp_path / "package"), + config, + source=str(source), + ) + serialized = Path(path).read_text(encoding="utf-8") + assert "&id" not in serialized + assert "*id" not in serialized + with open(path, encoding="utf-8") as handle: + workflow = yaml.safe_load(handle)["pipeline"]["workflow"] + decoder_invokes = [] + policy_invokes = {} + + def collect_decoder_invokes(node): + if isinstance(node, dict): + if node.get("kind") == "invoke" and node.get("component") == "decoder": + decoder_invokes.append(node) + elif node.get("kind") == "invoke": + policy_invokes[node["component"]] = node + for value in node.values(): + collect_decoder_invokes(value) + elif isinstance(node, list): + for value in node: + collect_decoder_invokes(value) + + collect_decoder_invokes(workflow["steps"]) + assert len(decoder.graph.inputs) == 106 + assert len(decoder.graph.outputs) == 105 + assert len(decoder_invokes) == 2 + assert all( + set(invoke["inputs"]) == {value.name for value in decoder.graph.inputs} + for invoke in decoder_invokes + ) + assert all("position_ids" not in invoke["inputs"] for invoke in decoder_invokes) + assert ( + len([name for name in workflow["state"] if name.removeprefix("cache_").isdigit()]) + == 104 + ) + assert workflow["inputs"]["package.max_context"]["default"] == 131072 + assert workflow["inputs"]["package.eos_ids"]["default"] == 200001 + assert workflow["state"]["cache_103"]["contract"] == { + "dtype": "bfloat16", + "rank": 4, + "shape": ["batch", 2, "past_sequence", 128], + "batch_layout": {"kind": "request_aligned", "axis": 0}, + } + assert workflow["state"]["logits"]["contract"] == { + "dtype": "float32", + "rank": 2, + "shape": ["batch", 202048], + "batch_layout": {"kind": "request_aligned", "axis": 0}, + } + # A capacity-preallocated cache binds a mask sized once, up front. + assert workflow["state"]["attention_mask"]["recurrence"] == {"kind": "invariant"} + assert workflow["state"]["cache_103"]["recurrence"] == { + "kind": "bounded", + "axis": 2, + "max": "package.max_context", + } + assert workflow["state"]["cache_lengths"]["initializer"] == "initializer.cache_lengths" + assert policy_invokes["decoder_state_initializer"]["inputs"] == { + "prompt_tokens": "request.prompt_tokens", + "prompt_lengths": "request.prompt_lengths", + "max_iterations": "request.max_iterations", + } + assert workflow["inputs"]["request.prompt_lengths"]["contract"]["shape"] == ["batch"] + assert workflow["inputs"]["request.max_iterations"]["contract"]["shape"] == [1] + assert workflow["steps"][0]["iteration"]["contract"] == { + "dtype": "int64", + "rank": 1, + "shape": [1], + } + assert workflow["inputs"]["package.one"]["contract"]["shape"] == ["batch"] + assert workflow["state"]["generated_lengths"]["initializer"] == ( + "initializer.generated_lengths" + ) + assert policy_invokes["token_sampler"]["inputs"]["active"] == "active" + assert policy_invokes["token_sampler"]["inputs"]["done"] == "done" + assert set(policy_invokes["token_state_update"]["inputs"]) == { + "current", + "update", + "active", + "done", + } + assert policy_invokes["token_state_update"]["outputs"] == {"next": "token.body"} + assert policy_invokes["generated_length_update"]["outputs"] == { + "total": "token.emitted_length" + } + + def collect_emits(node): + if isinstance(node, dict): + return ([node] if node.get("kind") == "emit" else []) + [ + emit for value in node.values() for emit in collect_emits(value) + ] + if isinstance(node, list): + return [emit for value in node for emit in collect_emits(value)] + return [] + + emit = next( + node for node in collect_emits(workflow["steps"]) if node["output"] == "tokens" + ) + assert emit["when"] == "active" + assert emit["valid_length"] == "token.emitted_length" + assert "row_ids" not in emit + assert "emit_row_identity" not in workflow["manifest"]["capabilities"] + assert policy_invokes["decoder_step_update"]["inputs"]["logical_length"] == "cache_lengths" + assert workflow["state"]["attention_mask"]["initializer"] == ("initializer.attention_mask") + assert any( + invoke["inputs"]["attention_mask"] == "decoder_step.body_attention_mask" + for invoke in decoder_invokes + ) + assert workflow["inputs"]["request.image"]["required"] is False + assert workflow["inputs"]["request.image"]["present_as"] == "request.image_present" + assert "request.has_media" not in workflow["inputs"] + assert "input_presence" in workflow["manifest"]["capabilities"] + assert workflow["steps"][0]["termination"] == "generation_eos" + media_branch = workflow["steps"][0]["setup"][0] + assert media_branch["kind"] == "branch" + assert media_branch["predicate"] == "request.image_present" + assert set(media_branch["cases"]) == {"true", "false"} + assert media_branch["cases"]["false"]["component"] == "empty_image_features" + state_service = workflow["serving"]["state_service"] + # Storage class, allocator, and compaction algorithm are the runtime's to + # choose; the package only declares semantics. + assert set(state_service) == {"groups"} + decoder_group = state_service["groups"]["decoder_cache"] + assert decoder_group["kind"] == "full_attention" + assert decoder_group["aliasing"] == "permitted" + assert decoder_group["reuse"] == { + "prefix_reusable": True, + "evictable_prefix": False, + } + assert "storage" not in decoder_group + carried = {item["cell"] for item in workflow["steps"][0]["carried"]} + assert not any("slot_ids" in cell for cell in carried) + assert { + "token", + "logits", + "generated_lengths", + "rng_counter", + "active", + "done", + "accepted_len", + "cache_lengths", + "attention_mask", + "cache_0", + } <= carried + decoder_cache = state_service["groups"]["decoder_cache"] + # Shared buffering is expressed by the admitted cache ports and runtime I/O + # binding, even when the graph has no node-level share-buffer attribute. + assert all("past_present_share_buffer" not in node.attributes for node in decoder.graph) + kv_ports = decoder_cache["ports"]["decoder"] + assert len(kv_ports) == 104 + assert kv_ports["cache_103"] == { + "input": "past_key_values.51.value", + "output": "present.51.value", + "role": "value", + "layer": 51, + } + assert (tmp_path / "package" / "policies" / "last_token_logits.onnx").is_file() + assert (tmp_path / "package" / "policies" / "empty_image_features.onnx").is_file() + + +def _masked_denoiser_package() -> ModelPackage: + input_ids = _value("input_ids", ir.DataType.INT64, ["batch", "sequence"]) + logits = _value("logits", ir.DataType.FLOAT, ["batch", "sequence", 128]) + proposed = _value("proposed_tokens", ir.DataType.INT64, ["batch", "sequence"]) + graph = ir.Graph( + inputs=[input_ids], + outputs=[logits, proposed], + nodes=[], + name="masked_denoiser", + opset_imports={"": 24}, + ) + return ModelPackage({"model": ir.Model(graph, ir_version=11)}) + + +def test_language_diffusion_uses_exclusive_ssa_workflow(): + metadata = build_language_diffusion_pipeline_metadata( + _masked_denoiser_package(), + num_inference_steps=8, + ) + pipeline = metadata["pipeline"] + assert set(pipeline) == {"workflow"} + + workflow = pipeline["workflow"] + assert workflow["state"]["tokens_state"]["recurrence"] == {"kind": "invariant"} + assert workflow["state"]["tokens_state"]["contract"]["shape"] == [ + "batch", + "sequence", + ] + assert workflow["state"]["rng_offset"]["contract"]["shape"] == ["batch"] + assert workflow["components"]["masked_update"]["contract"] == { + "id": "onnx-genai.masked-update", + "version": "1", + "bindings": { + "state": "current_tokens", + "proposal": "proposed_tokens", + "mask": "masked", + "step": "step", + "next_state": "next_state", + "next_mask": "next_mask", + "seed": "seed", + "offset": "offset", + "next_offset": "next_offset", + "continue": "continue", + }, + } + + graph = workflow["steps"][0] + assert graph["kind"] == "loop" + assert graph["continue_when"] == "loop_0_active" + assert graph["max_iterations"] == "request.max_iterations" + assert [node["component"] for node in graph["setup"]] == ["model"] + assert [node["kind"] for node in graph["steps"]] == [ + "invoke", + "emit", + "invoke", + ] + assert graph["iteration"]["value"] == "loop.iteration" + assert graph["steps"][0]["inputs"]["total_steps"] == "package.num_steps" + assert graph["steps"][1]["mode"] == "replace" + + +def _video_package(*, cache_ports: bool = True, latent_rank: int = 5) -> ModelPackage: + latent_shape: list[int | str] = ["batch", "num_frames", 4, "height", "width"] + if latent_rank == 4: + latent_shape = ["batch", 4, "height", "width"] + sample = _value("sample", ir.DataType.FLOAT, latent_shape) + timestep = _value("timestep", ir.DataType.INT64, ["batch"]) + conditioning = _value( + "encoder_hidden_states", ir.DataType.FLOAT, ["batch", "prompt_sequence", 32] + ) + noise_pred = _value("noise_pred", ir.DataType.FLOAT, latent_shape) + denoiser = ir.Graph( + inputs=[sample, timestep, conditioning], + outputs=[noise_pred], + nodes=[], + name="transformer", + opset_imports={"": 24}, + ) + + latent_sample = _value( + "latent_sample", + ir.DataType.FLOAT, + ["batch", 4, "latent_frames", "latent_height", "latent_width"], + ) + frames = _value( + "sample", + ir.DataType.FLOAT, + ["batch", 3, "frames", "2*latent_height", "2*latent_width"], + ) + vae_inputs = [latent_sample] + vae_outputs = [frames] + if cache_ports: + vae_inputs.append( + _value( + "conv_cache.conv_in", + ir.DataType.FLOAT, + ["batch", 4, "cache_frames", "latent_height", "latent_width"], + ) + ) + vae_outputs.append( + _value( + "conv_cache_out.conv_in", + ir.DataType.FLOAT, + ["batch", 4, "cache_frames", "latent_height", "latent_width"], + ) + ) + vae = ir.Graph( + inputs=vae_inputs, + outputs=vae_outputs, + nodes=[], + name="vae_decoder", + opset_imports={"": 24}, + ) + vae_model = ir.Model(vae, ir_version=11) + vae_model.metadata_props["mobius.conv_cache.spatial_scale.conv_cache.conv_in"] = "1" + return ModelPackage( + { + "transformer": ir.Model(denoiser, ir_version=11), + "vae_decoder": vae_model, + } + ) + + +def test_video_diffusion_keeps_the_temporal_axis_through_every_stage(): + metadata = build_video_diffusion_workflow_metadata( + _video_package(), num_inference_steps=2, solver="ddim" + ) + workflow = metadata["pipeline"]["workflow"] + + published = workflow["outputs"]["video"] + assert published["role"] == "video" + assert published["contract"]["rank"] == 5 + + latent = workflow["state"]["latent"] + assert latent["contract"]["rank"] == 5 + assert latent["contract"]["shape"][1] == "num_frames" + + # The scheduler trajectory is state, not telemetry: a multistep video solver + # reads it back. + history = workflow["state"]["scheduler_history"] + assert history["recurrence"]["kind"] == "growing" + assert history["recurrence"]["axis"] == 1 + + cache = workflow["state"]["conv_cache_conv_in"] + assert cache["management"] == "runtime" + assert cache["release_boundary"] == "invocation" + assert cache["recurrence"] == { + "kind": "bounded", + "axis": 2, + "max": "package.cache_frames", + } + + denoise, decode = (step for step in workflow["steps"] if step["kind"] == "loop") + assert denoise["max_iterations"] == "request.max_iterations" + # The decoder runs once per causal chunk, and the chunk count is computed at + # run time from the latent rather than fixed by the package. + assert decode["max_iterations"] == "decode.chunks" + emit = next(node for node in decode["steps"] if node["kind"] == "emit") + # Frames append along time. Appending on the last axis -- the token-sequence + # default -- would concatenate image columns instead. + assert emit["mode"] == "append" + assert emit["axis"] == 2 + + assert "bounded_state_recurrence" in workflow["manifest"]["capabilities"] + + +def test_video_diffusion_rejects_an_image_latent(): + with pytest.raises(ValueError, match="rank-5"): + build_video_diffusion_workflow_metadata( + _video_package(latent_rank=4), num_inference_steps=2 + ) + + +def test_video_diffusion_requires_paired_conv_caches(): + with pytest.raises(ValueError, match="conv_cache"): + build_video_diffusion_workflow_metadata( + _video_package(cache_ports=False), num_inference_steps=2 + ) + + +def test_language_diffusion_rejects_zero_steps(): + with pytest.raises(ValueError, match="num_inference_steps"): + build_language_diffusion_pipeline_metadata( + _masked_denoiser_package(), + num_inference_steps=0, + ) + + +def _graph_model( + name: str, + inputs: list[ir.Value], + outputs: list[ir.Value], +) -> ir.Model: + return ir.Model( + ir.Graph( + inputs=inputs, + outputs=outputs, + nodes=[], + name=name, + opset_imports={"": 24}, + ), + ir_version=11, + ) + + +def _speculative_package( + *, + adaptive: bool = False, + budget_dtype: ir.DataType = ir.DataType.INT64, +) -> ModelPackage: + proposer_inputs = [_value("tokens", ir.DataType.INT64, ["batch", 4])] + if adaptive: + proposer_inputs.append(_value("proposal_budget", budget_dtype, ["batch"])) + proposer = _graph_model( + "proposer", + proposer_inputs, + [ + _value("proposed_tokens", ir.DataType.INT64, ["batch", 4]), + _value("proposal_scores", ir.DataType.FLOAT, ["batch", 4, 32]), + ], + ) + verifier = _graph_model( + "verifier", + [ + _value("proposed_tokens", ir.DataType.INT64, ["batch", 4]), + _value( + "past_key_values.0.key", + ir.DataType.FLOAT, + ["batch", 2, "past_sequence", 8], + ), + ], + [ + _value("target_scores", ir.DataType.FLOAT, ["batch", 4, 32]), + _value( + "present.0.key", + ir.DataType.FLOAT, + ["batch", 2, "past_sequence + 4", 8], + ), + ], + ) + return ModelPackage({"proposer": proposer, "verifier": verifier}) + + +def test_adaptive_speculative_requires_int64_budget_port(): + with pytest.raises(ValueError, match="rank-1 proposer budget input"): + build_speculative_workflow_metadata( + _speculative_package(adaptive=True, budget_dtype=ir.DataType.INT32), + adaptive_k_max=4, + ) + + +def test_speculative_grammar_and_adaptive_k_use_typed_state_contracts(): + metadata = build_speculative_workflow_metadata( + _speculative_package(adaptive=True), + grammar_guidance=True, + adaptive_k_max=4, + ) + workflow = metadata["pipeline"]["workflow"] + assert workflow["components"]["grammar_commit"]["contract"] == { + "id": "onnx-genai.grammar-guidance", + "version": "1", + "bindings": { + "state": "state", + "tokens": "tokens", + "valid_length": "valid_length", + "transition_table": "transition_table", + "next_state": "next_state", + "consumed_length": "consumed_length", + "logits_mask": "logits_mask", + "forced_tokens": "forced_tokens", + "forced_length": "forced_length", + }, + "parameters": {"action": "commit"}, + } + assert workflow["components"]["adaptive_k"]["contract"]["id"] == ( + "onnx-genai.adaptive-proposal-budget" + ) + assert workflow["state"]["grammar"]["class"] == "semantic" + assert workflow["state"]["proposal_k"]["class"] == "advisory" + assert workflow["state"]["adaptive_estimates"]["class"] == "advisory" + assert all( + "ports" in component and "effects" not in component + for component in workflow["components"].values() + if component["implementation"]["kind"] == "onnx" and "contract" in component + ) + assert "ports" in workflow["components"]["grammar_commit"] + proposer = next( + node for node in workflow["steps"][0]["steps"] if node.get("component") == "proposer" + ) + assert proposer["inputs"]["proposal_budget"] == "proposal_k" + grammar_emit = next( + node + for node in workflow["steps"][0]["steps"] + if node.get("kind") == "emit" and node.get("value") == "grammar.token" + ) + assert grammar_emit["valid_length"] == "grammar.forced_length" + assert all("initial" not in carry for carry in workflow["steps"][0]["carried"]) + + +def test_speculative_workflow_uses_per_row_ragged_state_and_rng(): + workflow = build_speculative_workflow_metadata(_speculative_package())["pipeline"][ + "workflow" + ] + body = workflow["steps"][0]["steps"] + acceptance = body[2] + assert acceptance["inputs"]["offset"] == "rng_offset" + assert acceptance["outputs"]["next_offset"] == "rng_offset.body" + assert acceptance["outputs"]["accepted_len"] == "acceptance.length" + emit = next(node for node in body if node["kind"] == "emit") + assert emit["valid_length"] == "acceptance.length" + assert "row_ids" not in emit + assert not any(node["kind"] == "branch" for node in body) + assert workflow["serving"]["active"] == "active" + assert workflow["serving"]["done"] == "done" + assert workflow["serving"]["state_service"]["groups"]["verifier_cache"]["kind"] == ( + "full_attention" + ) + assert "slot_ids" not in workflow["serving"] + assert workflow["serving"]["state_service"]["groups"]["verifier_cache"]["ports"][ + "verifier" + ]["cache_0"] == { + "input": "past_key_values.0.key", + "output": "present.0.key", + "role": "key", + "layer": 0, + } + assert any(item["cell"].startswith("cache_") for item in workflow["steps"][0]["carried"]) + + +def _decoder_with_cache(domain: str, op_type: str) -> ir.Model: + """Minimal decoder whose KV cache is consumed by ``domain::op_type``.""" + past_key = _value("past_key_values.0.key", ir.DataType.FLOAT, ["batch", 2, "past", 8]) + past_value = _value("past_key_values.0.value", ir.DataType.FLOAT, ["batch", 2, "past", 8]) + hidden = _value("hidden", ir.DataType.FLOAT, ["batch", "seq", 16]) + present_key = _value("present.0.key", ir.DataType.FLOAT, ["batch", 2, "total", 8]) + present_value = _value("present.0.value", ir.DataType.FLOAT, ["batch", 2, "total", 8]) + logits = _value("logits", ir.DataType.FLOAT, ["batch", "seq", 32]) + attention = ir.Node( + domain, + op_type, + [hidden, past_key, past_value], + outputs=[logits, present_key, present_value], + name="attention", + ) + graph = ir.Graph( + inputs=[hidden, past_key, past_value], + outputs=[logits, present_key, present_value], + nodes=[attention], + name="decoder", + opset_imports={"": 21, "com.microsoft": 1}, + ) + return ir.Model(graph, ir_version=10) + + +def test_capacity_addressable_attention_declares_shared_buffer_storage(): + # GroupQueryAttention takes seqlens_k/total_sequence_length, so it can safely + # read a capacity-sized buffer whose tail is unwritten. + contract = _kv_storage_contract( + _decoder_with_cache("com.microsoft", "GroupQueryAttention") + ) + assert contract == {"paging": "none", "compaction": True, "storage": "shared_buffer"} + + +def test_standard_attention_declares_dynamic_storage(): + # The standard ONNX Attention operator derives the total sequence length from + # the past tensor's own shape, so a preallocated buffer would both attend over + # unwritten slots and contradict an exactly sized attention mask. + contract = _kv_storage_contract(_decoder_with_cache("", "Attention")) + assert contract == {"paging": "none", "compaction": True, "storage": "dynamic"} + + +def test_unconsumed_cache_ports_are_not_treated_as_capacity_addressable(): + # A graph that never reads its own cache cannot promise capacity-safe + # behaviour it does not exercise. + model = _decoder_with_cache("com.microsoft", "GroupQueryAttention") + model.graph.node("attention").replace_input_with(1, None) + model.graph.node("attention").replace_input_with(2, None) + assert _kv_storage_contract(model)["storage"] == "dynamic" + + +def test_paged_cache_inputs_take_precedence_over_operator_derivation(): + model = _decoder_with_cache("", "Attention") + model.graph.inputs.append(_value("block_tables", ir.DataType.INT32, ["batch", "blocks"])) + contract = _kv_storage_contract(model) + assert contract["paging"] == "paged" + assert contract["storage"] == "paged" + + +def _static_cache_model( + *, + capacities: list[int] | None = None, + scatter_axis: int | None = None, + paired: bool = True, + control_ports: bool = True, +) -> ir.Model: + """A minimal graph shaped like a mobius static-cache decoder export.""" + capacities = capacities or [32, 32] + inputs = [_value("input_ids", ir.DataType.INT64, ["batch", "sequence"])] + for layer, capacity in enumerate(capacities): + inputs.append(_value(f"key_cache.{layer}", ir.DataType.FLOAT, ["batch", capacity, 16])) + if control_ports: + inputs.append(_value("write_indices", ir.DataType.INT64, ["batch"])) + inputs.append(_value("nonpad_kv_seqlen", ir.DataType.INT64, ["batch"])) + outputs: list[tuple[str, ir.DataType, list[int | str]]] = [ + ("logits", ir.DataType.FLOAT, ["batch", "sequence", 128]) + ] + if paired: + for layer, capacity in enumerate(capacities): + outputs.append( + (f"updated_key_cache.{layer}", ir.DataType.FLOAT, ["batch", capacity, 16]) + ) + model = _model("decoder", inputs, outputs) + if scatter_axis is not None: + cache = model.graph.inputs[1] + scattered = _value("scattered", ir.DataType.FLOAT, ["batch", capacities[0], 16]) + model.graph.append( + ir.Node( + "", + "TensorScatter", + [cache, cache, model.graph.inputs[-2]], + outputs=[scattered], + attributes=[ir.AttrInt64("axis", scatter_axis)], + name="scatter", + ) + ) + return model + + +class TestStaticCachePortDiscovery: + """``_static_cache_ports`` reads the ABI from the graph or refuses to guess.""" + + def test_returns_none_without_the_control_ports(self): + # A dynamic decoder must not be mistaken for a fixed-capacity one. + assert _static_cache_ports(_static_cache_model(control_ports=False)) is None + + def test_discovers_buffers_control_ports_and_capacity(self): + ports = _static_cache_ports(_static_cache_model()) + assert ports["write_indices"] == "write_indices" + assert ports["kv_sequence_length"] == "nonpad_kv_seqlen" + assert ports["capacity"] == 32 + assert sorted(ports["buffers"]) == ["key_cache.0", "key_cache.1"] + + def test_rejects_control_ports_without_paired_buffers(self): + # Nothing to scatter into: the runtime would have no output to carry. + with pytest.raises(ValueError, match="no paired cache buffer"): + _static_cache_ports(_static_cache_model(paired=False)) + + def test_rejects_conflicting_capacities(self): + # One write cursor cannot address buffers of different lengths. + with pytest.raises(ValueError, match="conflicting capacities"): + _static_cache_ports(_static_cache_model(capacities=[32, 64])) + + def test_rejects_a_symbolic_capacity(self): + model = _static_cache_model() + model.graph.inputs[1].shape = ir.Shape(["batch", "capacity", 16]) + with pytest.raises(ValueError, match="symbolic extent"): + _static_cache_ports(model) + + def test_rejects_a_scatter_that_disagrees_with_the_declared_axis(self): + # The capacity axis published in the metadata has to be the axis the + # graph actually writes on, or a runtime sizes the wrong dimension. + with pytest.raises(ValueError, match="declared capacity axis"): + _static_cache_ports(_static_cache_model(scatter_axis=2)) + + def test_accepts_a_scatter_on_the_declared_axis(self): + assert _static_cache_ports(_static_cache_model(scatter_axis=1))["capacity"] == 32 diff --git a/src/mobius/integrations/transformers/_config_resolver.py b/src/mobius/integrations/transformers/_config_resolver.py index 610c0d842..e38e252fe 100644 --- a/src/mobius/integrations/transformers/_config_resolver.py +++ b/src/mobius/integrations/transformers/_config_resolver.py @@ -79,21 +79,29 @@ def _default_task_for_model(model_type: str) -> str: def _try_load_config_json(model_id: str, revision: str | None = None): """Try to load config.json directly for models not in transformers. - Returns a ``PretrainedConfig``-like object with attribute access, - or ``None`` if the file cannot be downloaded/parsed. + Accepts a Hugging Face repo id or a local directory. Returns a + ``PretrainedConfig``-like object with attribute access, or ``None`` if the + file cannot be downloaded/parsed. """ import json + import os from huggingface_hub import hf_hub_download - try: - kwargs = {"repo_id": model_id, "filename": "config.json"} - if revision is not None: - kwargs["revision"] = revision - path = hf_hub_download(**kwargs) - except (OSError, ValueError) as e: - logger.debug("Failed to download config.json for %s: %s", model_id, e) - return None + if os.path.isdir(model_id): + path = os.path.join(model_id, "config.json") + if not os.path.isfile(path): + logger.debug("No config.json in local directory %s", model_id) + return None + else: + try: + kwargs = {"repo_id": model_id, "filename": "config.json"} + if revision is not None: + kwargs["revision"] = revision + path = hf_hub_download(**kwargs) + except (OSError, ValueError) as e: + logger.debug("Failed to download config.json for %s: %s", model_id, e) + return None try: with open(path) as f: @@ -104,11 +112,45 @@ def _try_load_config_json(model_id: str, revision: str | None = None): model_type = config_dict.get("model_type") if not model_type: - return None + model_type = _model_type_from_architectures(config_dict.get("architectures")) + if not model_type: + return None + logger.info( + "config.json for %s declares no model_type; inferred '%s' from architectures=%s", + model_id, + model_type, + config_dict.get("architectures"), + ) + config_dict = {**config_dict, "model_type": model_type} return _dict_to_pretrained_config(config_dict) +def _model_type_from_architectures(architectures) -> str | None: + """Recover a HuggingFace ``model_type`` from a config's ``architectures``. + + Checkpoints published before ``model_type`` became mandatory omit the key + (``Rostlab/prot_bert`` is one), which makes ``AutoConfig`` refuse the repo + even though the architecture is a plain, fully supported one. The + architecture class name is not a guess in that situation: ``transformers`` + exports it, and the class states its own config class, which states the + model type. Anything that does not resolve through that chain returns + ``None`` so an unknown repo still fails loudly. + """ + if not architectures: + return None + + import transformers + + for architecture in architectures: + model_class = getattr(transformers, str(architecture), None) + config_class = getattr(model_class, "config_class", None) + model_type = getattr(config_class, "model_type", None) + if model_type: + return str(model_type) + return None + + def _dict_to_pretrained_config(d: dict): """Recursively convert a dict to a PretrainedConfig with attribute access. diff --git a/src/mobius/models/__init__.py b/src/mobius/models/__init__.py index 51b829c2f..99557e267 100644 --- a/src/mobius/models/__init__.py +++ b/src/mobius/models/__init__.py @@ -20,7 +20,9 @@ "FusedGateUpCausalLMModel", "ChatGLMCausalLMModel", "CodeGenCausalLMModel", + "AutoencoderKLCogVideoXModel", "CogVideoXTransformer3DModel", + "CogVideoXVAEConfig", "CohereCausalLMModel", "ControlNetModel", "Cosmos3EdgeTextModel", @@ -34,6 +36,8 @@ "DiTTransformer2DModel", "DiffLlamaCausalLMModel", "DistilBertModel", + "EsmConfig", + "EsmModel", "DogeCausalLMModel", "EncDecRNNTModel", "Ernie45MoECausalLMModel", @@ -188,6 +192,10 @@ from mobius.models.chatglm import ChatGLMCausalLMModel from mobius.models.clip import CLIPVisionModel, SigLIPVisionModel from mobius.models.cogvideox import CogVideoXTransformer3DModel +from mobius.models.cogvideox_vae import ( + AutoencoderKLCogVideoXModel, + CogVideoXVAEConfig, +) from mobius.models.cohere import CohereCausalLMModel from mobius.models.controlnet import ControlNetModel from mobius.models.cosmos import Cosmos3EdgeTextModel, Cosmos3EdgeVLModel @@ -203,6 +211,7 @@ from mobius.models.doge import DogeCausalLMModel from mobius.models.eagle3 import Eagle3DraftModel from mobius.models.ernie import ErnieCausalLMModel +from mobius.models.esm import EsmConfig, EsmModel from mobius.models.exaone4 import ExaOne4CausalLMModel from mobius.models.falcon import ( BloomCausalLMModel, diff --git a/src/mobius/models/bert.py b/src/mobius/models/bert.py index 13b227d59..7b841886e 100644 --- a/src/mobius/models/bert.py +++ b/src/mobius/models/bert.py @@ -24,7 +24,12 @@ from mobius._configs import ArchitectureConfig from mobius.components._activations import ACT2FN -from mobius.components._common import Embedding, LayerNorm, Linear +from mobius.components._common import ( + Embedding, + LayerNorm, + Linear, + create_padding_mask, +) if TYPE_CHECKING: import onnx_ir as ir @@ -73,6 +78,12 @@ def __init__(self, hidden_size: int, num_heads: int, eps: float, bias: bool): self.output = _BertAttentionOutput(hidden_size, eps, bias) def forward(self, op: OpBuilder, hidden_states: ir.Value, attention_mask: ir.Value): + """Bidirectional self-attention. + + ``attention_mask`` is the 4D bool padding mask + ``(batch, 1, seq, seq)`` prepared once per forward by + :meth:`BertModel.forward`, not the raw rank-2 request mask. + """ self_attn = self.self query = self_attn.query(op, hidden_states) key = self_attn.key(op, hidden_states) @@ -235,7 +246,11 @@ def forward( token_type_ids: ir.Value, ): hidden_states = self.embeddings(op, input_ids, token_type_ids) - hidden_states = self.encoder(op, hidden_states, attention_mask) + # The raw request mask is (batch, seq) int64; ONNX ``Attention`` needs a + # mask broadcastable to (batch, heads, q, kv). Building it once here + # keeps a single copy in the graph instead of one per layer. + padding_mask = create_padding_mask(op, input_ids, attention_mask) + hidden_states = self.encoder(op, hidden_states, padding_mask) return hidden_states def preprocess_weights( diff --git a/src/mobius/models/bert_test.py b/src/mobius/models/bert_test.py new file mode 100644 index 000000000..e9880889d --- /dev/null +++ b/src/mobius/models/bert_test.py @@ -0,0 +1,21 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT License. + +"""Shared tiny BERT-family configuration for model and metadata tests.""" + +from __future__ import annotations + +from mobius._configs import ArchitectureConfig + +PROTBERT_TINY_CONFIG = ArchitectureConfig( + hidden_size=64, + num_hidden_layers=2, + num_attention_heads=4, + num_key_value_heads=4, + intermediate_size=128, + vocab_size=30, + max_position_embeddings=512, + hidden_act="gelu", + rms_norm_eps=1e-12, + pad_token_id=0, +) diff --git a/src/mobius/models/cogvideox.py b/src/mobius/models/cogvideox.py index 0cb83eb86..043a18a46 100644 --- a/src/mobius/models/cogvideox.py +++ b/src/mobius/models/cogvideox.py @@ -120,6 +120,110 @@ def _get_3d_sincos_pos_embed( return np.concatenate([pos_temporal, pos_spatial], axis=-1) +def _sincos_omega(embed_dim: int) -> np.ndarray: + """Inverse frequency vector of a 1D sincos embedding, shape ``[embed_dim/2]``.""" + omega = np.arange(embed_dim // 2, dtype=np.float64) + omega /= embed_dim / 2.0 + return (1.0 / (10000**omega)).astype(np.float32) + + +def _sincos_in_graph( + op: OpBuilder, + embed_dim: int, + count: ir.Value, + scale: float, +) -> ir.Value: + """1D sincos embedding for ``count`` positions, evaluated inside the graph. + + ``count`` is a 1-D int64 tensor holding a single (possibly dynamic) extent, so + the table follows the real number of latent frames / patch rows / patch columns + instead of the sizes recorded in the checkpoint config. + + Returns: + Value of shape ``[count, embed_dim]``. + """ + positions = op.Range( + op.Constant(value_int=0), + op.Squeeze(count, op.Constant(value_ints=[0])), + op.Constant(value_int=1), + ) # (count,) int64 + positions = op.Cast(positions, to=ir.DataType.FLOAT) + if not math.isclose(scale, 1.0): + positions = op.Div(positions, op.Constant(value_float=float(scale))) + # (count, 1) * (embed_dim/2,) -> (count, embed_dim/2) + angles = op.Mul( + op.Unsqueeze(positions, op.Constant(value_ints=[-1])), + op.Constant(value_floats=_sincos_omega(embed_dim).tolist()), + ) + return op.Concat(op.Sin(angles), op.Cos(angles), axis=-1) # (count, embed_dim) + + +def _sincos_3d_in_graph( + op: OpBuilder, + embed_dim: int, + num_frames: ir.Value, + height: ir.Value, + width: ir.Value, + spatial_scale: float, + temporal_scale: float, +) -> ir.Value: + """3D sincos positional embedding built from runtime extents. + + Mirrors ``_get_3d_sincos_pos_embed`` (and diffusers' ``get_3d_sincos_pos_embed``) + but keeps the frame count and patch grid dynamic, so one exported graph serves + every video length and resolution. + + Returns: + Value of shape ``[1, num_frames*height*width, embed_dim]``. + """ + dim_spatial = 3 * embed_dim // 4 + dim_temporal = embed_dim // 4 + half_spatial = dim_spatial // 2 + + # diffusers meshgrid(grid_w, grid_h, indexing="xy") puts the width index in + # grid[0] and the height index in grid[1]; the 2D helper embeds grid[0] first. + w_table = _sincos_in_graph(op, half_spatial, width, spatial_scale) # (W, ds/2) + h_table = _sincos_in_graph(op, half_spatial, height, spatial_scale) # (H, ds/2) + grid_shape = op.Concat(height, width, op.Constant(value_ints=[half_spatial]), axis=0) + # (1, W, ds/2) -> (H, W, ds/2) and (H, 1, ds/2) -> (H, W, ds/2) + w_grid = op.Expand(op.Unsqueeze(w_table, op.Constant(value_ints=[0])), grid_shape) + h_grid = op.Expand(op.Unsqueeze(h_table, op.Constant(value_ints=[1])), grid_shape) + spatial = op.Concat(w_grid, h_grid, axis=-1) # (H, W, ds) + spatial = op.Reshape( + spatial, + op.Concat( + op.Constant(value_ints=[1]), + op.Mul(height, width), + op.Constant(value_ints=[dim_spatial]), + axis=0, + ), + ) # (1, H*W, ds) + + temporal = _sincos_in_graph(op, dim_temporal, num_frames, temporal_scale) # (T, dt) + temporal = op.Unsqueeze(temporal, op.Constant(value_ints=[1])) # (T, 1, dt) + + patches = op.Mul(height, width) + spatial = op.Expand( + spatial, + op.Concat(num_frames, patches, op.Constant(value_ints=[dim_spatial]), axis=0), + ) + temporal = op.Expand( + temporal, + op.Concat(num_frames, patches, op.Constant(value_ints=[dim_temporal]), axis=0), + ) + # Temporal channels come first, matching get_3d_sincos_pos_embed. + pos = op.Concat(temporal, spatial, axis=-1) # (T, H*W, D) + return op.Reshape( + pos, + op.Concat( + op.Constant(value_ints=[1]), + op.Mul(num_frames, patches), + op.Constant(value_ints=[embed_dim]), + axis=0, + ), + ) # (1, T*H*W, D) + + # --------------------------------------------------------------------------- # Model building blocks # --------------------------------------------------------------------------- @@ -393,34 +497,25 @@ def __init__(self, config: CogVideoXConfig, embed_dim: int): self._patch_size = p self._in_channels = in_ch + self._embed_dim = embed_dim + self._spatial_scale = config.spatial_interpolation_scale + self._temporal_scale = config.temporal_interpolation_scale + self._learned_pos = config.use_learned_positional_embeddings - # Pre-compute 3D sincos positional embedding as nn.Parameter post_patch_h = config.sample_height // p post_patch_w = config.sample_width // p num_time_patches = (config.sample_frames - 1) // config.temporal_compression_ratio + 1 - pos_embed = _get_3d_sincos_pos_embed( - embed_dim, - (post_patch_w, post_patch_h), - num_time_patches, - spatial_scale=config.spatial_interpolation_scale, - temporal_scale=config.temporal_interpolation_scale, - ) # [T, H'*W', D] - pos_embed = pos_embed.reshape(-1, embed_dim) # [T*H'*W', D] - - # Prepend zeros for text tokens (no positional embedding) - text_zeros = np.zeros((config.max_text_seq_length, embed_dim), dtype=np.float32) - # [max_text_seq + T*H'*W', D] - full_pos = np.concatenate([text_zeros, pos_embed], axis=0) - total_seq = full_pos.shape[0] - # [1, total_seq, D] for broadcasting with batch dim - full_pos = full_pos.reshape(1, total_seq, embed_dim) - - self.pos_embedding = nn.Parameter( - [1, total_seq, embed_dim], - name="patch_embed.pos_embedding.pos_embedding", - data=ir.tensor(full_pos), - ) + if self._learned_pos: + # I2V checkpoints ship a trained joint table, which is inherently tied + # to the sample resolution and frame count recorded in the config. + total_seq = ( + config.max_text_seq_length + num_time_patches * post_patch_h * post_patch_w + ) + self.pos_embedding = nn.Parameter( + [1, total_seq, embed_dim], + name="patch_embed.pos_embedding", + ) def forward( self, @@ -471,11 +566,25 @@ def forward( # Project text text = self.text_proj(op, text_embeds) # [B, seq, hidden] - # Concatenate text + video and add positional embedding - embeds = op.Concat(text, video, axis=1) # [B, total_seq, hidden] - embeds = op.Add(embeds, self.pos_embedding) + if self._learned_pos: + # Learned table covers text + video positions jointly. + embeds = op.Concat(text, video, axis=1) + return op.Add(embeds, self.pos_embedding) + + # Sincos positions are zero over the text prefix in diffusers, so adding + # them to the video stream alone is exact and keeps the text length free. + pos = _sincos_3d_in_graph( + op, + self._embed_dim, + num_frames, + h_patches, + w_patches, + self._spatial_scale, + self._temporal_scale, + ) # [1, T*H'*W', hidden] + video = op.Add(video, op.CastLike(pos, video)) - return embeds + return op.Concat(text, video, axis=1) # [B, total_seq, hidden] # --------------------------------------------------------------------------- @@ -648,11 +757,27 @@ def preprocess_weights( Renames: - ``ff.net.0.proj.*`` → ``ff.gelu_proj.*`` - ``ff.net.2.*`` → ``ff.linear_out.*`` + + Reshapes: + - ``patch_embed.proj.weight`` from the HuggingFace ``Conv2d`` layout + ``[embed_dim, C, p, p]`` to the flattened ``Linear`` layout + ``[embed_dim, C * p * p]``. A stride-``p`` ``Conv2d`` with a ``p x p`` + kernel is exactly a linear map over each ``(C, p, p)`` patch, and this + module patchifies to ``[B, T * H' * W', C * p * p]`` with the same + ``(C, p_h, p_w)`` element order, so a plain reshape is the correct + weight transform. """ - return rename_weight_keys( + state_dict = rename_weight_keys( state_dict, [ (".ff.net.0.proj.", ".ff.gelu_proj."), (".ff.net.2.", ".ff.linear_out."), ], ) + proj_weight = state_dict.get("patch_embed.proj.weight") + if proj_weight is not None and proj_weight.ndim == 4: + state_dict = dict(state_dict) + state_dict["patch_embed.proj.weight"] = proj_weight.reshape( + proj_weight.shape[0], -1 + ) + return state_dict diff --git a/src/mobius/models/cogvideox_vae.py b/src/mobius/models/cogvideox_vae.py new file mode 100644 index 000000000..d659dadfe --- /dev/null +++ b/src/mobius/models/cogvideox_vae.py @@ -0,0 +1,660 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT License. + +"""CogVideoX 3D causal video autoencoder. + +Replicates HuggingFace diffusers' ``AutoencoderKLCogVideoX`` decoder path +(``CogVideoXDecoder3D``), the autoencoder used by every CogVideoX text-to-video +and image-to-video pipeline. + +The decoder is *temporally causal*: each ``CogVideoXCausalConv3d`` pads only the +past side of the time axis (the first frame is replicated ``kernel_t - 1`` +times), so output frame ``t`` never depends on latent frames after ``t``. That +property is what lets the reference implementation decode a long video in +latent-frame chunks while carrying a small ``conv_cache`` across the chunk +boundary. + +Chunking is *not* transparent, however: the ``CogVideoXSpatialNorm3D`` group +normalizations reduce over the whole time axis of whatever is passed in, so the +statistics depend on the chunk. ``AutoencoderKLCogVideoX._decode`` always walks +the clip two latent frames at a time, and this module reproduces that boundary +exactly by exposing the same ``conv_cache`` tensors as graph inputs and outputs. +Decoding a clip in one call is only equivalent when it fits in a single +reference chunk (``T_latent <= 3``). + +Inputs / outputs +---------------- +- Decoder input ``latent_sample``: ``[B, latent_channels, T_latent, H, W]`` +- Decoder output ``sample``: ``[B, out_channels, T_pixels, H * s, W * s]`` + +with ``T_pixels = 2 ** log2(temporal_compression_ratio) * (T_latent - 1) + 1`` +for the odd latent-frame counts CogVideoX pipelines produce, and ``s`` the +spatial compression ratio. + +Temporal resampling in the reference model is ``F.interpolate(..., mode +="nearest")``, which is index arithmetic; this module reproduces it with +``Gather``/``Resize`` so the graph stays fully dynamic in ``T``, ``H`` and ``W`` +and never bakes in a frame count. +""" + +from __future__ import annotations + +import dataclasses +import math +import typing +from typing import TYPE_CHECKING + +from onnxscript import OpBuilder, nn + +from mobius.components import GroupNorm as _GroupNorm +from mobius.components import SiLU as _SiLU + +if TYPE_CHECKING: + import onnx_ir as ir + import torch + + +@dataclasses.dataclass +class CogVideoXVAEConfig: + """Configuration for ``AutoencoderKLCogVideoX``.""" + + in_channels: int = 3 + out_channels: int = 3 + latent_channels: int = 16 + block_out_channels: tuple[int, ...] = (128, 256, 256, 512) + layers_per_block: int = 3 + norm_num_groups: int = 32 + norm_eps: float = 1e-6 + temporal_compression_ratio: int = 4 + scaling_factor: float = 1.15258426 + use_post_quant_conv: bool = False + use_quant_conv: bool = False + + @classmethod + def from_diffusers(cls, config: dict) -> CogVideoXVAEConfig: + if hasattr(config, "to_dict"): + config = dict(config.items()) + return cls( + in_channels=config.get("in_channels", 3), + out_channels=config.get("out_channels", 3), + latent_channels=config.get("latent_channels", 16), + block_out_channels=tuple(config.get("block_out_channels", [128, 256, 256, 512])), + layers_per_block=config.get("layers_per_block", 3), + norm_num_groups=config.get("norm_num_groups", 32), + norm_eps=config.get("norm_eps", 1e-6), + temporal_compression_ratio=config.get("temporal_compression_ratio", 4), + scaling_factor=config.get("scaling_factor", 1.15258426), + use_post_quant_conv=config.get("use_post_quant_conv", False), + use_quant_conv=config.get("use_quant_conv", False), + ) + + @property + def spatial_compression_ratio(self) -> int: + """Spatial downsampling factor implied by the block count.""" + return 2 ** (len(self.block_out_channels) - 1) + + +# --------------------------------------------------------------------------- +# Shape / resampling helpers +# --------------------------------------------------------------------------- + +# ``Slice`` end sentinel for "to the end of the axis". +_INT64_MAX = 2**63 - 1 + + +def _dim(op: OpBuilder, value: ir.Value, axis: int) -> ir.Value: + """Single dimension of ``value`` as an int64 tensor of shape ``[1]``.""" + return op.Shape(value, start=axis, end=axis + 1) + + +def _nearest_indices(op: OpBuilder, out_len: ir.Value, in_len: ir.Value) -> ir.Value: + """``floor(i * in_len / out_len)`` for ``i`` in ``[0, out_len)``. + + This is exactly PyTorch's ``mode="nearest"`` source-index rule + (asymmetric coordinate transform with a floor rounding mode). + """ + zero = op.Constant(value_ints=[0]) + one = op.Constant(value_ints=[1]) + positions = op.Range(op.Squeeze(zero, [0]), op.Squeeze(out_len, [0]), op.Squeeze(one, [0])) + ratio = op.Div(op.Cast(in_len, to=1), op.Cast(out_len, to=1)) + source = op.Floor(op.Mul(op.Cast(positions, to=1), ratio)) + return op.Cast(source, to=7) + + +def _causal_nearest_indices(op: OpBuilder, out_len: ir.Value, in_len: ir.Value) -> ir.Value: + """Temporal source indices used by ``CogVideoXSpatialNorm3D``. + + When the conditioned feature map has an odd number of frames greater than + one, the reference implementation resamples the first frame on its own and + the remaining frames as a separate group:: + + idx[0] = 0 + idx[t] = 1 + floor((t - 1) * (T_in - 1) / (T_out - 1)) for t >= 1 + + Otherwise it resamples the whole clip uniformly. Both index programs are + computed and selected with ``Where`` so no ``If`` subgraph (and no static + frame count) is needed. + """ + zero = op.Constant(value_ints=[0]) + one = op.Constant(value_ints=[1]) + two = op.Constant(value_ints=[2]) + positions = op.Range(op.Squeeze(zero, [0]), op.Squeeze(out_len, [0]), op.Squeeze(one, [0])) + uniform = _nearest_indices(op, out_len, in_len) + + # Split-first-frame program. ``max(out_len - 1, 1)`` keeps the divisor + # well defined for the degenerate single-frame case, where only index 0 is + # ever materialized anyway. + out_rest = op.Max(op.Sub(out_len, one), one) + in_rest = op.Sub(in_len, one) + ratio = op.Div(op.Cast(in_rest, to=1), op.Cast(out_rest, to=1)) + shifted = op.Cast(op.Sub(positions, op.Squeeze(one, [0])), to=1) + split = op.Add( + op.Squeeze(one, [0]), + op.Cast(op.Floor(op.Mul(shifted, op.Squeeze(ratio, [0]))), to=7), + ) + split = op.Where(op.Equal(positions, op.Squeeze(zero, [0])), op.Squeeze(zero, [0]), split) + + is_odd = op.And( + op.Equal(op.Mod(out_len, two), one), + op.Greater(out_len, one), + ) + return op.Where(op.Squeeze(is_odd, [0]), split, uniform) + + +def _align_conditioning(op: OpBuilder, zq: ir.Value, target: ir.Value) -> ir.Value: + """Resample ``zq`` onto ``target``'s ``(T, H, W)`` grid, causally in time.""" + target_t = _dim(op, target, 2) + target_h = _dim(op, target, 3) + target_w = _dim(op, target, 4) + + time_indices = _causal_nearest_indices(op, target_t, _dim(op, zq, 2)) + zq = op.Gather(zq, time_indices, axis=2) + + # Spatial nearest resize; the temporal extent is already correct so the + # requested size keeps it unchanged. + sizes = op.Concat( + op.Shape(zq, start=0, end=2), + target_t, + target_h, + target_w, + axis=0, + ) + return op.Resize( + zq, + None, + None, + sizes, + mode="nearest", + coordinate_transformation_mode="asymmetric", + nearest_mode="floor", + ) + + +# --------------------------------------------------------------------------- +# Layers +# --------------------------------------------------------------------------- + + +class _SafeConv3d(nn.Module): + """Plain 3D convolution (HF ``CogVideoXSafeConv3d``, no temporal padding).""" + + def __init__( + self, + in_channels: int, + out_channels: int, + kernel_size: tuple[int, int, int] = (1, 1, 1), + padding: tuple[int, int, int] = (0, 0, 0), + ): + super().__init__() + self.weight = nn.Parameter((out_channels, in_channels, *kernel_size)) + self.bias = nn.Parameter((out_channels,)) + self._kernel_size = list(kernel_size) + self._padding = list(padding) + + def forward(self, op: OpBuilder, x: ir.Value) -> ir.Value: + return op.Conv( + x, + self.weight, + self.bias, + kernel_shape=self._kernel_size, + strides=[1, 1, 1], + pads=self._padding + self._padding, + ) + + +class _ConvCacheScope: + """Carries HuggingFace's ``conv_cache`` tensors through the decoder modules. + + A CogVideoX clip is decoded a few latent frames at a time. Each causal + convolution therefore has to start from the tail of the previous chunk + instead of replicating its own first frame, exactly as the reference + implementation's ``conv_cache`` dictionary does. + """ + + def __init__(self, inputs: dict[str, ir.Value]): + self.inputs = inputs + self.outputs: dict[str, ir.Value] = {} + + +class _CausalConv3d(nn.Module): + """Temporally causal 3D convolution. + + The ``conv`` attribute mirrors HuggingFace's ``CogVideoXCausalConv3d.conv`` + so parameter names match the checkpoint exactly. + """ + + def __init__(self, in_channels: int, out_channels: int, kernel_size: int = 3): + super().__init__() + spatial_pad = (kernel_size - 1) // 2 + self.conv = _SafeConv3d( + in_channels, + out_channels, + (kernel_size, kernel_size, kernel_size), + padding=(0, spatial_pad, spatial_pad), + ) + self._time_pad = kernel_size - 1 + + def forward( + self, + op: OpBuilder, + x: ir.Value, + scope: _ConvCacheScope | None = None, + path: str = "", + ) -> ir.Value: + if self._time_pad == 0: + return self.conv(op, x) + + previous = None if scope is None else scope.inputs.get(path) + if previous is None: + # Start of a clip: replicate the first latent frame ``k - 1`` times + # so frame ``t`` only ever sees frames ``<= t``. + first = op.Slice( + x, + op.Constant(value_ints=[0]), + op.Constant(value_ints=[1]), + op.Constant(value_ints=[2]), + ) + x = op.Concat(*([first] * self._time_pad), x, axis=2) + else: + # A zero-length cache means "no previous chunk"; edge padding then + # falls back to replicating frame 0, which is the same branch-free + # expression as the clip-start case above. + x = op.Concat(previous, x, axis=2) + deficit = op.Sub( + op.Constant(value_ints=[self._time_pad]), + op.Min(_dim(op, previous, 2), op.Constant(value_ints=[self._time_pad])), + ) + x = op.Pad( + x, + op.Concat(deficit, op.Constant(value_ints=[0]), axis=0), + None, + op.Constant(value_ints=[2]), + mode="edge", + ) + + if scope is not None: + # The reference caches the tail of the temporally padded input. + scope.outputs[path] = op.Slice( + x, + op.Constant(value_ints=[-self._time_pad]), + op.Constant(value_ints=[_INT64_MAX]), + op.Constant(value_ints=[2]), + ) + return self.conv(op, x) + + +class _SpatialNorm3D(nn.Module): + """Latent-conditioned normalization (HF ``CogVideoXSpatialNorm3D``). + + The reference implementation always builds its ``GroupNorm`` with + ``eps=1e-6``, independently of the autoencoder's ``norm_eps``. + """ + + def __init__(self, f_channels: int, zq_channels: int, groups: int): + super().__init__() + self.norm_layer = _GroupNorm(groups, f_channels, eps=1e-6) + self.conv_y = _CausalConv3d(zq_channels, f_channels, kernel_size=1) + self.conv_b = _CausalConv3d(zq_channels, f_channels, kernel_size=1) + + def forward(self, op: OpBuilder, f: ir.Value, zq: ir.Value) -> ir.Value: + zq = _align_conditioning(op, zq, f) + scale = self.conv_y(op, zq) + shift = self.conv_b(op, zq) + return op.Add(op.Mul(self.norm_layer(op, f), scale), shift) + + +class _ResnetBlock3D(nn.Module): + """Latent-conditioned causal 3D residual block (HF ``CogVideoXResnetBlock3D``).""" + + def __init__( + self, + in_channels: int, + out_channels: int, + zq_channels: int, + groups: int, + ): + super().__init__() + self.norm1 = _SpatialNorm3D(in_channels, zq_channels, groups) + self.conv1 = _CausalConv3d(in_channels, out_channels, kernel_size=3) + self.norm2 = _SpatialNorm3D(out_channels, zq_channels, groups) + self.conv2 = _CausalConv3d(out_channels, out_channels, kernel_size=3) + self.conv_shortcut = ( + _SafeConv3d(in_channels, out_channels, (1, 1, 1)) + if in_channels != out_channels + else None + ) + self._silu = _SiLU() + + def forward( + self, + op: OpBuilder, + hidden_states: ir.Value, + zq: ir.Value, + scope: _ConvCacheScope | None = None, + path: str = "", + ) -> ir.Value: + residual = hidden_states + hidden_states = self.norm1(op, hidden_states, zq) + hidden_states = self._silu(op, hidden_states) + hidden_states = self.conv1(op, hidden_states, scope, f"{path}conv1") + hidden_states = self.norm2(op, hidden_states, zq) + hidden_states = self._silu(op, hidden_states) + hidden_states = self.conv2(op, hidden_states, scope, f"{path}conv2") + if self.conv_shortcut is not None: + residual = self.conv_shortcut(op, residual) + return op.Add(hidden_states, residual) + + +class _Upsample3D(nn.Module): + """Spatial (and optionally temporal) 2x upsample (HF ``CogVideoXUpsample3D``). + + The trailing convolution is a per-frame ``Conv2d`` in HuggingFace. It is + emitted here as a ``(1, 3, 3)`` 3D convolution, which is the same linear + map without the reshape round trip; ``preprocess_weights`` inserts the + singleton temporal axis into the checkpoint weight. + """ + + def __init__(self, channels: int, compress_time: bool): + super().__init__() + self.conv = _SafeConv3d(channels, channels, (1, 3, 3), padding=(0, 1, 1)) + self._compress_time = compress_time + + def forward(self, op: OpBuilder, hidden_states: ir.Value) -> ir.Value: + if self._compress_time: + frames = _dim(op, hidden_states, 2) + one = op.Constant(value_ints=[1]) + two = op.Constant(value_ints=[2]) + is_odd = op.Equal(op.Mod(frames, two), one) + # Odd clips keep frame 0 fixed and double the remainder, matching + # the reference implementation's split-first-frame branch. + odd_len = op.Sub(op.Mul(frames, two), one) + even_len = op.Mul(frames, two) + out_len = op.Where(is_odd, odd_len, even_len) + + positions = op.Range( + op.Squeeze(op.Constant(value_ints=[0]), [0]), + op.Squeeze(out_len, [0]), + op.Squeeze(one, [0]), + ) + zero_scalar = op.Squeeze(op.Constant(value_ints=[0]), [0]) + one_scalar = op.Squeeze(one, [0]) + two_scalar = op.Squeeze(two, [0]) + odd_indices = op.Where( + op.Equal(positions, zero_scalar), + zero_scalar, + op.Add(one_scalar, op.Div(op.Sub(positions, one_scalar), two_scalar)), + ) + even_indices = op.Div(positions, two_scalar) + indices = op.Where(op.Squeeze(is_odd, [0]), odd_indices, even_indices) + hidden_states = op.Gather(hidden_states, indices, axis=2) + + hidden_states = op.Resize( + hidden_states, + None, + op.Constant(value_floats=[1.0, 1.0, 1.0, 2.0, 2.0]), + None, + mode="nearest", + coordinate_transformation_mode="asymmetric", + nearest_mode="floor", + ) + return self.conv(op, hidden_states) + + +class _MidBlock3D(nn.Module): + """Decoder mid block: two latent-conditioned residual blocks.""" + + def __init__(self, channels: int, zq_channels: int, groups: int, layers: int): + super().__init__() + self.resnets = nn.ModuleList( + [_ResnetBlock3D(channels, channels, zq_channels, groups) for _ in range(layers)] + ) + + def forward( + self, + op: OpBuilder, + hidden_states: ir.Value, + zq: ir.Value, + scope: _ConvCacheScope | None = None, + path: str = "", + ) -> ir.Value: + for index, resnet in enumerate(self.resnets): + hidden_states = resnet(op, hidden_states, zq, scope, f"{path}resnets.{index}.") + return hidden_states + + +class _UpBlock3D(nn.Module): + """Decoder up block: residual blocks plus an optional upsampler.""" + + def __init__( + self, + in_channels: int, + out_channels: int, + zq_channels: int, + groups: int, + layers: int, + add_upsample: bool, + compress_time: bool, + ): + super().__init__() + self.resnets = nn.ModuleList( + [ + _ResnetBlock3D( + in_channels if index == 0 else out_channels, + out_channels, + zq_channels, + groups, + ) + for index in range(layers) + ] + ) + self.upsamplers = ( + nn.ModuleList([_Upsample3D(out_channels, compress_time)]) if add_upsample else None + ) + + def forward( + self, + op: OpBuilder, + hidden_states: ir.Value, + zq: ir.Value, + scope: _ConvCacheScope | None = None, + path: str = "", + ) -> ir.Value: + for index, resnet in enumerate(self.resnets): + hidden_states = resnet(op, hidden_states, zq, scope, f"{path}resnets.{index}.") + if self.upsamplers is not None: + for upsampler in self.upsamplers: + hidden_states = upsampler(op, hidden_states) + return hidden_states + + +class _CogVideoXDecoder3D(nn.Module): + """Causal 3D decoder (HF ``CogVideoXDecoder3D``).""" + + def __init__(self, config: CogVideoXVAEConfig): + super().__init__() + reversed_channels = list(reversed(config.block_out_channels)) + zq_channels = config.latent_channels + groups = config.norm_num_groups + + self.conv_in = _CausalConv3d(zq_channels, reversed_channels[0], kernel_size=3) + self.mid_block = _MidBlock3D(reversed_channels[0], zq_channels, groups, layers=2) + + temporal_levels = int(math.log2(config.temporal_compression_ratio)) + self.up_blocks = nn.ModuleList() + output_channel = reversed_channels[0] + for index, channels in enumerate(reversed_channels): + prev_output_channel = output_channel + output_channel = channels + self.up_blocks.append( + _UpBlock3D( + prev_output_channel, + output_channel, + zq_channels, + groups, + layers=config.layers_per_block + 1, + add_upsample=index != len(reversed_channels) - 1, + compress_time=index < temporal_levels, + ) + ) + + self.norm_out = _SpatialNorm3D(reversed_channels[-1], zq_channels, groups) + self.conv_out = _CausalConv3d( + reversed_channels[-1], config.out_channels, kernel_size=3 + ) + self._silu = _SiLU() + + def forward( + self, + op: OpBuilder, + latent_sample: ir.Value, + scope: _ConvCacheScope | None = None, + ) -> ir.Value: + # The un-scaled latent conditions every normalization layer, so it is + # threaded through the whole decoder as ``zq``. + zq = latent_sample + hidden_states = self.conv_in(op, latent_sample, scope, "conv_in") + hidden_states = self.mid_block(op, hidden_states, zq, scope, "mid_block.") + for index, up_block in enumerate(self.up_blocks): + hidden_states = up_block(op, hidden_states, zq, scope, f"up_blocks.{index}.") + hidden_states = self.norm_out(op, hidden_states, zq) + hidden_states = self._silu(op, hidden_states) + return self.conv_out(op, hidden_states, scope, "conv_out") + + +class ConvCacheEntry(typing.NamedTuple): + """Shape contract of one carried ``conv_cache`` tensor. + + ``channels`` and ``spatial_scale`` are relative to the latent grid, so the + tensor is ``[batch, channels, frames, latent_height * spatial_scale, + latent_width * spatial_scale]``. + """ + + name: str + channels: int + frames: int + spatial_scale: int + + +class AutoencoderKLCogVideoXModel(nn.Module): + """CogVideoX 3D causal video autoencoder (decode path). + + Decodes ``[B, latent_channels, T_latent, H, W]`` latents into + ``[B, out_channels, T_pixels, H * s, W * s]`` video frames. + + Replicates HuggingFace diffusers' ``AutoencoderKLCogVideoX``. + """ + + default_task: str = "video-vae" + config_class = CogVideoXVAEConfig + category: str = "Diffusion" + + def __init__(self, config: CogVideoXVAEConfig): + super().__init__() + self.config = config + if config.use_post_quant_conv or config.use_quant_conv: + raise NotImplementedError( + "AutoencoderKLCogVideoX quant/post-quant convolutions are not built; " + "every published CogVideoX checkpoint sets use_quant_conv=False and " + "use_post_quant_conv=False." + ) + self.decoder = _CogVideoXDecoder3D(config) + + def conv_cache_spec(self) -> list[ConvCacheEntry]: + """Ordered contract of the carried ``conv_cache`` tensors. + + Mirrors the traversal order of :class:`_CogVideoXDecoder3D.forward`, so + callers can size the state cells for a chunked decode without inspecting + the graph. Only ``kernel_t > 1`` convolutions carry state; the ``k=1`` + convolutions inside the spatial norms consume no temporal context. + """ + config = self.config + reversed_channels = list(reversed(config.block_out_channels)) + frames = 2 # kernel_t - 1 for every cached convolution + entries = [ConvCacheEntry("conv_in", config.latent_channels, frames, 1)] + + head = reversed_channels[0] + for index in range(2): + entries.append(ConvCacheEntry(f"mid_block.resnets.{index}.conv1", head, frames, 1)) + entries.append(ConvCacheEntry(f"mid_block.resnets.{index}.conv2", head, frames, 1)) + + scale = 1 + output_channel = head + for block, channels in enumerate(reversed_channels): + prev_output_channel = output_channel + output_channel = channels + for layer in range(config.layers_per_block + 1): + in_channels = prev_output_channel if layer == 0 else output_channel + prefix = f"up_blocks.{block}.resnets.{layer}" + entries.append(ConvCacheEntry(f"{prefix}.conv1", in_channels, frames, scale)) + entries.append( + ConvCacheEntry(f"{prefix}.conv2", output_channel, frames, scale) + ) + if block != len(reversed_channels) - 1: + # The upsampler doubles the spatial grid for every later block. + scale *= 2 + + entries.append(ConvCacheEntry("conv_out", reversed_channels[-1], frames, scale)) + return entries + + def forward( + self, + op: OpBuilder, + latent_sample: ir.Value, + conv_cache: dict[str, ir.Value] | None = None, + ): + """Decode one chunk of latent frames. + + Args: + op: ONNX op builder. + latent_sample: ``[B, latent_channels, T_latent, H, W]``. + conv_cache: Carried temporal context per cached convolution. Pass + zero-length (``frames == 0``) tensors for the first chunk of a + clip; omit entirely to decode a clip in one call. + + Returns: + The decoded frames, or ``(frames, updated_conv_cache)`` when + ``conv_cache`` is supplied. + """ + if conv_cache is None: + return self.decoder(op, latent_sample) + scope = _ConvCacheScope(conv_cache) + sample = self.decoder(op, latent_sample, scope) + return sample, scope.outputs + + def preprocess_weights( + self, state_dict: dict[str, torch.Tensor] + ) -> dict[str, torch.Tensor]: + """Adapt the diffusers checkpoint to this module's parameter names. + + - Encoder parameters are dropped; only the decode path is built. + - ``upsamplers.*.conv.weight`` gains a singleton temporal axis because + the per-frame ``Conv2d`` is emitted as a ``(1, 3, 3)`` 3D convolution. + """ + processed: dict[str, torch.Tensor] = {} + for name, tensor in state_dict.items(): + if name.startswith("encoder."): + continue + if ".upsamplers." in name and name.endswith(".conv.weight") and tensor.ndim == 4: + tensor = tensor.unsqueeze(2) + processed[name] = tensor + return processed diff --git a/src/mobius/models/distilbert.py b/src/mobius/models/distilbert.py index 3130698fd..d7e814ff4 100644 --- a/src/mobius/models/distilbert.py +++ b/src/mobius/models/distilbert.py @@ -12,7 +12,7 @@ from mobius._configs import ArchitectureConfig from mobius.components import FCMLP -from mobius.components._common import Embedding, LayerNorm +from mobius.components._common import Embedding, LayerNorm, create_padding_mask from mobius.components._encoder import EncoderAttention if TYPE_CHECKING: @@ -61,9 +61,20 @@ def __init__(self, config: ArchitectureConfig): ) self.output_layer_norm = LayerNorm(config.hidden_size, eps=config.rms_norm_eps) - def forward(self, op: OpBuilder, hidden_states: ir.Value): + def forward( + self, + op: OpBuilder, + hidden_states: ir.Value, + attention_mask: ir.Value | None = None, + ): + """One encoder block. + + ``attention_mask`` is the 4D bool padding mask ``(batch, 1, seq, seq)`` + prepared once per forward by :meth:`DistilBertModel.forward`, not the + raw rank-2 request mask. + """ # Self-attention with post-norm - attn_output = self.attention(op, hidden_states) + attn_output = self.attention(op, hidden_states, attention_mask) hidden_states = self.sa_layer_norm(op, op.Add(hidden_states, attn_output)) # FFN with post-norm @@ -98,7 +109,13 @@ def forward( token_type_ids: ir.Value | None = None, ): hidden_states = self.embeddings(op, input_ids) - hidden_states = self.transformer(op, hidden_states) + # Build the 4D bool padding mask once and share it across layers: the + # raw rank-2 request mask is INT64 and does not broadcast to + # ``(batch, heads, q, kv)``, so feeding it to ``op.Attention`` fails + # outright for batch > 1 and adds a 0/1 bias -- not a mask -- at + # batch 1. + padding_mask = create_padding_mask(op, input_ids, attention_mask) + hidden_states = self.transformer(op, hidden_states, padding_mask) return hidden_states def preprocess_weights( @@ -121,9 +138,14 @@ def __init__(self, config: ArchitectureConfig): [_DistilBertEncoderLayer(config) for _ in range(config.num_hidden_layers)] ) - def forward(self, op: OpBuilder, hidden_states: ir.Value): + def forward( + self, + op: OpBuilder, + hidden_states: ir.Value, + attention_mask: ir.Value | None = None, + ): for layer in self.layer: - hidden_states = layer(op, hidden_states) + hidden_states = layer(op, hidden_states, attention_mask) return hidden_states diff --git a/src/mobius/models/esm.py b/src/mobius/models/esm.py new file mode 100644 index 000000000..f5f2342ec --- /dev/null +++ b/src/mobius/models/esm.py @@ -0,0 +1,407 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT License. + +"""ESM-2 protein language model encoder. + +Replicates HuggingFace's ``EsmModel`` (``transformers.models.esm``) for the +``feature-extraction`` task: amino-acid token ids in, per-residue contextual +embeddings out. + +ESM-2 is *not* a BERT clone, and the differences all change the numbers: + +* **Rotary position embeddings.** ``config.position_embedding_type`` is + ``"rotary"``; the learned ``embeddings.position_embeddings`` table shipped in + the checkpoint is dead weight and is dropped. Rotary is applied to Q and K + inside every layer, using positions ``0..seq_len-1``. +* **Pre-LayerNorm blocks.** ``attention.LayerNorm`` runs *before* self-attention + and the layer's own ``LayerNorm`` runs *before* the feed-forward, with plain + residual adds after each. BERT is post-norm, so reusing the BERT block would + put every norm in the wrong place. +* **No token-type embeddings and no embedding LayerNorm.** ``token_type_ids`` + is accepted for task-signature compatibility and ignored; + ``emb_layer_norm_before`` is ``False`` for the released ESM-2 checkpoints. +* **A final ``emb_layer_norm_after``** closes the encoder stack. +* **Token dropout.** When ``config.token_dropout`` is set, masked positions are + zeroed and the whole embedding is rescaled by + ``(1 - 0.15*0.8) / (1 - observed_mask_ratio)`` — a factor of ``0.88`` even + when no ```` token is present, so it cannot be skipped. +* **Padding is zeroed in the embedding**, not only masked in attention. + +Inputs: + input_ids: ``(batch, sequence_len)`` INT64 amino-acid token ids. + attention_mask: ``(batch, sequence_len)`` INT64, 1 = residue, 0 = padding. + token_type_ids: ``(batch, sequence_len)`` INT64, unused (task signature). + +Outputs: + last_hidden_state: ``(batch, sequence_len, hidden_size)`` per-residue + embeddings. +""" + +from __future__ import annotations + +import dataclasses +from typing import TYPE_CHECKING + +import torch +from onnxscript import OpBuilder, nn + +from mobius._configs import ArchitectureConfig +from mobius.components._activations import ACT2FN +from mobius.components._common import ( + Embedding, + LayerNorm, + Linear, + create_padding_mask, +) +from mobius.components._rotary_embedding import DefaultRope, apply_rotary_pos_emb + +if TYPE_CHECKING: + import onnx_ir as ir + +#: HuggingFace's ``EsmEmbeddings`` hard-codes the training-time masking rate as +#: ``0.15 * 0.8`` (15% of positions selected, 80% of those replaced by +#: ````). It is a property of how ESM-2 was trained, not a config value. +_MASK_RATIO_TRAIN = 0.15 * 0.8 + + +@dataclasses.dataclass +class EsmConfig(ArchitectureConfig): + """ESM-2 architecture config. + + Adds the four ESM-specific switches that decide *which* graph is built. + They are read from the HuggingFace config rather than assumed, because the + ESM family ships checkpoints on both sides of each one (ESM-1b uses + absolute positions and a pre-encoder LayerNorm; ESM-2 uses rotary and none). + """ + + position_embedding_type: str = "rotary" + emb_layer_norm_before: bool = False + token_dropout: bool = True + mask_token_id: int = 32 + + @classmethod + def from_transformers(cls, config, parent_config=None) -> EsmConfig: + base = super().from_transformers(config, parent_config=parent_config) + # ESM-2 declares rotary through ``position_embedding_type`` alone: its + # HuggingFace config carries no ``rope_theta`` / ``rope_scaling``, so + # the generic RoPE extractor sees no signal and leaves the fields unset. + # HF's ``EsmRotaryEmbedding`` hard-codes base 10000 over the full head + # dimension, which is what these values restate. + return dataclasses.replace( + base, + position_embedding_type=getattr(config, "position_embedding_type", "rotary"), + emb_layer_norm_before=bool(getattr(config, "emb_layer_norm_before", False)), + token_dropout=bool(getattr(config, "token_dropout", False)), + mask_token_id=int(getattr(config, "mask_token_id", 32) or 32), + rope_type="default", + rope_theta=10000.0, + partial_rotary_factor=1.0, + ) + + +class _EsmEmbeddings(nn.Module): + """Word embeddings + ESM token dropout + padding zeroing. + + Mirrors ``EsmEmbeddings``. The learned ``position_embeddings`` table is not + instantiated: for ``position_embedding_type == "rotary"`` HuggingFace never + reads it, and materializing it would add an initializer the graph never + uses. + """ + + def __init__(self, config: EsmConfig): + super().__init__() + self.word_embeddings = Embedding( + config.vocab_size, config.hidden_size, config.pad_token_id or 0 + ) + self.token_dropout = config.token_dropout + self.mask_token_id = config.mask_token_id + self.layer_norm = ( + LayerNorm(config.hidden_size, eps=config.rms_norm_eps) + if config.emb_layer_norm_before + else None + ) + + def forward(self, op: OpBuilder, input_ids: ir.Value, attention_mask: ir.Value): + # (batch, seq) -> (batch, seq, hidden) + embeddings = self.word_embeddings(op, input_ids) + + if self.token_dropout: + # Zero the rows, then rescale so the expected embedding + # magnitude matches training. With no present this is still + # a 0.88 scale, so it is not an inference-time no-op. + is_mask = op.Equal( + input_ids, op.Constant(value_int=self.mask_token_id) + ) # (batch, seq) + zero = op.CastLike(op.Constant(value_float=0.0), embeddings) + embeddings = op.Where(op.Unsqueeze(is_mask, [-1]), zero, embeddings) + + mask_float = op.CastLike(is_mask, embeddings) + valid_float = op.CastLike(attention_mask, embeddings) + # (batch,) counts of masked residues and of real residues + masked_count = op.ReduceSum(mask_float, [-1], keepdims=0) + src_lengths = op.ReduceSum(valid_float, [-1], keepdims=0) + observed = op.Div(masked_count, src_lengths) + one = op.CastLike(op.Constant(value_float=1.0), embeddings) + kept = op.CastLike(op.Constant(value_float=1.0 - _MASK_RATIO_TRAIN), embeddings) + scale = op.Div(kept, op.Sub(one, observed)) # (batch,) + # (batch,) -> (batch, 1, 1) so it scales every residue of a row + embeddings = op.Mul(embeddings, op.Unsqueeze(scale, [-1, -2])) + + if self.layer_norm is not None: + embeddings = self.layer_norm(op, embeddings) + + # ESM zeroes padded positions in the embedding itself, in addition to + # masking them in attention. + pad_scale = op.Unsqueeze(op.CastLike(attention_mask, embeddings), [-1]) + return op.Mul(embeddings, pad_scale) + + +class _EsmSelfAttention(nn.Module): + """Rotary bidirectional self-attention (HF ``EsmSelfAttention``).""" + + def __init__(self, hidden_size: int, num_heads: int): + super().__init__() + self.num_heads = num_heads + self.head_dim = hidden_size // num_heads + self.query = Linear(hidden_size, hidden_size, bias=True) + self.key = Linear(hidden_size, hidden_size, bias=True) + self.value = Linear(hidden_size, hidden_size, bias=True) + + def forward( + self, + op: OpBuilder, + hidden_states: ir.Value, + attention_mask: ir.Value, + position_embeddings: tuple, + ): + query = self.query(op, hidden_states) + key = self.key(op, hidden_states) + value = self.value(op, hidden_states) + + # Rotary acts within each head; HF scales Q by head_dim**-0.5 before + # the rotation, which commutes with it, so the equivalent `scale` + # attribute below reproduces the same scores. + query = apply_rotary_pos_emb( + op, x=query, position_embeddings=position_embeddings, num_heads=self.num_heads + ) + key = apply_rotary_pos_emb( + op, x=key, position_embeddings=position_embeddings, num_heads=self.num_heads + ) + + return op.Attention( + query, + key, + value, + attention_mask, + q_num_heads=self.num_heads, + kv_num_heads=self.num_heads, + scale=float(self.head_dim**-0.5), + ) + + +class _EsmSelfOutput(nn.Module): + """Attention output projection + residual (no LayerNorm — ESM is pre-norm).""" + + def __init__(self, hidden_size: int): + super().__init__() + self.dense = Linear(hidden_size, hidden_size, bias=True) + + def forward(self, op: OpBuilder, hidden_states: ir.Value, input_tensor: ir.Value): + return op.Add(self.dense(op, hidden_states), input_tensor) + + +class _EsmAttention(nn.Module): + """Pre-norm self-attention block. + + Parameter paths match HuggingFace: + ``attention.LayerNorm`` / ``attention.self.query`` / + ``attention.output.dense``. + """ + + def __init__(self, hidden_size: int, num_heads: int, eps: float): + super().__init__() + self.self = _EsmSelfAttention(hidden_size, num_heads) + self.output = _EsmSelfOutput(hidden_size) + # Capital 'LayerNorm' matches HF ESM naming; it runs *before* attention. + self.LayerNorm = LayerNorm(hidden_size, eps=eps) + + def forward( + self, + op: OpBuilder, + hidden_states: ir.Value, + attention_mask: ir.Value, + position_embeddings: tuple, + ): + normed = self.LayerNorm(op, hidden_states) + attn_out = self.self(op, normed, attention_mask, position_embeddings) + return self.output(op, attn_out, hidden_states) + + +class _EsmIntermediate(nn.Module): + """Feed-forward up-projection + activation (HF naming).""" + + def __init__(self, hidden_size: int, intermediate_size: int, hidden_act: str): + super().__init__() + self.dense = Linear(hidden_size, intermediate_size, bias=True) + self._act_fn = ACT2FN[hidden_act] + + def forward(self, op: OpBuilder, hidden_states: ir.Value): + return self._act_fn(op, self.dense(op, hidden_states)) + + +class _EsmOutput(nn.Module): + """Feed-forward down-projection + residual (no LayerNorm — ESM is pre-norm).""" + + def __init__(self, intermediate_size: int, hidden_size: int): + super().__init__() + self.dense = Linear(intermediate_size, hidden_size, bias=True) + + def forward(self, op: OpBuilder, hidden_states: ir.Value, input_tensor: ir.Value): + return op.Add(self.dense(op, hidden_states), input_tensor) + + +class _EsmLayer(nn.Module): + """One pre-norm ESM encoder layer. + + Parameter paths match HuggingFace: + ``layer.N.attention.*`` / ``layer.N.LayerNorm`` / + ``layer.N.intermediate.dense`` / ``layer.N.output.dense``. + """ + + def __init__(self, config: EsmConfig): + super().__init__() + self.attention = _EsmAttention( + config.hidden_size, config.num_attention_heads, config.rms_norm_eps + ) + self.intermediate = _EsmIntermediate( + config.hidden_size, config.intermediate_size, config.hidden_act + ) + self.output = _EsmOutput(config.intermediate_size, config.hidden_size) + # Pre-feed-forward norm. + self.LayerNorm = LayerNorm(config.hidden_size, eps=config.rms_norm_eps) + + def forward( + self, + op: OpBuilder, + hidden_states: ir.Value, + attention_mask: ir.Value, + position_embeddings: tuple, + ): + hidden_states = self.attention(op, hidden_states, attention_mask, position_embeddings) + normed = self.LayerNorm(op, hidden_states) + return self.output(op, self.intermediate(op, normed), hidden_states) + + +class _EsmEncoder(nn.Module): + """Stack of pre-norm layers closed by ``emb_layer_norm_after``.""" + + def __init__(self, config: EsmConfig): + super().__init__() + self.layer = nn.ModuleList( + [_EsmLayer(config) for _ in range(config.num_hidden_layers)] + ) + self.emb_layer_norm_after = LayerNorm(config.hidden_size, eps=config.rms_norm_eps) + + def forward( + self, + op: OpBuilder, + hidden_states: ir.Value, + attention_mask: ir.Value, + position_embeddings: tuple, + ): + for layer in self.layer: + hidden_states = layer(op, hidden_states, attention_mask, position_embeddings) + return self.emb_layer_norm_after(op, hidden_states) + + +class EsmModel(nn.Module): + """ESM-2 protein encoder for per-residue feature extraction. + + Replicates HuggingFace's ``EsmModel``; the output is ``last_hidden_state`` + (the pooler and the contact head are not part of the embedding contract and + are dropped). + """ + + default_task = "feature-extraction" + category = "encoder" + config_class = EsmConfig + + def __init__(self, config: EsmConfig): + super().__init__() + if config.position_embedding_type != "rotary": + raise ValueError( + "EsmModel currently builds the rotary ESM-2 variant; got " + f"position_embedding_type={config.position_embedding_type!r}" + ) + self.config = config + self.embeddings = _EsmEmbeddings(config) + self.encoder = _EsmEncoder(config) + # One shared rotary table: HF instantiates a RotaryEmbedding per layer, + # but every copy holds the same inv_freq, so a single cache is emitted. + self.rotary_emb = DefaultRope(config) + + def forward( + self, + op: OpBuilder, + input_ids: ir.Value, + attention_mask: ir.Value, + token_type_ids: ir.Value, # Unused: ESM has no token-type embeddings. + ): + del token_type_ids + hidden_states = self.embeddings(op, input_ids, attention_mask) + + # Rotary positions are 0..seq_len-1 for every row, matching HF, which + # derives them from the tensor shape rather than the mask. The ONNX + # ``RotaryEmbedding`` op requires cos/sin to carry the same batch extent + # as ``x``, so the row vector is expanded to (batch, seq) rather than + # left at (1, seq). + batch = op.Shape(input_ids, start=0, end=1) + seq_len = op.Shape(input_ids, start=1, end=2) + position_ids = op.Range( + op.Constant(value_int=0), + op.Squeeze(seq_len), + op.Constant(value_int=1), + ) + position_ids = op.Cast(position_ids, to=7) # INT64 + position_ids = op.Unsqueeze(position_ids, [0]) # (1, seq) + position_ids = op.Expand(position_ids, op.Concat(batch, seq_len, axis=0)) + position_embeddings = self.rotary_emb(op, position_ids) + + # (batch, 1, seq, seq) bool mask; a rank-2 int mask cannot broadcast + # onto the (batch, heads, q, kv) score tensor once batch > 1. + padding_mask = create_padding_mask(op, input_ids, attention_mask) + return self.encoder(op, hidden_states, padding_mask, position_embeddings) + + def preprocess_weights( + self, state_dict: dict[str, torch.Tensor] + ) -> dict[str, torch.Tensor]: + """Map HuggingFace ESM weight names onto this module tree.""" + new_state_dict: dict[str, torch.Tensor] = {} + for name, tensor in state_dict.items(): + new_name = _rename_esm_weight(name) + if new_name is not None: + new_state_dict[new_name] = tensor + return new_state_dict + + +def _rename_esm_weight(name: str) -> str | None: + """Rename one HuggingFace ESM weight, or drop it. + + Dropped: the masked-LM head, the contact head, the pooler, the unused + absolute ``position_embeddings`` table, the ``position_ids`` buffer, and the + per-layer ``inv_freq`` rotary buffers (this module emits one shared cos/sin + cache computed from ``rope_theta`` instead). + """ + if name.startswith("esm."): + name = name[4:] + + if name.startswith(("lm_head.", "contact_head.", "pooler.", "cls.")): + return None + if name in ("embeddings.position_ids", "embeddings.position_embeddings.weight"): + return None + if name.endswith(".rotary_embeddings.inv_freq"): + return None + + # attention.self.query -> attention.self.query (kept: HF nesting is mirrored + # by the module tree), attention.output.dense -> attention.output.dense. + return name diff --git a/src/mobius/models/esm_test.py b/src/mobius/models/esm_test.py new file mode 100644 index 000000000..c58018cec --- /dev/null +++ b/src/mobius/models/esm_test.py @@ -0,0 +1,125 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT License. + +"""Unit tests for the ESM-2 protein encoder. + +These build the ONNX graph from a tiny config -- no weights, no network -- and +assert the structural facts that separate ESM-2 from a BERT clone: rotary +positions instead of a learned position table, no token-type embedding, and a +final ``emb_layer_norm_after``. +""" + +from __future__ import annotations + +import onnx_ir as ir +import pytest + +from mobius.models.esm import EsmConfig, EsmModel, _rename_esm_weight +from mobius.tasks import FeatureExtractionTask + +#: Shape-faithful miniature of ``facebook/esm2_t6_8M_UR50D``: same vocabulary, +#: same special-token ids and the same architectural switches, with the widths +#: shrunk so the graph builds in well under a second. +TINY_CONFIG = EsmConfig( + hidden_size=64, + num_hidden_layers=2, + num_attention_heads=4, + num_key_value_heads=4, + intermediate_size=128, + vocab_size=33, + max_position_embeddings=1026, + hidden_act="gelu", + rms_norm_eps=1e-5, + position_embedding_type="rotary", + emb_layer_norm_before=False, + token_dropout=True, + mask_token_id=32, + pad_token_id=1, + rope_type="default", + rope_theta=10000.0, + partial_rotary_factor=1.0, +) + + +@pytest.fixture(scope="module") +def package(): + return FeatureExtractionTask().build(EsmModel(TINY_CONFIG), TINY_CONFIG) + + +class TestEsmGraph: + def test_emits_per_residue_embeddings(self, package) -> None: + model = package["model"] + outputs = {str(value.name): value for value in model.graph.outputs} + assert set(outputs) == {"last_hidden_state"} + shape = [str(dim) for dim in outputs["last_hidden_state"].shape] + assert shape[-1] == str(TINY_CONFIG.hidden_size) + assert outputs["last_hidden_state"].dtype == ir.DataType.FLOAT + + def test_declares_no_learned_position_table(self, package) -> None: + """ESM-2 is rotary; a learned position table would be dead weight.""" + names = set(package["model"].graph.initializers) + assert not any("position_embeddings" in name for name in names) + + def test_declares_no_token_type_embedding(self, package) -> None: + names = set(package["model"].graph.initializers) + assert not any("token_type_embeddings" in name for name in names) + + def test_closes_the_stack_with_a_final_layer_norm(self, package) -> None: + names = set(package["model"].graph.initializers) + assert "encoder.emb_layer_norm_after.weight" in names + + def test_shares_one_rotary_cache_across_layers(self, package) -> None: + """Every ESM-2 layer uses identical ``inv_freq``, so one cache suffices.""" + names = set(package["model"].graph.initializers) + caches = {name for name in names if "cos_cache" in name or "sin_cache" in name} + assert caches == {"rotary_emb.cos_cache", "rotary_emb.sin_cache"} + + def test_initializer_names_match_huggingface(self, package) -> None: + """The renamer only strips a prefix, so the scopes must already agree.""" + names = set(package["model"].graph.initializers) + for expected in ( + "embeddings.word_embeddings.weight", + "encoder.layer.0.attention.self.query.weight", + "encoder.layer.0.attention.output.dense.weight", + "encoder.layer.0.intermediate.dense.weight", + "encoder.layer.0.output.dense.weight", + ): + assert expected in names + + +class TestRenameEsmWeight: + @pytest.mark.parametrize( + "hf_name,expected", + [ + ("esm.embeddings.word_embeddings.weight", "embeddings.word_embeddings.weight"), + ( + "esm.encoder.layer.0.attention.self.query.weight", + "encoder.layer.0.attention.self.query.weight", + ), + ( + "esm.encoder.layer.0.attention.output.dense.bias", + "encoder.layer.0.attention.output.dense.bias", + ), + ("esm.encoder.emb_layer_norm_after.weight", "encoder.emb_layer_norm_after.weight"), + # A checkpoint saved without the task prefix passes through. + ("encoder.layer.1.output.dense.weight", "encoder.layer.1.output.dense.weight"), + ], + ) + def test_renames(self, hf_name: str, expected: str) -> None: + assert _rename_esm_weight(hf_name) == expected + + @pytest.mark.parametrize( + "hf_name", + [ + # Rotary is recomputed as a cache, absolute positions are unused, + # and the heads belong to other tasks. + "esm.embeddings.position_embeddings.weight", + "esm.embeddings.position_ids", + "esm.encoder.layer.0.attention.self.rotary_embeddings.inv_freq", + "esm.pooler.dense.weight", + "lm_head.decoder.weight", + "esm.contact_head.regression.weight", + ], + ) + def test_skipped_weights_return_none(self, hf_name: str) -> None: + assert _rename_esm_weight(hf_name) is None diff --git a/src/mobius/models/gemma4.py b/src/mobius/models/gemma4.py index 54d7de2c2..d69ce7a4e 100644 --- a/src/mobius/models/gemma4.py +++ b/src/mobius/models/gemma4.py @@ -32,7 +32,7 @@ import torch from onnxscript import OpBuilder, nn -from mobius._build_context import ep_capabilities +from mobius._build_context import ep_capabilities, is_prefill_prefix_pruning_enabled from mobius._configs import ArchitectureConfig, Gemma4Config from mobius._weight_utils import vlm_decoder_weights, vlm_embedding_weights from mobius.components import ( @@ -101,6 +101,39 @@ def _typed_scalar_constant(op: OpBuilder, value: float, dtype: ir.DataType) -> i return op.Constant(value=ir.tensor(np.asarray(value, dtype=dtype.numpy()))) +def _retain_last_position_embedding( + op: OpBuilder, position_embeddings: tuple[ir.Value, ir.Value] | None +) -> tuple[ir.Value, ir.Value] | None: + """Narrow a RoPE ``(cos, sin)`` pair to its final sequence position. + + Both caches are ``[B, S, rot_dim]`` and are indexed by query position, so + they must shrink in lockstep with the hidden states when prefill-prefix + pruning drops every position but the last. ``RotaryEmbedding`` requires + the cache sequence length to equal the query sequence length exactly. + """ + if position_embeddings is None: + return None + narrowed = [] + for cache in position_embeddings: + # (B, S, rot_dim) -> gather last position -> (B, rot_dim) -> (B, 1, rot_dim) + last = op.Gather(cache, op.Constant(value_int=-1), axis=1) + narrowed.append(op.Unsqueeze(last, op.Constant(value_ints=[1]))) + return (narrowed[0], narrowed[1]) + + +def _retain_last_bias_query_row(op: OpBuilder, bias: ir.Value | None) -> ir.Value | None: + """Narrow a ``[B, 1, S_q, S_kv]`` additive attention bias to its last query row. + + The key axis is left untouched: a pruned query still attends over the whole + key/value prefix, it just contributes a single query row. + """ + if bias is None: + return None + # (B, 1, S_q, S_kv) -> gather last query row -> (B, 1, S_kv) -> (B, 1, 1, S_kv) + last = op.Gather(bias, op.Constant(value_int=-1), axis=2) + return op.Unsqueeze(last, op.Constant(value_ints=[2])) + + def _text_quantization_config(config: Gemma4Config): """Return the active weight-quantization config, or ``None`` when off.""" quantization_config = getattr(config, "quantization", None) @@ -953,8 +986,15 @@ def forward( # ``o_proj``'s MatMul fail shape inference at model-load time. The # source shares the same KV-head/head-dim configuration as this # layer, so pin the known 4D BNSH shape to restore inference. + # + # Only an UNKNOWN shape may be pinned. A static-cache source hands + # over its rank-3 ``[batch, capacity, kv_heads * head_dim]`` scatter + # buffer, which is fully known; overwriting that with a 4D BNSH + # guess both mislabels the graph's declared cache output and defeats + # the rank-3 test below, which would then transpose a rank-3 tensor + # as if it were BNSH. for _kv in (src_key, src_value): - if _kv.shape is None or len(_kv.shape) != 4: + if _kv.shape is None: _kv.shape = ir.Shape( [ "batch", @@ -2210,6 +2250,30 @@ def forward( ): if i == self._first_kv_shared_layer and i < len(self.layers): hidden_states = _retain_last_sequence_token(op, hidden_states) + if is_prefill_prefix_pruning_enabled(): + # The stack has just narrowed to a single query position. + # Every per-layer tensor indexed by query position must + # narrow with it or the KV-shared layers receive an S-row + # RoPE cache / attention bias for a 1-row query. The GQA + # path is exempt: GroupQueryAttention derives its rotary + # offset from total_seq_len - q_len and takes a full-length + # cos/sin cache, so it already handles the narrowed query. + shared_pos_dict = fallback_pos_dict is position_embeddings_dict + position_embeddings_dict = { + key: _retain_last_position_embedding(op, value) + for key, value in position_embeddings_dict.items() + } + if shared_pos_dict: + fallback_pos_dict = position_embeddings_dict + else: + fallback_pos_dict = { + key: _retain_last_position_embedding(op, value) + for key, value in fallback_pos_dict.items() + } + fallback_bias_dict = { + key: _retain_last_bias_query_row(op, value) + for key, value in fallback_bias_dict.items() + } per_layer_input = per_layer_list[i] if per_layer_list is not None else None # Per-layer cache/attention dispatch: diff --git a/src/mobius/models/llada.py b/src/mobius/models/llada.py index b10f35a90..b11c6655e 100644 --- a/src/mobius/models/llada.py +++ b/src/mobius/models/llada.py @@ -6,9 +6,9 @@ LLaDA is a Llama backbone (RMSNorm, RoPE, SwiGLU MLP, full multi-head attention with no bias) run **bidirectionally**: every position attends to every other position, exactly like a BERT encoder and unlike a Llama decoder. -It is a mask predictor for discrete (masked) diffusion — onnx-genai's -``masked_diffusion`` scheduler drives the reverse process by repeatedly -feeding the predicted logits back as ``input_ids``. +It is a mask predictor for discrete (masked) diffusion. The onnx-genai generic +SSA workflow repeatedly invokes the denoiser and the packaged masked-update +policy artifact while carrying token, mask, and RNG state. Two properties distinguish it from the standard :class:`CausalLMModel`: @@ -177,8 +177,8 @@ class LLaDAModel(nn.Module): A Llama backbone run bidirectionally with a separate (untied) language modelling head. The forward pass maps ``input_ids [batch, sequence_len]`` (int64) to ``logits [batch, sequence_len, vocab_size]`` (float) in a - single full-sequence pass — the mask-predictor contract consumed by - onnx-genai's ``masked_diffusion`` scheduler. + single full-sequence pass. The task also exposes greedy token proposals for + the onnx-genai generic masked-update workflow. Replicates the ``LLaDALlamaBlock`` architecture of HuggingFace's ``LLaDAModelLM`` (``block_type: "llama"``, ``layer_norm_type: "rms"``, diff --git a/src/mobius/models/llada_test.py b/src/mobius/models/llada_test.py index 1a9afddea..89fe02641 100644 --- a/src/mobius/models/llada_test.py +++ b/src/mobius/models/llada_test.py @@ -187,10 +187,11 @@ def test_llada_matches_torch_reference(): expected = _reference_logits(config, state, input_ids).numpy() session = _build_onnx_session(config, state) - actual = session.run(None, {"input_ids": input_ids.numpy()})[0] + actual, proposed = session.run(None, {"input_ids": input_ids.numpy()}) max_delta = np.abs(actual - expected).max() assert max_delta < 1e-4, f"max|Δ|={max_delta}" + np.testing.assert_array_equal(proposed, np.argmax(actual, axis=-1)) def test_llada_attention_is_bidirectional(): @@ -216,23 +217,21 @@ def test_llada_attention_is_bidirectional(): def test_llada_export_signature_matches_masked_diffusion_metadata(): - """The exported ONNX I/O matches the onnx-genai masked-diffusion contract. + """The exported ONNX I/O matches the onnx-genai masked workflow contract. - Builds a tiny LLaDA package and asserts the graph exposes exactly - ``input_ids [B, S]`` int64 in and ``logits [B, S, V]`` f32 out with no - past/present KV, then checks that - :func:`build_language_diffusion_pipeline_metadata` emits a pipeline whose - denoiser self-edge references those same ports. + Builds a tiny LLaDA package and checks that the graph exposes the logits and + executable proposal ports consumed by the generic SSA workflow. """ import onnx_ir as ir - from mobius.integrations.onnx_genai.inference_metadata import ( + from mobius.integrations.onnx_genai.workflow_metadata import ( build_language_diffusion_pipeline_metadata, ) config = _make_config() module = LLaDAModel(config) - model = MaskedDiffusionTask().build(module, config)["model"] + package = MaskedDiffusionTask().build(module, config) + model = package["model"] graph = model.graph # Exactly one input: input_ids [B, S] int64. @@ -241,12 +240,14 @@ def test_llada_export_signature_matches_masked_diffusion_metadata(): assert input_ids.dtype == ir.DataType.INT64 assert len(input_ids.shape) == 2 - # Exactly one output: logits [B, S, V] float. - assert [value.name for value in graph.outputs] == ["logits"] + assert [value.name for value in graph.outputs] == ["logits", "proposed_tokens"] logits = graph.outputs[0] assert logits.dtype == ir.DataType.FLOAT assert len(logits.shape) == 3 assert logits.shape[2] == config.vocab_size + proposed = graph.outputs[1] + assert proposed.dtype == ir.DataType.INT64 + assert len(proposed.shape) == 2 # No KV-cache ports on either side. io_names = [value.name for value in (*graph.inputs, *graph.outputs)] @@ -254,22 +255,13 @@ def test_llada_export_signature_matches_masked_diffusion_metadata(): token in name for name in io_names for token in ("past", "present", "cache") ) - # The emitted metadata must wire those exact ports into a masked-diffusion - # iterative loop with a logits -> input_ids self-edge. meta = build_language_diffusion_pipeline_metadata( - mask_token_id=126336, + package, num_inference_steps=8, - input_ids_port="input_ids", - logits_port="logits", ) - pipeline = meta["pipeline"] - assert pipeline["models"]["denoiser"]["type"] == "denoiser" - assert pipeline["dataflow"] == [{"from": "denoiser.logits", "to": "denoiser.input_ids"}] - strategy = pipeline["strategy"] - assert strategy["kind"] == "iterative" - assert strategy["denoiser"] == "denoiser" - assert strategy["num_steps"] == 8 - assert strategy["scheduler_config"] == { - "kind": "masked_diffusion", - "mask_token_id": 126336, - } + workflow = meta["pipeline"]["workflow"] + assert "strategy" not in meta["pipeline"] + assert workflow["steps"][0]["kind"] == "loop" + assert workflow["components"]["masked_update"]["contract"]["id"] == ( + "onnx-genai.masked-update" + ) diff --git a/src/mobius/models/qwen3_tts.py b/src/mobius/models/qwen3_tts.py index 55746aff2..f24cc3a78 100644 --- a/src/mobius/models/qwen3_tts.py +++ b/src/mobius/models/qwen3_tts.py @@ -713,6 +713,12 @@ class Qwen3TTSForConditionalGeneration(nn.Module): "talker.model.codec_embedding", ), "speaker_encoder": ("speaker_encoder",), + # Loop-wiring graphs mobius authors itself; no HuggingFace sub-module + # backs them, so there is no runtime path to report. + "code_predictor_prefill": (), + "code_predictor_step_embedder": (), + "code_predictor_indices": (), + "talker_text_step": (), } def __init__(self, config: ArchitectureConfig): diff --git a/src/mobius/models/qwen3_tts_test.py b/src/mobius/models/qwen3_tts_test.py index 4851867c9..cec831214 100644 --- a/src/mobius/models/qwen3_tts_test.py +++ b/src/mobius/models/qwen3_tts_test.py @@ -35,8 +35,12 @@ num_key_value_heads=1, head_dim=4, vocab_size=_CODEC_VOCAB, + max_position_embeddings=128, rms_norm_eps=1e-6, hidden_act="silu", + mrope_section=[1, 1, 0], + mrope_interleaved=True, + rope_type="default", tts=TTSConfig( num_code_groups=_NUM_CODE_GROUPS, text_hidden_size=_TEXT_HIDDEN, @@ -122,6 +126,73 @@ def test_step_embedder_batched(): np.testing.assert_allclose(got, expected, atol=1e-5, rtol=1e-5) +def test_code_predictor_transition_components_use_exported_embeddings(): + task = TTSTask() + prefill = OnnxModelSession(task._build_code_predictor_prefill(_TINY_CONFIG)) + step = OnnxModelSession(task._build_code_predictor_step_embedder(_TINY_CONFIG)) + talker_hidden = np.arange(8, dtype=np.float32).reshape(1, 1, 8) + group_0_embed = talker_hidden + 10 + got_prefill = prefill.run( + { + "talker_hidden": talker_hidden, + "group_0_embed": group_0_embed, + } + )["inputs_embeds"] + np.testing.assert_array_equal( + got_prefill, + np.concatenate([talker_hidden, group_0_embed], axis=1), + ) + + tables = np.arange( + (_NUM_CODE_GROUPS - 1) * _CP_VOCAB * _HIDDEN, + dtype=np.float32, + ).reshape(_NUM_CODE_GROUPS - 1, _CP_VOCAB, _HIDDEN) + got_step = step.run( + { + "codec_embeddings": tables, + "token": np.array([4], np.int64), + "embedding_index": np.array(1, np.int64), + } + )["inputs_embeds"] + np.testing.assert_array_equal(got_step, tables[1, 4].reshape(1, 1, _HIDDEN)) + + +def test_talker_text_step_clamps_to_last_trailing_embedding(): + session = OnnxModelSession(TTSTask()._build_talker_text_step(_TINY_CONFIG)) + trailing = np.arange(3 * _HIDDEN, dtype=np.float32).reshape(1, 3, _HIDDEN) + got = session.run( + { + "trailing_text_embeds": trailing, + "iteration": np.array([7], np.int64), + } + )["text_embed"] + np.testing.assert_array_equal(got, trailing[:, 2:3]) + + +def test_every_built_component_declares_a_role(): + """model_roles must cover every key the task puts in the package. + + ``build_from_module`` looks each component up in ``model_roles`` and falls + back to ``"decoder"`` when it is absent, which would hand the parameter-free + loop-wiring graphs the GQA / QKV-packing passes meant for attention stacks. + ``inspect_components`` reports exactly ``model_roles``, so an undeclared + component would also be invisible to callers planning per-component work. + """ + module = Qwen3TTSForConditionalGeneration(_TINY_CONFIG) + package = TTSTask().build(module, _TINY_CONFIG) + + # speaker_encoder is optional, so the package is a subset of the declared + # roles; nothing may be built that is not declared. + assert set(package) <= set(TTSTask.model_roles) + assert set(TTSTask.model_roles) - set(package) == {"speaker_encoder"} + + # Glue components are pure wiring: every tensor they read is a graph input, + # so preprocess_weights never routes a parameter to one. + glue = {name for name, role in TTSTask.model_roles.items() if role == "glue"} + weights = module.preprocess_weights({}) + assert not [name for name in weights if name.split(".")[0] in glue] + + def test_step_embedder_weights_shared_with_existing_tables(): """preprocess_weights routes the same codec tables to the step embedder.""" model = Qwen3TTSForConditionalGeneration(_TINY_CONFIG) diff --git a/src/mobius/models/qwen_image_test.py b/src/mobius/models/qwen_image_test.py index fe185adbd..642c4b147 100644 --- a/src/mobius/models/qwen_image_test.py +++ b/src/mobius/models/qwen_image_test.py @@ -365,6 +365,118 @@ def test_edit_vae_matches_diffusers_on_real_source_image(): np.testing.assert_allclose(actual_image, expected_image, rtol=1e-4, atol=1e-4) +def _temporal_vae(dtype: ir.DataType = ir.DataType.FLOAT): + """Build a tiny edit VAE whose resamplers own temporal convolutions. + + ``temperal_downsample=(True,)`` is what the real Qwen Image Edit checkpoint + uses, and it is the only configuration that instantiates ``time_conv`` in the + down/upsample blocks -- the code path the image (single frame) case must skip. + """ + from mobius.integrations.diffusers._configs import QwenImageVAEConfig + from mobius.models.qwen_image_vae import AutoencoderKLQwenImageModel + from mobius.tasks import QwenImageEditVAETask + + config = QwenImageVAEConfig( + base_dim=8, + z_dim=4, + dim_mult=(1, 2), + num_res_blocks=1, + temperal_downsample=(True,), + latents_mean=(-0.2, -0.1, 0.1, 0.2), + latents_std=(1.1, 1.2, 1.3, 1.4), + dtype=dtype, + ) + module = AutoencoderKLQwenImageModel(config) + return config, module, QwenImageEditVAETask().build(module, config) + + +def test_edit_vae_skips_temporal_convolutions_for_single_frame_images(): + """A T=1 image must not run the temporal resampling branch. + + ``QwenImageResample.forward`` only applies ``time_conv`` (and the temporal + pixel shuffle) from the *second* cached chunk onward, so a lone image chunk + keeps its frame count. Applying it unconditionally produced a graph whose + Conv output dimension collapsed to zero and could not even be loaded, so + this pins both the graph shape and the numerical agreement with diffusers. + """ + ort = pytest.importorskip("onnxruntime") + torch = pytest.importorskip("torch") + diffusers = pytest.importorskip("diffusers") + + means = [-0.2, -0.1, 0.1, 0.2] + stds = [1.1, 1.2, 1.3, 1.4] + torch.manual_seed(17) + hf_vae = diffusers.AutoencoderKLQwenImage( + base_dim=8, + z_dim=4, + dim_mult=[1, 2], + num_res_blocks=1, + temperal_downsample=[True], + latents_mean=means, + latents_std=stds, + ).eval() + _, module, package = _temporal_vae() + weights = module.preprocess_weights(dict(hf_vae.state_dict())) + apply_weights(package["encoder"], weights) + apply_weights(package["decoder"], weights) + + generator = torch.Generator().manual_seed(19) + pixels = torch.randn((1, 3, 1, 16, 16), generator=generator) + with torch.no_grad(): + moments = hf_vae._encode(pixels) + mean = moments.chunk(2, dim=1)[0] + scale = torch.tensor(stds)[None, :, None, None, None] + offset = torch.tensor(means)[None, :, None, None, None] + expected_latents = (mean - offset) / scale + expected_image = hf_vae.decode(expected_latents * scale + offset).sample.numpy() + + with tempfile.TemporaryDirectory() as directory: + encoder_path = os.path.join(directory, "encoder.onnx") + decoder_path = os.path.join(directory, "decoder.onnx") + ir.save(package["encoder"], encoder_path) + ir.save(package["decoder"], decoder_path) + actual_latents = ort.InferenceSession( + encoder_path, providers=["CPUExecutionProvider"] + ).run(None, {"sample": pixels.numpy()})[0] + actual_image = ort.InferenceSession( + decoder_path, providers=["CPUExecutionProvider"] + ).run(None, {"latent_sample": actual_latents})[0] + + # The frame axis survives encode and decode: a single image stays a single image. + assert actual_latents.shape[2] == 1 + assert actual_image.shape == (1, 3, 1, 16, 16) + np.testing.assert_allclose(actual_latents, expected_latents.numpy(), rtol=1e-4, atol=1e-4) + np.testing.assert_allclose(actual_image, expected_image, rtol=1e-4, atol=1e-4) + + +@pytest.mark.parametrize("part", ["encoder", "decoder"]) +def test_edit_vae_bfloat16_graph_has_only_kernel_backed_ops(part): + """bfloat16 VAE graphs must avoid ops onnxruntime has no bfloat16 kernel for. + + onnxruntime ships no bfloat16 ``ReduceL2``, ``Resize`` or ``Clip`` kernel on + any execution provider, and a single unassignable node aborts session + creation outright. The RMS norm reduces in float32, ``Resize`` is sandwiched + between casts, and the bfloat16 lowering pass rewrites ``Clip`` into + ``Min``/``Max``, so no bfloat16-typed instance of those ops may survive. + """ + from mobius._optimizations import optimize_model + + package = _temporal_vae(ir.DataType.BFLOAT16)[2] + model = package[part] + optimize_model(model, ep="cuda", dtype=ir.DataType.BFLOAT16, model_role="encoder") + unsupported = {"ReduceL2", "Resize", "Clip"} + offenders = [ + node.op_type + for node in ir.traversal.RecursiveGraphIterator(model.graph) + if node.op_type in unsupported + and any( + value is not None and value.dtype == ir.DataType.BFLOAT16 + for value in (*node.inputs, *node.outputs) + ) + ] + assert offenders == [] + + @pytest.mark.integration @pytest.mark.integration_fast def test_deterministic_l4_l5_image_edit_golden(): diff --git a/src/mobius/models/qwen_image_vae.py b/src/mobius/models/qwen_image_vae.py index 8d7751210..9bb774e92 100644 --- a/src/mobius/models/qwen_image_vae.py +++ b/src/mobius/models/qwen_image_vae.py @@ -13,17 +13,14 @@ from __future__ import annotations import itertools -from typing import TYPE_CHECKING +import onnx_ir as ir from onnxscript import OpBuilder, nn from mobius.components import Conv2d as _Conv2d from mobius.components import SiLU as _SiLU from mobius.integrations.diffusers._configs import QwenImageVAEConfig -if TYPE_CHECKING: - import onnx_ir as ir - # --------------------------------------------------------------------------- # Building blocks # --------------------------------------------------------------------------- @@ -110,11 +107,15 @@ def __init__(self, dim: int, *, images: bool = False): def forward(self, op: OpBuilder, x: ir.Value): # F.normalize(x, dim=1) * scale * gamma - # L2 normalize along channel dimension - norm = op.ReduceL2(x, [1], keepdims=True) + # The L2 reduction runs in float32: this norm consumes raw convolution + # output, so squaring activations > 256 overflows float16, and + # onnxruntime ships no bfloat16 ReduceL2 kernel at all (a bfloat16 graph + # would fail to load because the node cannot be assigned to a provider). + x_f32 = op.Cast(x, to=ir.DataType.FLOAT) + norm = op.ReduceL2(x_f32, [1], keepdims=True) eps = 1e-12 norm = op.Max(norm, eps) - x_normalized = op.Div(x, norm) + x_normalized = op.CastLike(op.Div(x_f32, norm), x) scale = self._scale return op.Mul(op.Mul(x_normalized, scale), self.gamma) @@ -266,8 +267,18 @@ def forward(self, op: OpBuilder, x: ir.Value): class _Resample(nn.Module): """2D or 3D resampling for encoder (downsample) or decoder (upsample). - For ONNX export, temporal resampling uses CausalConv3d (downsample3d) - or CausalConv3d + pixel shuffle (upsample3d). + Temporal resampling in ``QwenImageResample`` is driven by the causal feature + cache: ``time_conv`` only runs for the *second and later* chunks of a video + (``feat_cache[idx] is not None``). Image pipelines encode/decode exactly one + ``T=1`` chunk, so the cache is always empty and both the temporal convolution + and the temporal pixel-shuffle are skipped entirely — ``downsample3d`` and + ``upsample3d`` then behave like their 2D counterparts. + + Mobius exports that single-chunk image path, so ``time_conv`` is declared to + keep the module hierarchy aligned with diffusers but is never invoked. + Applying it would be incorrect: a ``(3, 1, 1)`` stride-2 temporal convolution + over ``T = 1`` yields ``T = 0`` in the encoder, and would double ``T`` in the + decoder. """ def __init__(self, dim: int, mode: str): @@ -297,40 +308,13 @@ def __init__(self, dim: int, mode: str): self.time_conv = _CausalConv3d(dim, dim * 2, (3, 1, 1), padding=(1, 0, 0)) def forward(self, op: OpBuilder, x: ir.Value): - # x: (B, C, T, H, W) + # x: (B, C, T, H, W). Single-chunk image path: no temporal resampling. b_shape = op.Shape(x, start=0, end=1) c_shape = op.Shape(x, start=1, end=2) t_shape = op.Shape(x, start=2, end=3) h_shape = op.Shape(x, start=3, end=4) w_shape = op.Shape(x, start=4, end=5) - if self._mode == "upsample3d": - # Temporal upsample via CausalConv3d → pixel shuffle along T - x = self.time_conv(op, x) - # time_conv output: (B, 2C, T, H, W) → reshape to (B, 2, C, T, H, W) - two = op.Constant(value_ints=[2]) - x = op.Reshape( - x, op.Concat(b_shape, two, c_shape, t_shape, h_shape, w_shape, axis=0) - ) - # Interleave: stack dim=0 and dim=1 along temporal → (B, C, T*2, H, W) - x0 = op.Gather(x, op.Constant(value_ints=[0]), axis=1) - x1 = op.Gather(x, op.Constant(value_ints=[1]), axis=1) - t2 = op.Mul(t_shape, two) - # Interleave by stacking at dim 3, then reshaping - # (B, C, T, H, W) each → stack → (B, C, T, 2, H, W) → (B, C, T*2, H, W) - stacked = op.Concat( - op.Unsqueeze(x0, [3]), - op.Unsqueeze(x1, [3]), - axis=3, - ) - x = op.Reshape(stacked, op.Concat(b_shape, c_shape, t2, h_shape, w_shape, axis=0)) - - if self._mode == "downsample3d": - # Temporal downsample - x = self.time_conv(op, x) - # Update T after temporal conv - t_shape = op.Shape(x, start=2, end=3) - # Reshape to (B*T, C, H, W) for spatial operation bt = op.Mul(b_shape, t_shape) x_2d = op.Reshape( @@ -339,12 +323,17 @@ def forward(self, op: OpBuilder, x: ir.Value): ) if self._mode in ("upsample2d", "upsample3d"): - # Nearest-neighbor 2x upsample - x_2d = op.Resize( + # Nearest-neighbor 2x upsample. QwenImageUpsample runs the + # interpolation in float32 and casts back, which also keeps the + # graph off ORT's missing bfloat16 Resize kernel. + x_2d = op.CastLike( + op.Resize( + op.Cast(x_2d, to=ir.DataType.FLOAT), + None, + op.Constant(value_floats=[1.0, 1.0, 2.0, 2.0]), + mode="nearest", + ), x_2d, - None, - op.Constant(value_floats=[1.0, 1.0, 2.0, 2.0]), - mode="nearest", ) x_2d = self.resample(op, x_2d) diff --git a/src/mobius/models/unet_parity_test.py b/src/mobius/models/unet_parity_test.py index b42599bb5..52c2ab84c 100644 --- a/src/mobius/models/unet_parity_test.py +++ b/src/mobius/models/unet_parity_test.py @@ -129,7 +129,7 @@ def test_unet_matches_diffusers(): model, { "sample": sample.numpy(), - "timestep": timestep.numpy().astype(np.int64), + "timestep": timestep.numpy().astype(np.float32), "encoder_hidden_states": encoder_hidden_states.numpy(), }, ) @@ -194,7 +194,7 @@ def test_unet_sd1x_mixed_block_types_matches_diffusers(): model, { "sample": sample.numpy(), - "timestep": timestep.numpy().astype(np.int64), + "timestep": timestep.numpy().astype(np.float32), "encoder_hidden_states": encoder_hidden_states.numpy(), }, ) @@ -273,7 +273,7 @@ def test_unet_lora_gate_parity(): feed = { "sample": sample.numpy(), - "timestep": timestep.numpy().astype(np.int64), + "timestep": timestep.numpy().astype(np.float32), "encoder_hidden_states": encoder_hidden_states.numpy(), } off, on = _run_onnx( diff --git a/src/mobius/models/wav2vec2.py b/src/mobius/models/wav2vec2.py index b2821d2fe..64cece07a 100644 --- a/src/mobius/models/wav2vec2.py +++ b/src/mobius/models/wav2vec2.py @@ -3,17 +3,39 @@ """Wav2Vec2 encoder-only audio model. -Supports: wav2vec2, hubert, wavlm (all share similar architecture). +Supports: wav2vec2, hubert, wavlm (all share the same architecture). -Architecture: -1. CNN feature extractor: multiple Conv1d layers with group norm -2. Feature projection: Linear + LayerNorm -3. Transformer encoder: standard self-attention + FFN layers +Architecture (HF ``Wav2Vec2Model``): -HF weight naming: -- wav2vec2.feature_extractor.conv_layers.N.conv.weight → feature_extractor.conv_layers.N.conv.weight -- wav2vec2.feature_projection.projection.weight → feature_projection.projection.weight -- wav2vec2.encoder.layers.N.* → encoder.layers.N.* +1. ``feature_extractor`` — a stack of strided Conv1d layers that turn a raw + 16 kHz waveform into frames. The stride product fixes the downsampling + ratio (320 for the ``*-base-*`` checkpoints, i.e. one frame per 20 ms). +2. ``feature_projection`` — LayerNorm over the conv channel dim + Linear into + ``hidden_size``. +3. ``encoder`` — a relative positional conv embedding added to the frames, + followed by transformer layers. + +Two encoder variants exist and are selected by ``config.do_stable_layer_norm``: + +- ``False`` (e.g. ``facebook/wav2vec2-base-960h``): ``layer_norm`` runs *before* + the transformer stack and each layer is **post-norm**. +- ``True`` (e.g. ``facebook/wav2vec2-large-960h-lv60-self``, MMS): each layer is + **pre-norm** and ``layer_norm`` runs *after* the stack. + +Getting this ordering wrong silently produces plausible-looking but incorrect +logits, so the variant is chosen from the checkpoint rather than assumed. + +Inputs: + ``input_values`` — (batch, num_samples) float waveform + ``attention_mask`` — (batch, num_samples) int64 padding mask, optional + +Output: + ``last_hidden_state`` — (batch, num_frames, hidden_size) + +HF weight naming is mirrored one-to-one except for the FFN (HF +``intermediate_dense``/``output_dense`` vs the shared ``FCMLP`` +``up_proj``/``down_proj``) and the weight-normalized positional conv, which is +materialized in :meth:`Wav2Vec2Model.preprocess_weights`. """ from __future__ import annotations @@ -23,69 +45,152 @@ import torch from onnxscript import OpBuilder, nn -from mobius._configs import ArchitectureConfig +from mobius._configs import ArchitectureConfig, MMSConfig from mobius.components import FCMLP -from mobius.components._common import LayerNorm, Linear +from mobius.components._common import GroupNorm, LayerNorm, Linear +from mobius.components._whisper import Conv1d if TYPE_CHECKING: import onnx_ir as ir -class _Conv1dFeatureExtractor(nn.Module): - """CNN feature extractor: extracts features from raw audio waveform.""" +def _conv_geometry(config: ArchitectureConfig) -> tuple[tuple[int, ...], ...]: + """Return ``(conv_dim, conv_kernel, conv_stride)`` for *config*. - def __init__(self, conv_channels: list[int], conv_kernel_sizes: list[int]): - super().__init__() - self.conv_layers = nn.ModuleList() - for i in range(len(conv_channels) - 1): - in_ch = conv_channels[i] - out_ch = conv_channels[i + 1] - kernel = conv_kernel_sizes[i] - layer = _ConvLayerBlock(in_ch, out_ch, kernel, use_group_norm=(i == 0)) - self.conv_layers.append(layer) - - def forward(self, op: OpBuilder, input_values: ir.Value): - # input_values: [batch, time] → [batch, 1, time] - hidden_states = op.Unsqueeze(input_values, [1]) - for layer in self.conv_layers: - hidden_states = layer(op, hidden_states) - # Output: [batch, channels, time'] → transpose to [batch, time', channels] - hidden_states = op.Transpose(hidden_states, perm=[0, 2, 1]) - return hidden_states + Falls back to the wav2vec2-base geometry when a caller supplies a bare + :class:`ArchitectureConfig` (the audio-feature-extraction path). + """ + conv_dim = tuple(getattr(config, "conv_dim", None) or (512,) * 7) + conv_kernel = tuple(getattr(config, "conv_kernel", None) or (10, 3, 3, 3, 3, 2, 2)) + conv_stride = tuple(getattr(config, "conv_stride", None) or (5, 2, 2, 2, 2, 2, 2)) + return conv_dim, conv_kernel, conv_stride class _ConvLayerBlock(nn.Module): - """Single Conv1d + optional GroupNorm + GELU.""" + """One feature-encoder block: strided Conv1d → optional norm → GELU. + + ``norm`` mirrors HF's three conv-layer classes: + ``"group"`` (``Wav2Vec2GroupNormConvLayer``), ``"layer"`` + (``Wav2Vec2LayerNormConvLayer``) and ``"none"`` + (``Wav2Vec2NoLayerNormConvLayer``). + """ def __init__( self, in_channels: int, out_channels: int, kernel_size: int, - use_group_norm: bool = False, + stride: int, + bias: bool, + norm: str = "none", ): super().__init__() - self.conv = nn.Parameter((out_channels, in_channels, kernel_size)) - self.conv_bias = nn.Parameter((out_channels,)) - self.use_group_norm = use_group_norm - self.out_channels = out_channels - if use_group_norm: - self.layer_norm = nn.Parameter((out_channels,)) - self.layer_norm_bias = nn.Parameter((out_channels,)) + # Named ``conv`` with ``.weight``/``.bias`` so HF names map through + # unchanged. ``bias=False`` must not materialize a bias parameter at + # all: an unset initializer makes the exported graph unloadable. + self.conv = Conv1d(in_channels, out_channels, kernel_size, stride=stride, bias=bias) + self._norm = norm + if norm == "group": + # HF uses nn.GroupNorm(num_groups=C, num_channels=C) — one group per + # channel — with torch's default eps of 1e-5. + self.layer_norm = GroupNorm(out_channels, out_channels, eps=1e-5) + elif norm == "layer": + self.layer_norm = LayerNorm(out_channels, eps=1e-5) + + def forward(self, op: OpBuilder, hidden_states: ir.Value) -> ir.Value: + # (batch, in_channels, time) -> (batch, out_channels, time') + hidden_states = self.conv(op, hidden_states) + if self._norm == "group": + hidden_states = self.layer_norm(op, hidden_states) + elif self._norm == "layer": + # HF normalizes over the channel axis, which is axis 1 here, so the + # tensor is moved to channels-last for the LayerNormalization op. + hidden_states = op.Transpose(hidden_states, perm=[0, 2, 1]) + hidden_states = self.layer_norm(op, hidden_states) + hidden_states = op.Transpose(hidden_states, perm=[0, 2, 1]) + return op.Gelu(hidden_states) - def forward(self, op: OpBuilder, hidden_states: ir.Value): - hidden_states = op.Conv(hidden_states, self.conv, self.conv_bias) - if self.use_group_norm: - hidden_states = op.GroupNormalization( - hidden_states, - self.layer_norm, - self.layer_norm_bias, - num_groups=self.out_channels, + +class _Conv1dFeatureExtractor(nn.Module): + """CNN feature encoder: raw waveform → strided frame features. + + Matches HF ``Wav2Vec2FeatureEncoder``. With ``feat_extract_norm="group"`` + only the first layer is normalized; with ``"layer"`` every layer is. + """ + + def __init__( + self, + conv_dim: tuple[int, ...], + conv_kernel: tuple[int, ...], + conv_stride: tuple[int, ...], + conv_bias: bool, + feat_extract_norm: str, + ): + super().__init__() + self.conv_layers = nn.ModuleList() + in_channels = 1 + for i, out_channels in enumerate(conv_dim): + if feat_extract_norm == "group": + norm = "group" if i == 0 else "none" + else: + norm = "layer" + self.conv_layers.append( + _ConvLayerBlock( + in_channels, + out_channels, + conv_kernel[i], + conv_stride[i], + conv_bias, + norm, + ) ) - hidden_states = op.Gelu(hidden_states) + in_channels = out_channels + + def forward(self, op: OpBuilder, input_values: ir.Value) -> ir.Value: + # (batch, time) -> (batch, 1, time) + hidden_states = op.Unsqueeze(input_values, [1]) + for layer in self.conv_layers: + hidden_states = layer(op, hidden_states) + # Channels-first (batch, channels, frames); the caller transposes. return hidden_states +class _PositionalConvEmbedding(nn.Module): + """Relative position embedding via a grouped Conv1d + GELU. + + Matches HF ``Wav2Vec2PositionalConvEmbedding``. The convolution is + weight-normalized in the checkpoint (``weight_g``/``weight_v``); the product + is materialized during weight preprocessing so the graph holds one dense + kernel. + + An even ``kernel_size`` with ``padding = kernel_size // 2`` emits one frame + too many, which HF trims via ``Wav2Vec2SamePadLayer``. + """ + + def __init__(self, hidden_size: int, kernel_size: int, groups: int): + super().__init__() + self.conv = Conv1d( + hidden_size, + hidden_size, + kernel_size, + stride=1, + padding=kernel_size // 2, + bias=True, + groups=groups, + ) + self._num_pad_remove = 1 if kernel_size % 2 == 0 else 0 + + def forward(self, op: OpBuilder, hidden_states: ir.Value) -> ir.Value: + # (batch, frames, hidden) -> (batch, hidden, frames) + hidden_states = op.Transpose(hidden_states, perm=[0, 2, 1]) + hidden_states = self.conv(op, hidden_states) + if self._num_pad_remove: + # Drop the trailing frame produced by the even-kernel "same" padding. + hidden_states = op.Slice(hidden_states, [0], [-1], [2]) + hidden_states = op.Gelu(hidden_states) + return op.Transpose(hidden_states, perm=[0, 2, 1]) + + class _FeatureProjection(nn.Module): """Projects CNN features to hidden size with LayerNorm.""" @@ -100,8 +205,49 @@ def forward(self, op: OpBuilder, hidden_states: ir.Value): return hidden_states +class _Wav2Vec2Attention(nn.Module): + """Bidirectional self-attention for the Wav2Vec2 encoder.""" + + def __init__(self, hidden_size: int, num_heads: int, head_dim: int): + super().__init__() + self.num_heads = num_heads + self.head_dim = head_dim + self.q_proj = Linear(hidden_size, hidden_size, bias=True) + self.k_proj = Linear(hidden_size, hidden_size, bias=True) + self.v_proj = Linear(hidden_size, hidden_size, bias=True) + self.out_proj = Linear(hidden_size, hidden_size, bias=True) + + def forward( + self, + op: OpBuilder, + hidden_states: ir.Value, + attention_mask: ir.Value | None = None, + ): + q = self.q_proj(op, hidden_states) + k = self.k_proj(op, hidden_states) + v = self.v_proj(op, hidden_states) + # ``attention_mask`` is a BOOL keep-mask broadcastable to + # (batch, heads, q_frames, kv_frames); padded frames must not be + # attended to or a padded batch would not match an unpadded run. + attn_out = op.Attention( + q, + k, + v, + attention_mask, + q_num_heads=self.num_heads, + kv_num_heads=self.num_heads, + is_causal=0, + scale=float(self.head_dim**-0.5), + ) + return self.out_proj(op, attn_out) + + class _Wav2Vec2EncoderLayer(nn.Module): - """Standard transformer encoder layer: self-attn + FFN with pre-norm.""" + """Post-norm transformer layer (HF ``Wav2Vec2EncoderLayer``). + + Normalization runs *after* each residual add, which is the layout used by + ``do_stable_layer_norm=False`` checkpoints. + """ def __init__( self, hidden_size: int, intermediate_size: int, num_heads: int, eps: float = 1e-5 @@ -120,28 +266,20 @@ def forward( attention_mask: ir.Value | None = None, ): residual = hidden_states - hidden_states = self.layer_norm(op, hidden_states) hidden_states = self.attention(op, hidden_states, attention_mask) hidden_states = op.Add(residual, hidden_states) + hidden_states = self.layer_norm(op, hidden_states) - residual = hidden_states - hidden_states = self.final_layer_norm(op, hidden_states) - hidden_states = self.feed_forward(op, hidden_states) - hidden_states = op.Add(residual, hidden_states) - return hidden_states + hidden_states = op.Add(hidden_states, self.feed_forward(op, hidden_states)) + return self.final_layer_norm(op, hidden_states) -class _Wav2Vec2Attention(nn.Module): - """Self-attention for Wav2Vec2 encoder (bidirectional).""" +class _Wav2Vec2EncoderLayerStableLayerNorm(_Wav2Vec2EncoderLayer): + """Pre-norm transformer layer (HF ``Wav2Vec2EncoderLayerStableLayerNorm``). - def __init__(self, hidden_size: int, num_heads: int, head_dim: int): - super().__init__() - self.num_heads = num_heads - self.head_dim = head_dim - self.q_proj = Linear(hidden_size, hidden_size, bias=True) - self.k_proj = Linear(hidden_size, hidden_size, bias=True) - self.v_proj = Linear(hidden_size, hidden_size, bias=True) - self.out_proj = Linear(hidden_size, hidden_size, bias=True) + Same parameters as the post-norm layer; only the application order differs, + so the weight names stay identical. + """ def forward( self, @@ -149,35 +287,41 @@ def forward( hidden_states: ir.Value, attention_mask: ir.Value | None = None, ): - q = self.q_proj(op, hidden_states) - k = self.k_proj(op, hidden_states) - v = self.v_proj(op, hidden_states) - attn_out = op.Attention( - q, - k, - v, - q_num_heads=self.num_heads, - kv_num_heads=self.num_heads, - is_causal=0, - scale=float(self.head_dim**-0.5), - ) - return self.out_proj(op, attn_out) + residual = hidden_states + hidden_states = self.layer_norm(op, hidden_states) + hidden_states = self.attention(op, hidden_states, attention_mask) + hidden_states = op.Add(residual, hidden_states) + + normed = self.final_layer_norm(op, hidden_states) + return op.Add(hidden_states, self.feed_forward(op, normed)) class _Wav2Vec2Encoder(nn.Module): - """Wrapper matching HF encoder.layers.{i} nesting.""" + """Post-norm encoder: pos-conv → layer_norm → layers (HF ``Wav2Vec2Encoder``).""" + + layer_class: type[_Wav2Vec2EncoderLayer] = _Wav2Vec2EncoderLayer def __init__(self, config: ArchitectureConfig, eps: float = 1e-5): super().__init__() + self.pos_conv_embed = _PositionalConvEmbedding( + config.hidden_size, + getattr(config, "num_conv_pos_embeddings", 128), + getattr(config, "num_conv_pos_embedding_groups", 16), + ) + self.layer_norm = LayerNorm(config.hidden_size, eps=eps) self.layers = nn.ModuleList() for _ in range(config.num_hidden_layers): - layer = _Wav2Vec2EncoderLayer( - config.hidden_size, - config.intermediate_size, - config.num_attention_heads, - eps=eps, + self.layers.append( + self.layer_class( + config.hidden_size, + config.intermediate_size, + config.num_attention_heads, + eps=eps, + ) ) - self.layers.append(layer) + + def _add_positions(self, op: OpBuilder, hidden_states: ir.Value) -> ir.Value: + return op.Add(hidden_states, self.pos_conv_embed(op, hidden_states)) def forward( self, @@ -185,11 +329,34 @@ def forward( hidden_states: ir.Value, attention_mask: ir.Value | None = None, ): + hidden_states = self._add_positions(op, hidden_states) + hidden_states = self.layer_norm(op, hidden_states) for layer in self.layers: hidden_states = layer(op, hidden_states, attention_mask) return hidden_states +class _Wav2Vec2EncoderStableLayerNorm(_Wav2Vec2Encoder): + """Pre-norm encoder: pos-conv → layers → layer_norm. + + Matches HF ``Wav2Vec2EncoderStableLayerNorm``; the final normalization moves + to the end of the stack. + """ + + layer_class: type[_Wav2Vec2EncoderLayer] = _Wav2Vec2EncoderLayerStableLayerNorm + + def forward( + self, + op: OpBuilder, + hidden_states: ir.Value, + attention_mask: ir.Value | None = None, + ): + hidden_states = self._add_positions(op, hidden_states) + for layer in self.layers: + hidden_states = layer(op, hidden_states, attention_mask) + return self.layer_norm(op, hidden_states) + + class Wav2Vec2Model(nn.Module): """Wav2Vec2 encoder-only audio model. @@ -199,36 +366,93 @@ class Wav2Vec2Model(nn.Module): default_task: str = "audio-feature-extraction" category: str = "Audio" + # Declared on the module so config resolution picks the wav2vec2-shaped + # config even when the registry entry is reached by architecture re-routing + # rather than by ``model_type``. Without it the convolution geometry and + # normalization placement silently fall back to wav2vec2-base defaults. + config_class = MMSConfig def __init__(self, config: ArchitectureConfig): super().__init__() self.config = config - # CNN feature extractor - conv_channels = getattr( - config, "conv_channels", [1, 512, 512, 512, 512, 512, 512, 512] - ) - conv_kernel_sizes = getattr(config, "conv_kernel_sizes", [10, 3, 3, 3, 3, 2, 2]) - self.feature_extractor = _Conv1dFeatureExtractor(conv_channels, conv_kernel_sizes) - - # Feature projection - conv_dim = conv_channels[-1] - self.feature_projection = _FeatureProjection( + conv_dim, conv_kernel, conv_stride = _conv_geometry(config) + self._conv_kernel = conv_kernel + self._conv_stride = conv_stride + self.feature_extractor = _Conv1dFeatureExtractor( conv_dim, - config.hidden_size, - eps=getattr(config, "layer_norm_eps", 1e-5), + conv_kernel, + conv_stride, + bool(getattr(config, "conv_bias", False)), + getattr(config, "feat_extract_norm", None) or "group", ) - # Transformer encoder - self.encoder = _Wav2Vec2Encoder( - config, - eps=getattr(config, "layer_norm_eps", 1e-5), + eps = getattr(config, "layer_norm_eps", None) or 1e-5 + self.feature_projection = _FeatureProjection(conv_dim[-1], config.hidden_size, eps=eps) + + encoder_class = ( + _Wav2Vec2EncoderStableLayerNorm + if getattr(config, "do_stable_layer_norm", False) + else _Wav2Vec2Encoder ) + self.encoder = encoder_class(config, eps=eps) + + @property + def batch_padding_sensitive(self) -> bool: + """Whether a row's outputs depend on the padded width of its batch. + + ``feat_extract_norm="group"`` puts a ``GroupNorm`` with one group per + channel on the first convolution, which reduces over the *time* axis. + Its statistics therefore include whatever padding was appended to reach + the batch width, so co-batching rows of unequal length perturbs every + frame of the shorter rows. ``"layer"`` reduces over channels instead + and leaves rows independent. + """ + return getattr(self.config, "feat_extract_norm", None) == "group" - self.layer_norm = LayerNorm( - config.hidden_size, - eps=getattr(config, "layer_norm_eps", 1e-5), + def frame_lengths(self, op: OpBuilder, attention_mask: ir.Value) -> ir.Value: + """Return the per-row valid frame count for a sample-level mask. + + Mirrors ``_get_feat_extract_output_lengths``: every conv contributes + ``floor((L - kernel) / stride) + 1``. The result lets a caller segment + a padded batch back into per-row outputs. + + Args: + attention_mask: (batch, num_samples) INT64, 1 for valid samples. + + Returns: + (batch,) INT64 frame counts. + """ + lengths = op.ReduceSum( + op.Cast(attention_mask, to=7), # INT64 + [1], + keepdims=0, ) + for kernel, stride in zip(self._conv_kernel, self._conv_stride): + # Integer division truncates toward zero in ONNX; lengths are + # non-negative here so that matches Python floor division. + lengths = op.Add( + op.Div( + op.Sub(lengths, op.Constant(value_int=int(kernel))), + op.Constant(value_int=int(stride)), + ), + op.Constant(value_int=1), + ) + return lengths + + def frame_mask( + self, op: OpBuilder, attention_mask: ir.Value, hidden_states: ir.Value + ) -> ir.Value: + """Build a (batch, num_frames) BOOL keep-mask from a sample-level mask.""" + lengths = self.frame_lengths(op, attention_mask) # (batch,) + num_frames = op.Shape(hidden_states, start=1, end=2) # (1,) + positions = op.Range( + op.Constant(value_int=0), + op.Squeeze(num_frames, [0]), + op.Constant(value_int=1), + ) # (num_frames,) + # (1, num_frames) < (batch, 1) -> (batch, num_frames) + return op.Less(op.Unsqueeze(positions, [0]), op.Unsqueeze(lengths, [1])) def forward( self, @@ -240,40 +464,115 @@ def forward( Args: op: ONNX op builder. - input_values: [batch, time] raw audio waveform - attention_mask: [batch, time] optional mask + input_values: (batch, num_samples) raw waveform. + attention_mask: (batch, num_samples) INT64 sample-level padding mask. Returns: - last_hidden_state: [batch, time', hidden_size] + last_hidden_state: (batch, num_frames, hidden_size) """ + # (batch, samples) -> (batch, channels, frames) -> (batch, frames, channels) hidden_states = self.feature_extractor(op, input_values) + hidden_states = op.Transpose(hidden_states, perm=[0, 2, 1]) hidden_states = self.feature_projection(op, hidden_states) - hidden_states = self.encoder(op, hidden_states, attention_mask) + attention_bias = None + if attention_mask is not None: + keep = self.frame_mask(op, attention_mask, hidden_states) # (batch, frames) + # Zero the padded frames exactly as HF does before the pos-conv, so + # the convolution never mixes real frames with padding residue. + hidden_states = op.Where( + op.Unsqueeze(keep, [-1]), + hidden_states, + op.CastLike(op.Constant(value_float=0.0), hidden_states), + ) + # The ONNX Attention op does not broadcast the query axis of + # ``attn_mask``, so the (batch, frames) keep-mask is expanded to + # (batch, 1, q_frames, kv_frames). Masking only the key axis + # matches HF ``create_bidirectional_mask``: a padded *query* row + # still attends to the valid keys, so no row is fully masked and + # softmax never sees an all -inf row. + num_frames = op.Shape(hidden_states, start=1, end=2) # (1,) + batch = op.Shape(hidden_states, start=0, end=1) # (1,) + mask_shape = op.Concat( + batch, + op.Constant(value_ints=[1]), + num_frames, + num_frames, + axis=0, + ) + attention_bias = op.Expand(op.Unsqueeze(keep, [1, 2]), mask_shape) - hidden_states = self.layer_norm(op, hidden_states) - return hidden_states + return self.encoder(op, hidden_states, attention_bias) + + @staticmethod + def _materialize_weight_norm( + state_dict: dict[str, torch.Tensor], prefix: str + ) -> torch.Tensor | None: + """Recombine a weight-normalized conv kernel into a dense tensor. + + ``torch.nn.utils.weight_norm(conv, name="weight", dim=2)`` stores a + magnitude ``g`` of shape ``(1, 1, K)`` and a direction ``v`` of shape + ``(C_out, C_in/groups, K)``. The effective kernel is + ``g * v / ||v||`` with the norm taken over every axis except ``dim``. + + Checkpoints saved by newer torch use the ``parametrizations`` names, so + both spellings are accepted. + """ + pairs = ( + (f"{prefix}.weight_g", f"{prefix}.weight_v"), + ( + f"{prefix}.parametrizations.weight.original0", + f"{prefix}.parametrizations.weight.original1", + ), + ) + for g_key, v_key in pairs: + if g_key in state_dict and v_key in state_dict: + g = state_dict[g_key].float() + v = state_dict[v_key].float() + norm = v.pow(2).sum(dim=(0, 1), keepdim=True).sqrt() + return g * v / norm + return None def preprocess_weights( self, state_dict: dict[str, torch.Tensor] ) -> dict[str, torch.Tensor]: """Map HF Wav2Vec2 weight names to our names. - HF prefix: wav2vec2.* → strip it. - Attribute names are aligned with HF (out_proj, encoder.layers). - FFN renames: intermediate_dense → up_proj, output_dense → down_proj (FCMLP naming). + Only three families of names actually differ: + + - the ``wav2vec2.``/``hubert.``/``wavlm.`` prefix is stripped, + - HF's ``intermediate_dense``/``output_dense`` become ``FCMLP``'s + ``up_proj``/``down_proj``, + - the weight-normalized positional conv is collapsed to a dense kernel. + + ``masked_spec_embed`` is a SpecAugment training parameter and has no + inference effect, so it is dropped rather than exported. """ - new_state_dict = {} + stripped: dict[str, torch.Tensor] = {} for key, value in state_dict.items(): new_key = key - # Strip wav2vec2. / hubert. prefix for prefix in ("wav2vec2.", "hubert.", "wavlm."): if new_key.startswith(prefix): new_key = new_key[len(prefix) :] break - # FFN: intermediate_dense → up_proj, output_dense → down_proj - new_key = new_key.replace(".intermediate_dense.", ".up_proj.").replace( + stripped[new_key] = value + + pos_conv = "encoder.pos_conv_embed.conv" + dense_pos_conv = self._materialize_weight_norm(stripped, pos_conv) + + new_state_dict: dict[str, torch.Tensor] = {} + for key, value in stripped.items(): + if key == "masked_spec_embed": + continue + if key.startswith((f"{pos_conv}.weight_g", f"{pos_conv}.weight_v")): + continue + if key.startswith(f"{pos_conv}.parametrizations."): + continue + new_key = key.replace(".intermediate_dense.", ".up_proj.").replace( ".output_dense.", ".down_proj." ) new_state_dict[new_key] = value + + if dense_pos_conv is not None: + new_state_dict[f"{pos_conv}.weight"] = dense_pos_conv return new_state_dict diff --git a/src/mobius/models/wav2vec2_ctc.py b/src/mobius/models/wav2vec2_ctc.py index bdf91f9a4..61f2133dd 100644 --- a/src/mobius/models/wav2vec2_ctc.py +++ b/src/mobius/models/wav2vec2_ctc.py @@ -29,10 +29,11 @@ import re import onnx_ir as ir +import torch from onnxscript import nn from onnxscript._internal import builder -from mobius._configs import ArchitectureConfig +from mobius._configs import ArchitectureConfig, MMSConfig from mobius.components._common import LayerNorm, Linear from mobius.models.wav2vec2 import Wav2Vec2Model @@ -166,6 +167,7 @@ class Wav2Vec2ForCTCModel(Wav2Vec2Model): default_task: str = "ctc-asr" category: str = "Speech-to-Text" + config_class = MMSConfig def __init__(self, config: ArchitectureConfig): super().__init__(config) @@ -184,6 +186,8 @@ def __init__(self, config: ArchitectureConfig): adapter_kernel_size=adapter_kernel_size, adapter_stride=adapter_stride, ) + self._adapter_stride = adapter_stride + self._num_adapter_layers = num_adapter_layers # CTC head: projects hidden states to per-frame vocabulary logits self.lm_head = Linear(output_hidden_size, config.vocab_size, bias=True) @@ -216,67 +220,43 @@ def forward( # CTC head: (B, T'', H) → (B, T'', vocab_size) return self.lm_head(op, hidden_states) - def preprocess_weights(self, state_dict: dict[str, object]) -> dict[str, object]: - """Map HuggingFace weight names to ONNX module names. - - HF layout (Wav2Vec2ForCTC): - wav2vec2.feature_extractor.conv_layers.N.conv.weight → feature_extractor.conv_layers.N.conv - wav2vec2.feature_extractor.conv_layers.N.conv.bias → feature_extractor.conv_layers.N.conv_bias - wav2vec2.feature_extractor.conv_layers.0.layer_norm.* → feature_extractor.conv_layers.0.layer_norm[_bias] - wav2vec2.encoder.layer_norm.weight → layer_norm.weight (top-level post-encoder norm) - wav2vec2.encoder.pos_conv_embed.* → (dropped — not in our model) - wav2vec2.encoder.layers.N.* → encoder.layers.N.* - wav2vec2.adapter.layers.N.conv.weight → adapter.layers.N.conv (bare param) - wav2vec2.feature_projection.* → feature_projection.* - lm_head.weight / lm_head.bias → lm_head.weight / lm_head.bias - """ - result: dict[str, object] = {} - for key, value in state_dict.items(): - k = key - - # Strip outer model prefix (wav2vec2.*, hubert.*, wavlm.*) - for prefix in ("wav2vec2.", "hubert.", "wavlm."): - if k.startswith(prefix): - k = k[len(prefix) :] - break - - # FFN weight renames: HF uses intermediate_dense/output_dense, we use up_proj/down_proj - k = k.replace(".intermediate_dense.", ".up_proj.").replace( - ".output_dense.", ".down_proj." - ) + def frame_lengths(self, op: builder.OpBuilder, attention_mask: ir.Value) -> ir.Value: + """Per-row valid frame count, including the MMS adapter downsampling. - # Feature extractor CNN: nn.Conv1d has .weight/.bias sub-attributes; - # our _ConvLayerBlock uses bare nn.Parameter named 'conv' and 'conv_bias'. - if k.startswith("feature_extractor.conv_layers."): - # conv.weight → conv (bare param) - k = re.sub(r"\.conv\.weight$", ".conv", k) - # conv.bias → conv_bias (bare param named differently) - k = re.sub(r"\.conv\.bias$", ".conv_bias", k) - # GroupNorm weight/bias → bare params layer_norm / layer_norm_bias - k = re.sub(r"\.layer_norm\.weight$", ".layer_norm", k) - k = re.sub(r"\.layer_norm\.bias$", ".layer_norm_bias", k) - result[k] = value - continue - - # Stable encoder layer_norm lives at HF's encoder.layer_norm.*; - # our Wav2Vec2Model puts it at top-level as self.layer_norm. - if k == "encoder.layer_norm.weight": - result["layer_norm.weight"] = value - continue - if k == "encoder.layer_norm.bias": - result["layer_norm.bias"] = value - continue - - # Positional conv embedding — not present in our simplified encoder. - if k.startswith("encoder.pos_conv_embed."): - continue - - # Adapter conv weights: HF uses nn.Conv1d → .conv.weight / .conv.bias; - # our _AdapterLayer uses bare params named 'conv' / 'conv_bias'. - if k.startswith("adapter.layers."): - k = re.sub(r"\.conv\.weight$", ".conv", k) - k = re.sub(r"\.conv\.bias$", ".conv_bias", k) - - result[k] = value - - return result + The adapter's strided convolutions shrink the frame axis further, so the + logits' time axis is shorter than the encoder's. Reporting the encoder + count here would over-run the logits when segmenting a padded batch. + """ + lengths = super().frame_lengths(op, attention_mask) + if not hasattr(self, "adapter"): + return lengths + stride = op.Constant(value_int=int(self._adapter_stride)) + one = op.Constant(value_int=1) + for _ in range(self._num_adapter_layers): + # Adapter convs use kernel 1 semantics for length purposes: HF's + # _get_feat_extract_output_lengths passes kernel_size=1. + lengths = op.Add(op.Div(op.Sub(lengths, one), stride), one) + return lengths + + def preprocess_weights( + self, state_dict: dict[str, torch.Tensor] + ) -> dict[str, torch.Tensor]: + """Map HuggingFace ``Wav2Vec2ForCTC`` weight names to ONNX module names. + + The encoder names are handled by :meth:`Wav2Vec2Model.preprocess_weights` + (prefix strip, FFN rename, positional-conv weight-norm materialization). + Only the MMS adapter needs extra work: HF wraps its convolution in an + ``nn.Conv1d`` (``*.conv.weight``/``*.conv.bias``) while ``_AdapterLayer`` + holds bare parameters named ``conv`` and ``conv_bias``. + + ``lm_head.weight``/``lm_head.bias`` already match. + """ + result = super().preprocess_weights(state_dict) + + adapted: dict[str, torch.Tensor] = {} + for key, value in result.items(): + if key.startswith("adapter.layers."): + key = re.sub(r"\.conv\.weight$", ".conv", key) + key = re.sub(r"\.conv\.bias$", ".conv_bias", key) + adapted[key] = value + return adapted diff --git a/src/mobius/rewrite_rules/__init__.py b/src/mobius/rewrite_rules/__init__.py index 2d270e7a9..015e18d29 100644 --- a/src/mobius/rewrite_rules/__init__.py +++ b/src/mobius/rewrite_rules/__init__.py @@ -30,6 +30,7 @@ __all__ = [ "bias_gelu_rules", + "clip_to_min_max_rules", "decompose_attention_pass", "decompose_rope_rules", "fuse_dense_moe_to_qmoe", @@ -48,6 +49,7 @@ ] from mobius.rewrite_rules._bias_gelu import bias_gelu_rules +from mobius.rewrite_rules._clip_to_min_max import clip_to_min_max_rules from mobius.rewrite_rules._decompose_attention import decompose_attention_pass from mobius.rewrite_rules._decompose_rope import decompose_rope_rules from mobius.rewrite_rules._gelu_fusion import gelu_fusion_rules diff --git a/src/mobius/rewrite_rules/_clip_to_min_max.py b/src/mobius/rewrite_rules/_clip_to_min_max.py new file mode 100644 index 000000000..c63ce9ef8 --- /dev/null +++ b/src/mobius/rewrite_rules/_clip_to_min_max.py @@ -0,0 +1,81 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT License. + +"""Lower bfloat16 ``Clip`` to ``Min``/``Max`` primitives. + +ONNX Runtime has no ``Clip`` kernel for ``BFLOAT16`` on any execution provider. +Because ``Clip`` carries an ONNX function body, ORT silently expands it into +comparison primitives (``Less``/``Where``) that also lack bfloat16 kernels, so +the expanded nodes cannot be assigned to any provider and **session creation +fails outright**:: + + FAIL : Exception during initialization: transformer_memcpy.cc:253 + Provider type for Less node with name '' is not set. + +``Min`` and ``Max`` do have bfloat16 kernels, and ``Clip(x, lo, hi)`` is exactly +``Min(Max(x, lo), hi)`` — including for NaN inputs, since both lowerings inherit +the same propagation behaviour from the underlying comparisons. Rewriting the op +is therefore numerically exact, not an approximation. + +The rule only fires for bfloat16 inputs; integer clamps (index clamping) and +float32/float16 clamps keep the compact ``Clip`` op and its native kernel. + +These rules are applied automatically by +:func:`~mobius._optimizations.optimize_model` for bfloat16 models. They can also +be applied manually:: + + from mobius.rewrite_rules import clip_to_min_max_rules + from onnxscript.rewriter import rewrite + + rewrite(model, pattern_rewrite_rules=clip_to_min_max_rules()) +""" + +from __future__ import annotations + +import onnx_ir as ir +from onnxscript.rewriter._basics import MatchResult +from onnxscript.rewriter._rewrite_rule import RewriteRuleClassBase, RewriteRuleSet + + +class _ClipToMinMaxBase(RewriteRuleClassBase): + """Shared bfloat16 check for the ``Clip`` lowering variants.""" + + def check(self, context, x, **_): + result = MatchResult() + if x.dtype != ir.DataType.BFLOAT16: + return result.fail("Clip input is not bfloat16") + return result + + +class ClipBothToMinMax(_ClipToMinMaxBase): + """Rewrite ``Clip(x, lo, hi)`` → ``Min(Max(x, lo), hi)`` for bfloat16.""" + + def pattern(self, op, x, lo, hi): + return op.Clip(x, lo, hi) + + def rewrite(self, op, x, lo, hi): + return op.Min(op.Max(x, lo), hi) + + +class ClipMinToMax(_ClipToMinMaxBase): + """Rewrite the lower-bound-only ``Clip(x, lo)`` → ``Max(x, lo)`` for bfloat16.""" + + def pattern(self, op, x, lo): + return op.Clip(x, lo) + + def rewrite(self, op, x, lo): + return op.Max(x, lo) + + +def clip_to_min_max_rules() -> RewriteRuleSet: + """Return rules lowering bfloat16 ``Clip`` to ``Min``/``Max``. + + ONNX Runtime cannot execute ``Clip`` in bfloat16 on any execution provider; + the ONNX function expansion produces ``Less``/``Where`` nodes with no + bfloat16 kernel, which aborts session creation. ``Min``/``Max`` have + bfloat16 kernels and are numerically identical. + + Returns: + :class:`RewriteRuleSet` with the two-bound and lower-bound-only rules. + """ + return RewriteRuleSet([ClipBothToMinMax().rule(), ClipMinToMax().rule()]) diff --git a/src/mobius/rewrite_rules/_clip_to_min_max_test.py b/src/mobius/rewrite_rules/_clip_to_min_max_test.py new file mode 100644 index 000000000..c2598f995 --- /dev/null +++ b/src/mobius/rewrite_rules/_clip_to_min_max_test.py @@ -0,0 +1,90 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT License. + +"""Tests for the bfloat16 ``Clip`` → ``Min``/``Max`` lowering rule.""" + +from __future__ import annotations + +from collections import Counter + +import ml_dtypes +import numpy as np +import onnx_ir as ir +from onnxscript.rewriter import rewrite + +from mobius._constants import OPSET_VERSION +from mobius.rewrite_rules import clip_to_min_max_rules + +_NUMPY_DTYPE = { + ir.DataType.BFLOAT16: ml_dtypes.bfloat16, + ir.DataType.FLOAT: np.float32, +} + + +def _clip_model(dtype: ir.DataType, *, both_bounds: bool = True) -> ir.Model: + """Build a single-node ``Clip`` graph of the given dtype.""" + np_dtype = _NUMPY_DTYPE[dtype] + x = ir.Value(name="x", type=ir.TensorType(dtype), shape=ir.Shape(["batch", 4])) + lo = ir.Value( + name="lo", + const_value=ir.tensor(np.array(-1.0, dtype=np_dtype)), + type=ir.TensorType(dtype), + ) + inputs = [x, lo] + initializers = [lo] + if both_bounds: + hi = ir.Value( + name="hi", + const_value=ir.tensor(np.array(1.0, dtype=np_dtype)), + type=ir.TensorType(dtype), + ) + inputs.append(hi) + initializers.append(hi) + y = ir.Value(name="y", type=ir.TensorType(dtype), shape=ir.Shape(["batch", 4])) + node = ir.Node("", "Clip", inputs=inputs, outputs=[y], name="clip") + graph = ir.Graph( + inputs=[x], + outputs=[y], + nodes=[node], + initializers=initializers, + opset_imports={"": OPSET_VERSION}, + name="clip_graph", + ) + return ir.Model(graph, ir_version=10) + + +def _op_counts(model: ir.Model) -> Counter: + return Counter(node.op_type for node in ir.traversal.RecursiveGraphIterator(model.graph)) + + +def test_bfloat16_clip_lowers_to_min_max() -> None: + model = _clip_model(ir.DataType.BFLOAT16) + rewrite(model, pattern_rewrite_rules=clip_to_min_max_rules()) + counts = _op_counts(model) + assert counts["Clip"] == 0 + assert counts["Min"] == 1 + assert counts["Max"] == 1 + + +def test_bfloat16_single_bound_clip_lowers_to_max() -> None: + model = _clip_model(ir.DataType.BFLOAT16, both_bounds=False) + rewrite(model, pattern_rewrite_rules=clip_to_min_max_rules()) + counts = _op_counts(model) + assert counts["Clip"] == 0 + assert counts["Max"] == 1 + assert counts["Min"] == 0 + + +def test_float32_clip_is_untouched() -> None: + model = _clip_model(ir.DataType.FLOAT) + rewrite(model, pattern_rewrite_rules=clip_to_min_max_rules()) + counts = _op_counts(model) + assert counts["Clip"] == 1 + assert counts["Min"] == 0 + assert counts["Max"] == 0 + + +def test_lowering_is_numerically_exact() -> None: + """``Min(Max(x, lo), hi)`` reproduces ``Clip(x, lo, hi)`` elementwise.""" + x = np.array([-3.0, -1.0, -0.25, 0.0, 0.5, 1.0, 7.0], dtype=np.float32) + np.testing.assert_array_equal(np.minimum(np.maximum(x, -1.0), 1.0), np.clip(x, -1.0, 1.0)) diff --git a/src/mobius/tasks/__init__.py b/src/mobius/tasks/__init__.py index d2e32bca0..fc53ddce7 100644 --- a/src/mobius/tasks/__init__.py +++ b/src/mobius/tasks/__init__.py @@ -80,6 +80,7 @@ "TTSTask", "VAETask", "VideoDenoisingTask", + "VideoVAETask", "VisionLanguageTask", "VisionEncoderDecoderTask", "WorldModelTask", @@ -146,6 +147,7 @@ from mobius.tasks._tts import TTSTask from mobius.tasks._vae import VAETask from mobius.tasks._video_denoising import VideoDenoisingTask +from mobius.tasks._video_vae import VideoVAETask from mobius.tasks._vision_encoder_decoder import VisionEncoderDecoderTask from mobius.tasks._vision_language import Qwen3VLVisionLanguageTask from mobius.tasks._vision_language_3model import ( @@ -228,6 +230,7 @@ "ssm2-text-generation": SSM2CausalLMTask, "tts": TTSTask, "video-denoising": VideoDenoisingTask, + "video-vae": VideoVAETask, "world-model": WorldModelTask, } diff --git a/src/mobius/tasks/_causal_lm.py b/src/mobius/tasks/_causal_lm.py index 5ae9c5230..41758389e 100644 --- a/src/mobius/tasks/_causal_lm.py +++ b/src/mobius/tasks/_causal_lm.py @@ -10,6 +10,10 @@ from mobius._build_context import prefill_prefix_pruning from mobius._configs import ArchitectureConfig +from mobius._constants import ( + STATIC_CACHE_KV_SEQUENCE_LENGTH, + STATIC_CACHE_WRITE_INDICES, +) from mobius._model_package import ModelPackage from mobius.components._attention import StaticCacheState from mobius.tasks._base import ( @@ -425,12 +429,12 @@ def _make_static_cache_inputs( # Shared inputs across all layers write_indices = builder.input( - "write_indices", + STATIC_CACHE_WRITE_INDICES, dtype=ir.DataType.INT64, shape=[batch], ) nonpad_kv_seqlen = builder.input( - "nonpad_kv_seqlen", + STATIC_CACHE_KV_SEQUENCE_LENGTH, dtype=ir.DataType.INT64, shape=[batch], ) diff --git a/src/mobius/tasks/_controlnet.py b/src/mobius/tasks/_controlnet.py index 4354250fc..e797e52d8 100644 --- a/src/mobius/tasks/_controlnet.py +++ b/src/mobius/tasks/_controlnet.py @@ -36,7 +36,7 @@ def build( dtype=ir.DataType.FLOAT, shape=["batch", config.in_channels, "height", "width"], ) - timestep = builder.input("timestep", dtype=ir.DataType.INT64, shape=["batch"]) + timestep = builder.input("timestep", dtype=ir.DataType.FLOAT, shape=["batch"]) encoder_hidden_states = builder.input( "encoder_hidden_states", dtype=ir.DataType.FLOAT, diff --git a/src/mobius/tasks/_ctc_asr.py b/src/mobius/tasks/_ctc_asr.py index 87d756b2d..b94290063 100644 --- a/src/mobius/tasks/_ctc_asr.py +++ b/src/mobius/tasks/_ctc_asr.py @@ -19,6 +19,23 @@ from mobius._model_package import ModelPackage from mobius.tasks._base import ModelTask, _make_graph, _make_model +#: Graph metadata key recording whether a row's outputs depend on the padded +#: width of its batch. Carried on the ONNX model so a metadata producer can +#: publish the fact without knowing anything about the architecture, and so it +#: survives a round trip through the serialized file. +BATCH_PADDING_SENSITIVE_KEY = "mobius.batch_padding_sensitive" + + +def _record_batch_padding_sensitivity(model: ir.Model, module: object) -> None: + """Copy a module's batch-padding sensitivity onto the built ONNX model. + + Modules that do not answer the question leave the key absent, which readers + must treat as unstated rather than as "rows are independent". + """ + sensitive = getattr(module, "batch_padding_sensitive", None) + if isinstance(sensitive, bool): + model.metadata_props[BATCH_PADDING_SENSITIVE_KEY] = "true" if sensitive else "false" + class CTCAsrTask(ModelTask): """Build ONNX graph for CTC-based ASR (raw waveform → frame logits). @@ -30,8 +47,14 @@ class CTCAsrTask(ModelTask): required graph input; callers with no padding should pass an all-ones mask. - Output: - ``logits`` — (batch, num_frames, vocab_size) CTC logit scores FLOAT + Outputs: + ``logits`` — (batch, num_frames, vocab_size) CTC logit scores + ``frame_lengths`` — (batch,) INT64 count of non-padded frames per row + + ``frame_lengths`` is emitted so a padded batch can be segmented back into + per-row transcripts without the caller re-deriving the convolutional + downsampling ratio. It is only emitted when the module knows how to compute + it, keeping the task usable for encoders with a different contract. """ name = "ctc-asr" @@ -60,7 +83,14 @@ def build( logits = module(builder.op, input_values=input_values, attention_mask=attention_mask) builder.add_output(logits, "logits") - return ModelPackage({"model": _make_model(graph)}, config=config) + frame_lengths_fn = getattr(module, "frame_lengths", None) + if callable(frame_lengths_fn): + frame_lengths = frame_lengths_fn(builder.op, attention_mask) + builder.add_output(frame_lengths, "frame_lengths") + + model = _make_model(graph) + _record_batch_padding_sensitivity(model, module) + return ModelPackage({"model": model}, config=config) class FeatureCTCAsrTask(ModelTask): diff --git a/src/mobius/tasks/_denoising.py b/src/mobius/tasks/_denoising.py index df406d410..46e4bcdd4 100644 --- a/src/mobius/tasks/_denoising.py +++ b/src/mobius/tasks/_denoising.py @@ -36,7 +36,7 @@ def build( dtype=ir.DataType.FLOAT, shape=["batch", config.in_channels, "height", "width"], ) - timestep = builder.input("timestep", dtype=ir.DataType.INT64, shape=["batch"]) + timestep = builder.input("timestep", dtype=ir.DataType.FLOAT, shape=["batch"]) encoder_hidden_states = builder.input( "encoder_hidden_states", dtype=ir.DataType.FLOAT, @@ -62,6 +62,16 @@ def build( **extra_kwargs, ) + # The denoiser is spatially shape preserving, so republish the latent's named + # dimensions on the estimate instead of leaving the anonymous symbols shape + # inference produces, which no consumer can relate back to the inputs. A + # denoiser that learns the variance emits twice the latent channels. + out_channels = getattr(config, "out_channels", None) + if out_channels is None: + out_channels = config.in_channels * ( + 2 if getattr(config, "learn_sigma", False) else 1 + ) + noise_pred.shape = ir.Shape(["batch", out_channels, "height", "width"]) builder.add_output(noise_pred, "noise_pred") return ModelPackage({"model": _make_model(graph)}, config=config) diff --git a/src/mobius/tasks/_gemma4.py b/src/mobius/tasks/_gemma4.py index fcbe329fe..458a0d17f 100644 --- a/src/mobius/tasks/_gemma4.py +++ b/src/mobius/tasks/_gemma4.py @@ -28,6 +28,10 @@ from mobius._build_context import ep_capabilities, prefill_prefix_pruning from mobius._configs import Gemma4Config +from mobius._constants import ( + STATIC_CACHE_KV_SEQUENCE_LENGTH, + STATIC_CACHE_WRITE_INDICES, +) from mobius._model_package import ModelPackage from mobius._pipeline_contract import ( declare_component_presence, @@ -43,6 +47,21 @@ ) +def _has_dynamic_cache_layer(config: Gemma4Config) -> bool: + """Whether the static-cache layout still leaves a layer on a dynamic cache. + + Only full-attention layers are addressable by a fixed-capacity scatter; a + sliding layer keeps a growing ``past_key_values.N.*`` pair, and the last + ``num_kv_shared_layers`` layers borrow KV and own no cache at all. + """ + layer_types = config.layer_types or (["sliding_attention"] * config.num_hidden_layers) + num_kv_layers = config.num_hidden_layers - (config.num_kv_shared_layers or 0) + return any( + (layer_types[i] if i < len(layer_types) else "sliding_attention") != "full_attention" + for i in range(num_kv_layers) + ) + + def _register_hybrid_cache_outputs( builder: GraphBuilder, present_key_values: list[tuple[ir.Value, ir.Value]], @@ -146,12 +165,12 @@ def _make_gemma4_static_cache_inputs( write_indices = nonpad_kv_seqlen = None if has_static: write_indices = builder.input( - "write_indices", + STATIC_CACHE_WRITE_INDICES, dtype=ir.DataType.INT64, shape=[batch], ) nonpad_kv_seqlen = builder.input( - "nonpad_kv_seqlen", + STATIC_CACHE_KV_SEQUENCE_LENGTH, dtype=ir.DataType.INT64, shape=[batch], ) @@ -504,17 +523,21 @@ def _build_decoder( shape=[batch, seq_len, config.hidden_size], ) - if not static: - past_seq_len = ir.SymbolicDim("past_sequence_len") - attention_mask = builder.input( + past_seq_len = ir.SymbolicDim("past_sequence_len") + # A static-cache layer masks itself from ``write_indices`` and + # ``nonpad_kv_seqlen`` and takes no bias, but a sliding layer keeps a + # dynamic cache even in static mode and still needs the mask to know + # where each row's real tokens are. Mint the mask exactly when such a + # layer survives, so a fully static decoder does not carry a port + # nothing reads and a hybrid one does not lose its padding information. + if not static or _has_dynamic_cache_layer(config): + attention_mask: ir.Value | None = builder.input( "attention_mask", dtype=ir.DataType.INT64, shape=[batch, "past_seq_len + seq_len"], ) else: - # Static cache still needs past_seq_len for sliding-window layers - # that use dynamic cache within the hybrid static/dynamic scheme. - past_seq_len = ir.SymbolicDim("past_sequence_len") + attention_mask = None position_ids = builder.input( "position_ids", @@ -558,7 +581,6 @@ def _build_decoder( max_seq_len, past_seq_len, ) - attention_mask = None # Static cache uses position-based attention else: past_key_values = _make_gemma4_kv_cache_inputs( builder, diff --git a/src/mobius/tasks/_masked_diffusion.py b/src/mobius/tasks/_masked_diffusion.py index 46bd417b8..87ccfa314 100644 --- a/src/mobius/tasks/_masked_diffusion.py +++ b/src/mobius/tasks/_masked_diffusion.py @@ -6,9 +6,9 @@ Builds an ONNX graph for a discrete masked-diffusion *mask predictor*: it maps an int64 token sequence to per-position vocabulary logits in a single full-sequence pass. There is no KV cache and no attention-mask input — the -model attends bidirectionally over the whole sequence. onnx-genai's -``masked_diffusion`` scheduler drives the reverse process, feeding the emitted -logits back to ``input_ids`` via a loop-carried self-edge. +model attends bidirectionally over the whole sequence. The graph also exposes +the greedy full-sequence proposal consumed by the generic workflow's +``masked_update`` policy component. """ from __future__ import annotations @@ -31,6 +31,7 @@ class MaskedDiffusionTask(ModelTask): Outputs: - logits: [batch, sequence_len, vocab_size] FLOAT + - proposed_tokens: [batch, sequence_len] INT64 """ # ``encoder`` role: attention is bidirectional, so the decoder-only GQA / @@ -51,7 +52,9 @@ def build( input_ids = builder.input("input_ids", dtype=ir.DataType.INT64, shape=[batch, seq_len]) logits = module(op, input_ids=input_ids) + proposed_tokens = op.ArgMax(logits, axis=-1, keepdims=0) builder.add_output(logits, "logits") + builder.add_output(proposed_tokens, "proposed_tokens") 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 46b35fb37..00036eb45 100644 --- a/src/mobius/tasks/_task_test.py +++ b/src/mobius/tasks/_task_test.py @@ -410,7 +410,7 @@ def test_input_types(self): model = pkg["model"] inputs_by_name = {v.name: v for v in model.graph.inputs} assert inputs_by_name["sample"].dtype == ir.DataType.FLOAT - assert inputs_by_name["timestep"].dtype == ir.DataType.INT64 + assert inputs_by_name["timestep"].dtype == ir.DataType.FLOAT assert inputs_by_name["encoder_hidden_states"].dtype == ir.DataType.FLOAT def test_outputs(self): diff --git a/src/mobius/tasks/_tts.py b/src/mobius/tasks/_tts.py index 331f2c59a..73ec89600 100644 --- a/src/mobius/tasks/_tts.py +++ b/src/mobius/tasks/_tts.py @@ -10,6 +10,10 @@ 4. **talker_step_embedder**: frame_codes + text_embed → inputs_embeds 5. **talker_prefill_embedder**: text_ids → prefill_embeds + trailing_text_embeds 6. **speaker_encoder**: mel_input → speaker_embedding +7. **code_predictor_prefill**: talker hidden + group-0 embedding → predictor input +8. **code_predictor_step_embedder**: prior code + predictor tables → next input +9. **talker_text_step**: trailing text embeddings + loop index → one text embedding +10. **code_predictor_indices**: inner induction → embedding/head/frame indices Used by Qwen3TTSForConditionalGeneration. """ @@ -53,6 +57,10 @@ class TTSTask(ModelTask): Each sub-module is wired into its own ONNX graph. """ + # Every key the package produces must appear here: ``model_roles`` is what + # ``inspect_components`` reports and what ``build_from_module`` uses to pick + # optimization passes. An undeclared component silently falls back to the + # ``"decoder"`` role and would be handed GQA / QKV-packing fusion. model_roles: ClassVar[dict[str, str]] = { "talker": "decoder", "code_predictor": "decoder", @@ -60,6 +68,13 @@ class TTSTask(ModelTask): "talker_step_embedder": "embedding", "talker_prefill_embedder": "embedding", "speaker_encoder": "encoder", + # Parameter-free graphs that wire the generation loop. They read every + # tensor they use from their own graph inputs, so they carry no weights + # and no fusion pass applies to them. + "code_predictor_prefill": "glue", + "code_predictor_step_embedder": "glue", + "code_predictor_indices": "glue", + "talker_text_step": "glue", } components: ClassVar[ComponentSpec] = ComponentSpec( talker="talker", @@ -86,6 +101,12 @@ def build( models["talker_prefill_embedder"] = self._build_talker_prefill_embedder( module.talker_prefill_embedder, config ) + models["code_predictor_prefill"] = self._build_code_predictor_prefill(config) + models["code_predictor_step_embedder"] = self._build_code_predictor_step_embedder( + config + ) + models["code_predictor_indices"] = self._build_code_predictor_indices() + models["talker_text_step"] = self._build_talker_text_step(config) if module.speaker_encoder is not None: models["speaker_encoder"] = self._build_speaker_encoder( module.speaker_encoder, config @@ -93,6 +114,79 @@ def build( return ModelPackage(models, config=config) + def _build_code_predictor_prefill(self, config: ArchitectureConfig) -> ir.Model: + """Build the trained group-0 transition input for code-predictor prefill.""" + graph, builder = _make_graph(name="code_predictor_prefill") + talker_hidden = builder.input( + "talker_hidden", + dtype=config.dtype, + shape=["batch", 1, config.hidden_size], + ) + group_0_embed = builder.input( + "group_0_embed", + dtype=config.dtype, + shape=["batch", 1, config.hidden_size], + ) + inputs_embeds = builder.op.Concat(talker_hidden, group_0_embed, axis=1) + inputs_embeds.shape = ir.Shape(["batch", 2, config.hidden_size]) + builder.add_output(inputs_embeds, "inputs_embeds") + return _make_model(graph) + + def _build_code_predictor_step_embedder(self, config: ArchitectureConfig) -> ir.Model: + """Gather the trained predictor embedding for the previously sampled code.""" + tts = config.tts + cp = tts.code_predictor if tts else None + num_groups = tts.num_code_groups if tts else 16 + cp_vocab = cp.vocab_size if cp else 2048 + graph, builder = _make_graph(name="code_predictor_step_embedder") + codec_embeddings = builder.input( + "codec_embeddings", + dtype=config.dtype, + shape=[num_groups - 1, cp_vocab, config.hidden_size], + ) + token = builder.input("token", dtype=ir.DataType.INT64, shape=["batch"]) + embedding_index = builder.input("embedding_index", dtype=ir.DataType.INT64, shape=[]) + table = builder.op.Gather(codec_embeddings, embedding_index, axis=0) + inputs_embeds = builder.op.Gather(table, token, axis=0) + inputs_embeds = builder.op.Unsqueeze(inputs_embeds, [1]) + inputs_embeds.shape = ir.Shape(["batch", 1, config.hidden_size]) + builder.add_output(inputs_embeds, "inputs_embeds") + return _make_model(graph) + + def _build_talker_text_step(self, config: ArchitectureConfig) -> ir.Model: + """Select one trailing-text embedding for the outer talker iteration.""" + graph, builder = _make_graph(name="talker_text_step") + trailing = builder.input( + "trailing_text_embeds", + dtype=config.dtype, + shape=["batch", "trailing_sequence", config.hidden_size], + ) + iteration = builder.input("iteration", dtype=ir.DataType.INT64, shape=["batch"]) + sequence_length = builder.op.Shape(trailing, start=1, end=2) + index = builder.op.Min( + builder.op.Gather(iteration, 0, axis=0), + builder.op.Sub( + builder.op.Squeeze(sequence_length, [0]), + builder.op.Constant(value_int=1), + ), + ) + text_embed = builder.op.Gather(trailing, index, axis=1) + text_embed = builder.op.Unsqueeze(text_embed, [1]) + text_embed.shape = ir.Shape(["batch", 1, config.hidden_size]) + builder.add_output(text_embed, "text_embed") + return _make_model(graph) + + def _build_code_predictor_indices(self) -> ir.Model: + """Derive predictor/table/frame indices from the zero-based inner loop.""" + graph, builder = _make_graph(name="code_predictor_indices") + iteration = builder.input("iteration", dtype=ir.DataType.INT64, shape=[]) + step_index = builder.op.Add(iteration, builder.op.Constant(value_int=1)) + frame_index = builder.op.Add(iteration, builder.op.Constant(value_int=2)) + builder.add_output(builder.op.Identity(iteration), "embedding_index") + builder.add_output(step_index, "step_index") + builder.add_output(frame_index, "frame_index") + return _make_model(graph) + def _build_talker( self, talker: nn.Module, @@ -142,6 +236,13 @@ def _build_talker( position_ids=position_ids, past_key_values=past_key_values, ) + for present in present_key_values: + present[0].shape = ir.Shape( + [batch, config.num_key_value_heads, "total_sequence_len", config.head_dim] + ) + present[1].shape = ir.Shape( + [batch, config.num_key_value_heads, "total_sequence_len", config.head_dim] + ) builder.add_output(logits, "logits") builder.add_output(last_hidden_state, "last_hidden_state") @@ -224,6 +325,13 @@ def _build_code_predictor( position_ids=position_ids, past_key_values=past_key_values, ) + for present in present_key_values: + present[0].shape = ir.Shape( + [batch, cp_num_key_value_heads, "total_sequence_len", cp_head_dim] + ) + present[1].shape = ir.Shape( + [batch, cp_num_key_value_heads, "total_sequence_len", cp_head_dim] + ) builder.add_output(logits, "logits") # Expose stacked codec embeddings for generation loop to extract. @@ -340,6 +448,10 @@ def _build_talker_prefill_embedder( builder.op, text_ids=text_ids, ) + prefill_embeds.shape = ir.Shape([batch, "prefill_sequence_len", config.hidden_size]) + trailing_text_embeds.shape = ir.Shape( + [batch, "trailing_sequence_len", config.hidden_size] + ) builder.add_output(prefill_embeds, "prefill_embeds") builder.add_output(trailing_text_embeds, "trailing_text_embeds") diff --git a/src/mobius/tasks/_vae.py b/src/mobius/tasks/_vae.py index 007497994..4d222c400 100644 --- a/src/mobius/tasks/_vae.py +++ b/src/mobius/tasks/_vae.py @@ -76,6 +76,12 @@ def _build_decoder_graph( hidden_states = module.post_quant_conv(op, hidden_states) hidden_states = module.decoder(op, latent_sample=hidden_states) + # The decoder upsamples by a fixed factor, but the factor is a property of the + # configured block stack rather than of the graph, so publish named image dims + # instead of anonymous inferred symbols. + hidden_states.shape = ir.Shape( + ["batch", config.out_channels, "image_height", "image_width"] + ) builder.add_output(hidden_states, "sample") return _make_model(graph) diff --git a/src/mobius/tasks/_video_denoising.py b/src/mobius/tasks/_video_denoising.py index 95657810e..c2601e79e 100644 --- a/src/mobius/tasks/_video_denoising.py +++ b/src/mobius/tasks/_video_denoising.py @@ -50,6 +50,13 @@ def build( timestep=timestep, encoder_hidden_states=encoder_hidden_states, ) + # Unpatchify recomputes the spatial extents from Shape ops, which mints + # fresh symbolic dimensions. The prediction is elementwise with the + # latent, so the declared contract has to say so: a solver that carries + # the latent across steps needs both to unify. + noise_pred.shape = ir.Shape( + ["batch", "num_frames", config.out_channels, "height", "width"] + ) builder.add_output(noise_pred, "noise_pred") diff --git a/src/mobius/tasks/_video_vae.py b/src/mobius/tasks/_video_vae.py new file mode 100644 index 000000000..ccf440595 --- /dev/null +++ b/src/mobius/tasks/_video_vae.py @@ -0,0 +1,101 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT License. + +"""Video VAE decode task for 3D causal autoencoders. + +Builds a decoder graph whose latent input and frame output are rank-5 so the +temporal axis stays explicit end to end: + +- ``latent_sample``: ``[batch, latent_channels, latent_frames, height, width]`` +- ``sample``: ``[batch, out_channels, frames, out_height, out_width]`` + +The frame count is a free dimension: a causal video decoder expands +``latent_frames`` by the temporal compression ratio, so nothing about the graph +may assume a single frame or a fixed clip length. + +The decoder additionally exposes the reference implementation's ``conv_cache`` +as paired ``conv_cache.`` inputs and ``conv_cache_out.`` outputs. +Long clips are decoded a few latent frames at a time, and those tensors are the +only state that crosses a chunk boundary; a caller that decodes a whole clip in +one call passes zero-length caches and drops the outputs. +""" + +from __future__ import annotations + +from typing import ClassVar + +import onnx_ir as ir + +from mobius._model_package import ModelPackage +from mobius.models.cogvideox_vae import CogVideoXVAEConfig +from mobius.tasks._base import ModelTask, _make_graph, _make_model + +CONV_CACHE_INPUT_PREFIX = "conv_cache." +CONV_CACHE_SCALE_METADATA = "mobius.conv_cache.spatial_scale." +CONV_CACHE_OUTPUT_PREFIX = "conv_cache_out." + + +class VideoVAETask(ModelTask): + """Build the decode graph of a 3D causal video autoencoder.""" + + model_roles: ClassVar[dict[str, str]] = {"decoder": "decoder"} + + def build( + self, + module, + config: CogVideoXVAEConfig, + ) -> ModelPackage: + graph, builder = _make_graph(name="video_vae_decoder") + op = builder.op + + latent_sample = builder.input( + "latent_sample", + dtype=ir.DataType.FLOAT, + shape=[ + "batch", + config.latent_channels, + "latent_frames", + "latent_height", + "latent_width", + ], + ) + + conv_cache = {} + for entry in module.conv_cache_spec(): + scale = entry.spatial_scale + height = "latent_height" if scale == 1 else f"{scale}*latent_height" + width = "latent_width" if scale == 1 else f"{scale}*latent_width" + conv_cache[entry.name] = builder.input( + f"{CONV_CACHE_INPUT_PREFIX}{entry.name}", + dtype=ir.DataType.FLOAT, + shape=["batch", entry.channels, "cache_frames", height, width], + ) + + sample, updated_cache = module(op, latent_sample=latent_sample, conv_cache=conv_cache) + # The temporal and spatial extents are recovered from Shape ops inside + # the decoder, which mints anonymous symbolic dimensions. Name them so + # the published contract states the compression ratios instead of + # leaking solver-internal dimension identities. + spatial = 2 ** (len(config.block_out_channels) - 1) + sample.shape = ir.Shape( + [ + "batch", + config.out_channels, + "frames", + f"{spatial}*latent_height" if spatial != 1 else "latent_height", + f"{spatial}*latent_width" if spatial != 1 else "latent_width", + ] + ) + builder.add_output(sample, "sample") + for name, value in updated_cache.items(): + builder.add_output(value, f"{CONV_CACHE_OUTPUT_PREFIX}{name}") + + model = _make_model(graph) + # Record each cache port's resolution relative to the latent. A consumer + # that has to allocate the empty first-chunk caches cannot recover this + # from the port's symbolic dimensions, which shape inference may rename + # when it unifies a declared name with an internal value. + for entry in module.conv_cache_spec(): + key = f"{CONV_CACHE_SCALE_METADATA}{CONV_CACHE_INPUT_PREFIX}{entry.name}" + model.metadata_props[key] = str(entry.spatial_scale) + return ModelPackage({"decoder": model}, config=config) diff --git a/testdata/cases/diffusion/qwen-image-edit-2509.yaml b/testdata/cases/diffusion/qwen-image-edit-2509.yaml index 587355096..0f0a0e332 100644 --- a/testdata/cases/diffusion/qwen-image-edit-2509.yaml +++ b/testdata/cases/diffusion/qwen-image-edit-2509.yaml @@ -9,5 +9,22 @@ inputs: - "Make the cat wear a bright red scarf while preserving its pose." images: - "pipeline-cat-chonk.jpeg" -skip_reason: "The official 20.4B pipeline requires over 57 GB of weights; genuine reduced-config diffusion L4/L5 goldens are covered by src/mobius/models/qwen_image_test.py and are explicitly not official-weight outputs." +skip_reason: "The checked-in golden JSON is a reduced-config reference (seeded random weights), so this runner cannot be pointed at the 57.7 GB official checkpoint without comparing against the wrong reference. Reduced-config L4/L5 goldens live in src/mobius/models/qwen_image_test.py; official-weight validation is recorded under real_weight_validation below." notes: "QwenImageEditPlusPipeline. The reduced golden encodes pipeline-cat-chonk.jpeg through a seeded diffusers Qwen Image VAE, packs the normalized posterior mode as source conditioning, runs the diffusers transformer and three FlowMatch Euler steps, and decodes the final edit latent. Mobius independently replays the same VAE-to-denoiser-to-decoder chain." +real_weight_validation: + status: "passed" + checkpoint: "Qwen/Qwen-Image-Edit-2509 @ d3968ef930e841f4c73640fb8afa3b306a78167e (public, ungated, 57.7 GB)" + hardware: "NVIDIA H200 143 GB, CUDA 13.3, driver 580.105.08; 58 GiB peak GPU" + upstream: "diffusers 0.39.0 QwenImageEditPlusPipeline, seed 42, 8 steps, true_cfg_scale 4.0, 1216x864 output" + export: "mobius bfloat16 / CUDA EP; vae_encoder + vae_decoder + transformer + 3-stage text conditioning" + component_parity: + vae_encoder: "cos 0.99997, max abs 3.44e-2" + vae_decoder: "53.31 dB PSNR" + transformer_noise_pred: "cos 0.999974 (conditional) / 0.999981 (unconditional)" + vision_encoder: "cos 0.99771" + text_encoder: "cos 0.9636 isolated (bfloat16 cancellation in decoder layer 27; not a Mobius defect)" + end_to_end: + python_onnx_replay: "38.17 dB PSNR vs upstream image" + onnx_genai_runtime: "37.10 dB PSNR / cos 0.99971 vs upstream, executed from Mobius-emitted inference_metadata.yaml" + onnx_genai_runtime_onnx_conditioning: "33.49 dB PSNR / cos 0.99937 with ONNX-produced prompt embeds" + batch_of_two: "row 0 34.79 dB / cos 0.99952; row 1 25.54 dB / cos 0.97753 (28.02 dB unpadded at batch 1)" diff --git a/testdata/cases/encoder/esm2-8m.yaml b/testdata/cases/encoder/esm2-8m.yaml index 68a73f86c..b8ed81470 100644 --- a/testdata/cases/encoder/esm2-8m.yaml +++ b/testdata/cases/encoder/esm2-8m.yaml @@ -10,5 +10,8 @@ inputs: level: "L4" -skip_reason: "Architecture-specific differences from generic BERT (extra embeddings, missing params)." -notes: "ESM-2 8M. Protein language model." +notes: >- + ESM-2 8M protein language model (MIT licence). Not a BERT variant: rotary + positions, pre-LayerNorm blocks, a final emb_layer_norm_after, no token-type + embedding, and a token-dropout rescale that applies even with no mask token. + Handled by mobius.models.esm.EsmModel. diff --git a/testdata/cases/encoder/protbert.yaml b/testdata/cases/encoder/protbert.yaml new file mode 100644 index 000000000..252aabddf --- /dev/null +++ b/testdata/cases/encoder/protbert.yaml @@ -0,0 +1,23 @@ +model_id: "Rostlab/prot_bert" +model_type: "bert" +revision: "7a894481acdc12202f0a415dd567f6cfdb698908" +task_type: "feature-extraction" +dtype: "float32" + +inputs: + prompts: + - "M K T V R Q E R L K S I V R I L E R S K E P V S G A Q L A E E L S V S R Q V I V Q D I A Y L R S L G Y N I V A T P R G Y V L A G G" + +level: "L4" + +skip_reason: >- + Upstream repo ships pytorch_model.bin only (no safetensors), which mobius + does not load, and its config.json predates the mandatory model_type key. + Both are resolvable locally -- convert the checkpoint to safetensors and + build with --config -- but neither is reproducible from the hub id + alone, so this case cannot run unattended. +notes: >- + ProtBert: BERT-large over a 30-token amino-acid vocabulary, whitespace + separated. No licence is declared on the model card. Verified by hand + against HuggingFace at revision 7a89448: relative L2 error 1.7e-5 and + per-residue cosine > 0.999999 on a padded batch of two proteins. diff --git a/testdata/cases/schema.json b/testdata/cases/schema.json index 6ce504866..3f53cdcf1 100644 --- a/testdata/cases/schema.json +++ b/testdata/cases/schema.json @@ -210,6 +210,55 @@ "type": "string", "description": "Human-readable description of the test case. Not parsed by the test runner." }, + "real_weight_validation": { + "type": "object", + "description": "Record of an out-of-band validation run against the official checkpoint, for cases whose checked-in golden is a reduced-config reference and whose real weights are too large for CI. Not parsed by the test runner; it documents evidence a reviewer would otherwise have to take on trust.", + "required": [ + "status", + "checkpoint" + ], + "properties": { + "status": { + "enum": [ + "passed", + "failed", + "partial" + ], + "description": "Outcome of the official-weight run." + }, + "checkpoint": { + "type": "string", + "description": "Checkpoint and revision the run used, including its size." + }, + "hardware": { + "type": "string", + "description": "Accelerator, driver, and peak memory the run needed." + }, + "upstream": { + "type": "string", + "description": "Reference implementation and the exact generation settings it ran with." + }, + "export": { + "type": "string", + "description": "Mobius export configuration under test (dtype, execution provider, sub-models)." + }, + "component_parity": { + "type": "object", + "additionalProperties": { + "type": "string" + }, + "description": "Per-component agreement with the reference, keyed by sub-model name." + }, + "end_to_end": { + "type": "object", + "additionalProperties": { + "type": "string" + }, + "description": "Whole-pipeline agreement with the reference, keyed by the execution path measured." + } + }, + "additionalProperties": false + }, "architecture": { "type": "string", "description": "Optional registry architecture key (e.g. Qwen35MtpModel) forcing a specific module class + task at build time. Needed for auxiliary heads (DFlash, MTP) that share a base checkpoint whose architectures field would otherwise auto-route to the base model." diff --git a/tests/_test_configs.py b/tests/_test_configs.py index 776e59e4d..c1a8c08c7 100644 --- a/tests/_test_configs.py +++ b/tests/_test_configs.py @@ -52,6 +52,7 @@ YolosConfig, Zamba2Config, ) +from mobius.models import EsmConfig # --------------------------------------------------------------------------- # Tiny model dimensions shared by all configs @@ -1465,7 +1466,18 @@ def _base_config(config_cls=None, **overrides) -> ArchitectureConfig: ("electra", {"hidden_act": "gelu", "type_vocab_size": 2}, False), ("ernie", {"hidden_act": "gelu", "type_vocab_size": 2}, False), ("ernie_m", {"hidden_act": "gelu", "type_vocab_size": 2}, False), - ("esm", {"hidden_act": "gelu", "type_vocab_size": 2}, False), + ( + "esm", + { + "_config_cls": EsmConfig, + "hidden_act": "gelu", + "type_vocab_size": 2, + # ESM-2 rotates the whole head dimension with base 10000 and + # declares it through ``position_embedding_type`` alone. + "partial_rotary_factor": 1.0, + }, + False, + ), ("flaubert", {"hidden_act": "gelu", "type_vocab_size": 2}, False), ("ibert", {"hidden_act": "gelu", "type_vocab_size": 2}, False), ( diff --git a/tests/arch_validation_test.py b/tests/arch_validation_test.py index d0abc65f7..6aefdbd67 100644 --- a/tests/arch_validation_test.py +++ b/tests/arch_validation_test.py @@ -141,7 +141,7 @@ def _resolve_hf_config(hf_config): def _build_graph(model_type: str, model_id: str): - """Download config, build ONNX graph, return ModelPackage. + """Download config, build ONNX graph, return ``(ModelPackage, task)``. Uses get_task().build() directly (same pattern as build_graph_test.py) to bypass ArchitectureConfig.validate() which rejects non-LM configs @@ -164,7 +164,7 @@ def _build_graph(model_type: str, model_id: str): module = registration.module_class(config) task_name = registration.task or _default_task_for_model(model_type) task = get_task(task_name) - return task.build(module, config) + return task.build(module, config), task @pytest.mark.arch_validation @@ -195,7 +195,7 @@ def test_full_graph_builds(self, model_type: str, model_id: str): shape mismatches, missing fields, or initialization errors, this test will catch them. """ - pkg = _build_graph(model_type, model_id) + pkg, task = _build_graph(model_type, model_id) # Validate: every component has a non-empty graph assert len(pkg) > 0, "ModelPackage is empty" @@ -206,6 +206,16 @@ def test_full_graph_builds(self, model_type: str, model_id: str): assert len(model.graph.inputs) > 0, f"{component_name} has no inputs" assert len(model.graph.outputs) > 0, f"{component_name} has no outputs" + # Every component the task builds must declare a role. An undeclared + # component falls back to the "decoder" role in build_from_module and + # would be handed fusion passes meant for attention stacks, and + # inspect_components would not report it at all. + undeclared = sorted(set(pkg) - set(task.model_roles or {})) + assert not undeclared, ( + f"{model_type}: components {undeclared} are built but missing from " + f"{type(task).__name__}.model_roles" + ) + del pkg @pytest.mark.parametrize("model_type,model_id", _GRAPH_PARAMS) @@ -216,10 +226,16 @@ def test_graph_shapes_consistent(self, model_type: str, model_id: str): inputs/outputs are defined. Note: output type info may not be available for all outputs when building without shape inference. """ - pkg = _build_graph(model_type, model_id) + pkg, task = _build_graph(model_type, model_id) + roles = task.model_roles or {} for component_name, model in pkg.items(): - # Model should have initializers (parameters) + # Model should have initializers (parameters). "glue" components are + # parameter-free loop wiring — they read every tensor they use from + # a graph input, so they hold only hoisted constants (if anything) + # and the parameter check does not apply. + if roles.get(component_name) == "glue": + continue initializers = list(model.graph.initializers) assert len(initializers) > 0, ( f"{component_name} has no initializers — graph may be missing parameters" diff --git a/tests/build_graph_test.py b/tests/build_graph_test.py index c3076e436..c64d5a5b4 100644 --- a/tests/build_graph_test.py +++ b/tests/build_graph_test.py @@ -4459,6 +4459,77 @@ def test_cogvideox_graph_builds(self): assert len(sample_input.shape) == 5 # [B, T, C, H, W] +class TestBuildCogVideoXVAEGraph: + """Verify the CogVideoX causal 3D VAE decoder graph construction.""" + + @staticmethod + def _config(): + from mobius.models.cogvideox_vae import CogVideoXVAEConfig + + return CogVideoXVAEConfig( + in_channels=3, + out_channels=3, + latent_channels=4, + block_out_channels=(8, 8, 8, 8), + layers_per_block=1, + norm_num_groups=2, + temporal_compression_ratio=4, + scaling_factor=1.15258426, + ) + + def test_video_vae_graph_builds_with_paired_conv_caches(self): + from mobius.models.cogvideox_vae import AutoencoderKLCogVideoXModel + from mobius.tasks import VideoVAETask + from mobius.tasks._video_vae import ( + CONV_CACHE_INPUT_PREFIX, + CONV_CACHE_OUTPUT_PREFIX, + CONV_CACHE_SCALE_METADATA, + ) + + config = self._config() + model = VideoVAETask().build(AutoencoderKLCogVideoXModel(config), config)["decoder"] + + assert model.graph is not None + latent = next(v for v in model.graph.inputs if v.name == "latent_sample") + # [B, C, T, H, W]: the temporal axis is explicit, and the frame count is + # a free dimension rather than a baked clip length. + assert len(latent.shape) == 5 + assert str(latent.shape[2]) == "latent_frames" + + sample = next(v for v in model.graph.outputs if v.name == "sample") + assert len(sample.shape) == 5 + assert int(sample.shape[1]) == config.out_channels + + cache_inputs = { + v.name[len(CONV_CACHE_INPUT_PREFIX) :] + for v in model.graph.inputs + if v.name.startswith(CONV_CACHE_INPUT_PREFIX) + } + cache_outputs = { + v.name[len(CONV_CACHE_OUTPUT_PREFIX) :] + for v in model.graph.outputs + if v.name.startswith(CONV_CACHE_OUTPUT_PREFIX) + } + # Every cached convolution has to be readable and writable, or a clip + # decoded in chunks would silently lose the frames before each chunk. + assert cache_inputs + assert cache_inputs == cache_outputs + for name in cache_inputs: + key = f"{CONV_CACHE_SCALE_METADATA}{CONV_CACHE_INPUT_PREFIX}{name}" + assert key in model.metadata_props + + def test_video_vae_cache_spec_matches_upsampled_resolutions(self): + from mobius.models.cogvideox_vae import AutoencoderKLCogVideoXModel + + config = self._config() + module = AutoencoderKLCogVideoXModel(config) + scales = {entry.name: entry.spatial_scale for entry in module.conv_cache_spec()} + assert scales["conv_in"] == 1 + # Three upsampling stages for four blocks, so the last cached + # convolutions live at the full frame resolution. + assert scales["conv_out"] == 2 ** (len(config.block_out_channels) - 1) + + class TestBuildAdapterGraph: """Verify T2I-Adapter and IP-Adapter graph construction.""" diff --git a/tests/canonical_workflow_contract_test.py b/tests/canonical_workflow_contract_test.py new file mode 100644 index 000000000..2e9bea22d --- /dev/null +++ b/tests/canonical_workflow_contract_test.py @@ -0,0 +1,947 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT License. + +"""One serialized contract shape, for every package this producer can emit. + +A package describes itself in exactly one place: ``pipeline.workflow``. That is +true of a three-graph vision-language package, and it is equally true of a bare +single-file decoder — the single-file case is a one-component workflow, not a +different kind of document with its own keys. A runtime is free to *lower* that +one component onto an optimized single-graph path; what it may not do is read a +second, independently writable statement of the same facts, because nothing +would force the two to agree and a reader of one never learns that the other +said something else. + +These tests exist because that property is only worth anything if it cannot +quietly lapse. It would be easy for a feature added to one export shape — a +fixed-capacity cache, an FP8 buffer, a heterogeneous decoder — to grow its own +top-level block "just for this case", and easy for that to go unnoticed while +every feature-specific test kept passing. So the assertions here are +deliberately shape-agnostic: they are asked of dynamic, static-cache, FP8, +heterogeneous and composite packages through the same code path, and of every +checked-in fixture package, and none of them mentions a feature by name. +""" + +from __future__ import annotations + +import copy +import glob +import os +import re +from typing import Any + +import onnx_ir as ir +import pytest +import yaml + +from mobius import registry +from mobius._configs import ArchitectureConfig +from mobius._optimizations import optimize_model +from mobius.integrations.onnx_genai.workflow_metadata import ( + build_decoder_workflow_metadata, + build_vlm_workflow_metadata, +) +from mobius.tasks import CausalLMTask + +CAPACITY = 64 + +# Enough layers that a cell label's lexicographic order stops agreeing with its +# numeric one: with ten or more cells, ``cache_10`` sorts between ``cache_1`` +# and ``cache_2``. Below that threshold every ordering rule looks correct. +DEEP_LAYERS = 12 + +FIXTURE_ROOT = os.path.join(os.path.dirname(__file__), "fixtures", "onnx_genai_workflows") + + +def _text_config(**overrides: Any) -> ArchitectureConfig: + params: dict[str, Any] = { + "num_hidden_layers": 2, + "hidden_size": 64, + "intermediate_size": 128, + "num_attention_heads": 4, + "num_key_value_heads": 2, + "head_dim": 16, + "vocab_size": 256, + "rms_norm_eps": 1e-6, + "hidden_act": "silu", + "max_position_embeddings": 512, + } + params.update(overrides) + return ArchitectureConfig(**params) + + +def _dynamic_decoder() -> tuple[Any, dict[str, Any]]: + """One ONNX file, an appending cache: the simplest package there is.""" + config = _text_config() + pkg = CausalLMTask().build(registry.get("qwen2")(config), config) + return pkg, build_decoder_workflow_metadata(pkg, config) + + +def _static_cache_decoder() -> tuple[Any, dict[str, Any]]: + """One ONNX file whose cache is scattered into fixed-capacity buffers.""" + config = _text_config() + task = CausalLMTask(static_cache=True, max_seq_len=CAPACITY) + pkg = task.build(registry.get("qwen2")(config), config) + return pkg, build_decoder_workflow_metadata(pkg, config) + + +def _fp8_decoder() -> tuple[Any, dict[str, Any]]: + """One ONNX file whose cache buffers are FP8 rather than the compute dtype.""" + config = _text_config() + pkg = CausalLMTask().build(registry.get("qwen2")(config), config) + optimize_model( + pkg["model"], + ep="cuda", + dtype=ir.DataType.FLOAT16, + model_role="decoder", + fp8_kv_cache=True, + ) + return pkg, build_decoder_workflow_metadata(pkg, config) + + +def _heterogeneous_decoder() -> tuple[Any, dict[str, Any]]: + """One ONNX file that publishes two state groups instead of one. + + A hybrid decoder interleaves sliding and full attention, so its cells land + in different groups with different eviction rules. The contract shape must + not depend on how many groups a package happens to need. + """ + config = _text_config( + sliding_window=8, layer_types=["sliding_attention", "full_attention"] + ) + pkg = CausalLMTask().build(registry.get("qwen2")(config), config) + return pkg, build_decoder_workflow_metadata(pkg, config) + + +def _composite_vision_language() -> tuple[Any, dict[str, Any]]: + """Three ONNX files driven by one workflow. + + This is the same package the runtime conformance suite executes, built by + the same helper, so the composite arm of these tests and the composite arm + of the executed fixtures cannot describe different things. + """ + from generate_onnx_genai_validation_packages import _executable_vlm_package + + pkg = _executable_vlm_package() + return pkg, build_vlm_workflow_metadata(pkg, pkg.config) + + +_PACKAGES = { + "dynamic": _dynamic_decoder, + "static_cache": _static_cache_decoder, + "fp8": _fp8_decoder, + "heterogeneous": _heterogeneous_decoder, + "composite": _composite_vision_language, +} + + +@pytest.fixture(scope="module") +def built() -> dict[str, tuple[Any, dict[str, Any]]]: + return {name: build() for name, build in _PACKAGES.items()} + + +@pytest.fixture(params=sorted(_PACKAGES), scope="module") +def package(request, built): + return built[request.param] + + +def _onnx_components(workflow: dict[str, Any]) -> dict[str, dict[str, Any]]: + return { + name: component + for name, component in workflow["components"].items() + if component.get("implementation", {}).get("kind") == "onnx" + } + + +def _graph_ports(model: ir.Model) -> tuple[set[str], set[str]]: + """The ports an artifact actually exposes, which is the only authority.""" + return ( + {str(value.name) for value in model.graph.inputs}, + {str(value.name) for value in model.graph.outputs}, + ) + + +def _artifacts(pkg: Any, workflow: dict[str, Any]) -> dict[str, ir.Model]: + """Every ONNX component of *workflow* paired with the graph it references. + + A component names an artifact, and the artifact is what a runtime binds + against. Resolving a declaration through this map is what makes these + assertions checks of the package rather than checks of the metadata's + internal consistency with itself. + """ + policies = { + name: component.model + for name, component in getattr(pkg, "policy_components", {}).items() + } + graphs = {**dict(pkg.items()), **policies} + return {name: graphs[name] for name in _onnx_components(workflow) if name in graphs} + + +def _walk_steps(steps: list[dict[str, Any]]): + """Every step of a workflow, including the ones nested in loops and branches.""" + for step in steps: + yield step + for key in ("setup", "steps", "nodes"): + if isinstance(step.get(key), list): + yield from _walk_steps(step[key]) + for case in (step.get("cases") or {}).values(): + yield from _walk_steps([case]) + if isinstance(step.get("default"), dict): + yield from _walk_steps([step["default"]]) + + +def _groups(workflow: dict[str, Any]) -> dict[str, Any]: + return (workflow.get("serving") or {}).get("state_service", {}).get("groups", {}) or {} + + +def _scatter_bound_components(workflow: dict[str, Any]) -> set[str]: + """Components a fixed-capacity group hands its write cursor and valid length to. + + These two ports are consumed by exactly one thing: the driver that writes + into a preallocated cache at an index. That driver binds its ports from the + resolved decode ABI, so naming a component here is a claim that the + component is resolvable as a decoder. + """ + bound: set[str] = set() + for group in _groups(workflow).values(): + update = group.get("update") or {} + if update.get("kind") != "indexed_scatter": + continue + bound |= set(update.get("write_indices_ports") or {}) + bound |= set(update.get("kv_length_ports") or {}) + return bound + + +def _declares_a_sequence_role(component: dict[str, Any]) -> bool: + roles = ((component.get("ports") or {}).get("roles")) or {} + return bool({"token_ids", "inputs_embeds"} & set(roles.values())) + + +class TestOneSerializedContract: + """Facts that hold for every package shape, asserted through one code path.""" + + def test_the_workflow_is_where_a_package_describes_itself(self, package): + _, metadata = package + assert "workflow" in metadata["pipeline"] + + def test_no_package_states_its_graph_abi_a_second_time(self, package): + """The point of one representation is that there is no other one. + + ``model`` may still carry package-wide geometry, but the moment it + carries an ``io`` block the package has two writable answers to "what + does the decode step look like", and a reader of either never learns the + other exists. + """ + _, metadata = package + assert "io" not in (metadata.get("model") or {}) + + def test_an_exported_graph_is_not_transcribed_into_the_workflow(self, package): + """A model component declares its roles and nothing the artifact says. + + The ``.onnx`` file ships inside the package and is authoritative for + which ports exist and what each one's dtype, rank and shape is. Copying + that into YAML would create a second statement of the same fact with + nothing to keep the two in agreement, which is the failure this whole + module exists to prevent — the copy just happens to sit one level down + from ``model.io`` rather than beside it. + """ + pkg, metadata = package + components = _onnx_components(metadata["pipeline"]["workflow"]) + exported = [name for name in components if name in pkg] + assert exported, "a package with no exported graph describes nothing" + for name in exported: + ports = components[name].get("ports", {}) + assert not ports.get("inputs") + assert not ports.get("outputs") + + def test_a_synthesized_policy_graph_still_types_the_dataflow(self, package): + """The producer's own control graphs are the workflow's type annotations. + + A policy graph is not an external interface this producer describes; it + is a graph this producer emits to realize the workflow's control flow, + and the value an invocation produces takes its dtype, rank and request + axis from the port that produced it. A validator reads metadata without + the artifacts, so dropping these would leave every derived value + untyped — the contract here is load-bearing, not a transcription. + """ + pkg, metadata = package + workflow = metadata["pipeline"]["workflow"] + components = _onnx_components(workflow) + policies = [ + components[name] + for name in getattr(pkg, "policy_components", {}) + if name in components + ] + for component in policies: + ports = component["ports"] + assert ports["inputs"] or ports["outputs"] + for contract in (*ports["inputs"].values(), *ports["outputs"].values()): + assert contract["rank"] == len(contract["shape"]) + + def test_every_declared_role_names_a_port_the_graph_exposes(self, package): + """A role is only meaningful if it resolves in the artifact. + + This is what replaces the transcription: rather than restating the + graph and checking the restatement against itself, the role is checked + against the graph it claims to describe. A role naming a port the file + does not expose binds nothing, and is caught here exactly as the + runtime would catch it against a live session. + """ + pkg, metadata = package + workflow = metadata["pipeline"]["workflow"] + components = _onnx_components(workflow) + roled = 0 + for name, model in _artifacts(pkg, workflow).items(): + inputs, outputs = _graph_ports(model) + for port in (components[name].get("ports") or {}).get("roles", {}): + assert port in inputs or port in outputs + roled += 1 + assert roled, "no component says what it does with any of its ports" + + def test_every_invocation_binds_a_port_of_the_artifact(self, package): + """A binding to a port the graph lacks names nothing a runtime can feed.""" + pkg, metadata = package + workflow = metadata["pipeline"]["workflow"] + graphs = _artifacts(pkg, workflow) + for step in _walk_steps(workflow["steps"]): + if step.get("kind") != "invoke" or step["component"] not in graphs: + continue + inputs, outputs = _graph_ports(graphs[step["component"]]) + assert set(step.get("inputs", {})) <= inputs + assert set(step.get("outputs", {})) <= outputs + + def test_every_state_pair_names_ports_of_the_artifact(self, package): + """State is carried through ports, so both halves have to exist.""" + pkg, metadata = package + workflow = metadata["pipeline"]["workflow"] + graphs = _artifacts(pkg, workflow) + for group in _groups(workflow).values(): + for component, aliases in (group.get("ports") or {}).items(): + if component not in graphs: + continue + inputs, outputs = _graph_ports(graphs[component]) + for alias in aliases.values(): + assert alias["input"] in inputs + assert alias["output"] in outputs + + def test_split_cache_halves_and_layers_are_stated_not_positional(self, package): + """A layer's key and value buffers are indistinguishable once listed. + + They are the same dtype and the same shape, and a cell's label sorts + lexicographically (``cache_10`` before ``cache_2``), so a consumer that + paired them positionally would transpose two layers' caches without + anything failing. When a pair carries a half, it says so, and it says + which layer it belongs to. + """ + _, metadata = package + for group in _groups(metadata["pipeline"]["workflow"]).values(): + for aliases in (group.get("ports") or {}).values(): + halves = [alias.get("role") for alias in aliases.values()] + if not any(halves): + continue + assert all(half in {"key", "value", "combined"} for half in halves) + assert all("layer" in alias for alias in aliases.values()) + keys = [alias for alias in aliases.values() if alias["role"] == "key"] + values = [alias for alias in aliases.values() if alias["role"] == "value"] + assert len(keys) == len(values) + assert {alias["layer"] for alias in keys} == { + alias["layer"] for alias in values + } + + def test_control_ports_of_a_fixed_capacity_cache_exist_in_the_artifact(self, package): + """A scatter's two control vectors are rank-1 integers, so they must be named. + + Nothing distinguishes the write cursor from the valid length by shape. + A package that scatters therefore names both against a component whose + graph exposes both; a package that appends declares no scatter at all, + and this assertion is vacuous for it — which is the point, since the + same test runs over every shape. + """ + pkg, metadata = package + workflow = metadata["pipeline"]["workflow"] + graphs = _artifacts(pkg, workflow) + for group in _groups(workflow).values(): + update = group.get("update") or {} + if update.get("kind") != "indexed_scatter": + continue + bound = set(group["ports"]) + assert set(update["write_indices_ports"]) == bound + assert set(update["kv_length_ports"]) == bound + for component in bound: + inputs, _ = _graph_ports(graphs[component]) + assert update["write_indices_ports"][component] in inputs + assert update["kv_length_ports"][component] in inputs + + def test_the_decode_step_is_recoverable_from_the_workflow_alone(self, package): + """Reconstruct the decode ABI the way a runtime lowering would. + + This is the assertion that makes removing the second copy safe: if the + sequence input, the logits output and the per-layer cache pairs can all + be read off the workflow, then a separate block stating them again was + never carrying information — only risk. Roles are the only thing read + here, and every name they yield is checked against the graph, so the + reconstruction never falls back to recognizing a spelling. + """ + pkg, metadata = package + workflow = metadata["pipeline"]["workflow"] + graphs = _artifacts(pkg, workflow) + decoders = [] + for name, component in _onnx_components(workflow).items(): + roles = (component.get("ports") or {}).get("roles", {}) + consumes = {"token_ids", "inputs_embeds"} & set(roles.values()) + owns_state = any( + name in (group.get("ports") or {}) for group in _groups(workflow).values() + ) + if consumes and ("logits" in roles.values() or owns_state): + decoders.append((name, roles)) + assert decoders, "no component declares what it does with the sequence" + for name, roles in decoders: + inputs, outputs = _graph_ports(graphs[name]) + sequence = [ + port + for port, role in roles.items() + if role in {"token_ids", "inputs_embeds"} and port in inputs + ] + assert len(sequence) == 1 + assert all(port in outputs for port, role in roles.items() if role == "logits") + pairs = [ + alias + for group in _groups(workflow).values() + for alias in (group.get("ports") or {}).get(name, {}).values() + ] + assert all( + alias["input"] in inputs and alias["output"] in outputs for alias in pairs + ) + + +def _cell_order(alias: dict[str, Any], label: str) -> tuple[Any, ...]: + """Canonical order of a state pair: its layer, then its half.""" + return (alias.get("layer", 0), alias.get("role", ""), label) + + +def _resolve_decode_abi(workflow: dict[str, Any], component: str) -> dict[str, Any]: + """Bind the decode ABI the way a consumer with no port vocabulary must. + + Nothing here looks at how a port is spelled. The sequence and logits ports + come from declared roles, and every cache port comes from a state-service + alias, ordered by the layer and half the alias states. A consumer that took + any other route would be recognizing names — which is the failure mode this + resolution exists to make impossible. + """ + declaration = workflow["components"][component] + roles = (declaration.get("ports") or {}).get("roles", {}) + caches: list[tuple[str, str]] = [] + write_indices: str | None = None + kv_length: str | None = None + for group in _groups(workflow).values(): + aliases = (group.get("ports") or {}).get(component) or {} + for label in sorted(aliases, key=lambda label: _cell_order(aliases[label], label)): + caches.append((aliases[label]["input"], aliases[label]["output"])) + update = group.get("update") or {} + if update.get("kind") == "indexed_scatter": + write_indices = update["write_indices_ports"][component] + kv_length = update["kv_length_ports"][component] + return { + "token_ids": [port for port, role in roles.items() if role == "token_ids"], + "inputs_embeds": [port for port, role in roles.items() if role == "inputs_embeds"], + "logits": [port for port, role in roles.items() if role == "logits"], + "cache_inputs": [pair[0] for pair in caches], + "cache_outputs": [pair[1] for pair in caches], + "write_indices": write_indices, + "kv_sequence_length": kv_length, + } + + +def _rename_component_ports( + workflow: dict[str, Any], component: str, renames: dict[str, str] +) -> None: + """Rewrite every place the workflow names a port of *component*. + + The four places are the whole surface: the role table, the invocations that + bind the ports, the state pairs that carry buffers through them, and the + scatter's two control ports. + """ + declaration = workflow["components"][component] + ports = declaration.get("ports") or {} + if "roles" in ports: + ports["roles"] = {renames.get(k, k): v for k, v in ports["roles"].items()} + for side in ("inputs", "outputs"): + if ports.get(side): + ports[side] = {renames.get(k, k): v for k, v in ports[side].items()} + for step in _walk_steps(workflow["steps"]): + if step.get("kind") != "invoke" or step.get("component") != component: + continue + for side in ("inputs", "outputs"): + if step.get(side): + step[side] = {renames.get(k, k): v for k, v in step[side].items()} + for group in _groups(workflow).values(): + for alias in ((group.get("ports") or {}).get(component) or {}).values(): + alias["input"] = renames.get(alias["input"], alias["input"]) + alias["output"] = renames.get(alias["output"], alias["output"]) + update = group.get("update") or {} + for key in ("write_indices_ports", "kv_length_ports"): + bindings = update.get(key) or {} + if component in bindings: + bindings[component] = renames.get(bindings[component], bindings[component]) + + +class TestRolesAloneCarryTheDecodeAbi: + """Omitting the port contracts must not push a consumer back to guessing. + + A component that ships an artifact declares no port contracts, so the only + thing left saying what a port *means* is its role. The risk that creates is + specific: a consumer that cannot find a role does not fail loudly, it falls + back to matching the spelling ``input_ids`` — and then a package that + spells it differently binds the wrong tensor with the right shape. + + These tests hold the producer to the side of that contract it owns. Every + export shape must declare the roles a consumer needs, and the whole decode + ABI must survive renaming every port to an opaque label: if any part of the + binding still resolved after that, it was resolving by name. + """ + + @staticmethod + def _roled(workflow: dict[str, Any]) -> list[str]: + """Every component that says what any of its ports means.""" + return [ + name + for name, component in _onnx_components(workflow).items() + if ((component.get("ports") or {}).get("roles") or {}) + ] + + @staticmethod + def _decoders(workflow: dict[str, Any]) -> list[str]: + """Components that drive a decode step. + + A decoder consumes the sequence and either produces logits or owns + cache state. An embedding component consumes tokens too, but it is not + the step a runtime specializes. + """ + decoders = [] + for name, component in _onnx_components(workflow).items(): + roles = set(((component.get("ports") or {}).get("roles") or {}).values()) + owns_state = any( + name in (group.get("ports") or {}) for group in _groups(workflow).values() + ) + if {"token_ids", "inputs_embeds"} & roles and ("logits" in roles or owns_state): + decoders.append(name) + return decoders + + def test_a_shipped_artifact_declares_roles_and_no_contracts(self, package): + """The omission and the role are one decision, not two. + + Dropping the contracts is only safe because the role is there; a + component with neither would be a graph a consumer can only guess at. + """ + pkg, metadata = package + workflow = metadata["pipeline"]["workflow"] + for name in _artifacts(pkg, workflow): + if name not in pkg: + continue + ports = workflow["components"][name].get("ports") or {} + assert not ports.get("inputs") and not ports.get("outputs") + assert self._decoders(workflow), "no exported graph says what consumes the sequence" + + def test_the_sequence_port_is_declared_not_spelled(self, package): + """``input_ids`` and ``position_ids`` are both rank-2 int64. + + Nothing in a graph distinguishes them, so the producer states which one + is the autoregressive sequence rather than leaving a runtime to infer + it from a name it has no right to assume. A decode step consumes + exactly one sequence — tokens or embeddings, never both. + """ + _, metadata = package + workflow = metadata["pipeline"]["workflow"] + for name in self._decoders(workflow): + abi = _resolve_decode_abi(workflow, name) + assert len(abi["token_ids"]) + len(abi["inputs_embeds"]) == 1 + + def test_the_whole_abi_survives_renaming_every_port(self, package): + """Rename the ports to opaque labels; the ABI must resolve identically. + + This is the assertion with teeth. The renamed package contains no port + called ``input_ids``, ``logits`` or ``key_cache.0``, so a resolution + that still returns the right ports cannot have been reading names. The + answers are compared as positions, because the names deliberately no + longer match. + """ + pkg, metadata = package + workflow = copy.deepcopy(metadata)["pipeline"]["workflow"] + roled = self._roled(workflow) + assert roled + original = {name: _resolve_decode_abi(workflow, name) for name in roled} + renames = {} + for index, name in enumerate(sorted(original)): + inputs, outputs = _graph_ports(_artifacts(pkg, workflow)[name]) + mapping = { + port: f"c{index}.p{position}" + for position, port in enumerate(sorted(inputs | outputs)) + } + renames[name] = mapping + _rename_component_ports(workflow, name, mapping) + recognizable = {"input_ids", "inputs_embeds", "logits"} + for name, before in original.items(): + after = _resolve_decode_abi(workflow, name) + mapping = renames[name] + for key, value in before.items(): + expected = ( + [mapping[port] for port in value] + if isinstance(value, list) + else (mapping[value] if value is not None else None) + ) + assert after[key] == expected + # Nothing recognizable is left to have matched on. + resolved = (*after["token_ids"], *after["inputs_embeds"], *after["logits"]) + assert not recognizable & set(resolved) + + def test_deleting_the_roles_is_what_breaks_it(self, package): + """The mutation that proves the role is doing the work. + + If the sequence port were still recoverable with the role table gone, + then something else — a position, a shape, a spelling — was carrying + the fact, and these tests would be asserting nothing. + """ + _, metadata = package + workflow = copy.deepcopy(metadata)["pipeline"]["workflow"] + roled = self._roled(workflow) + assert roled + for name in roled: + workflow["components"][name]["ports"].pop("roles") + for name in roled: + abi = _resolve_decode_abi(workflow, name) + assert abi["token_ids"] == [] + assert abi["inputs_embeds"] == [] + assert abi["logits"] == [] + + def test_state_pairs_still_bind_without_a_role_table(self, package): + """The cache half of the ABI is the state service's, not the role table's. + + Deleting the roles must not disturb it: a runtime that resolved caches + through the roles would break on any component whose cache ports carry + no role, and the two halves are deliberately independent. + """ + _, metadata = package + workflow = copy.deepcopy(metadata)["pipeline"]["workflow"] + roled = self._roled(workflow) + before = {name: _resolve_decode_abi(workflow, name) for name in roled} + for name in roled: + workflow["components"][name]["ports"].pop("roles") + for name in roled: + after = _resolve_decode_abi(workflow, name) + for key in ( + "cache_inputs", + "cache_outputs", + "write_indices", + "kv_sequence_length", + ): + assert after[key] == before[name][key] + + def test_a_scatter_bound_component_is_always_resolvable_as_a_decoder(self, package): + """Naming a component in an ``indexed_scatter`` group obliges it to say so. + + The write cursor and valid length exist for one consumer: the driver + that writes into a preallocated buffer at an index. That driver binds + its ports from the resolved decode ABI, and every field of that ABI is + found by role — so a component handed those two control ports while + declaring no sequence role cannot be resolved as a decoder at all. + + Nothing rejects that combination for us. Identifying the decoder + requires a sequence role, so a component that omits one is invisible to + the very check that would have caught it, and the package validates, + loads, and quietly falls back to inferring ports from shapes. The + producer is the only place the contradiction is visible, which is why + it is asserted here. + """ + _, metadata = package + workflow = metadata["pipeline"]["workflow"] + for name in _scatter_bound_components(workflow): + component = workflow["components"][name] + assert _declares_a_sequence_role(component), ( + f"{name} is handed a fixed-capacity write cursor but declares no " + "sequence role, so no decode ABI can be resolved for it" + ) + assert _resolve_decode_abi(workflow, name)["write_indices"] + + +def _deep_dynamic() -> dict[str, Any]: + """An appending cache with enough layers that labels sort out of order.""" + config = _text_config(num_hidden_layers=DEEP_LAYERS) + pkg = CausalLMTask().build(registry.get("qwen2")(config), config) + return build_decoder_workflow_metadata(pkg, config) + + +def _deep_static_cache() -> dict[str, Any]: + """The same depth, scattered into fixed-capacity buffers.""" + config = _text_config(num_hidden_layers=DEEP_LAYERS) + task = CausalLMTask(static_cache=True, max_seq_len=CAPACITY) + pkg = task.build(registry.get("qwen2")(config), config) + return build_decoder_workflow_metadata(pkg, config) + + +def _deep_heterogeneous() -> dict[str, Any]: + """A deep hybrid whose cache-owning layers are a non-contiguous subset. + + Alternating the layer types makes the full-attention group own layers + 1, 3, 5, ... only, so a cell's position in its group is never its layer. + """ + config = _text_config( + num_hidden_layers=DEEP_LAYERS, + sliding_window=8, + layer_types=["sliding_attention", "full_attention"] * (DEEP_LAYERS // 2), + ) + pkg = CausalLMTask().build(registry.get("qwen2")(config), config) + return build_decoder_workflow_metadata(pkg, config) + + +_DEEP_PACKAGES = { + "dynamic": _deep_dynamic, + "static_cache": _deep_static_cache, + "heterogeneous": _deep_heterogeneous, +} + + +@pytest.fixture(scope="module") +def deep_built() -> dict[str, dict[str, Any]]: + return {name: build() for name, build in _DEEP_PACKAGES.items()} + + +@pytest.fixture(params=sorted(_DEEP_PACKAGES), scope="module") +def deep_package(request, deep_built): + return deep_built[request.param] + + +def _roled_aliases(metadata: dict[str, Any]) -> list[dict[str, dict[str, Any]]]: + """Every per-component alias map that carries key/value halves.""" + maps = [] + for group in _groups(metadata["pipeline"]["workflow"]).values(): + for aliases in (group.get("ports") or {}).values(): + if any(alias.get("role") for alias in aliases.values()): + maps.append(aliases) + return maps + + +class TestDeepDecodersDeclareLayersRatherThanPositions: + """The layer annotation only earns its place above nine layers. + + Every other package in this file has two layers, and with two layers a + cell's label sorts the same way whichever rule is used: lexicographic, + numeric, or insertion order all agree. So do the alternatives a producer + could accidentally implement — ``layer`` taken from the enumeration index + rather than parsed from the port name is indistinguishable from the + correct value until some layer's cells outnumber a single digit. + + Real decoders have twenty to eighty layers, so that region is the normal + case in production and the unreachable case in this suite. These tests put + a package there. They are written to fail loudly if the annotation ever + degrades into a restatement of position, because the failure it prevents + is silent: two transposed caches have identical shapes and identical + dtypes, and the only symptom is that generated text is subtly wrong. + """ + + def test_labels_really_do_sort_out_of_order_at_this_depth(self, deep_package): + # Guards the premise of every other test in this class: if the labels + # happened to sort numerically, the rest would prove nothing. + for aliases in _roled_aliases(deep_package): + by_label = [aliases[cell]["layer"] for cell in sorted(aliases)] + assert by_label != sorted(by_label), ( + "labels sort into layer order, so this package cannot " + "distinguish a declared layer from a positional one" + ) + + def test_the_declared_layer_is_the_one_the_port_name_states(self, deep_package): + """The annotation restates the exporter's own port name, not an index.""" + for aliases in _roled_aliases(deep_package): + for alias in aliases.values(): + stated = re.search( + r"\.(\d+)\.(?:key|value)$|(?:key|value)_cache\.(\d+)$", alias["input"] + ) + assert stated is not None, alias["input"] + assert alias["layer"] == int(stated.group(1) or stated.group(2)) + + def test_ordering_by_the_declared_layer_recovers_the_buffer_lists(self, deep_package): + """Sorting by (layer, half) is what a consumer does; it must be numeric. + + This mirrors how the runtime collects a group's ports, so the assertion + fails here rather than as transposed caches at inference time. + """ + for aliases in _roled_aliases(deep_package): + ordered = sorted( + aliases.values(), key=lambda alias: (alias["layer"], alias["role"]) + ) + layers = [alias["layer"] for alias in ordered] + assert layers == sorted(layers) + keys = [alias["input"] for alias in ordered if alias["role"] == "key"] + values = [alias["input"] for alias in ordered if alias["role"] == "value"] + assert len(keys) == len(values) + # Each layer contributes exactly one key and one value, in step. + assert [alias["layer"] for alias in ordered if alias["role"] == "key"] == [ + alias["layer"] for alias in ordered if alias["role"] == "value" + ] + + def test_a_cells_position_in_its_group_is_not_its_layer(self, deep_built): + """Each group of a hybrid owns an alternating half of the layers. + + This is the case that separates a declared layer from a positional one + even for a producer that sorts numerically: the sliding group owns + layers 0, 2, 4, ... and the full-attention group owns 1, 3, 5, ..., so + a cell's position within its own group is never its layer. + """ + aliases_by_group = _roled_aliases(deep_built["heterogeneous"]) + assert len(aliases_by_group) == 2, "expected one group per attention type" + owned = [ + frozenset(alias["layer"] for alias in aliases.values() if alias["role"] == "key") + for aliases in aliases_by_group + ] + evens = frozenset(index for index in range(DEEP_LAYERS) if index % 2 == 0) + odds = frozenset(index for index in range(DEEP_LAYERS) if index % 2 == 1) + assert set(owned) == {evens, odds} + # Neither group's layers are its own positions, which is the property a + # positional annotation would have satisfied by construction. + for layers in owned: + assert layers != frozenset(range(len(layers))) + + +def _fixture_packages() -> list[str]: + return sorted( + os.path.dirname(path) + for path in glob.glob(os.path.join(FIXTURE_ROOT, "*", "inference_metadata.yaml")) + ) + + +@pytest.mark.parametrize( + "directory", _fixture_packages(), ids=lambda path: os.path.basename(path) +) +class TestCheckedInPackagesShareTheShape: + """The same contract, asserted against every package checked into the tree. + + The built packages above cover the decoder shapes this producer can + construct in a unit test. The fixtures cover the ones it cannot — audio + codecs, diffusion, video, speculative decoding, text-to-speech — and they + are the exact bytes the runtime conformance suite executes, so a shape that + drifts here drifts in something already proven to run. + """ + + @staticmethod + def _metadata(directory: str) -> dict[str, Any]: + with open( + os.path.join(directory, "inference_metadata.yaml"), encoding="utf-8" + ) as handle: + return yaml.safe_load(handle) + + @staticmethod + def _artifact_root(directory: str, materialized: str) -> str: + """Where this package's graphs live once they have been generated.""" + return os.path.join(materialized, os.path.basename(directory)) + + @classmethod + def _ports( + cls, directory: str, workflow: dict[str, Any], materialized: str + ) -> dict[str, tuple[set, set]]: + """Read the ports each graph really exposes. + + The metadata under review is the committed one; the graph it describes + is generated, because it is a deterministic function of the producer + and committing it would store bytes no reviewer reads. Resolving the + committed declaration against the generated graph is what makes this an + assertion about the artifact rather than about a copy of itself. + """ + root = cls._artifact_root(directory, materialized) + ports: dict[str, tuple[set, set]] = {} + for name, component in _onnx_components(workflow).items(): + artifact = os.path.join(root, component["implementation"]["artifact"]) + assert os.path.exists(artifact), ( + f"{os.path.basename(directory)} declares component {name!r} backed by " + f"{component['implementation']['artifact']!r}, but the generator emitted " + f"no such file. Two of the assertions below skip components they cannot " + f"resolve, so a missing graph would quietly turn them into no-ops." + ) + ports[name] = _graph_ports(ir.load(artifact)) + return ports + + def test_describes_itself_only_through_the_workflow(self, directory): + """One serialized representation of the executable ABI, not two. + + The runtime resolves the workflow first and never consults a second + declaration when one is present, so a surviving ``model.io`` or + ``pipeline.models`` would not be a redundant copy that merely risks + drifting: it would be inert, and nothing would say so. ONNX GenAI + rejects a package that carries both rather than silently picking one. + """ + metadata = self._metadata(directory) + assert "workflow" in metadata["pipeline"] + assert "model" not in metadata or "io" not in metadata["model"], ( + "model.io declares the same executable ABI as the workflow; the " + "workflow is canonical, so this one would be discarded at load" + ) + assert "models" not in metadata["pipeline"], ( + "pipeline.models is the legacy composite ABI, superseded by " + "pipeline.workflow's components and invoke bindings" + ) + + def test_no_component_transcribes_the_ports_of_its_own_artifact( + self, directory, materialized_workflow_packages + ): + """Whatever a component declares must not be a restatement of its graph. + + A policy graph declares contracts because they type the workflow's + dataflow, and a workflow value has no other source for its dtype and + request axis. What no component may do is declare a port the artifact + does not have, or declare only some of them: either turns the + declaration into a partial second truth that drifts silently. + """ + workflow = self._metadata(directory)["pipeline"]["workflow"] + graphs = self._ports(directory, workflow, materialized_workflow_packages) + assert graphs, "a fixture that ships no artifact proves nothing" + for name, component in _onnx_components(workflow).items(): + ports = component.get("ports") or {} + if not (ports.get("inputs") or ports.get("outputs")): + continue + inputs, outputs = graphs[name] + assert set(ports.get("inputs", {})) == inputs + assert set(ports.get("outputs", {})) == outputs + + def test_every_binding_and_state_pair_resolves_in_the_artifact( + self, directory, materialized_workflow_packages + ): + workflow = self._metadata(directory)["pipeline"]["workflow"] + graphs = self._ports(directory, workflow, materialized_workflow_packages) + for step in _walk_steps(workflow["steps"]): + if step.get("kind") != "invoke" or step["component"] not in graphs: + continue + inputs, outputs = graphs[step["component"]] + assert set(step.get("inputs", {})) <= inputs + assert set(step.get("outputs", {})) <= outputs + for group in _groups(workflow).values(): + for component, aliases in (group.get("ports") or {}).items(): + if component not in graphs: + continue + inputs, outputs = graphs[component] + for alias in aliases.values(): + assert alias["input"] in inputs + assert alias["output"] in outputs + + def test_every_declared_role_resolves_in_the_artifact( + self, directory, materialized_workflow_packages + ): + workflow = self._metadata(directory)["pipeline"]["workflow"] + graphs = self._ports(directory, workflow, materialized_workflow_packages) + for name, component in _onnx_components(workflow).items(): + if name not in graphs: + continue + inputs, outputs = graphs[name] + for port in (component.get("ports") or {}).get("roles", {}): + assert port in inputs or port in outputs + + def test_a_scatter_bound_component_is_always_resolvable_as_a_decoder(self, directory): + """The shipped bytes carry the same obligation the built packages do. + + A fixture is edited and regenerated by hand far more often than the + producer is changed, so this is the copy of the rule that catches a + package whose sequence role was dropped on the way to disk. + """ + workflow = self._metadata(directory)["pipeline"]["workflow"] + for name in _scatter_bound_components(workflow): + assert _declares_a_sequence_role(workflow["components"][name]), ( + f"{name} is handed a fixed-capacity write cursor but declares no " + "sequence role, so no decode ABI can be resolved for it" + ) diff --git a/tests/cli_test.py b/tests/cli_test.py index 17ea88b6f..59ec850dd 100644 --- a/tests/cli_test.py +++ b/tests/cli_test.py @@ -222,6 +222,120 @@ def test_revision_is_forwarded_to_detection_and_build(self): ) assert mock_build.call_args.kwargs["revision"] == revision + def test_static_cache_with_onnx_genai_runtime_emits_scatter_abi(self): + """A static-cache export is describable, so the CLI must describe it. + + The two control ports are rank-1 integer vectors and are therefore + shape-indistinguishable from one another, which is exactly why the ABI + is *declared* rather than inferred. It is declared once, in the + workflow: the state group that scatters into the buffers names the port + carrying the write cursor and the port carrying the non-pad length, and + the component those ports belong to declares both. + """ + import onnx_ir as ir + import yaml + + with tempfile.TemporaryDirectory() as tmpdir: + main( + [ + "build", + "--model", + "Qwen/Qwen2.5-0.5B", + tmpdir, + "--no-weights", + "--features", + "static-cache", + "--max-seq-len", + "128", + "--runtime", + "onnx-genai", + ] + ) + with open( + os.path.join(tmpdir, "inference_metadata.yaml"), encoding="utf-8" + ) as handle: + metadata = yaml.safe_load(handle) + # The exported graph is the authority on which ports exist and what + # rank they have, so read it here rather than trusting a copy of it + # in the metadata — a copy is what this contract exists to avoid. + artifact = metadata["pipeline"]["workflow"]["components"]["model"][ + "implementation" + ]["artifact"] + exported = { + str(value.name): len(value.shape) + for value in ir.load(os.path.join(tmpdir, artifact)).graph.inputs + } + + # One canonical description: no second copy of the port ABI outside it. + assert "io" not in metadata.get("model", {}) + + workflow = metadata["pipeline"]["workflow"] + assert workflow["inputs"]["package.cache_capacity"]["default"] == 128 + groups = workflow["serving"]["state_service"]["groups"] + update = next(group["update"] for group in groups.values() if "update" in group) + assert update["kind"] == "indexed_scatter" + assert update["capacity"] == "package.cache_capacity" + # The write cursor and the logical length are the same quantity. + assert update["write_indices"] == "cache_lengths" + assert update["write_indices_ports"] == {"model": "write_indices"} + assert update["kv_length_ports"] == {"model": "nonpad_kv_seqlen"} + + # The component transcribes none of this: it declares the roles a graph + # cannot state, and the scatter's control ports are named by the state + # group. Both names have to resolve in the artifact. + assert not (workflow["components"]["model"]["ports"].get("inputs")) + assert workflow["components"]["model"]["ports"]["roles"]["input_ids"] == "token_ids" + assert exported["write_indices"] == 1 + assert exported["nonpad_kv_seqlen"] == 1 + + group = next(group for group in groups.values() if "update" in group) + pairs = group["ports"]["model"] + assert {alias["output"] for alias in pairs.values()} == { + f"updated_{alias['input']}" for alias in pairs.values() + } + assert pairs["cache_0"] == { + "input": "key_cache.0", + "output": "updated_key_cache.0", + "role": "key", + "layer": 0, + } + + def test_static_cache_task_follows_text_only_substitution(self): + """``text-only`` + ``static-cache`` must resolve the *text* task. + + ``build()`` swaps a multimodal ``model_type`` for its text-only + registry sibling, so the deferred static-cache task has to be resolved + against the substituted type. Resolving against the raw checkpoint type + pairs the text-only module with the multimodal task, which then fails + looking for sub-modules a text-only module does not have. + """ + from mobius.tasks._gemma4 import Gemma4TextCausalLMTask + + hf_config = mock.MagicMock() + hf_config.model_type = "gemma4" + with ( + tempfile.TemporaryDirectory() as tmpdir, + mock.patch("transformers.AutoConfig.from_pretrained", return_value=hf_config), + mock.patch("mobius.__main__.build", return_value=mock.MagicMock()) as mock_build, + mock.patch("mobius.__main__._save_package"), + ): + main( + [ + "build", + "--model", + "google/gemma-4-E2B-it", + tmpdir, + "--no-weights", + "--features", + "text-only,static-cache", + "--max-seq-len", + "128", + ] + ) + + task = mock_build.call_args.kwargs["task"] + assert isinstance(task, Gemma4TextCausalLMTask) + def test_build_static_cache(self): with tempfile.TemporaryDirectory() as tmpdir: main( @@ -612,10 +726,11 @@ def test_runtime_ort_genai_rejects_mage_vl_before_saving(self): save.assert_not_called() config_writer.assert_not_called() - def test_runtime_onnx_genai_uses_native_vlm_emitter(self): + def test_runtime_onnx_genai_routes_vlm_through_workflow_emitter(self): + """A VLM package emits the workflow IR, not a legacy composite pipeline.""" pkg = mock.MagicMock() pkg.items.return_value = [] - pkg.__iter__.return_value = iter(()) + pkg.__iter__.return_value = iter(("vision_encoder", "embedding", "decoder")) pkg.config = object() args = SimpleNamespace( max_shard_size=None, @@ -630,27 +745,19 @@ def test_runtime_onnx_genai_uses_native_vlm_emitter(self): with ( tempfile.TemporaryDirectory() as tmpdir, mock.patch( - "mobius.integrations.onnx_genai.inference_metadata.is_native_vlm_package", - return_value=True, - ), - mock.patch( - "mobius.integrations.onnx_genai.inference_metadata." - "write_native_vlm_package_metadata", + "mobius.integrations.onnx_genai.write_onnx_genai_config", return_value={}, - ) as native_writer, - mock.patch( - "mobius.integrations.onnx_genai.write_onnx_genai_config" - ) as generic_writer, + ) as writer, ): _save_package(pkg, tmpdir, args, None, None) - native_writer.assert_called_once_with( + writer.assert_called_once_with( pkg, tmpdir, config=pkg.config, source="/models/vlm", + revision=None, ) - generic_writer.assert_not_called() def test_runtime_onnx_genai_does_not_fallback_for_unsupported_vlm(self): pkg = mock.MagicMock() @@ -670,25 +777,16 @@ def test_runtime_onnx_genai_does_not_fallback_for_unsupported_vlm(self): with ( tempfile.TemporaryDirectory() as tmpdir, mock.patch( - "mobius.integrations.onnx_genai.inference_metadata.is_native_vlm_package", - return_value=True, - ), - mock.patch( - "mobius.integrations.onnx_genai.inference_metadata." - "write_native_vlm_package_metadata", + "mobius.integrations.onnx_genai.write_onnx_genai_config", side_effect=ValueError( "unsupported VLM signature; regenerate processor assets or register it" ), - ) as native_writer, - mock.patch( - "mobius.integrations.onnx_genai.write_onnx_genai_config" - ) as generic_writer, + ) as writer, pytest.raises(SystemExit, match=r"regenerate.*register"), ): _save_package(pkg, tmpdir, args, None, None) - native_writer.assert_called_once() - generic_writer.assert_not_called() + writer.assert_called_once() def test_no_runtime_does_not_call_write_ort_genai_config(self): """Omitting --runtime does NOT call write_ort_genai_config().""" diff --git a/tests/compare_onnx_genai_validation_packages.py b/tests/compare_onnx_genai_validation_packages.py new file mode 100644 index 000000000..4ff252196 --- /dev/null +++ b/tests/compare_onnx_genai_validation_packages.py @@ -0,0 +1,130 @@ +from __future__ import annotations + +import argparse +import hashlib +import json +from pathlib import Path +from typing import Any + +import onnx_ir as ir +import yaml + + +def _dimension(dimension: Any) -> int | str | None: + return getattr(dimension, "value", dimension) + + +def _value(value: ir.Value | None) -> Any: + if value is None: + return None + return { + "name": value.name, + "dtype": str(value.dtype), + "shape": ( + [_dimension(dimension) for dimension in value.shape] + if value.shape is not None + else None + ), + } + + +def _tensor(tensor: ir.TensorProtocol) -> dict[str, Any]: + return { + "dtype": str(tensor.dtype), + "shape": [_dimension(dimension) for dimension in tensor.shape], + "sha256": hashlib.sha256(tensor.tobytes()).hexdigest(), + } + + +def _graph(graph: ir.Graph) -> dict[str, Any]: + return { + "inputs": [_value(value) for value in graph.inputs], + "outputs": [_value(value) for value in graph.outputs], + "initializers": { + name: _tensor(value.const_value) + for name, value in sorted(graph.initializers.items()) + if value.const_value is not None + }, + } + + +def _model(path: Path) -> dict[str, Any]: + model = ir.load(path) + return _graph(model.graph) + + +# Bytes the generator produces but the tree does not carry. They are a +# deterministic function of the generator, so a committed copy would restate +# it in a form no reviewer can read; ``_artifacts_exist`` below checks that +# every one the metadata names was actually produced. +_GENERATED_SUFFIXES = {".onnx", ".data", ".safetensors"} + + +def _reviewable_files(directory: Path) -> set[Path]: + """The files that are committed, and so are the reviewable contract.""" + return { + path.relative_to(directory) + for path in directory.rglob("*") + if path.is_file() and path.suffix not in _GENERATED_SUFFIXES + } + + +def _artifacts_exist(committed: Path, generated: Path) -> list[str]: + """Every artifact the committed metadata names must have been generated. + + The graphs are no longer committed, so nothing else would notice if the + generator stopped emitting one: the metadata would still compare equal and + the package would still look complete. + """ + missing: list[str] = [] + for metadata_path in sorted(committed.rglob("inference_metadata.yaml")): + package = metadata_path.parent.relative_to(committed) + metadata = yaml.safe_load(metadata_path.read_text(encoding="utf-8")) + workflow = ((metadata.get("pipeline") or {}).get("workflow")) or {} + for name, component in (workflow.get("components") or {}).items(): + implementation = component.get("implementation") or {} + artifact = implementation.get("artifact") + if implementation.get("kind") != "onnx" or not artifact: + continue + if not (generated / package / artifact).exists(): + missing.append(f"{package}: component {name!r} -> {artifact}") + return missing + + +def _content(path: Path) -> Any: + if path.suffix == ".onnx": + return _model(path) + if path.suffix in {".yaml", ".yml"}: + return yaml.safe_load(path.read_text(encoding="utf-8")) + if path.suffix == ".json": + return json.loads(path.read_text(encoding="utf-8")) + return path.read_bytes() + + +def main() -> None: + parser = argparse.ArgumentParser() + parser.add_argument("expected", type=Path) + parser.add_argument("actual", type=Path) + args = parser.parse_args() + + expected_files = _reviewable_files(args.expected) + actual_files = _reviewable_files(args.actual) + if expected_files != actual_files: + missing = sorted(str(path) for path in expected_files - actual_files) + extra = sorted(str(path) for path in actual_files - expected_files) + raise SystemExit(f"package file mismatch: missing={missing}, extra={extra}") + + if absent := _artifacts_exist(args.expected, args.actual): + raise SystemExit(f"declared artifact was not generated: {absent}") + + changed = [ + str(relative) + for relative in sorted(expected_files) + if _content(args.expected / relative) != _content(args.actual / relative) + ] + if changed: + raise SystemExit(f"package semantic mismatch: {changed}") + + +if __name__ == "__main__": + main() diff --git a/tests/conftest.py b/tests/conftest.py index 8e81c779d..a961cf510 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -156,3 +156,27 @@ def deterministic_seed(): random.seed(seed) np.random.seed(seed) return seed + + +@pytest.fixture(scope="session") +def materialized_workflow_packages(tmp_path_factory) -> str: + """Build the conformance packages once, with their ONNX artifacts. + + Only the metadata of these packages is committed: it is the contract under + review, and a diff of it is readable. The graphs are a deterministic + function of the generator, so committing them would store ~14 MB of bytes + no one reads to restate what the script already says — and one of them is a + 5 MB weight blob. + + Tests that need to resolve a declaration against the graph that really + exposes a port get the artifacts from here. Generation takes a few seconds + and is shared across the whole session. + """ + import sys + + sys.path.insert(0, os.path.dirname(__file__)) + from generate_onnx_genai_validation_packages import generate_packages + + directory = tmp_path_factory.mktemp("workflow_packages") + generate_packages(directory) + return str(directory) diff --git a/tests/fixtures/onnx_genai_workflows/README.md b/tests/fixtures/onnx_genai_workflows/README.md new file mode 100644 index 000000000..dbf9c5df2 --- /dev/null +++ b/tests/fixtures/onnx_genai_workflows/README.md @@ -0,0 +1,23 @@ +# ONNX GenAI workflow conformance fixtures + +Generated by `tests/generate_onnx_genai_validation_packages.py` for semantic +validation and runtime conformance. The ONNX GenAI revision these are checked +against is pinned in `.github/workflows/main.yml`; it is not restated here, +because a second copy of a SHA only ever drifts from the first. + +Only the metadata is committed under `tests/fixtures/onnx_genai_workflows`. +The graphs and adapter weights are a deterministic function of this script, so +committing them would store megabytes of bytes no reviewer can read. CI +regenerates them, compares the metadata against the committed copy, and runs +validation and conformance against the regenerated tree. Tests that need the +graphs use the `materialized_workflow_packages` fixture. + +The decoder, VLM, diffusion, masked diffusion, speculative, and codec packages +contain executable synthetic models. The TTS fixture uses the real tiny +Qwen3-TTS producer graphs with deterministic synthetic weights. The two +protein-encoder fixtures use the real tiny ESM-2 and ProtBert producer graphs, +also with deterministic synthetic weights; they cover the non-generative +embedding shape, which has one invocation, no carried state and no sampler. +No downloaded model weights are included. The adapter fixture covers +authoritative target metadata, portable artifacts, ordered heterogeneous +composition, inactive rows, compaction, and request-epoch slot reuse. diff --git a/tests/fixtures/onnx_genai_workflows/adapter/adapters/blue/adapter.json b/tests/fixtures/onnx_genai_workflows/adapter/adapters/blue/adapter.json new file mode 100644 index 000000000..8fe084a30 --- /dev/null +++ b/tests/fixtures/onnx_genai_workflows/adapter/adapters/blue/adapter.json @@ -0,0 +1 @@ +{"targets":{"projection":{"a":[0,1],"b":[3,4]}}} \ No newline at end of file diff --git a/tests/fixtures/onnx_genai_workflows/adapter/adapters/green/adapter.json b/tests/fixtures/onnx_genai_workflows/adapter/adapters/green/adapter.json new file mode 100644 index 000000000..bc5d36d1a --- /dev/null +++ b/tests/fixtures/onnx_genai_workflows/adapter/adapters/green/adapter.json @@ -0,0 +1 @@ +{"targets":{"projection":{"a":[1,1],"b":[1,1]}}} \ No newline at end of file diff --git a/tests/fixtures/onnx_genai_workflows/adapter/adapters/peft/adapter.json b/tests/fixtures/onnx_genai_workflows/adapter/adapters/peft/adapter.json new file mode 100644 index 000000000..5dbbb0b9b --- /dev/null +++ b/tests/fixtures/onnx_genai_workflows/adapter/adapters/peft/adapter.json @@ -0,0 +1 @@ +{"targets":{"projection":{"a":[1,-1],"b":[0.5,0.25]}}} \ No newline at end of file diff --git a/tests/fixtures/onnx_genai_workflows/adapter/adapters/peft/adapter_config.json b/tests/fixtures/onnx_genai_workflows/adapter/adapters/peft/adapter_config.json new file mode 100644 index 000000000..c55844d20 --- /dev/null +++ b/tests/fixtures/onnx_genai_workflows/adapter/adapters/peft/adapter_config.json @@ -0,0 +1 @@ +{"base_model_name_or_path":"synthetic/adapter-base","lora_alpha":1.0,"r":1,"revision":"synthetic-revision","target_modules":["projection"]} \ No newline at end of file diff --git a/tests/fixtures/onnx_genai_workflows/adapter/adapters/red/adapter.json b/tests/fixtures/onnx_genai_workflows/adapter/adapters/red/adapter.json new file mode 100644 index 000000000..bd9db6cbe --- /dev/null +++ b/tests/fixtures/onnx_genai_workflows/adapter/adapters/red/adapter.json @@ -0,0 +1 @@ +{"targets":{"projection":{"a":[1,0],"b":[1,2]}}} \ No newline at end of file diff --git a/tests/fixtures/onnx_genai_workflows/adapter/inference_metadata.yaml b/tests/fixtures/onnx_genai_workflows/adapter/inference_metadata.yaml new file mode 100644 index 000000000..b8e5e5a82 --- /dev/null +++ b/tests/fixtures/onnx_genai_workflows/adapter/inference_metadata.yaml @@ -0,0 +1,276 @@ +schema_version: v1 +pipeline: + workflow: + manifest: + adapter_abis: + onnx-genai.parameter-overlay: '1' + capabilities: + - workflow_ssa + - typed_emit + - parameter_adapters + - heterogeneous_adapter_batching + inputs: + request.active: + contract: + dtype: bool + rank: 1 + shape: + - batch + batch_layout: + kind: request_aligned + axis: 0 + role: + kind: runtime + version: '1.0' + role: adapter_active + source: + kind: request + activations: + contract: + dtype: float32 + rank: 2 + shape: + - batch + - 2 + batch_layout: + kind: request_aligned + axis: 0 + role: + kind: opaque + source: + kind: application + name: activations + request.adapter_segments: + contract: + dtype: int64 + rank: 2 + shape: + - batch + - 2 + batch_layout: + kind: request_aligned + axis: 0 + role: + kind: runtime + version: '1.0' + role: adapter_segments + source: + kind: request + request.adapter_counts: + contract: + dtype: int64 + rank: 1 + shape: + - batch + batch_layout: + kind: request_aligned + axis: 0 + role: + kind: runtime + version: '1.0' + role: adapter_counts + source: + kind: request + request.adapter_scales: + contract: + dtype: float32 + rank: 2 + shape: + - batch + - 2 + batch_layout: + kind: request_aligned + axis: 0 + role: + kind: runtime + version: '1.0' + role: adapter_scales + source: + kind: request + outputs: + result: + contract: + dtype: float32 + rank: 2 + shape: + - batch + - 2 + batch_layout: + kind: request_aligned + axis: 0 + role: tensor + stage: pre_adapter + components: + decoder: + implementation: + kind: onnx + artifact: model.onnx + overlay: + implementation: + kind: adapter + abi: onnx-genai.parameter-overlay + version: '1' + ports: + inputs: + input: + dtype: float32 + rank: 2 + shape: + - batch + - 2 + batch_layout: + kind: request_aligned + axis: 0 + outputs: + output: + dtype: float32 + rank: 2 + shape: + - batch + - 2 + batch_layout: + kind: request_aligned + axis: 0 + contract: + id: onnx-genai.parameter-overlay + version: '1' + bindings: + input: input + output: output + parameters: + action: apply + component: decoder + parameter: projection + state: {} + steps: + - kind: invoke + component: overlay + inputs: + input: activations + outputs: + output: adapted + - kind: emit + value: adapted + output: result + mode: replace +adapters: + base_model_fingerprint: onnx-genai-targeted-base-v1:sha256:702704827be590661fd77676445ddaf1557e4f4d446eaa7bcfb78754466cccdb + target_manifest: + targets: + - id: projection + component: decoder + initializer: projection + node_name: projection + output_name: projection.output + activation_dtype: float32 + input_features: 2 + output_features: 2 + layer_index: 0 + rank: 1 + alpha: 1.0 + graph_inputs: + a: lora.projection.a + b: lora.projection.b + scale: lora.projection.scale + discovery_fallback: disabled + selection: + segments: request.adapter_segments + adapter_counts: request.adapter_counts + scales: request.adapter_scales + active: request.active + max_adapters: 2 + application_capability: onnx-genai.adapters@1 + portable_fallback: true + cache: + max_entries: 2 + eviction: lru + planning: + bucket_by_adapter_set: true + stable_buffers: true + invalidate_capture_on_eviction: true + artifacts: + blue: + index: 0 + identity: blue + version: '1' + base_model_fingerprint: onnx-genai-targeted-base-v1:sha256:702704827be590661fd77676445ddaf1557e4f4d446eaa7bcfb78754466cccdb + rank: 1 + alpha: 1.0 + dtype: float32 + provenance: + producer: mobius + weights: + - location: adapters/blue/adapter.json + loader_capability: onnx-genai.adapters.json@1 + sha256: d9e67499ad74d4c2d62c45a79ceeea8dfa94893afde9615ea6030b79d052ddbe + scale_encoding: alpha_over_rank + format: json + bindings: + - target: projection + weight_key: projection + green: + index: 1 + identity: green + version: '1' + base_model_fingerprint: onnx-genai-targeted-base-v1:sha256:702704827be590661fd77676445ddaf1557e4f4d446eaa7bcfb78754466cccdb + rank: 1 + alpha: 1.0 + dtype: float32 + provenance: + producer: mobius + weights: + - location: adapters/green/adapter.json + loader_capability: onnx-genai.adapters.json@1 + sha256: d49ae9a3336884fa1ed8503caca1380699788add893870a670f6290c473bc935 + scale_encoding: alpha_over_rank + format: json + bindings: + - target: projection + weight_key: projection + peft: + index: 2 + identity: peft + version: '1' + base_model_fingerprint: onnx-genai-targeted-base-v1:sha256:702704827be590661fd77676445ddaf1557e4f4d446eaa7bcfb78754466cccdb + rank: 1 + alpha: 1.0 + dtype: float32 + provenance: + producer: mobius + source: synthetic/adapter-base + revision: synthetic-revision + weights: + - location: adapters/peft/adapter.json + loader_capability: onnx-genai.adapters.json@1 + sha256: 4351d2c682e0d1666ca3fdd6e967baaab6ab93c33cf5ace7bf3d6269713e85ce + scale_encoding: alpha_over_rank + format: json + - location: adapters/peft/adapter_model.safetensors + loader_capability: onnx-genai.adapters.hf-peft@1 + sha256: c927c136719514dc03fd9f2e11b49ecbca0b4f28a87925a801fb069e29e078d8 + config_location: adapters/peft/adapter_config.json + config_sha256: b8419ce55415ad1be8844bd6eaef9a873881f172f427694b1d8e46049a8784df + scale_encoding: alpha_over_rank + format: hf_peft + bindings: + - target: projection + weight_key: projection + red: + index: 3 + identity: red + version: '1' + base_model_fingerprint: onnx-genai-targeted-base-v1:sha256:702704827be590661fd77676445ddaf1557e4f4d446eaa7bcfb78754466cccdb + rank: 1 + alpha: 1.0 + dtype: float32 + provenance: + producer: mobius + weights: + - location: adapters/red/adapter.json + loader_capability: onnx-genai.adapters.json@1 + sha256: b4bd656305aa2fea1d518df351af5fd3bf99b3d0a6c8bd8501dfc5107ac278a1 + scale_encoding: alpha_over_rank + format: json + bindings: + - target: projection + weight_key: projection diff --git a/tests/fixtures/onnx_genai_workflows/codec/inference_metadata.yaml b/tests/fixtures/onnx_genai_workflows/codec/inference_metadata.yaml new file mode 100644 index 000000000..09801175b --- /dev/null +++ b/tests/fixtures/onnx_genai_workflows/codec/inference_metadata.yaml @@ -0,0 +1,67 @@ +schema_version: v1 +pipeline: + workflow: + manifest: + capabilities: + - workflow_ssa + - linear_effects + - typed_emit + inputs: + request.waveform: + contract: + dtype: float32 + rank: 3 + shape: + - batch + - 1 + - audio_samples + batch_layout: + kind: request_aligned + axis: 0 + role: + kind: runtime + version: '1.0' + role: media + source: + kind: request + required: true + outputs: + waveform: + contract: + dtype: float32 + rank: 3 + shape: + - batch + - 1 + - audio_samples + batch_layout: + kind: request_aligned + axis: 0 + role: audio + stage: post_adapter + components: + encoder: + implementation: + kind: onnx + artifact: encoder/model.onnx + decoder: + implementation: + kind: onnx + artifact: decoder/model.onnx + steps: + - kind: invoke + component: encoder + inputs: + waveform: request.waveform + outputs: + codes: codec.codes + - kind: invoke + component: decoder + inputs: + codes: codec.codes + outputs: + waveform: codec.waveform + - kind: emit + value: codec.waveform + output: waveform + mode: replace diff --git a/tests/fixtures/onnx_genai_workflows/decoder/inference_metadata.yaml b/tests/fixtures/onnx_genai_workflows/decoder/inference_metadata.yaml new file mode 100644 index 000000000..7253fb474 --- /dev/null +++ b/tests/fixtures/onnx_genai_workflows/decoder/inference_metadata.yaml @@ -0,0 +1,1331 @@ +schema_version: '1.0' +pipeline: + workflow: + manifest: + capabilities: + - workflow_ssa + - linear_effects + - nested_control_flow + - typed_emit + - emit_valid_length + - loop_induction_values + - serving_service_contract + - bounded_state_recurrence + inputs: + request.input_ids: + contract: + dtype: int64 + rank: 2 + shape: + - batch + - sequence + batch_layout: + kind: request_aligned + axis: 0 + role: + kind: runtime + version: '1.0' + role: prompt_tokens + source: + kind: request + required: true + request.max_iterations: + contract: + dtype: int64 + rank: 1 + shape: + - 1 + role: + kind: runtime + version: '1.0' + role: max_output_tokens + source: + kind: request + required: true + package.eos_ids: + contract: + dtype: int64 + rank: 1 + shape: + - E + role: + kind: opaque + source: + kind: literal + required: true + default: 127 + package.one_token: + contract: + dtype: int64 + rank: 1 + shape: + - batch + batch_layout: + kind: request_aligned + axis: 0 + role: + kind: opaque + source: + kind: literal + required: false + default: 1 + package.one_step: + contract: + dtype: int64 + rank: 1 + shape: + - 1 + role: + kind: opaque + source: + kind: literal + required: false + default: 1 + package.max_context: + contract: + dtype: int64 + rank: 1 + shape: + - 1 + role: + kind: opaque + source: + kind: literal + required: false + default: 8192 + request.prompt_lengths: + contract: + dtype: int64 + rank: 1 + shape: + - batch + batch_layout: + kind: request_aligned + axis: 0 + role: + kind: opaque + source: + kind: application + name: prompt_lengths + required: false + default: -1 + request.eos_ids: + contract: + dtype: int64 + rank: 2 + shape: + - batch + - num_eos + batch_layout: + kind: request_aligned + axis: 0 + role: + kind: opaque + source: + kind: application + name: eos_ids + required: false + default: 127 + request.eos_lengths: + contract: + dtype: int64 + rank: 1 + shape: + - batch + batch_layout: + kind: request_aligned + axis: 0 + role: + kind: opaque + source: + kind: application + name: eos_lengths + required: false + default: 1 + request.row_max_iterations: + contract: + dtype: int64 + rank: 1 + shape: + - batch + batch_layout: + kind: request_aligned + axis: 0 + role: + kind: opaque + source: + kind: application + name: row_max_iterations + required: false + default: -1 + request.temperature: + contract: + dtype: float32 + rank: 1 + shape: + - batch + batch_layout: + kind: request_aligned + axis: 0 + role: + kind: runtime + version: '1.0' + role: sampling_temperature + source: + kind: request + required: false + default: 1.0 + request.top_k: + contract: + dtype: int64 + rank: 1 + shape: + - batch + batch_layout: + kind: request_aligned + axis: 0 + role: + kind: runtime + version: '1.0' + role: sampling_top_k + source: + kind: request + required: false + default: 1 + request.top_p: + contract: + dtype: float32 + rank: 1 + shape: + - batch + batch_layout: + kind: request_aligned + axis: 0 + role: + kind: runtime + version: '1.0' + role: sampling_top_p + source: + kind: request + required: false + default: 1.0 + request.min_p: + contract: + dtype: float32 + rank: 1 + shape: + - batch + batch_layout: + kind: request_aligned + axis: 0 + role: + kind: runtime + version: '1.0' + role: sampling_min_p + source: + kind: request + required: false + default: 0.0 + request.seed: + contract: + dtype: int64 + rank: 1 + shape: + - batch + batch_layout: + kind: request_aligned + axis: 0 + role: + kind: runtime + version: '1.0' + role: seed + source: + kind: request + required: false + default: 0 + request.rng_counter: + contract: + dtype: int64 + rank: 1 + shape: + - batch + batch_layout: + kind: request_aligned + axis: 0 + role: + kind: opaque + source: + kind: application + name: rng_counter + required: false + default: 0 + package.active: + contract: + dtype: bool + rank: 1 + shape: + - batch + batch_layout: + kind: request_aligned + axis: 0 + role: + kind: opaque + source: + kind: literal + required: false + default: true + package.not_done: + contract: + dtype: bool + rank: 1 + shape: + - batch + batch_layout: + kind: request_aligned + axis: 0 + role: + kind: opaque + source: + kind: literal + required: false + default: false + package.cache_lengths: + contract: + dtype: int64 + rank: 1 + shape: + - batch + batch_layout: + kind: request_aligned + axis: 0 + role: + kind: opaque + source: + kind: literal + required: false + default: 0 + package.zero_batch: + contract: + dtype: int64 + rank: 1 + shape: + - batch + batch_layout: + kind: request_aligned + axis: 0 + role: + kind: opaque + source: + kind: literal + required: false + default: 0 + outputs: + tokens: + contract: + dtype: int64 + rank: 2 + shape: + - batch + - generated_sequence + batch_layout: + kind: request_aligned + axis: 0 + role: tokens + stage: pre_adapter + components: + model: + implementation: + kind: onnx + artifact: model.onnx + ports: + roles: + input_ids: token_ids + attention_mask: attention_mask + position_ids: position_ids + logits: logits + token_sampler: + implementation: + kind: onnx + artifact: policies/token_sampler.onnx + ports: + inputs: + logits: + dtype: float32 + rank: 2 + shape: + - batch + - vocabulary + batch_layout: + kind: request_aligned + axis: 0 + temperature: + dtype: float32 + rank: 1 + shape: + - batch + batch_layout: + kind: request_aligned + axis: 0 + top_k: + dtype: int64 + rank: 1 + shape: + - batch + batch_layout: + kind: request_aligned + axis: 0 + top_p: + dtype: float32 + rank: 1 + shape: + - batch + batch_layout: + kind: request_aligned + axis: 0 + min_p: + dtype: float32 + rank: 1 + shape: + - batch + batch_layout: + kind: request_aligned + axis: 0 + seed: + dtype: int64 + rank: 1 + shape: + - batch + batch_layout: + kind: request_aligned + axis: 0 + counter: + dtype: int64 + rank: 1 + shape: + - batch + batch_layout: + kind: request_aligned + axis: 0 + active: + dtype: bool + rank: 1 + shape: + - batch + batch_layout: + kind: request_aligned + axis: 0 + done: + dtype: bool + rank: 1 + shape: + - batch + batch_layout: + kind: request_aligned + axis: 0 + outputs: + token: + dtype: int64 + rank: 1 + shape: + - batch + batch_layout: + kind: request_aligned + axis: 0 + next_counter: + dtype: int64 + rank: 1 + shape: + - batch + batch_layout: + kind: request_aligned + axis: 0 + contract: + id: onnx-genai.token-sampler + version: '2' + bindings: + logits: logits + token: token + temperature: temperature + top_k: top_k + top_p: top_p + min_p: min_p + active: active + done: done + seed: seed + counter: counter + next_counter: next_counter + parameters: + mode: seeded_stochastic + batching: per_row + inactive_rows: preserve + application_overridable: true + termination: + implementation: + kind: onnx + artifact: policies/termination.onnx + ports: + inputs: + tokens: + dtype: int64 + rank: 1 + shape: + - batch + batch_layout: + kind: request_aligned + axis: 0 + eos_ids: + dtype: int64 + rank: 2 + shape: + - batch + - num_eos + batch_layout: + kind: request_aligned + axis: 0 + eos_lengths: + dtype: int64 + rank: 1 + shape: + - batch + batch_layout: + kind: request_aligned + axis: 0 + iteration: + dtype: int64 + rank: 1 + shape: + - 1 + max_iterations: + dtype: int64 + rank: 1 + shape: + - batch + batch_layout: + kind: request_aligned + axis: 0 + active: + dtype: bool + rank: 1 + shape: + - batch + batch_layout: + kind: request_aligned + axis: 0 + outputs: + done: + dtype: bool + rank: 1 + shape: + - batch + batch_layout: + kind: request_aligned + axis: 0 + next_active: + dtype: bool + rank: 1 + shape: + - batch + batch_layout: + kind: request_aligned + axis: 0 + continue: + dtype: bool + rank: 1 + shape: + - 1 + contract: + id: onnx-genai.termination-predicate + version: '2' + bindings: + tokens: tokens + eos_ids: eos_ids + iteration: iteration + max_iterations: max_iterations + eos_lengths: eos_lengths + active: active + done: done + next_active: next_active + continue: continue + parameters: + batching: per_row + inactive_rows: preserve + token_state_update: + implementation: + kind: onnx + artifact: policies/token_state_update.onnx + ports: + inputs: + current: + dtype: int64 + rank: 2 + shape: + - batch + - 1 + batch_layout: + kind: request_aligned + axis: 0 + update: + dtype: int64 + rank: 2 + shape: + - batch + - 1 + batch_layout: + kind: request_aligned + axis: 0 + active: + dtype: bool + rank: 1 + shape: + - batch + batch_layout: + kind: request_aligned + axis: 0 + done: + dtype: bool + rank: 1 + shape: + - batch + batch_layout: + kind: request_aligned + axis: 0 + outputs: + next: + dtype: int64 + rank: 2 + shape: + - batch + - 1 + batch_layout: + kind: request_aligned + axis: 0 + contract: + id: onnx-genai.state-update + version: '2' + bindings: + current: current + update: update + active: active + done: done + next: next + parameters: + batching: per_row + inactive_rows: preserve + last_token_logits: + implementation: + kind: onnx + artifact: policies/last_token_logits.onnx + ports: + inputs: + logits: + dtype: float32 + rank: 3 + shape: + - batch + - sequence + - vocabulary + batch_layout: + kind: request_aligned + axis: 0 + outputs: + last_logits: + dtype: float32 + rank: 2 + shape: + - batch + - vocabulary + batch_layout: + kind: request_aligned + axis: 0 + decoder_state_initializer: + implementation: + kind: onnx + artifact: policies/decoder_state_initializer.onnx + ports: + inputs: + prompt_tokens: + dtype: int64 + rank: 2 + shape: + - batch + - prompt_sequence + batch_layout: + kind: request_aligned + axis: 0 + prompt_lengths: + dtype: int64 + rank: 1 + shape: + - batch + batch_layout: + kind: request_aligned + axis: 0 + outputs: + attention_mask: + dtype: int64 + rank: 2 + shape: + - batch + - total_sequence + batch_layout: + kind: request_aligned + axis: 0 + position_ids: + dtype: int64 + rank: 2 + shape: + - batch + - prompt_sequence + batch_layout: + kind: request_aligned + axis: 0 + body_attention_mask: + dtype: int64 + rank: 2 + shape: + - batch + - prompt_sequence + 1 + batch_layout: + kind: request_aligned + axis: 0 + body_position_ids: + dtype: int64 + rank: 2 + shape: + - batch + - 1 + batch_layout: + kind: request_aligned + axis: 0 + token_slot: + dtype: int64 + rank: 2 + shape: + - batch + - 1 + batch_layout: + kind: request_aligned + axis: 0 + generated_lengths: + dtype: int64 + rank: 1 + shape: + - batch + batch_layout: + kind: request_aligned + axis: 0 + past_key_values.0.key: + dtype: float32 + rank: 4 + shape: + - batch + - 2 + - past_sequence + - 8 + batch_layout: + kind: request_aligned + axis: 0 + decoder_step_update: + implementation: + kind: onnx + artifact: policies/decoder_step_update.onnx + ports: + inputs: + attention_mask: + dtype: int64 + rank: 2 + shape: + - batch + - context + batch_layout: + kind: request_aligned + axis: 0 + position_ids: + dtype: int64 + rank: 2 + shape: + - batch + - 1 + batch_layout: + kind: request_aligned + axis: 0 + outputs: + next_attention_mask: + dtype: int64 + rank: 2 + shape: + - batch + - context + 1 + batch_layout: + kind: request_aligned + axis: 0 + next_position_ids: + dtype: int64 + rank: 2 + shape: + - batch + - 1 + batch_layout: + kind: request_aligned + axis: 0 + cache_length_update: + implementation: + kind: onnx + artifact: policies/cache_length_update.onnx + ports: + inputs: + left: + dtype: int64 + rank: 1 + shape: + - batch + batch_layout: + kind: request_aligned + axis: 0 + right: + dtype: int64 + rank: 1 + shape: + - batch + batch_layout: + kind: request_aligned + axis: 0 + active: + dtype: bool + rank: 1 + shape: + - batch + batch_layout: + kind: request_aligned + axis: 0 + done: + dtype: bool + rank: 1 + shape: + - batch + batch_layout: + kind: request_aligned + axis: 0 + outputs: + total: + dtype: int64 + rank: 1 + shape: + - batch + batch_layout: + kind: request_aligned + axis: 0 + termination_batch_initializer: + implementation: + kind: onnx + artifact: policies/termination_batch_initializer.onnx + ports: + inputs: + input_eos_ids: + dtype: int64 + rank: 2 + shape: + - batch + - num_eos + batch_layout: + kind: request_aligned + axis: 0 + input_eos_lengths: + dtype: int64 + rank: 1 + shape: + - batch + batch_layout: + kind: request_aligned + axis: 0 + input_max_iterations: + dtype: int64 + rank: 1 + shape: + - batch + batch_layout: + kind: request_aligned + axis: 0 + fallback_max_iterations: + dtype: int64 + rank: 1 + shape: + - 1 + active: + dtype: bool + rank: 1 + shape: + - batch + batch_layout: + kind: request_aligned + axis: 0 + outputs: + row_eos_ids: + dtype: int64 + rank: 2 + shape: + - batch + - num_eos + batch_layout: + kind: request_aligned + axis: 0 + eos_lengths: + dtype: int64 + rank: 1 + shape: + - batch + batch_layout: + kind: request_aligned + axis: 0 + max_iterations: + dtype: int64 + rank: 1 + shape: + - batch + batch_layout: + kind: request_aligned + axis: 0 + token_to_slot: + implementation: + kind: onnx + artifact: policies/token_to_slot.onnx + ports: + inputs: + token: + dtype: int64 + rank: 1 + shape: + - batch + batch_layout: + kind: request_aligned + axis: 0 + outputs: + slot: + dtype: int64 + rank: 2 + shape: + - batch + - 1 + batch_layout: + kind: request_aligned + axis: 0 + generated_length_update: + implementation: + kind: onnx + artifact: policies/generated_length_update.onnx + ports: + inputs: + left: + dtype: int64 + rank: 1 + shape: + - batch + batch_layout: + kind: request_aligned + axis: 0 + right: + dtype: int64 + rank: 1 + shape: + - batch + batch_layout: + kind: request_aligned + axis: 0 + active: + dtype: bool + rank: 1 + shape: + - batch + batch_layout: + kind: request_aligned + axis: 0 + done: + dtype: bool + rank: 1 + shape: + - batch + batch_layout: + kind: request_aligned + axis: 0 + outputs: + total: + dtype: int64 + rank: 1 + shape: + - batch + batch_layout: + kind: request_aligned + axis: 0 + state: + token: + contract: + dtype: int64 + rank: 2 + shape: + - batch + - 1 + batch_layout: + kind: request_aligned + axis: 0 + scope: invocation + initializer: initializer.token_slot + recurrence: + kind: invariant + logits: + contract: + dtype: float32 + rank: 2 + shape: + - batch + - 128 + batch_layout: + kind: request_aligned + axis: 0 + scope: invocation + initializer: decoder.setup.last_logits + recurrence: + kind: invariant + generated_lengths: + contract: + dtype: int64 + rank: 1 + shape: + - batch + batch_layout: + kind: request_aligned + axis: 0 + class: semantic + scope: invocation + initializer: initializer.generated_lengths + recurrence: + kind: invariant + active: + contract: + dtype: bool + rank: 1 + shape: + - batch + batch_layout: + kind: request_aligned + axis: 0 + class: semantic + scope: invocation + initializer: package.active + recurrence: + kind: invariant + done: + contract: + dtype: bool + rank: 1 + shape: + - batch + batch_layout: + kind: request_aligned + axis: 0 + class: semantic + scope: invocation + initializer: package.not_done + recurrence: + kind: invariant + accepted_len: + contract: + dtype: int64 + rank: 1 + shape: + - batch + batch_layout: + kind: request_aligned + axis: 0 + class: semantic + scope: invocation + initializer: package.zero_batch + recurrence: + kind: invariant + cache_lengths: + contract: + dtype: int64 + rank: 1 + shape: + - batch + batch_layout: + kind: request_aligned + axis: 0 + class: semantic + scope: invocation + initializer: package.cache_lengths + recurrence: + kind: invariant + rng_counter: + contract: + dtype: int64 + rank: 1 + shape: + - batch + batch_layout: + kind: request_aligned + axis: 0 + scope: invocation + class: semantic + initializer: request.rng_counter + recurrence: + kind: invariant + attention_mask: + contract: + dtype: int64 + rank: 2 + shape: + - batch + - context + batch_layout: + kind: request_aligned + axis: 0 + scope: invocation + initializer: initializer.body_attention_mask + recurrence: + kind: growing + axis: 1 + increment: package.one_step + max: package.max_context + position_ids: + contract: + dtype: int64 + rank: 2 + shape: + - batch + - 1 + batch_layout: + kind: request_aligned + axis: 0 + scope: invocation + initializer: initializer.body_position_ids + recurrence: + kind: invariant + cache_0: + contract: + dtype: float32 + rank: 4 + shape: + - batch + - 2 + - past_sequence + - 8 + batch_layout: + kind: request_aligned + axis: 0 + scope: invocation + initializer: decoder.setup.present.0.key + recurrence: + kind: bounded + axis: 2 + max: package.max_context + management: runtime + release_boundary: invocation + service_group: decoder_cache + serving: + active: active + done: done + accepted_len: accepted_len + state_service: + groups: + decoder_cache: + kind: full_attention + sequence_axis: 2 + layout: bnsh + aliasing: forbidden + reuse: + prefix_reusable: true + evictable_prefix: false + ports: + model: + cache_0: + input: past_key_values.0.key + output: present.0.key + role: key + layer: 0 + steps: + - kind: loop + setup: + - kind: invoke + component: decoder_state_initializer + inputs: + prompt_tokens: request.input_ids + prompt_lengths: request.prompt_lengths + outputs: + attention_mask: initializer.attention_mask + body_attention_mask: initializer.body_attention_mask + token_slot: initializer.token_slot + generated_lengths: initializer.generated_lengths + position_ids: initializer.position_ids + body_position_ids: initializer.body_position_ids + past_key_values.0.key: initializer.past_key_values.0.key + - kind: invoke + component: model + inputs: + input_ids: request.input_ids + past_key_values.0.key: initializer.past_key_values.0.key + attention_mask: initializer.attention_mask + position_ids: initializer.position_ids + outputs: + logits: decoder.setup.logits + present.0.key: decoder.setup.present.0.key + - kind: invoke + component: termination_batch_initializer + inputs: + input_eos_ids: request.eos_ids + input_eos_lengths: request.eos_lengths + input_max_iterations: request.row_max_iterations + fallback_max_iterations: request.max_iterations + active: package.active + outputs: + row_eos_ids: termination.eos_ids + eos_lengths: termination.eos_lengths + max_iterations: termination.max_iterations + - kind: invoke + component: last_token_logits + inputs: + logits: decoder.setup.logits + outputs: + last_logits: decoder.setup.last_logits + steps: + - kind: invoke + component: token_sampler + inputs: + logits: logits + temperature: request.temperature + top_k: request.top_k + top_p: request.top_p + min_p: request.min_p + seed: request.seed + counter: rng_counter + active: active + done: done + outputs: + token: sample.body + next_counter: sample.next_counter + - kind: invoke + component: token_to_slot + inputs: + token: sample.body + outputs: + slot: sample.slot + - kind: invoke + component: generated_length_update + inputs: + left: generated_lengths + right: package.one_token + active: active + done: done + outputs: + total: token.next_lengths + - kind: invoke + component: generated_length_update + inputs: + left: package.zero_batch + right: package.one_token + active: active + done: done + outputs: + total: token.emitted_length + - kind: invoke + component: token_state_update + inputs: + current: token + update: sample.slot + active: active + done: done + outputs: + next: token.body + - kind: invoke + component: termination + inputs: + tokens: sample.body + eos_ids: termination.eos_ids + eos_lengths: termination.eos_lengths + iteration: loop.iteration + max_iterations: termination.max_iterations + active: active + outputs: + done: loop.done + continue: loop.continue + next_active: loop.next_active + - kind: invoke + component: cache_length_update + inputs: + left: cache_lengths + right: package.one_token + active: active + done: done + outputs: + total: cache_lengths.next + - kind: invoke + component: cache_length_update + inputs: + left: package.zero_batch + right: package.one_token + active: active + done: done + outputs: + total: accepted_len.next + - kind: emit + value: token.body + output: tokens + mode: append + valid_length: token.emitted_length + when: active + - kind: invoke + component: model + inputs: + input_ids: token.body + past_key_values.0.key: cache_0 + attention_mask: attention_mask + position_ids: position_ids + outputs: + logits: decoder.body.logits + present.0.key: decoder.body.present.0.key + - kind: invoke + component: last_token_logits + inputs: + logits: decoder.body.logits + outputs: + last_logits: decoder.body.last_logits + - kind: invoke + component: decoder_step_update + inputs: + attention_mask: attention_mask + position_ids: position_ids + outputs: + next_attention_mask: decoder_step.body_attention_mask + next_position_ids: decoder_step.body_position_ids + continue_when: active + max_iterations: request.max_iterations + carried: + - cell: token + next: token.body + - cell: logits + next: decoder.body.last_logits + - cell: generated_lengths + next: token.next_lengths + - cell: active + next: loop.next_active + - cell: done + next: loop.done + - cell: cache_lengths + next: cache_lengths.next + - cell: accepted_len + next: accepted_len.next + - cell: rng_counter + next: sample.next_counter + - cell: attention_mask + next: decoder_step.body_attention_mask + - cell: position_ids + next: decoder_step.body_position_ids + - cell: cache_0 + next: decoder.body.present.0.key + termination: generation_eos + iteration: + value: loop.iteration + contract: + dtype: int64 + rank: 1 + shape: + - 1 diff --git a/tests/fixtures/onnx_genai_workflows/diffusion/inference_metadata.yaml b/tests/fixtures/onnx_genai_workflows/diffusion/inference_metadata.yaml new file mode 100644 index 000000000..638decb07 --- /dev/null +++ b/tests/fixtures/onnx_genai_workflows/diffusion/inference_metadata.yaml @@ -0,0 +1,530 @@ +schema_version: v1 +pipeline: + workflow: + manifest: + capabilities: + - workflow_ssa + - linear_effects + - nested_control_flow + - loop_induction_values + - typed_emit + inputs: + request.max_iterations: + contract: + dtype: int64 + rank: 1 + shape: + - 1 + role: + kind: runtime + version: '1.0' + role: max_iterations + source: + kind: request + required: false + default: 30 + package.false: + contract: + dtype: bool + rank: 1 + shape: + - batch + batch_layout: + kind: request_aligned + axis: 0 + role: + kind: opaque + source: + kind: literal + required: false + default: false + request.noise: + contract: + dtype: float32 + rank: 4 + shape: + - batch + - 4 + - height + - width + batch_layout: + kind: request_aligned + axis: 0 + role: + kind: opaque + source: + kind: application + name: noise + required: true + externally_suppliable: true + request.input_ids: + contract: + dtype: int64 + rank: 2 + shape: + - batch + - prompt_sequence + batch_layout: + kind: request_aligned + axis: 0 + role: + kind: runtime + version: '1.0' + role: prompt_tokens + source: + kind: request + required: true + externally_suppliable: true + package.loop_0_active: + contract: + dtype: bool + rank: 1 + shape: + - 1 + role: + kind: opaque + source: + kind: literal + required: false + default: true + outputs: + image: + contract: + dtype: float32 + rank: 4 + shape: + - batch + - 3 + - height + - width + batch_layout: + kind: request_aligned + axis: 0 + role: image + stage: pre_adapter + latent: + contract: + dtype: float32 + rank: 4 + shape: + - batch + - 4 + - height + - width + batch_layout: + kind: request_aligned + axis: 0 + role: tensor + stage: pre_adapter + noise_estimate: + contract: + dtype: float32 + rank: 4 + shape: + - batch + - 4 + - height + - noise_estimate_width + batch_layout: + kind: request_aligned + axis: 0 + role: tensor + stage: pre_adapter + latent_trajectory: + contract: + dtype: float32 + rank: 4 + shape: + - batch + - 4 + - height + - trajectory_width + batch_layout: + kind: request_aligned + axis: 0 + role: tensor + stage: pre_adapter + components: + text_encoder: + implementation: + kind: onnx + artifact: text_encoder/model.onnx + ports: + roles: + input_ids: token_ids + encoder_hidden_states: encoder_hidden_states + denoiser: + implementation: + kind: onnx + artifact: denoiser/model.onnx + ports: + roles: + encoder_hidden_states: encoder_hidden_states + vae_decoder: + implementation: + kind: onnx + artifact: vae_decoder/model.onnx + solver_step: + implementation: + kind: onnx + artifact: policies/solver_step.onnx + ports: + inputs: + sample: + dtype: float32 + rank: 4 + shape: + - batch + - channels + - height + - width + batch_layout: + kind: request_aligned + axis: 0 + derivative: + dtype: float32 + rank: 4 + shape: + - batch + - channels + - height + - width + batch_layout: + kind: request_aligned + axis: 0 + step: + dtype: int64 + rank: 1 + shape: + - batch + batch_layout: + kind: request_aligned + axis: 0 + schedule: + dtype: float32 + rank: 1 + shape: + - schedule_length + outputs: + next_state: + dtype: float32 + rank: 4 + shape: + - batch + - channels + - height + - width + batch_layout: + kind: request_aligned + axis: 0 + contract: + id: onnx-genai.solver-step + version: '1' + bindings: + state: sample + estimate: derivative + step: step + schedule: schedule + next_state: next_state + continue_predicate: + implementation: + kind: onnx + artifact: policies/continue_predicate.onnx + ports: + inputs: + done: + dtype: bool + rank: 1 + shape: + - batch + batch_layout: + kind: request_aligned + axis: 0 + outputs: + continue: + dtype: bool + rank: 1 + shape: + - 1 + model_input_scale: + implementation: + kind: onnx + artifact: policies/model_input_scale.onnx + ports: + inputs: + sample: + dtype: float32 + rank: 4 + shape: + - batch + - channels + - height + - width + batch_layout: + kind: request_aligned + axis: 0 + step: + dtype: int64 + rank: 1 + shape: + - batch + batch_layout: + kind: request_aligned + axis: 0 + schedule: + dtype: float32 + rank: 1 + shape: + - schedule_length + outputs: + model_input: + dtype: float32 + rank: 4 + shape: + - batch + - channels + - height + - width + batch_layout: + kind: request_aligned + axis: 0 + diffusion_schedule: + implementation: + kind: onnx + artifact: policies/diffusion_schedule.onnx + ports: + inputs: {} + outputs: + schedule: + dtype: float32 + rank: 1 + shape: + - 31 + diffusion_timesteps: + implementation: + kind: onnx + artifact: policies/diffusion_timesteps.onnx + ports: + inputs: {} + outputs: + schedule: + dtype: float32 + rank: 1 + shape: + - 30 + schedule_lookup: + implementation: + kind: onnx + artifact: policies/schedule_lookup.onnx + ports: + inputs: + schedule: + dtype: float32 + rank: 1 + shape: + - schedule_length + step: + dtype: int64 + rank: 1 + shape: + - batch + batch_layout: + kind: request_aligned + axis: 0 + outputs: + timestep: + dtype: float32 + rank: 1 + shape: + - batch + batch_layout: + kind: request_aligned + axis: 0 + tensor_scale: + implementation: + kind: onnx + artifact: policies/tensor_scale.onnx + ports: + inputs: + tensor: + dtype: float32 + rank: 4 + shape: + - batch + - channels + - height + - width + batch_layout: + kind: request_aligned + axis: 0 + scale: + dtype: float32 + rank: 1 + shape: + - 1 + outputs: + scaled: + dtype: float32 + rank: 4 + shape: + - batch + - channels + - height + - width + batch_layout: + kind: request_aligned + axis: 0 + initial_state_scale: + implementation: + kind: onnx + artifact: policies/initial_state_scale.onnx + ports: + inputs: {} + outputs: + value: + dtype: float32 + rank: 1 + shape: + - 1 + state: + latent_state: + contract: + dtype: float32 + rank: 4 + shape: + - batch + - 4 + - height + - width + batch_layout: + kind: request_aligned + axis: 0 + scope: invocation + initializer: diffusion.initial_state + recurrence: + kind: invariant + loop_0_active: + contract: + dtype: bool + rank: 1 + shape: + - 1 + scope: invocation + initializer: package.loop_0_active + recurrence: + kind: invariant + steps: + - kind: loop + setup: + - kind: invoke + component: diffusion_schedule + inputs: {} + outputs: + schedule: diffusion.schedule + - kind: invoke + component: diffusion_timesteps + inputs: {} + outputs: + schedule: diffusion.timesteps + - kind: invoke + component: initial_state_scale + inputs: {} + outputs: + value: diffusion.initial_scale + - kind: invoke + component: tensor_scale + inputs: + tensor: request.noise + scale: diffusion.initial_scale + outputs: + scaled: diffusion.initial_state + - kind: invoke + component: text_encoder + inputs: + input_ids: request.input_ids + outputs: + encoder_hidden_states: conditioning.hidden_states + - kind: invoke + component: continue_predicate + inputs: + done: package.false + outputs: + continue: setup.continue + steps: + - kind: invoke + component: schedule_lookup + inputs: + schedule: diffusion.timesteps + step: loop.iteration + outputs: + timestep: diffusion.timestep + - kind: invoke + component: model_input_scale + inputs: + sample: latent_state + step: loop.iteration + schedule: diffusion.schedule + outputs: + model_input: diffusion.model_input + - kind: invoke + component: denoiser + inputs: + sample: diffusion.model_input + timestep: diffusion.timestep + encoder_hidden_states: conditioning.hidden_states + outputs: + noise_pred: denoiser.estimate + - kind: invoke + component: solver_step + inputs: + sample: latent_state + step: loop.iteration + schedule: diffusion.schedule + derivative: denoiser.estimate + outputs: + next_state: latent.body + - kind: emit + value: denoiser.estimate + output: noise_estimate + mode: append + - kind: emit + value: latent.body + output: latent_trajectory + mode: append + - kind: invoke + component: continue_predicate + inputs: + done: package.false + outputs: + continue: loop.continue + continue_when: loop_0_active + max_iterations: request.max_iterations + carried: + - cell: latent_state + next: latent.body + - cell: loop_0_active + next: loop.continue + iteration: + value: loop.iteration + contract: + dtype: int64 + rank: 1 + shape: + - batch + batch_layout: + kind: request_aligned + axis: 0 + - kind: invoke + component: vae_decoder + inputs: + latent: latent_state + outputs: + image: vae.image + - kind: emit + value: latent_state + output: latent + mode: replace + - kind: emit + value: vae.image + output: image + mode: replace diff --git a/tests/fixtures/onnx_genai_workflows/diffusion_guided/inference_metadata.yaml b/tests/fixtures/onnx_genai_workflows/diffusion_guided/inference_metadata.yaml new file mode 100644 index 000000000..a40fb218d --- /dev/null +++ b/tests/fixtures/onnx_genai_workflows/diffusion_guided/inference_metadata.yaml @@ -0,0 +1,781 @@ +schema_version: v1 +pipeline: + workflow: + manifest: + capabilities: + - workflow_ssa + - linear_effects + - nested_control_flow + - loop_induction_values + - typed_emit + inputs: + request.max_iterations: + contract: + dtype: int64 + rank: 1 + shape: + - 1 + role: + kind: runtime + version: '1.0' + role: max_iterations + source: + kind: request + required: false + default: 3 + package.false: + contract: + dtype: bool + rank: 1 + shape: + - batch + batch_layout: + kind: request_aligned + axis: 0 + role: + kind: opaque + source: + kind: literal + required: false + default: false + request.seed: + contract: + dtype: int64 + rank: 1 + shape: + - batch + batch_layout: + kind: request_aligned + axis: 0 + role: + kind: runtime + version: '1.0' + role: seed + source: + kind: request + required: false + default: 0 + externally_suppliable: true + package.rng_offset: + contract: + dtype: int64 + rank: 1 + shape: + - batch + batch_layout: + kind: request_aligned + axis: 0 + role: + kind: opaque + source: + kind: literal + required: false + default: 0 + request.input_ids: + contract: + dtype: int64 + rank: 2 + shape: + - batch + - prompt_sequence + batch_layout: + kind: request_aligned + axis: 0 + role: + kind: runtime + version: '1.0' + role: prompt_tokens + source: + kind: request + required: true + externally_suppliable: true + request.negative_input_ids: + contract: + dtype: int64 + rank: 2 + shape: + - batch + - prompt_sequence + batch_layout: + kind: request_aligned + axis: 0 + role: + kind: opaque + source: + kind: application + name: negative_input_ids + required: true + externally_suppliable: true + request.guidance_scale: + contract: + dtype: float32 + rank: 1 + shape: + - batch + batch_layout: + kind: request_aligned + axis: 0 + role: + kind: opaque + source: + kind: application + name: guidance_scale + required: false + default: 7.5 + package.loop_0_active: + contract: + dtype: bool + rank: 1 + shape: + - 1 + role: + kind: opaque + source: + kind: literal + required: false + default: true + outputs: + image: + contract: + dtype: float32 + rank: 4 + shape: + - batch + - 3 + - height + - width + batch_layout: + kind: request_aligned + axis: 0 + role: image + stage: pre_adapter + latent: + contract: + dtype: float32 + rank: 4 + shape: + - batch + - 4 + - height + - width + batch_layout: + kind: request_aligned + axis: 0 + role: tensor + stage: pre_adapter + noise_estimate: + contract: + dtype: float32 + rank: 4 + shape: + - batch + - 4 + - height + - noise_estimate_width + batch_layout: + kind: request_aligned + axis: 0 + role: tensor + stage: pre_adapter + latent_trajectory: + contract: + dtype: float32 + rank: 4 + shape: + - batch + - 4 + - height + - trajectory_width + batch_layout: + kind: request_aligned + axis: 0 + role: tensor + stage: pre_adapter + rng_offset: + contract: + dtype: int64 + rank: 1 + shape: + - batch + batch_layout: + kind: request_aligned + axis: 0 + role: tensor + stage: pre_adapter + components: + text_encoder: + implementation: + kind: onnx + artifact: text_encoder/model.onnx + ports: + roles: + input_ids: token_ids + encoder_hidden_states: encoder_hidden_states + denoiser: + implementation: + kind: onnx + artifact: denoiser/model.onnx + ports: + roles: + encoder_hidden_states: encoder_hidden_states + vae_decoder: + implementation: + kind: onnx + artifact: vae_decoder/model.onnx + solver_step: + implementation: + kind: onnx + artifact: policies/solver_step.onnx + ports: + inputs: + sample: + dtype: float32 + rank: 4 + shape: + - batch + - channels + - height + - width + batch_layout: + kind: request_aligned + axis: 0 + estimate: + dtype: float32 + rank: 4 + shape: + - batch + - channels + - height + - width + batch_layout: + kind: request_aligned + axis: 0 + history: + dtype: float32 + rank: 4 + shape: + - batch + - channels + - height + - width + batch_layout: + kind: request_aligned + axis: 0 + step: + dtype: int64 + rank: 1 + shape: + - batch + batch_layout: + kind: request_aligned + axis: 0 + schedule: + dtype: float32 + rank: 1 + shape: + - schedule_length + outputs: + next_state: + dtype: float32 + rank: 4 + shape: + - batch + - channels + - height + - width + batch_layout: + kind: request_aligned + axis: 0 + next_history: + dtype: float32 + rank: 4 + shape: + - batch + - channels + - height + - width + batch_layout: + kind: request_aligned + axis: 0 + contract: + id: onnx-genai.solver-step + version: '1' + bindings: + state: sample + estimate: estimate + history: history + step: step + schedule: schedule + next_state: next_state + next_history: next_history + continue_predicate: + implementation: + kind: onnx + artifact: policies/continue_predicate.onnx + ports: + inputs: + done: + dtype: bool + rank: 1 + shape: + - batch + batch_layout: + kind: request_aligned + axis: 0 + outputs: + continue: + dtype: bool + rank: 1 + shape: + - 1 + diffusion_schedule: + implementation: + kind: onnx + artifact: policies/diffusion_schedule.onnx + ports: + inputs: {} + outputs: + schedule: + dtype: float32 + rank: 1 + shape: + - 4 + diffusion_timesteps: + implementation: + kind: onnx + artifact: policies/diffusion_timesteps.onnx + ports: + inputs: {} + outputs: + schedule: + dtype: float32 + rank: 1 + shape: + - 3 + schedule_lookup: + implementation: + kind: onnx + artifact: policies/schedule_lookup.onnx + ports: + inputs: + schedule: + dtype: float32 + rank: 1 + shape: + - schedule_length + step: + dtype: int64 + rank: 1 + shape: + - batch + batch_layout: + kind: request_aligned + axis: 0 + outputs: + timestep: + dtype: float32 + rank: 1 + shape: + - batch + batch_layout: + kind: request_aligned + axis: 0 + tensor_scale: + implementation: + kind: onnx + artifact: policies/tensor_scale.onnx + ports: + inputs: + tensor: + dtype: float32 + rank: 4 + shape: + - batch + - channels + - height + - width + batch_layout: + kind: request_aligned + axis: 0 + scale: + dtype: float32 + rank: 1 + shape: + - 1 + outputs: + scaled: + dtype: float32 + rank: 4 + shape: + - batch + - channels + - height + - width + batch_layout: + kind: request_aligned + axis: 0 + decoder_input_scale: + implementation: + kind: onnx + artifact: policies/decoder_input_scale.onnx + ports: + inputs: {} + outputs: + value: + dtype: float32 + rank: 1 + shape: + - 1 + history_initializer: + implementation: + kind: onnx + artifact: policies/history_initializer.onnx + ports: + inputs: + reference: + dtype: float32 + rank: 4 + shape: + - batch + - channels + - height + - width + batch_layout: + kind: request_aligned + axis: 0 + outputs: + zeros: + dtype: float32 + rank: 4 + shape: + - batch + - channels + - height + - width + batch_layout: + kind: request_aligned + axis: 0 + guidance_combine: + implementation: + kind: onnx + artifact: policies/guidance_combine.onnx + ports: + inputs: + unconditional: + dtype: float32 + rank: 4 + shape: + - batch + - channels + - height + - width + batch_layout: + kind: request_aligned + axis: 0 + conditional: + dtype: float32 + rank: 4 + shape: + - batch + - channels + - height + - width + batch_layout: + kind: request_aligned + axis: 0 + scale: + dtype: float32 + rank: 1 + shape: + - batch + batch_layout: + kind: request_aligned + axis: 0 + outputs: + estimate: + dtype: float32 + rank: 4 + shape: + - batch + - channels + - height + - width + batch_layout: + kind: request_aligned + axis: 0 + contract: + id: onnx-genai.guidance-combine + version: '1' + bindings: + unconditional: unconditional + conditional: conditional + scale: scale + estimate: estimate + latent_row_shape: + implementation: + kind: onnx + artifact: policies/latent_row_shape.onnx + ports: + inputs: {} + outputs: + shape: + dtype: int64 + rank: 1 + shape: + - 3 + latent_noise: + implementation: + kind: onnx + artifact: policies/latent_noise.onnx + ports: + inputs: + seed: + dtype: int64 + rank: 1 + shape: + - batch + batch_layout: + kind: request_aligned + axis: 0 + offset: + dtype: int64 + rank: 1 + shape: + - batch + batch_layout: + kind: request_aligned + axis: 0 + row_shape: + dtype: int64 + rank: 1 + shape: + - row_rank + outputs: + noise: + dtype: float32 + rank: 4 + shape: + - batch + - channels + - height + - width + batch_layout: + kind: request_aligned + axis: 0 + next_offset: + dtype: int64 + rank: 1 + shape: + - batch + batch_layout: + kind: request_aligned + axis: 0 + contract: + id: onnx-genai.counter-rng + version: '1' + bindings: + seed: seed + offset: offset + row_shape: row_shape + noise: noise + next_offset: next_offset + state: + latent_state: + contract: + dtype: float32 + rank: 4 + shape: + - batch + - 4 + - height + - width + batch_layout: + kind: request_aligned + axis: 0 + scope: invocation + initializer: diffusion.noise + recurrence: + kind: invariant + history: + contract: + dtype: float32 + rank: 4 + shape: + - batch + - 4 + - height + - width + batch_layout: + kind: request_aligned + axis: 0 + scope: invocation + initializer: diffusion.initial_history + recurrence: + kind: invariant + loop_0_active: + contract: + dtype: bool + rank: 1 + shape: + - 1 + scope: invocation + initializer: package.loop_0_active + recurrence: + kind: invariant + steps: + - kind: loop + setup: + - kind: invoke + component: diffusion_schedule + inputs: {} + outputs: + schedule: diffusion.schedule + - kind: invoke + component: diffusion_timesteps + inputs: {} + outputs: + schedule: diffusion.timesteps + - kind: invoke + component: decoder_input_scale + inputs: {} + outputs: + value: diffusion.decoder_scale + - kind: invoke + component: latent_row_shape + inputs: {} + outputs: + shape: diffusion.latent_row_shape + - kind: invoke + component: latent_noise + inputs: + seed: request.seed + offset: package.rng_offset + row_shape: diffusion.latent_row_shape + outputs: + noise: diffusion.noise + next_offset: diffusion.rng_offset + - kind: invoke + component: text_encoder + inputs: + input_ids: request.input_ids + outputs: + encoder_hidden_states: conditioning.hidden_states + - kind: invoke + component: text_encoder + inputs: + input_ids: request.negative_input_ids + outputs: + encoder_hidden_states: conditioning.unconditional + - kind: invoke + component: history_initializer + inputs: + reference: diffusion.noise + outputs: + zeros: diffusion.initial_history + - kind: invoke + component: continue_predicate + inputs: + done: package.false + outputs: + continue: setup.continue + steps: + - kind: invoke + component: schedule_lookup + inputs: + schedule: diffusion.timesteps + step: loop.iteration + outputs: + timestep: diffusion.timestep + - kind: invoke + component: denoiser + inputs: + sample: latent_state + timestep: diffusion.timestep + encoder_hidden_states: conditioning.unconditional + outputs: + noise_pred: denoiser.unconditional + - kind: invoke + component: denoiser + inputs: + sample: latent_state + timestep: diffusion.timestep + encoder_hidden_states: conditioning.hidden_states + outputs: + noise_pred: denoiser.conditional + - kind: invoke + component: guidance_combine + inputs: + unconditional: denoiser.unconditional + conditional: denoiser.conditional + scale: request.guidance_scale + outputs: + estimate: denoiser.estimate + - kind: invoke + component: solver_step + inputs: + sample: latent_state + step: loop.iteration + schedule: diffusion.schedule + estimate: denoiser.estimate + history: history + outputs: + next_state: latent.body + next_history: history.body + - kind: emit + value: denoiser.estimate + output: noise_estimate + mode: append + - kind: emit + value: latent.body + output: latent_trajectory + mode: append + - kind: invoke + component: continue_predicate + inputs: + done: package.false + outputs: + continue: loop.continue + continue_when: loop_0_active + max_iterations: request.max_iterations + carried: + - cell: latent_state + next: latent.body + - cell: history + next: history.body + - cell: loop_0_active + next: loop.continue + iteration: + value: loop.iteration + contract: + dtype: int64 + rank: 1 + shape: + - batch + batch_layout: + kind: request_aligned + axis: 0 + - kind: invoke + component: tensor_scale + inputs: + tensor: latent_state + scale: diffusion.decoder_scale + outputs: + scaled: diffusion.decoder_input + - kind: invoke + component: vae_decoder + inputs: + latent: diffusion.decoder_input + outputs: + image: vae.image + - kind: emit + value: latent_state + output: latent + mode: replace + - kind: emit + value: vae.image + output: image + mode: replace + - kind: emit + value: diffusion.rng_offset + output: rng_offset + mode: replace diff --git a/tests/fixtures/onnx_genai_workflows/esm2_protein_embeddings/inference_metadata.yaml b/tests/fixtures/onnx_genai_workflows/esm2_protein_embeddings/inference_metadata.yaml new file mode 100644 index 000000000..56ef7f831 --- /dev/null +++ b/tests/fixtures/onnx_genai_workflows/esm2_protein_embeddings/inference_metadata.yaml @@ -0,0 +1,95 @@ +schema_version: v1 +profiles: + embedding: + kind: embedding + version: '1.0' + requirement: required + outputs: + last_hidden_state: last_hidden_state + pooling: + kind: mean + axis: 1 + normalize: false + batch_invariance: row_independent +pipeline: + workflow: + manifest: + capabilities: + - workflow_ssa + - linear_effects + - typed_emit + effects: + encode: + retry: pure + speculation_safety: + kind: clonable + inputs: + request.input_ids: + contract: + dtype: int64 + rank: 2 + shape: + - batch + - sequence_len + batch_layout: + kind: request_aligned + axis: 0 + role: + kind: runtime + version: '1.0' + role: prompt_tokens + source: + kind: request + required: true + request.attention_mask: + contract: + dtype: int64 + rank: 2 + shape: + - batch + - sequence_len + batch_layout: + kind: request_aligned + axis: 0 + role: + kind: opaque + source: + kind: application + name: request.attention_mask + required: true + outputs: + last_hidden_state: + contract: + dtype: float32 + rank: 3 + shape: + - batch + - sequence_len + - 64 + batch_layout: + kind: request_aligned + axis: 0 + role: tensor + stage: post_adapter + components: + encoder: + implementation: + kind: onnx + artifact: model.onnx + ports: + roles: + input_ids: token_ids + attention_mask: attention_mask + last_hidden_state: hidden_states + steps: + - kind: invoke + component: encoder + inputs: + input_ids: request.input_ids + attention_mask: request.attention_mask + outputs: + last_hidden_state: encoder.last_hidden_state + - kind: emit + value: encoder.last_hidden_state + output: last_hidden_state + mode: replace diff --git a/tests/fixtures/onnx_genai_workflows/masked/inference_metadata.yaml b/tests/fixtures/onnx_genai_workflows/masked/inference_metadata.yaml new file mode 100644 index 000000000..51892ae4d --- /dev/null +++ b/tests/fixtures/onnx_genai_workflows/masked/inference_metadata.yaml @@ -0,0 +1,412 @@ +schema_version: '1.0' +pipeline: + workflow: + manifest: + capabilities: + - workflow_ssa + - linear_effects + - nested_control_flow + - typed_emit + - emit_valid_length + - loop_induction_values + inputs: + request.input_ids: + contract: + dtype: int64 + rank: 2 + shape: + - batch + - sequence + batch_layout: + kind: request_aligned + axis: 0 + role: + kind: runtime + version: '1.0' + role: prompt_tokens + source: + kind: request + required: true + request.mask: + contract: + dtype: bool + rank: 2 + shape: + - batch + - sequence + batch_layout: + kind: request_aligned + axis: 0 + role: + kind: opaque + source: + kind: application + name: masked_positions + required: true + request.seed: + contract: + dtype: int64 + rank: 1 + shape: + - batch + batch_layout: + kind: request_aligned + axis: 0 + role: + kind: runtime + version: '1.0' + role: seed + source: + kind: request + required: false + default: 0 + request.rng_offset: + contract: + dtype: int64 + rank: 1 + shape: + - batch + batch_layout: + kind: request_aligned + axis: 0 + role: + kind: opaque + source: + kind: application + name: rng_offset + required: false + default: 0 + request.max_iterations: + contract: + dtype: int64 + rank: 1 + shape: + - 1 + role: + kind: runtime + version: '1.0' + role: max_iterations + source: + kind: request + required: false + default: 8 + package.num_steps: + contract: + dtype: int64 + rank: 1 + shape: + - batch + batch_layout: + kind: request_aligned + axis: 0 + role: + kind: opaque + source: + kind: literal + required: false + default: 8 + package.loop_0_active: + contract: + dtype: bool + rank: 1 + shape: + - 1 + role: + kind: opaque + source: + kind: literal + required: false + default: true + outputs: + tokens: + contract: + dtype: int64 + rank: 2 + shape: + - batch + - sequence + batch_layout: + kind: request_aligned + axis: 0 + role: tokens + stage: pre_adapter + components: + model: + implementation: + kind: onnx + artifact: model.onnx + ports: + roles: + input_ids: token_ids + logits: logits + masked_update: + implementation: + kind: onnx + artifact: policies/masked_update.onnx + ports: + inputs: + current_tokens: + dtype: int64 + rank: 2 + shape: + - batch + - sequence + batch_layout: + kind: request_aligned + axis: 0 + proposed_tokens: + dtype: int64 + rank: 2 + shape: + - batch + - sequence + batch_layout: + kind: request_aligned + axis: 0 + logits: + dtype: float32 + rank: 3 + shape: + - batch + - sequence + - vocabulary + batch_layout: + kind: request_aligned + axis: 0 + masked: + dtype: bool + rank: 2 + shape: + - batch + - sequence + batch_layout: + kind: request_aligned + axis: 0 + step: + dtype: int64 + rank: 1 + shape: + - batch + batch_layout: + kind: request_aligned + axis: 0 + total_steps: + dtype: int64 + rank: 1 + shape: + - batch + batch_layout: + kind: request_aligned + axis: 0 + seed: + dtype: int64 + rank: 1 + shape: + - batch + batch_layout: + kind: request_aligned + axis: 0 + offset: + dtype: int64 + rank: 1 + shape: + - batch + batch_layout: + kind: request_aligned + axis: 0 + outputs: + next_state: + dtype: int64 + rank: 2 + shape: + - batch + - sequence + batch_layout: + kind: request_aligned + axis: 0 + next_mask: + dtype: bool + rank: 2 + shape: + - batch + - sequence + batch_layout: + kind: request_aligned + axis: 0 + next_offset: + dtype: int64 + rank: 1 + shape: + - batch + batch_layout: + kind: request_aligned + axis: 0 + done: + dtype: bool + rank: 1 + shape: + - batch + batch_layout: + kind: request_aligned + axis: 0 + continue: + dtype: bool + rank: 1 + shape: + - 1 + contract: + id: onnx-genai.masked-update + version: '1' + bindings: + state: current_tokens + proposal: proposed_tokens + mask: masked + step: step + next_state: next_state + next_mask: next_mask + continue: continue + seed: seed + offset: offset + next_offset: next_offset + state: + tokens_state: + contract: + dtype: int64 + rank: 2 + shape: + - batch + - sequence + batch_layout: + kind: request_aligned + axis: 0 + scope: invocation + initializer: request.input_ids + recurrence: + kind: invariant + mask: + contract: + dtype: bool + rank: 2 + shape: + - batch + - sequence + batch_layout: + kind: request_aligned + axis: 0 + scope: invocation + initializer: request.mask + recurrence: + kind: invariant + rng_offset: + contract: + dtype: int64 + rank: 1 + shape: + - batch + batch_layout: + kind: request_aligned + axis: 0 + scope: invocation + initializer: request.rng_offset + recurrence: + kind: invariant + logits: + contract: + dtype: float32 + rank: 3 + shape: + - batch + - sequence + - 128 + batch_layout: + kind: request_aligned + axis: 0 + scope: invocation + initializer: denoiser.setup.logits + recurrence: + kind: invariant + proposal: + contract: + dtype: int64 + rank: 2 + shape: + - batch + - sequence + batch_layout: + kind: request_aligned + axis: 0 + scope: invocation + initializer: denoiser.setup.proposal + recurrence: + kind: invariant + loop_0_active: + contract: + dtype: bool + rank: 1 + shape: + - 1 + scope: invocation + initializer: package.loop_0_active + recurrence: + kind: invariant + steps: + - kind: loop + setup: + - kind: invoke + component: model + inputs: + input_ids: request.input_ids + outputs: + logits: denoiser.setup.logits + proposed_tokens: denoiser.setup.proposal + steps: + - kind: invoke + component: masked_update + inputs: + current_tokens: tokens_state + proposed_tokens: proposal + logits: logits + masked: mask + step: loop.iteration + total_steps: package.num_steps + seed: request.seed + offset: rng_offset + outputs: + next_state: denoiser.body.tokens + next_mask: denoiser.body.mask + next_offset: denoiser.body.rng_offset + done: denoiser.body.done + continue: denoiser.body.continue + - kind: emit + value: denoiser.body.tokens + output: tokens + mode: replace + - kind: invoke + component: model + inputs: + input_ids: denoiser.body.tokens + outputs: + logits: denoiser.body.logits + proposed_tokens: denoiser.body.proposal + continue_when: loop_0_active + max_iterations: request.max_iterations + carried: + - cell: tokens_state + next: denoiser.body.tokens + - cell: mask + next: denoiser.body.mask + - cell: rng_offset + next: denoiser.body.rng_offset + - cell: logits + next: denoiser.body.logits + - cell: proposal + next: denoiser.body.proposal + - cell: loop_0_active + next: denoiser.body.continue + iteration: + value: loop.iteration + contract: + dtype: int64 + rank: 1 + shape: + - 1 diff --git a/tests/fixtures/onnx_genai_workflows/protbert_protein_embeddings/inference_metadata.yaml b/tests/fixtures/onnx_genai_workflows/protbert_protein_embeddings/inference_metadata.yaml new file mode 100644 index 000000000..28ff17289 --- /dev/null +++ b/tests/fixtures/onnx_genai_workflows/protbert_protein_embeddings/inference_metadata.yaml @@ -0,0 +1,112 @@ +schema_version: v1 +profiles: + embedding: + kind: embedding + version: '1.0' + requirement: required + outputs: + last_hidden_state: last_hidden_state + pooling: + kind: mean + axis: 1 + normalize: false + batch_invariance: row_independent +pipeline: + workflow: + manifest: + capabilities: + - workflow_ssa + - linear_effects + - typed_emit + effects: + encode: + retry: pure + speculation_safety: + kind: clonable + inputs: + request.input_ids: + contract: + dtype: int64 + rank: 2 + shape: + - batch + - sequence_len + batch_layout: + kind: request_aligned + axis: 0 + role: + kind: runtime + version: '1.0' + role: prompt_tokens + source: + kind: request + required: true + request.attention_mask: + contract: + dtype: int64 + rank: 2 + shape: + - batch + - sequence_len + batch_layout: + kind: request_aligned + axis: 0 + role: + kind: opaque + source: + kind: application + name: request.attention_mask + required: true + request.token_type_ids: + contract: + dtype: int64 + rank: 2 + shape: + - batch + - sequence_len + batch_layout: + kind: request_aligned + axis: 0 + role: + kind: opaque + source: + kind: application + name: request.token_type_ids + required: true + outputs: + last_hidden_state: + contract: + dtype: float32 + rank: 3 + shape: + - batch + - last_hidden_state_dim_1 + - 64 + batch_layout: + kind: request_aligned + axis: 0 + role: tensor + stage: post_adapter + components: + encoder: + implementation: + kind: onnx + artifact: model.onnx + ports: + roles: + input_ids: token_ids + attention_mask: attention_mask + last_hidden_state: hidden_states + steps: + - kind: invoke + component: encoder + inputs: + input_ids: request.input_ids + attention_mask: request.attention_mask + token_type_ids: request.token_type_ids + outputs: + last_hidden_state: encoder.last_hidden_state + - kind: emit + value: encoder.last_hidden_state + output: last_hidden_state + mode: replace diff --git a/tests/fixtures/onnx_genai_workflows/speculative/inference_metadata.yaml b/tests/fixtures/onnx_genai_workflows/speculative/inference_metadata.yaml new file mode 100644 index 000000000..5ca10d33f --- /dev/null +++ b/tests/fixtures/onnx_genai_workflows/speculative/inference_metadata.yaml @@ -0,0 +1,1322 @@ +schema_version: v1 +pipeline: + workflow: + manifest: + capabilities: + - workflow_ssa + - linear_effects + - nested_control_flow + - loop_induction_values + - typed_emit + - emit_valid_length + - bounded_state_recurrence + - serving_service_contract + - grammar_guidance_adapter + - adaptive_proposal_budget + - advisory_state + adapter_abis: + onnx-genai.grammar-guidance: '1' + inputs: + request.tokens: + contract: + dtype: int64 + rank: 2 + shape: + - batch + - 4 + batch_layout: + kind: request_aligned + axis: 0 + role: + kind: runtime + version: '1.0' + role: prompt_tokens + source: + kind: request + required: true + request.seed: + contract: + dtype: int64 + rank: 1 + shape: + - batch + batch_layout: + kind: request_aligned + axis: 0 + role: + kind: runtime + version: '1.0' + role: seed + source: + kind: request + required: false + default: 0 + request.max_iterations: + contract: + dtype: int64 + rank: 1 + shape: + - 1 + role: + kind: runtime + version: '1.0' + role: max_output_tokens + source: + kind: request + required: true + package.zero: + contract: + dtype: int64 + rank: 1 + shape: + - batch + batch_layout: + kind: request_aligned + axis: 0 + role: + kind: opaque + source: + kind: literal + required: false + default: 0 + package.one: + contract: + dtype: int64 + rank: 1 + shape: + - batch + batch_layout: + kind: request_aligned + axis: 0 + role: + kind: opaque + source: + kind: literal + required: false + default: 1 + package.max_context: + contract: + dtype: int64 + rank: 1 + shape: + - 1 + role: + kind: opaque + source: + kind: literal + required: false + default: 4096 + package.false: + contract: + dtype: bool + rank: 1 + shape: + - batch + batch_layout: + kind: request_aligned + axis: 0 + role: + kind: opaque + source: + kind: literal + required: false + default: false + package.active: + contract: + dtype: bool + rank: 1 + shape: + - batch + batch_layout: + kind: request_aligned + axis: 0 + role: + kind: opaque + source: + kind: literal + required: false + default: true + request.cache_lengths: + contract: + dtype: int64 + rank: 1 + shape: + - batch + batch_layout: + kind: request_aligned + axis: 0 + role: + kind: opaque + source: + kind: application + name: serving.cache_lengths + required: false + default: 0 + request.grammar_state: + contract: + dtype: int64 + rank: 1 + shape: + - batch + batch_layout: + kind: request_aligned + axis: 0 + role: + kind: opaque + source: + kind: application + name: grammar.initial_state + required: true + request.grammar_transition_table: + contract: + dtype: int64 + rank: 2 + shape: + - grammar_states + - vocabulary + role: + kind: opaque + source: + kind: application + name: grammar.transition_table + required: true + request.adaptive_k: + contract: + dtype: int64 + rank: 1 + shape: + - batch + batch_layout: + kind: request_aligned + axis: 0 + role: + kind: opaque + source: + kind: application + name: adaptive.current_k + required: false + default: 1 + request.adaptive_estimates: + contract: + dtype: float32 + rank: 2 + shape: + - batch + - 24 + batch_layout: + kind: request_aligned + axis: 0 + role: + kind: opaque + source: + kind: application + name: adaptive.estimates + required: false + default: 0.0 + request.draft_ms: + contract: + dtype: float32 + rank: 1 + shape: + - batch + batch_layout: + kind: request_aligned + axis: 0 + role: + kind: opaque + source: + kind: application + name: telemetry.draft_ms + required: true + request.target_ms: + contract: + dtype: float32 + rank: 1 + shape: + - batch + batch_layout: + kind: request_aligned + axis: 0 + role: + kind: opaque + source: + kind: application + name: telemetry.target_ms + required: true + request.verifier.past_key_values.0.key: + contract: + dtype: float32 + rank: 4 + shape: + - batch + - 2 + - past_sequence + - 8 + batch_layout: + kind: request_aligned + axis: 0 + role: + kind: opaque + source: + kind: application + name: verifier.past_key_values.0.key + required: true + outputs: + tokens: + contract: + dtype: int64 + rank: 2 + shape: + - batch + - accepted_sequence + batch_layout: + kind: request_aligned + axis: 0 + role: tokens + stage: pre_adapter + components: + proposer: + implementation: + kind: onnx + artifact: proposer/model.onnx + verifier: + implementation: + kind: onnx + artifact: verifier/model.onnx + grammar_clone: + implementation: + kind: adapter + abi: onnx-genai.grammar-guidance + version: '1' + ports: + inputs: + state: + dtype: int64 + rank: 1 + shape: + - batch + batch_layout: + kind: request_aligned + axis: 0 + tokens: + dtype: int64 + rank: 2 + shape: + - batch + - proposal + batch_layout: + kind: request_aligned + axis: 0 + valid_length: + dtype: int64 + rank: 1 + shape: + - batch + batch_layout: + kind: request_aligned + axis: 0 + transition_table: + dtype: int64 + rank: 2 + shape: + - grammar_states + - vocabulary + outputs: + next_state: + dtype: int64 + rank: 1 + shape: + - batch + batch_layout: + kind: request_aligned + axis: 0 + consumed_length: + dtype: int64 + rank: 1 + shape: + - batch + batch_layout: + kind: request_aligned + axis: 0 + logits_mask: + dtype: bool + rank: 2 + shape: + - batch + - vocabulary + batch_layout: + kind: request_aligned + axis: 0 + forced_tokens: + dtype: int64 + rank: 2 + shape: + - batch + - 1 + batch_layout: + kind: request_aligned + axis: 0 + forced_length: + dtype: int64 + rank: 1 + shape: + - batch + batch_layout: + kind: request_aligned + axis: 0 + contract: + id: onnx-genai.grammar-guidance + version: '1' + bindings: + state: state + tokens: tokens + valid_length: valid_length + transition_table: transition_table + next_state: next_state + consumed_length: consumed_length + logits_mask: logits_mask + forced_tokens: forced_tokens + forced_length: forced_length + parameters: + action: clone + effects: + - grammar + row_scope: + axis: 0 + stateful: true + grammar_lookahead: + implementation: + kind: adapter + abi: onnx-genai.grammar-guidance + version: '1' + ports: + inputs: + state: + dtype: int64 + rank: 1 + shape: + - batch + batch_layout: + kind: request_aligned + axis: 0 + tokens: + dtype: int64 + rank: 2 + shape: + - batch + - proposal + batch_layout: + kind: request_aligned + axis: 0 + valid_length: + dtype: int64 + rank: 1 + shape: + - batch + batch_layout: + kind: request_aligned + axis: 0 + transition_table: + dtype: int64 + rank: 2 + shape: + - grammar_states + - vocabulary + outputs: + next_state: + dtype: int64 + rank: 1 + shape: + - batch + batch_layout: + kind: request_aligned + axis: 0 + consumed_length: + dtype: int64 + rank: 1 + shape: + - batch + batch_layout: + kind: request_aligned + axis: 0 + logits_mask: + dtype: bool + rank: 2 + shape: + - batch + - vocabulary + batch_layout: + kind: request_aligned + axis: 0 + forced_tokens: + dtype: int64 + rank: 2 + shape: + - batch + - 1 + batch_layout: + kind: request_aligned + axis: 0 + forced_length: + dtype: int64 + rank: 1 + shape: + - batch + batch_layout: + kind: request_aligned + axis: 0 + contract: + id: onnx-genai.grammar-guidance + version: '1' + bindings: + state: state + tokens: tokens + valid_length: valid_length + transition_table: transition_table + next_state: next_state + consumed_length: consumed_length + logits_mask: logits_mask + forced_tokens: forced_tokens + forced_length: forced_length + parameters: + action: lookahead + effects: + - grammar + row_scope: + axis: 0 + stateful: true + grammar_commit: + implementation: + kind: adapter + abi: onnx-genai.grammar-guidance + version: '1' + ports: + inputs: + state: + dtype: int64 + rank: 1 + shape: + - batch + batch_layout: + kind: request_aligned + axis: 0 + tokens: + dtype: int64 + rank: 2 + shape: + - batch + - proposal + batch_layout: + kind: request_aligned + axis: 0 + valid_length: + dtype: int64 + rank: 1 + shape: + - batch + batch_layout: + kind: request_aligned + axis: 0 + transition_table: + dtype: int64 + rank: 2 + shape: + - grammar_states + - vocabulary + outputs: + next_state: + dtype: int64 + rank: 1 + shape: + - batch + batch_layout: + kind: request_aligned + axis: 0 + consumed_length: + dtype: int64 + rank: 1 + shape: + - batch + batch_layout: + kind: request_aligned + axis: 0 + logits_mask: + dtype: bool + rank: 2 + shape: + - batch + - vocabulary + batch_layout: + kind: request_aligned + axis: 0 + forced_tokens: + dtype: int64 + rank: 2 + shape: + - batch + - 1 + batch_layout: + kind: request_aligned + axis: 0 + forced_length: + dtype: int64 + rank: 1 + shape: + - batch + batch_layout: + kind: request_aligned + axis: 0 + contract: + id: onnx-genai.grammar-guidance + version: '1' + bindings: + state: state + tokens: tokens + valid_length: valid_length + transition_table: transition_table + next_state: next_state + consumed_length: consumed_length + logits_mask: logits_mask + forced_tokens: forced_tokens + forced_length: forced_length + parameters: + action: commit + effects: + - grammar + row_scope: + axis: 0 + stateful: true + speculative_acceptance: + implementation: + kind: onnx + artifact: policies/speculative_acceptance.onnx + ports: + inputs: + target_scores: + dtype: float32 + rank: 3 + shape: + - batch + - draft_sequence + - vocabulary + batch_layout: + kind: request_aligned + axis: 0 + proposed_tokens: + dtype: int64 + rank: 2 + shape: + - batch + - draft_sequence + batch_layout: + kind: request_aligned + axis: 0 + seed: + dtype: int64 + rank: 1 + shape: + - batch + batch_layout: + kind: request_aligned + axis: 0 + offset: + dtype: int64 + rank: 1 + shape: + - batch + batch_layout: + kind: request_aligned + axis: 0 + outputs: + accepted_tokens: + dtype: int64 + rank: 2 + shape: + - batch + - draft_sequence + batch_layout: + kind: request_aligned + axis: 0 + accepted_len: + dtype: int64 + rank: 1 + shape: + - batch + batch_layout: + kind: request_aligned + axis: 0 + done: + dtype: bool + rank: 1 + shape: + - batch + batch_layout: + kind: request_aligned + axis: 0 + next_offset: + dtype: int64 + rank: 1 + shape: + - batch + batch_layout: + kind: request_aligned + axis: 0 + rollback_len: + dtype: int64 + rank: 1 + shape: + - batch + batch_layout: + kind: request_aligned + axis: 0 + continue: + dtype: bool + rank: 1 + shape: + - batch + batch_layout: + kind: request_aligned + axis: 0 + contract: + id: onnx-genai.speculative-verifier + version: '1' + bindings: + target_scores: target_scores + proposed_tokens: proposed_tokens + accepted_tokens: accepted_tokens + accepted_len: accepted_len + done: done + continue: continue + seed: seed + offset: offset + next_offset: next_offset + grammar_guidance: + implementation: + kind: onnx + artifact: policies/grammar_guidance.onnx + ports: + inputs: + logits: + dtype: float32 + rank: 2 + shape: + - batch + - vocabulary + batch_layout: + kind: request_aligned + axis: 0 + logits_mask: + dtype: bool + rank: 2 + shape: + - batch + - vocabulary + batch_layout: + kind: request_aligned + axis: 0 + forced_tokens: + dtype: int64 + rank: 2 + shape: + - batch + - 1 + batch_layout: + kind: request_aligned + axis: 0 + forced_length: + dtype: int64 + rank: 1 + shape: + - batch + batch_layout: + kind: request_aligned + axis: 0 + outputs: + token: + dtype: int64 + rank: 2 + shape: + - batch + - 1 + batch_layout: + kind: request_aligned + axis: 0 + adaptive_k: + implementation: + kind: onnx + artifact: policies/adaptive_k.onnx + ports: + inputs: + current_k: + dtype: int64 + rank: 1 + shape: + - batch + batch_layout: + kind: request_aligned + axis: 0 + accepted: + dtype: int64 + rank: 1 + shape: + - batch + batch_layout: + kind: request_aligned + axis: 0 + evaluated: + dtype: int64 + rank: 1 + shape: + - batch + batch_layout: + kind: request_aligned + axis: 0 + committed_tokens: + dtype: int64 + rank: 1 + shape: + - batch + batch_layout: + kind: request_aligned + axis: 0 + filled_proposal_budget: + dtype: bool + rank: 1 + shape: + - batch + batch_layout: + kind: request_aligned + axis: 0 + draft_ms: + dtype: float32 + rank: 1 + shape: + - batch + batch_layout: + kind: request_aligned + axis: 0 + target_ms: + dtype: float32 + rank: 1 + shape: + - batch + batch_layout: + kind: request_aligned + axis: 0 + estimates: + dtype: float32 + rank: 2 + shape: + - batch + - 24 + batch_layout: + kind: request_aligned + axis: 0 + outputs: + next_k: + dtype: int64 + rank: 1 + shape: + - batch + batch_layout: + kind: request_aligned + axis: 0 + next_estimates: + dtype: float32 + rank: 2 + shape: + - batch + - 24 + batch_layout: + kind: request_aligned + axis: 0 + contract: + id: onnx-genai.adaptive-proposal-budget + version: '1' + bindings: + current_k: current_k + accepted: accepted + evaluated: evaluated + committed_tokens: committed_tokens + filled_proposal_budget: filled_proposal_budget + draft_ms: draft_ms + target_ms: target_ms + estimates: estimates + next_k: next_k + next_estimates: next_estimates + grammar_length: + implementation: + kind: onnx + artifact: policies/grammar_length.onnx + ports: + inputs: + left: + dtype: int64 + rank: 1 + shape: + - batch + batch_layout: + kind: request_aligned + axis: 0 + right: + dtype: int64 + rank: 1 + shape: + - batch + batch_layout: + kind: request_aligned + axis: 0 + outputs: + minimum: + dtype: int64 + rank: 1 + shape: + - batch + batch_layout: + kind: request_aligned + axis: 0 + grammar_sampler_logits: + implementation: + kind: onnx + artifact: policies/grammar_sampler_logits.onnx + ports: + inputs: + logits: + dtype: float32 + rank: 3 + shape: + - batch + - sequence + - vocabulary + batch_layout: + kind: request_aligned + axis: 0 + outputs: + last_logits: + dtype: float32 + rank: 2 + shape: + - batch + - vocabulary + batch_layout: + kind: request_aligned + axis: 0 + proposal_metrics: + implementation: + kind: onnx + artifact: policies/proposal_metrics.onnx + ports: + inputs: + proposed_tokens: + dtype: int64 + rank: 2 + shape: + - batch + - proposal + batch_layout: + kind: request_aligned + axis: 0 + requested_k: + dtype: int64 + rank: 1 + shape: + - batch + batch_layout: + kind: request_aligned + axis: 0 + outputs: + evaluated: + dtype: int64 + rank: 1 + shape: + - batch + batch_layout: + kind: request_aligned + axis: 0 + filled_proposal_budget: + dtype: bool + rank: 1 + shape: + - batch + batch_layout: + kind: request_aligned + axis: 0 + cache_length_update: + implementation: + kind: onnx + artifact: policies/cache_length_update.onnx + ports: + inputs: + left: + dtype: int64 + rank: 1 + shape: + - batch + batch_layout: + kind: request_aligned + axis: 0 + right: + dtype: int64 + rank: 1 + shape: + - batch + batch_layout: + kind: request_aligned + axis: 0 + outputs: + total: + dtype: int64 + rank: 1 + shape: + - batch + batch_layout: + kind: request_aligned + axis: 0 + state: + tokens_state: + contract: + dtype: int64 + rank: 2 + shape: + - batch + - 4 + batch_layout: + kind: request_aligned + axis: 0 + class: semantic + scope: invocation + initializer: request.tokens + recurrence: + kind: invariant + rng_offset: + contract: + dtype: int64 + rank: 1 + shape: + - batch + batch_layout: + kind: request_aligned + axis: 0 + class: semantic + scope: invocation + initializer: package.zero + recurrence: + kind: invariant + active: + contract: + dtype: bool + rank: 1 + shape: + - batch + batch_layout: + kind: request_aligned + axis: 0 + class: semantic + scope: invocation + initializer: package.active + recurrence: + kind: invariant + done: + contract: + dtype: bool + rank: 1 + shape: + - batch + batch_layout: + kind: request_aligned + axis: 0 + class: semantic + scope: invocation + initializer: package.false + recurrence: + kind: invariant + accepted_len: + contract: + dtype: int64 + rank: 1 + shape: + - batch + batch_layout: + kind: request_aligned + axis: 0 + class: semantic + scope: invocation + initializer: package.zero + recurrence: + kind: invariant + cache_lengths: + contract: + dtype: int64 + rank: 1 + shape: + - batch + batch_layout: + kind: request_aligned + axis: 0 + class: semantic + scope: invocation + initializer: request.cache_lengths + recurrence: + kind: invariant + grammar: + contract: + dtype: int64 + rank: 1 + shape: + - batch + batch_layout: + kind: request_aligned + axis: 0 + class: semantic + scope: invocation + initializer: request.grammar_state + recurrence: + kind: invariant + proposal_k: + contract: + dtype: int64 + rank: 1 + shape: + - batch + batch_layout: + kind: request_aligned + axis: 0 + class: advisory + scope: invocation + initializer: request.adaptive_k + recurrence: + kind: invariant + adaptive_estimates: + contract: + dtype: float32 + rank: 2 + shape: + - batch + - 24 + batch_layout: + kind: request_aligned + axis: 0 + class: advisory + scope: invocation + initializer: request.adaptive_estimates + recurrence: + kind: invariant + cache_0: + contract: + dtype: float32 + rank: 4 + shape: + - batch + - 2 + - past_sequence + - 8 + batch_layout: + kind: request_aligned + axis: 0 + class: semantic + scope: invocation + initializer: request.verifier.past_key_values.0.key + recurrence: + kind: bounded + axis: 2 + max: package.max_context + service_group: verifier_cache + management: runtime + release_boundary: invocation + serving: + active: active + done: done + accepted_len: accepted_len + state_service: + groups: + verifier_cache: + kind: full_attention + sequence_axis: 2 + layout: bnsh + aliasing: forbidden + reuse: + prefix_reusable: true + evictable_prefix: false + capabilities: + snapshot: true + fork: true + ports: + verifier: + cache_0: + input: past_key_values.0.key + output: present.0.key + role: key + layer: 0 + logical_lengths: cache_lengths + effects: + grammar: + retry: transactional + speculation_safety: + kind: clonable + steps: + - kind: loop + setup: [] + steps: + - kind: invoke + component: proposer + inputs: + tokens: tokens_state + proposal_budget: proposal_k + outputs: + proposed_tokens: proposal.tokens + proposal_scores: proposal.scores + - kind: invoke + component: proposal_metrics + inputs: + proposed_tokens: proposal.tokens + requested_k: proposal_k + outputs: + evaluated: proposal.evaluated + filled_proposal_budget: proposal.filled_budget + - kind: invoke + component: grammar_clone + inputs: + state: grammar + tokens: proposal.tokens + valid_length: package.zero + transition_table: request.grammar_transition_table + outputs: + next_state: grammar.clone.state + consumed_length: grammar.clone.consumed + logits_mask: grammar.clone.mask + forced_tokens: grammar.clone.forced + forced_length: grammar.clone.forced_length + - kind: invoke + component: grammar_lookahead + inputs: + state: grammar.clone.state + tokens: proposal.tokens + valid_length: proposal.evaluated + transition_table: request.grammar_transition_table + outputs: + next_state: grammar.lookahead.state + consumed_length: grammar.valid_length + logits_mask: grammar.lookahead.mask + forced_tokens: grammar.lookahead.forced + forced_length: grammar.lookahead.forced_length + - kind: invoke + component: verifier + inputs: + proposed_tokens: proposal.tokens + past_key_values.0.key: cache_0 + outputs: + target_scores: target.scores + present.0.key: verifier.present.0.key + - kind: invoke + component: speculative_acceptance + inputs: + target_scores: target.scores + proposed_tokens: proposal.tokens + seed: request.seed + offset: rng_offset + outputs: + accepted_tokens: acceptance.tokens + accepted_len: acceptance.length + done: acceptance.done + continue: acceptance.continue + next_offset: rng_offset.body + rollback_len: acceptance.rollback_length + - kind: invoke + component: grammar_length + inputs: + left: acceptance.length + right: grammar.valid_length + outputs: + minimum: grammar.committed_length + - kind: invoke + component: grammar_commit + inputs: + state: grammar + tokens: acceptance.tokens + valid_length: grammar.committed_length + transition_table: request.grammar_transition_table + outputs: + next_state: grammar.next + consumed_length: grammar.committed + logits_mask: grammar.mask + forced_tokens: grammar.forced + forced_length: grammar.forced_length + - kind: invoke + component: grammar_sampler_logits + inputs: + logits: target.scores + outputs: + last_logits: grammar.sampler_logits + - kind: invoke + component: grammar_guidance + inputs: + logits: grammar.sampler_logits + logits_mask: grammar.mask + forced_tokens: grammar.forced + forced_length: grammar.forced_length + outputs: + token: grammar.token + - kind: invoke + component: cache_length_update + inputs: + left: cache_lengths + right: grammar.committed_length + outputs: + total: cache_lengths.next + - kind: invoke + component: adaptive_k + inputs: + current_k: proposal_k + accepted: grammar.committed_length + evaluated: proposal.evaluated + committed_tokens: grammar.committed_length + filled_proposal_budget: proposal.filled_budget + draft_ms: request.draft_ms + target_ms: request.target_ms + estimates: adaptive_estimates + outputs: + next_k: adaptive.next_k + next_estimates: adaptive.next_estimates + - kind: emit + value: acceptance.tokens + output: tokens + mode: append + valid_length: grammar.committed_length + - kind: emit + value: grammar.token + output: tokens + mode: append + valid_length: grammar.forced_length + continue_when: active + max_iterations: request.max_iterations + carried: + - cell: tokens_state + next: acceptance.tokens + - cell: rng_offset + next: rng_offset.body + - cell: active + next: acceptance.continue + - cell: done + next: acceptance.done + - cell: accepted_len + next: grammar.committed_length + - cell: cache_lengths + next: cache_lengths.next + - cell: grammar + next: grammar.next + - cell: proposal_k + next: adaptive.next_k + - cell: adaptive_estimates + next: adaptive.next_estimates + - cell: cache_0 + next: verifier.present.0.key + iteration: + value: speculative.iteration + contract: + dtype: int64 + rank: 1 + shape: + - batch + batch_layout: + kind: request_aligned + axis: 0 diff --git a/tests/fixtures/onnx_genai_workflows/static_cache/inference_metadata.yaml b/tests/fixtures/onnx_genai_workflows/static_cache/inference_metadata.yaml new file mode 100644 index 000000000..534530704 --- /dev/null +++ b/tests/fixtures/onnx_genai_workflows/static_cache/inference_metadata.yaml @@ -0,0 +1,1348 @@ +schema_version: '1.0' +pipeline: + workflow: + manifest: + capabilities: + - workflow_ssa + - linear_effects + - nested_control_flow + - typed_emit + - emit_valid_length + - loop_induction_values + - serving_service_contract + - bounded_state_recurrence + inputs: + request.input_ids: + contract: + dtype: int64 + rank: 2 + shape: + - batch + - sequence + batch_layout: + kind: request_aligned + axis: 0 + role: + kind: runtime + version: '1.0' + role: prompt_tokens + source: + kind: request + required: true + request.max_iterations: + contract: + dtype: int64 + rank: 1 + shape: + - 1 + role: + kind: runtime + version: '1.0' + role: max_output_tokens + source: + kind: request + required: true + package.eos_ids: + contract: + dtype: int64 + rank: 1 + shape: + - E + role: + kind: opaque + source: + kind: literal + required: true + default: 127 + package.one_token: + contract: + dtype: int64 + rank: 1 + shape: + - batch + batch_layout: + kind: request_aligned + axis: 0 + role: + kind: opaque + source: + kind: literal + required: false + default: 1 + package.one_step: + contract: + dtype: int64 + rank: 1 + shape: + - 1 + role: + kind: opaque + source: + kind: literal + required: false + default: 1 + package.max_context: + contract: + dtype: int64 + rank: 1 + shape: + - 1 + role: + kind: opaque + source: + kind: literal + required: false + default: 8192 + package.cache_capacity: + contract: + dtype: int64 + rank: 1 + shape: + - 1 + role: + kind: opaque + source: + kind: literal + required: false + default: 16 + request.prompt_lengths: + contract: + dtype: int64 + rank: 1 + shape: + - batch + batch_layout: + kind: request_aligned + axis: 0 + role: + kind: opaque + source: + kind: application + name: prompt_lengths + required: false + default: -1 + request.eos_ids: + contract: + dtype: int64 + rank: 2 + shape: + - batch + - num_eos + batch_layout: + kind: request_aligned + axis: 0 + role: + kind: opaque + source: + kind: application + name: eos_ids + required: false + default: 127 + request.eos_lengths: + contract: + dtype: int64 + rank: 1 + shape: + - batch + batch_layout: + kind: request_aligned + axis: 0 + role: + kind: opaque + source: + kind: application + name: eos_lengths + required: false + default: 1 + request.row_max_iterations: + contract: + dtype: int64 + rank: 1 + shape: + - batch + batch_layout: + kind: request_aligned + axis: 0 + role: + kind: opaque + source: + kind: application + name: row_max_iterations + required: false + default: -1 + request.temperature: + contract: + dtype: float32 + rank: 1 + shape: + - batch + batch_layout: + kind: request_aligned + axis: 0 + role: + kind: runtime + version: '1.0' + role: sampling_temperature + source: + kind: request + required: false + default: 1.0 + request.top_k: + contract: + dtype: int64 + rank: 1 + shape: + - batch + batch_layout: + kind: request_aligned + axis: 0 + role: + kind: runtime + version: '1.0' + role: sampling_top_k + source: + kind: request + required: false + default: 1 + request.top_p: + contract: + dtype: float32 + rank: 1 + shape: + - batch + batch_layout: + kind: request_aligned + axis: 0 + role: + kind: runtime + version: '1.0' + role: sampling_top_p + source: + kind: request + required: false + default: 1.0 + request.min_p: + contract: + dtype: float32 + rank: 1 + shape: + - batch + batch_layout: + kind: request_aligned + axis: 0 + role: + kind: runtime + version: '1.0' + role: sampling_min_p + source: + kind: request + required: false + default: 0.0 + request.seed: + contract: + dtype: int64 + rank: 1 + shape: + - batch + batch_layout: + kind: request_aligned + axis: 0 + role: + kind: runtime + version: '1.0' + role: seed + source: + kind: request + required: false + default: 0 + request.rng_counter: + contract: + dtype: int64 + rank: 1 + shape: + - batch + batch_layout: + kind: request_aligned + axis: 0 + role: + kind: opaque + source: + kind: application + name: rng_counter + required: false + default: 0 + package.active: + contract: + dtype: bool + rank: 1 + shape: + - batch + batch_layout: + kind: request_aligned + axis: 0 + role: + kind: opaque + source: + kind: literal + required: false + default: true + package.not_done: + contract: + dtype: bool + rank: 1 + shape: + - batch + batch_layout: + kind: request_aligned + axis: 0 + role: + kind: opaque + source: + kind: literal + required: false + default: false + package.cache_lengths: + contract: + dtype: int64 + rank: 1 + shape: + - batch + batch_layout: + kind: request_aligned + axis: 0 + role: + kind: opaque + source: + kind: literal + required: false + default: 0 + package.zero_batch: + contract: + dtype: int64 + rank: 1 + shape: + - batch + batch_layout: + kind: request_aligned + axis: 0 + role: + kind: opaque + source: + kind: literal + required: false + default: 0 + outputs: + tokens: + contract: + dtype: int64 + rank: 2 + shape: + - batch + - generated_sequence + batch_layout: + kind: request_aligned + axis: 0 + role: tokens + stage: pre_adapter + components: + model: + implementation: + kind: onnx + artifact: model.onnx + ports: + roles: + input_ids: token_ids + position_ids: position_ids + logits: logits + token_sampler: + implementation: + kind: onnx + artifact: policies/token_sampler.onnx + ports: + inputs: + logits: + dtype: float32 + rank: 2 + shape: + - batch + - vocabulary + batch_layout: + kind: request_aligned + axis: 0 + temperature: + dtype: float32 + rank: 1 + shape: + - batch + batch_layout: + kind: request_aligned + axis: 0 + top_k: + dtype: int64 + rank: 1 + shape: + - batch + batch_layout: + kind: request_aligned + axis: 0 + top_p: + dtype: float32 + rank: 1 + shape: + - batch + batch_layout: + kind: request_aligned + axis: 0 + min_p: + dtype: float32 + rank: 1 + shape: + - batch + batch_layout: + kind: request_aligned + axis: 0 + seed: + dtype: int64 + rank: 1 + shape: + - batch + batch_layout: + kind: request_aligned + axis: 0 + counter: + dtype: int64 + rank: 1 + shape: + - batch + batch_layout: + kind: request_aligned + axis: 0 + active: + dtype: bool + rank: 1 + shape: + - batch + batch_layout: + kind: request_aligned + axis: 0 + done: + dtype: bool + rank: 1 + shape: + - batch + batch_layout: + kind: request_aligned + axis: 0 + outputs: + token: + dtype: int64 + rank: 1 + shape: + - batch + batch_layout: + kind: request_aligned + axis: 0 + next_counter: + dtype: int64 + rank: 1 + shape: + - batch + batch_layout: + kind: request_aligned + axis: 0 + contract: + id: onnx-genai.token-sampler + version: '2' + bindings: + logits: logits + token: token + temperature: temperature + top_k: top_k + top_p: top_p + min_p: min_p + active: active + done: done + seed: seed + counter: counter + next_counter: next_counter + parameters: + mode: seeded_stochastic + batching: per_row + inactive_rows: preserve + application_overridable: true + termination: + implementation: + kind: onnx + artifact: policies/termination.onnx + ports: + inputs: + tokens: + dtype: int64 + rank: 1 + shape: + - batch + batch_layout: + kind: request_aligned + axis: 0 + eos_ids: + dtype: int64 + rank: 2 + shape: + - batch + - num_eos + batch_layout: + kind: request_aligned + axis: 0 + eos_lengths: + dtype: int64 + rank: 1 + shape: + - batch + batch_layout: + kind: request_aligned + axis: 0 + iteration: + dtype: int64 + rank: 1 + shape: + - 1 + max_iterations: + dtype: int64 + rank: 1 + shape: + - batch + batch_layout: + kind: request_aligned + axis: 0 + active: + dtype: bool + rank: 1 + shape: + - batch + batch_layout: + kind: request_aligned + axis: 0 + outputs: + done: + dtype: bool + rank: 1 + shape: + - batch + batch_layout: + kind: request_aligned + axis: 0 + next_active: + dtype: bool + rank: 1 + shape: + - batch + batch_layout: + kind: request_aligned + axis: 0 + continue: + dtype: bool + rank: 1 + shape: + - 1 + contract: + id: onnx-genai.termination-predicate + version: '2' + bindings: + tokens: tokens + eos_ids: eos_ids + iteration: iteration + max_iterations: max_iterations + eos_lengths: eos_lengths + active: active + done: done + next_active: next_active + continue: continue + parameters: + batching: per_row + inactive_rows: preserve + token_state_update: + implementation: + kind: onnx + artifact: policies/token_state_update.onnx + ports: + inputs: + current: + dtype: int64 + rank: 2 + shape: + - batch + - 1 + batch_layout: + kind: request_aligned + axis: 0 + update: + dtype: int64 + rank: 2 + shape: + - batch + - 1 + batch_layout: + kind: request_aligned + axis: 0 + active: + dtype: bool + rank: 1 + shape: + - batch + batch_layout: + kind: request_aligned + axis: 0 + done: + dtype: bool + rank: 1 + shape: + - batch + batch_layout: + kind: request_aligned + axis: 0 + outputs: + next: + dtype: int64 + rank: 2 + shape: + - batch + - 1 + batch_layout: + kind: request_aligned + axis: 0 + contract: + id: onnx-genai.state-update + version: '2' + bindings: + current: current + update: update + active: active + done: done + next: next + parameters: + batching: per_row + inactive_rows: preserve + last_token_logits: + implementation: + kind: onnx + artifact: policies/last_token_logits.onnx + ports: + inputs: + logits: + dtype: float32 + rank: 3 + shape: + - batch + - sequence + - vocabulary + batch_layout: + kind: request_aligned + axis: 0 + outputs: + last_logits: + dtype: float32 + rank: 2 + shape: + - batch + - vocabulary + batch_layout: + kind: request_aligned + axis: 0 + decoder_state_initializer: + implementation: + kind: onnx + artifact: policies/decoder_state_initializer.onnx + ports: + inputs: + prompt_tokens: + dtype: int64 + rank: 2 + shape: + - batch + - prompt_sequence + batch_layout: + kind: request_aligned + axis: 0 + prompt_lengths: + dtype: int64 + rank: 1 + shape: + - batch + batch_layout: + kind: request_aligned + axis: 0 + outputs: + position_ids: + dtype: int64 + rank: 2 + shape: + - batch + - prompt_sequence + batch_layout: + kind: request_aligned + axis: 0 + body_position_ids: + dtype: int64 + rank: 2 + shape: + - batch + - 1 + batch_layout: + kind: request_aligned + axis: 0 + token_slot: + dtype: int64 + rank: 2 + shape: + - batch + - 1 + batch_layout: + kind: request_aligned + axis: 0 + generated_lengths: + dtype: int64 + rank: 1 + shape: + - batch + batch_layout: + kind: request_aligned + axis: 0 + cache_lengths: + dtype: int64 + rank: 1 + shape: + - batch + batch_layout: + kind: request_aligned + axis: 0 + write_indices: + dtype: int64 + rank: 1 + shape: + - batch + batch_layout: + kind: request_aligned + axis: 0 + key_cache.0: + dtype: float32 + rank: 3 + shape: + - batch + - 16 + - 8 + batch_layout: + kind: request_aligned + axis: 0 + value_cache.0: + dtype: float32 + rank: 3 + shape: + - batch + - 16 + - 8 + batch_layout: + kind: request_aligned + axis: 0 + decoder_step_update: + implementation: + kind: onnx + artifact: policies/decoder_step_update.onnx + ports: + inputs: + position_ids: + dtype: int64 + rank: 2 + shape: + - batch + - 1 + batch_layout: + kind: request_aligned + axis: 0 + outputs: + next_position_ids: + dtype: int64 + rank: 2 + shape: + - batch + - 1 + batch_layout: + kind: request_aligned + axis: 0 + cache_length_update: + implementation: + kind: onnx + artifact: policies/cache_length_update.onnx + ports: + inputs: + left: + dtype: int64 + rank: 1 + shape: + - batch + batch_layout: + kind: request_aligned + axis: 0 + right: + dtype: int64 + rank: 1 + shape: + - batch + batch_layout: + kind: request_aligned + axis: 0 + active: + dtype: bool + rank: 1 + shape: + - batch + batch_layout: + kind: request_aligned + axis: 0 + done: + dtype: bool + rank: 1 + shape: + - batch + batch_layout: + kind: request_aligned + axis: 0 + outputs: + total: + dtype: int64 + rank: 1 + shape: + - batch + batch_layout: + kind: request_aligned + axis: 0 + termination_batch_initializer: + implementation: + kind: onnx + artifact: policies/termination_batch_initializer.onnx + ports: + inputs: + input_eos_ids: + dtype: int64 + rank: 2 + shape: + - batch + - num_eos + batch_layout: + kind: request_aligned + axis: 0 + input_eos_lengths: + dtype: int64 + rank: 1 + shape: + - batch + batch_layout: + kind: request_aligned + axis: 0 + input_max_iterations: + dtype: int64 + rank: 1 + shape: + - batch + batch_layout: + kind: request_aligned + axis: 0 + fallback_max_iterations: + dtype: int64 + rank: 1 + shape: + - 1 + active: + dtype: bool + rank: 1 + shape: + - batch + batch_layout: + kind: request_aligned + axis: 0 + outputs: + row_eos_ids: + dtype: int64 + rank: 2 + shape: + - batch + - num_eos + batch_layout: + kind: request_aligned + axis: 0 + eos_lengths: + dtype: int64 + rank: 1 + shape: + - batch + batch_layout: + kind: request_aligned + axis: 0 + max_iterations: + dtype: int64 + rank: 1 + shape: + - batch + batch_layout: + kind: request_aligned + axis: 0 + token_to_slot: + implementation: + kind: onnx + artifact: policies/token_to_slot.onnx + ports: + inputs: + token: + dtype: int64 + rank: 1 + shape: + - batch + batch_layout: + kind: request_aligned + axis: 0 + outputs: + slot: + dtype: int64 + rank: 2 + shape: + - batch + - 1 + batch_layout: + kind: request_aligned + axis: 0 + generated_length_update: + implementation: + kind: onnx + artifact: policies/generated_length_update.onnx + ports: + inputs: + left: + dtype: int64 + rank: 1 + shape: + - batch + batch_layout: + kind: request_aligned + axis: 0 + right: + dtype: int64 + rank: 1 + shape: + - batch + batch_layout: + kind: request_aligned + axis: 0 + active: + dtype: bool + rank: 1 + shape: + - batch + batch_layout: + kind: request_aligned + axis: 0 + done: + dtype: bool + rank: 1 + shape: + - batch + batch_layout: + kind: request_aligned + axis: 0 + outputs: + total: + dtype: int64 + rank: 1 + shape: + - batch + batch_layout: + kind: request_aligned + axis: 0 + state: + token: + contract: + dtype: int64 + rank: 2 + shape: + - batch + - 1 + batch_layout: + kind: request_aligned + axis: 0 + scope: invocation + initializer: initializer.token_slot + recurrence: + kind: invariant + logits: + contract: + dtype: float32 + rank: 2 + shape: + - batch + - 128 + batch_layout: + kind: request_aligned + axis: 0 + scope: invocation + initializer: decoder.setup.last_logits + recurrence: + kind: invariant + generated_lengths: + contract: + dtype: int64 + rank: 1 + shape: + - batch + batch_layout: + kind: request_aligned + axis: 0 + class: semantic + scope: invocation + initializer: initializer.generated_lengths + recurrence: + kind: invariant + active: + contract: + dtype: bool + rank: 1 + shape: + - batch + batch_layout: + kind: request_aligned + axis: 0 + class: semantic + scope: invocation + initializer: package.active + recurrence: + kind: invariant + done: + contract: + dtype: bool + rank: 1 + shape: + - batch + batch_layout: + kind: request_aligned + axis: 0 + class: semantic + scope: invocation + initializer: package.not_done + recurrence: + kind: invariant + accepted_len: + contract: + dtype: int64 + rank: 1 + shape: + - batch + batch_layout: + kind: request_aligned + axis: 0 + class: semantic + scope: invocation + initializer: package.zero_batch + recurrence: + kind: invariant + cache_lengths: + contract: + dtype: int64 + rank: 1 + shape: + - batch + batch_layout: + kind: request_aligned + axis: 0 + class: semantic + scope: invocation + initializer: initializer.cache_lengths + recurrence: + kind: invariant + rng_counter: + contract: + dtype: int64 + rank: 1 + shape: + - batch + batch_layout: + kind: request_aligned + axis: 0 + scope: invocation + class: semantic + initializer: request.rng_counter + recurrence: + kind: invariant + position_ids: + contract: + dtype: int64 + rank: 2 + shape: + - batch + - 1 + batch_layout: + kind: request_aligned + axis: 0 + scope: invocation + initializer: initializer.body_position_ids + recurrence: + kind: invariant + cache_0: + contract: + dtype: float32 + rank: 3 + shape: + - batch + - 16 + - 8 + batch_layout: + kind: request_aligned + axis: 0 + scope: invocation + initializer: decoder.setup.updated_key_cache.0 + recurrence: + kind: invariant + management: runtime + release_boundary: invocation + service_group: decoder_cache + cache_1: + contract: + dtype: float32 + rank: 3 + shape: + - batch + - 16 + - 8 + batch_layout: + kind: request_aligned + axis: 0 + scope: invocation + initializer: decoder.setup.updated_value_cache.0 + recurrence: + kind: invariant + management: runtime + release_boundary: invocation + service_group: decoder_cache + serving: + active: active + done: done + accepted_len: accepted_len + state_service: + groups: + decoder_cache: + kind: full_attention + sequence_axis: 1 + layout: bsh + logical_lengths: cache_lengths + update: + kind: indexed_scatter + write_indices: cache_lengths + capacity: package.cache_capacity + write_indices_ports: + model: write_indices + kv_length_ports: + model: nonpad_kv_seqlen + aliasing: permitted + reuse: + prefix_reusable: true + evictable_prefix: false + ports: + model: + cache_0: + input: key_cache.0 + output: updated_key_cache.0 + role: key + layer: 0 + cache_1: + input: value_cache.0 + output: updated_value_cache.0 + role: value + layer: 0 + steps: + - kind: loop + setup: + - kind: invoke + component: decoder_state_initializer + inputs: + prompt_tokens: request.input_ids + prompt_lengths: request.prompt_lengths + outputs: + token_slot: initializer.token_slot + generated_lengths: initializer.generated_lengths + cache_lengths: initializer.cache_lengths + write_indices: initializer.write_indices + position_ids: initializer.position_ids + body_position_ids: initializer.body_position_ids + key_cache.0: initializer.key_cache.0 + value_cache.0: initializer.value_cache.0 + - kind: invoke + component: model + inputs: + input_ids: request.input_ids + key_cache.0: initializer.key_cache.0 + value_cache.0: initializer.value_cache.0 + position_ids: initializer.position_ids + write_indices: initializer.write_indices + nonpad_kv_seqlen: initializer.cache_lengths + outputs: + logits: decoder.setup.logits + updated_key_cache.0: decoder.setup.updated_key_cache.0 + updated_value_cache.0: decoder.setup.updated_value_cache.0 + - kind: invoke + component: termination_batch_initializer + inputs: + input_eos_ids: request.eos_ids + input_eos_lengths: request.eos_lengths + input_max_iterations: request.row_max_iterations + fallback_max_iterations: request.max_iterations + active: package.active + outputs: + row_eos_ids: termination.eos_ids + eos_lengths: termination.eos_lengths + max_iterations: termination.max_iterations + - kind: invoke + component: last_token_logits + inputs: + logits: decoder.setup.logits + outputs: + last_logits: decoder.setup.last_logits + steps: + - kind: invoke + component: token_sampler + inputs: + logits: logits + temperature: request.temperature + top_k: request.top_k + top_p: request.top_p + min_p: request.min_p + seed: request.seed + counter: rng_counter + active: active + done: done + outputs: + token: sample.body + next_counter: sample.next_counter + - kind: invoke + component: token_to_slot + inputs: + token: sample.body + outputs: + slot: sample.slot + - kind: invoke + component: generated_length_update + inputs: + left: generated_lengths + right: package.one_token + active: active + done: done + outputs: + total: token.next_lengths + - kind: invoke + component: generated_length_update + inputs: + left: package.zero_batch + right: package.one_token + active: active + done: done + outputs: + total: token.emitted_length + - kind: invoke + component: token_state_update + inputs: + current: token + update: sample.slot + active: active + done: done + outputs: + next: token.body + - kind: invoke + component: termination + inputs: + tokens: sample.body + eos_ids: termination.eos_ids + eos_lengths: termination.eos_lengths + iteration: loop.iteration + max_iterations: termination.max_iterations + active: active + outputs: + done: loop.done + continue: loop.continue + next_active: loop.next_active + - kind: invoke + component: cache_length_update + inputs: + left: cache_lengths + right: package.one_token + active: active + done: done + outputs: + total: cache_lengths.next + - kind: invoke + component: cache_length_update + inputs: + left: package.zero_batch + right: package.one_token + active: active + done: done + outputs: + total: accepted_len.next + - kind: emit + value: token.body + output: tokens + mode: append + valid_length: token.emitted_length + when: active + - kind: invoke + component: model + inputs: + input_ids: token.body + key_cache.0: cache_0 + value_cache.0: cache_1 + position_ids: position_ids + write_indices: cache_lengths + nonpad_kv_seqlen: cache_lengths.next + outputs: + logits: decoder.body.logits + updated_key_cache.0: decoder.body.updated_key_cache.0 + updated_value_cache.0: decoder.body.updated_value_cache.0 + - kind: invoke + component: last_token_logits + inputs: + logits: decoder.body.logits + outputs: + last_logits: decoder.body.last_logits + - kind: invoke + component: decoder_step_update + inputs: + position_ids: position_ids + outputs: + next_position_ids: decoder_step.body_position_ids + continue_when: active + max_iterations: request.max_iterations + carried: + - cell: token + next: token.body + - cell: logits + next: decoder.body.last_logits + - cell: generated_lengths + next: token.next_lengths + - cell: active + next: loop.next_active + - cell: done + next: loop.done + - cell: cache_lengths + next: cache_lengths.next + - cell: accepted_len + next: accepted_len.next + - cell: rng_counter + next: sample.next_counter + - cell: position_ids + next: decoder_step.body_position_ids + - cell: cache_0 + next: decoder.body.updated_key_cache.0 + - cell: cache_1 + next: decoder.body.updated_value_cache.0 + termination: generation_eos + iteration: + value: loop.iteration + contract: + dtype: int64 + rank: 1 + shape: + - 1 diff --git a/tests/fixtures/onnx_genai_workflows/tts/inference_metadata.yaml b/tests/fixtures/onnx_genai_workflows/tts/inference_metadata.yaml new file mode 100644 index 000000000..113cbbe00 --- /dev/null +++ b/tests/fixtures/onnx_genai_workflows/tts/inference_metadata.yaml @@ -0,0 +1,2277 @@ +schema_version: v1 +pipeline: + workflow: + manifest: + capabilities: + - workflow_ssa + - linear_effects + - nested_control_flow + - loop_induction_values + - typed_emit + - serving_service_contract + - bounded_state_recurrence + inputs: + request.prompt_tokens: + contract: + dtype: int64 + rank: 2 + shape: + - batch + - text_sequence_len + batch_layout: + kind: request_aligned + axis: 0 + role: + kind: runtime + version: '1.0' + role: prompt_tokens + source: + kind: request + required: true + request.max_iterations: + contract: + dtype: int64 + rank: 1 + shape: + - 1 + role: + kind: runtime + version: '1.0' + role: max_output_tokens + source: + kind: request + required: true + package.false: + contract: + dtype: bool + rank: 1 + shape: + - batch + batch_layout: + kind: request_aligned + axis: 0 + role: + kind: opaque + source: + kind: literal + required: false + default: false + package.zero_scalar: + contract: + dtype: int64 + rank: 0 + shape: [] + role: + kind: opaque + source: + kind: literal + required: false + default: 0 + package.one_scalar: + contract: + dtype: int64 + rank: 0 + shape: [] + role: + kind: opaque + source: + kind: literal + required: false + default: 1 + package.remaining_groups: + contract: + dtype: int64 + rank: 1 + shape: + - 1 + role: + kind: opaque + source: + kind: literal + required: false + default: 2 + package.predictor_context_limit: + contract: + dtype: int64 + rank: 1 + shape: + - 1 + role: + kind: opaque + source: + kind: literal + required: false + default: 4 + package.predictor_mask_limit: + contract: + dtype: int64 + rank: 1 + shape: + - 1 + role: + kind: opaque + source: + kind: literal + required: false + default: 5 + package.talker_context_limit: + contract: + dtype: int64 + rank: 1 + shape: + - 1 + role: + kind: opaque + source: + kind: literal + required: false + default: 128 + package.one_control: + contract: + dtype: int64 + rank: 1 + shape: + - 1 + role: + kind: opaque + source: + kind: literal + required: false + default: 1 + package.zero_batch: + contract: + dtype: int64 + rank: 1 + shape: + - batch + batch_layout: + kind: request_aligned + axis: 0 + role: + kind: opaque + source: + kind: literal + required: false + default: 0 + package.one_batch: + contract: + dtype: int64 + rank: 1 + shape: + - batch + batch_layout: + kind: request_aligned + axis: 0 + role: + kind: opaque + source: + kind: literal + required: false + default: 1 + package.true: + contract: + dtype: bool + rank: 1 + shape: + - batch + batch_layout: + kind: request_aligned + axis: 0 + role: + kind: opaque + source: + kind: literal + required: false + default: true + package.setup_predictor_iteration_0: + contract: + dtype: int64 + rank: 0 + shape: [] + role: + kind: opaque + source: + kind: literal + required: false + default: 0 + package.setup_predictor_iteration_1: + contract: + dtype: int64 + rank: 0 + shape: [] + role: + kind: opaque + source: + kind: literal + required: false + default: 1 + package.loop_1_active: + contract: + dtype: bool + rank: 1 + shape: + - 1 + role: + kind: opaque + source: + kind: literal + required: false + default: true + package.loop_0_active: + contract: + dtype: bool + rank: 1 + shape: + - 1 + role: + kind: opaque + source: + kind: literal + required: false + default: true + outputs: + waveform: + contract: + dtype: float32 + rank: 3 + shape: + - batch + - 1 + - frames + batch_layout: + kind: request_aligned + axis: 0 + role: audio + stage: post_adapter + components: + talker: + implementation: + kind: onnx + artifact: talker/model.onnx + ports: + roles: + inputs_embeds: inputs_embeds + attention_mask: attention_mask + position_ids: position_ids + logits: logits + last_hidden_state: hidden_states + code_predictor: + implementation: + kind: onnx + artifact: code_predictor/model.onnx + ports: + roles: + inputs_embeds: inputs_embeds + attention_mask: attention_mask + position_ids: position_ids + logits: logits + embedding: + implementation: + kind: onnx + artifact: embedding/model.onnx + talker_step_embedder: + implementation: + kind: onnx + artifact: talker_step_embedder/model.onnx + ports: + roles: + inputs_embeds: inputs_embeds + talker_prefill_embedder: + implementation: + kind: onnx + artifact: talker_prefill_embedder/model.onnx + code_predictor_prefill: + implementation: + kind: onnx + artifact: code_predictor_prefill/model.onnx + ports: + roles: + inputs_embeds: inputs_embeds + code_predictor_step_embedder: + implementation: + kind: onnx + artifact: code_predictor_step_embedder/model.onnx + ports: + roles: + inputs_embeds: inputs_embeds + code_predictor_indices: + implementation: + kind: onnx + artifact: code_predictor_indices/model.onnx + talker_text_step: + implementation: + kind: onnx + artifact: talker_text_step/model.onnx + codec: + implementation: + kind: onnx + artifact: codec/model.onnx + last_token_logits: + implementation: + kind: onnx + artifact: policies/last_token_logits.onnx + ports: + inputs: + logits: + dtype: float32 + rank: 3 + shape: + - batch + - sequence + - vocabulary + batch_layout: + kind: request_aligned + axis: 0 + outputs: + last_logits: + dtype: float32 + rank: 2 + shape: + - batch + - vocabulary + batch_layout: + kind: request_aligned + axis: 0 + setup_talker_sampler: + implementation: + kind: onnx + artifact: policies/setup_talker_sampler.onnx + ports: + inputs: + logits: + dtype: float32 + rank: 2 + shape: + - batch + - vocabulary + batch_layout: + kind: request_aligned + axis: 0 + outputs: + token: + dtype: int64 + rank: 1 + shape: + - batch + batch_layout: + kind: request_aligned + axis: 0 + contract: + id: onnx-genai.token-sampler + version: '1' + bindings: + logits: logits + token: token + parameters: + mode: greedy + application_overridable: true + setup_predictor_sampler: + implementation: + kind: onnx + artifact: policies/setup_predictor_sampler.onnx + ports: + inputs: + logits: + dtype: float32 + rank: 2 + shape: + - batch + - vocabulary + batch_layout: + kind: request_aligned + axis: 0 + outputs: + token: + dtype: int64 + rank: 1 + shape: + - batch + batch_layout: + kind: request_aligned + axis: 0 + contract: + id: onnx-genai.token-sampler + version: '1' + bindings: + logits: logits + token: token + parameters: + mode: greedy + application_overridable: true + talker_sampler: + implementation: + kind: onnx + artifact: policies/talker_sampler.onnx + ports: + inputs: + logits: + dtype: float32 + rank: 2 + shape: + - batch + - vocabulary + batch_layout: + kind: request_aligned + axis: 0 + outputs: + token: + dtype: int64 + rank: 1 + shape: + - batch + batch_layout: + kind: request_aligned + axis: 0 + contract: + id: onnx-genai.token-sampler + version: '1' + bindings: + logits: logits + token: token + parameters: + mode: greedy + application_overridable: true + predictor_prefill_sampler: + implementation: + kind: onnx + artifact: policies/predictor_prefill_sampler.onnx + ports: + inputs: + logits: + dtype: float32 + rank: 2 + shape: + - batch + - vocabulary + batch_layout: + kind: request_aligned + axis: 0 + outputs: + token: + dtype: int64 + rank: 1 + shape: + - batch + batch_layout: + kind: request_aligned + axis: 0 + contract: + id: onnx-genai.token-sampler + version: '1' + bindings: + logits: logits + token: token + parameters: + mode: greedy + application_overridable: true + predictor_body_sampler: + implementation: + kind: onnx + artifact: policies/predictor_body_sampler.onnx + ports: + inputs: + logits: + dtype: float32 + rank: 2 + shape: + - batch + - vocabulary + batch_layout: + kind: request_aligned + axis: 0 + outputs: + token: + dtype: int64 + rank: 1 + shape: + - batch + batch_layout: + kind: request_aligned + axis: 0 + contract: + id: onnx-genai.token-sampler + version: '1' + bindings: + logits: logits + token: token + parameters: + mode: greedy + application_overridable: true + continue_predicate: + implementation: + kind: onnx + artifact: policies/continue_predicate.onnx + ports: + inputs: + done: + dtype: bool + rank: 1 + shape: + - batch + batch_layout: + kind: request_aligned + axis: 0 + outputs: + continue: + dtype: bool + rank: 1 + shape: + - 1 + tts_state_initializer: + implementation: + kind: onnx + artifact: policies/tts_state_initializer.onnx + ports: + inputs: + prompt_tokens: + dtype: int64 + rank: 2 + shape: + - batch + - sequence + batch_layout: + kind: request_aligned + axis: 0 + outputs: + frame_codes: + dtype: int64 + rank: 2 + shape: + - batch + - 4 + batch_layout: + kind: request_aligned + axis: 0 + token_slot: + dtype: int64 + rank: 2 + shape: + - batch + - 1 + batch_layout: + kind: request_aligned + axis: 0 + code_history: + dtype: int64 + rank: 3 + shape: + - batch + - 0 + - 4 + batch_layout: + kind: request_aligned + axis: 0 + token_to_slot: + implementation: + kind: onnx + artifact: policies/token_to_slot.onnx + ports: + inputs: + token: + dtype: int64 + rank: 1 + shape: + - batch + batch_layout: + kind: request_aligned + axis: 0 + outputs: + slot: + dtype: int64 + rank: 2 + shape: + - batch + - 1 + batch_layout: + kind: request_aligned + axis: 0 + code_frame_update: + implementation: + kind: onnx + artifact: policies/code_frame_update.onnx + ports: + inputs: + frame_codes: + dtype: int64 + rank: 2 + shape: + - batch + - 4 + batch_layout: + kind: request_aligned + axis: 0 + token: + dtype: int64 + rank: 1 + shape: + - batch + batch_layout: + kind: request_aligned + axis: 0 + index: + dtype: int64 + rank: 0 + shape: [] + outputs: + next_frame: + dtype: int64 + rank: 2 + shape: + - batch + - 4 + batch_layout: + kind: request_aligned + axis: 0 + code_history_append: + implementation: + kind: onnx + artifact: policies/code_history_append.onnx + ports: + inputs: + history: + dtype: int64 + rank: 3 + shape: + - batch + - frames + - 4 + batch_layout: + kind: request_aligned + axis: 0 + frame: + dtype: int64 + rank: 2 + shape: + - batch + - 4 + batch_layout: + kind: request_aligned + axis: 0 + outputs: + next_history: + dtype: int64 + rank: 3 + shape: + - batch + - frames + 1 + - 4 + batch_layout: + kind: request_aligned + axis: 0 + cache_length_update: + implementation: + kind: onnx + artifact: policies/cache_length_update.onnx + ports: + inputs: + left: + dtype: int64 + rank: 1 + shape: + - batch + batch_layout: + kind: request_aligned + axis: 0 + right: + dtype: int64 + rank: 1 + shape: + - batch + batch_layout: + kind: request_aligned + axis: 0 + outputs: + total: + dtype: int64 + rank: 1 + shape: + - batch + batch_layout: + kind: request_aligned + axis: 0 + talker_state_initializer: + implementation: + kind: onnx + artifact: policies/talker_state_initializer.onnx + ports: + inputs: + prefill_embeds: + dtype: float32 + rank: 3 + shape: + - batch + - prefill_sequence + - 8 + batch_layout: + kind: request_aligned + axis: 0 + outputs: + attention_mask: + dtype: int64 + rank: 2 + shape: + - batch + - past_seq_len + seq_len + batch_layout: + kind: request_aligned + axis: 0 + position_ids: + dtype: int64 + rank: 3 + shape: + - 3 + - batch + - sequence_len + body_attention_mask: + dtype: int64 + rank: 2 + shape: + - batch + - prefill_sequence + 1 + batch_layout: + kind: request_aligned + axis: 0 + body_position_ids: + dtype: int64 + rank: 3 + shape: + - 3 + - batch + - 1 + past_key_values.0.key: + dtype: float32 + rank: 4 + shape: + - batch + - 1 + - past_sequence_len + - 4 + batch_layout: + kind: request_aligned + axis: 0 + past_key_values.0.value: + dtype: float32 + rank: 4 + shape: + - batch + - 1 + - past_sequence_len + - 4 + batch_layout: + kind: request_aligned + axis: 0 + predictor_state_initializer: + implementation: + kind: onnx + artifact: policies/predictor_state_initializer.onnx + ports: + inputs: + prefill_embeds: + dtype: float32 + rank: 3 + shape: + - batch + - prefill_sequence + - 8 + batch_layout: + kind: request_aligned + axis: 0 + outputs: + attention_mask: + dtype: int64 + rank: 2 + shape: + - batch + - past_seq_len + seq_len + batch_layout: + kind: request_aligned + axis: 0 + position_ids: + dtype: int64 + rank: 2 + shape: + - batch + - sequence_len + batch_layout: + kind: request_aligned + axis: 0 + body_attention_mask: + dtype: int64 + rank: 2 + shape: + - batch + - prefill_sequence + 1 + batch_layout: + kind: request_aligned + axis: 0 + body_position_ids: + dtype: int64 + rank: 2 + shape: + - batch + - 1 + batch_layout: + kind: request_aligned + axis: 0 + past_key_values.0.key: + dtype: float32 + rank: 4 + shape: + - batch + - 8 + - past_sequence_len + - 128 + batch_layout: + kind: request_aligned + axis: 0 + past_key_values.0.value: + dtype: float32 + rank: 4 + shape: + - batch + - 8 + - past_sequence_len + - 128 + batch_layout: + kind: request_aligned + axis: 0 + past_key_values.1.key: + dtype: float32 + rank: 4 + shape: + - batch + - 8 + - past_sequence_len + - 128 + batch_layout: + kind: request_aligned + axis: 0 + past_key_values.1.value: + dtype: float32 + rank: 4 + shape: + - batch + - 8 + - past_sequence_len + - 128 + batch_layout: + kind: request_aligned + axis: 0 + past_key_values.2.key: + dtype: float32 + rank: 4 + shape: + - batch + - 8 + - past_sequence_len + - 128 + batch_layout: + kind: request_aligned + axis: 0 + past_key_values.2.value: + dtype: float32 + rank: 4 + shape: + - batch + - 8 + - past_sequence_len + - 128 + batch_layout: + kind: request_aligned + axis: 0 + past_key_values.3.key: + dtype: float32 + rank: 4 + shape: + - batch + - 8 + - past_sequence_len + - 128 + batch_layout: + kind: request_aligned + axis: 0 + past_key_values.3.value: + dtype: float32 + rank: 4 + shape: + - batch + - 8 + - past_sequence_len + - 128 + batch_layout: + kind: request_aligned + axis: 0 + past_key_values.4.key: + dtype: float32 + rank: 4 + shape: + - batch + - 8 + - past_sequence_len + - 128 + batch_layout: + kind: request_aligned + axis: 0 + past_key_values.4.value: + dtype: float32 + rank: 4 + shape: + - batch + - 8 + - past_sequence_len + - 128 + batch_layout: + kind: request_aligned + axis: 0 + talker_step_update: + implementation: + kind: onnx + artifact: policies/talker_step_update.onnx + ports: + inputs: + attention_mask: + dtype: int64 + rank: 2 + shape: + - batch + - context + batch_layout: + kind: request_aligned + axis: 0 + position_ids: + dtype: int64 + rank: 3 + shape: + - 3 + - batch + - 1 + outputs: + next_attention_mask: + dtype: int64 + rank: 2 + shape: + - batch + - context + 1 + batch_layout: + kind: request_aligned + axis: 0 + next_position_ids: + dtype: int64 + rank: 3 + shape: + - 3 + - batch + - 1 + predictor_step_update: + implementation: + kind: onnx + artifact: policies/predictor_step_update.onnx + ports: + inputs: + attention_mask: + dtype: int64 + rank: 2 + shape: + - batch + - context + batch_layout: + kind: request_aligned + axis: 0 + position_ids: + dtype: int64 + rank: 2 + shape: + - batch + - 1 + batch_layout: + kind: request_aligned + axis: 0 + outputs: + next_attention_mask: + dtype: int64 + rank: 2 + shape: + - batch + - context + 1 + batch_layout: + kind: request_aligned + axis: 0 + next_position_ids: + dtype: int64 + rank: 2 + shape: + - batch + - 1 + batch_layout: + kind: request_aligned + axis: 0 + codec_layout: + implementation: + kind: onnx + artifact: policies/codec_layout.onnx + ports: + inputs: + history: + dtype: int64 + rank: 3 + shape: + - batch + - frames + - 4 + batch_layout: + kind: request_aligned + axis: 0 + outputs: + codes: + dtype: int64 + rank: 3 + shape: + - batch + - 4 + - frames + batch_layout: + kind: request_aligned + axis: 0 + state: + last_frame: + contract: + dtype: int64 + rank: 2 + shape: + - batch + - 4 + batch_layout: + kind: request_aligned + axis: 0 + scope: invocation + initializer: setup.predictor.remaining_1.frame + recurrence: + kind: invariant + history: + contract: + dtype: int64 + rank: 3 + shape: + - batch + - frames + - 4 + batch_layout: + kind: request_aligned + axis: 0 + scope: invocation + initializer: history.setup + recurrence: + kind: growing + axis: 1 + increment: package.one_control + max: package.talker_context_limit + talker_mask: + contract: + dtype: int64 + rank: 2 + shape: + - batch + - talker_context + batch_layout: + kind: request_aligned + axis: 0 + scope: invocation + initializer: talker.initializer.body_attention_mask + recurrence: + kind: growing + axis: 1 + increment: package.one_control + max: package.talker_context_limit + talker_position: + contract: + dtype: int64 + rank: 3 + shape: + - 3 + - batch + - 1 + scope: invocation + initializer: talker.initializer.body_position_ids + recurrence: + kind: invariant + frame: + contract: + dtype: int64 + rank: 2 + shape: + - batch + - 4 + batch_layout: + kind: request_aligned + axis: 0 + scope: invocation + initializer: frame.frame_prefill + recurrence: + kind: invariant + code_token: + contract: + dtype: int64 + rank: 1 + shape: + - batch + batch_layout: + kind: request_aligned + axis: 0 + scope: invocation + initializer: frame.group1 + recurrence: + kind: invariant + predictor_mask: + contract: + dtype: int64 + rank: 2 + shape: + - batch + - predictor_context + batch_layout: + kind: request_aligned + axis: 0 + scope: invocation + initializer: frame.predictor.initializer.body_attention_mask + recurrence: + kind: growing + axis: 1 + increment: package.one_control + max: package.predictor_mask_limit + predictor_position: + contract: + dtype: int64 + rank: 2 + shape: + - batch + - 1 + batch_layout: + kind: request_aligned + axis: 0 + scope: invocation + initializer: frame.predictor.initializer.body_position_ids + recurrence: + kind: invariant + active: + contract: + dtype: bool + rank: 1 + shape: + - batch + batch_layout: + kind: request_aligned + axis: 0 + class: semantic + scope: invocation + initializer: package.true + recurrence: + kind: invariant + done: + contract: + dtype: bool + rank: 1 + shape: + - batch + batch_layout: + kind: request_aligned + axis: 0 + class: semantic + scope: invocation + initializer: package.false + recurrence: + kind: invariant + accepted_len: + contract: + dtype: int64 + rank: 1 + shape: + - batch + batch_layout: + kind: request_aligned + axis: 0 + class: semantic + scope: invocation + initializer: package.zero_batch + recurrence: + kind: invariant + talker_cache_lengths: + contract: + dtype: int64 + rank: 1 + shape: + - batch + batch_layout: + kind: request_aligned + axis: 0 + class: semantic + scope: invocation + initializer: package.zero_batch + recurrence: + kind: invariant + predictor_cache_lengths: + contract: + dtype: int64 + rank: 1 + shape: + - batch + batch_layout: + kind: request_aligned + axis: 0 + class: semantic + scope: invocation + initializer: package.zero_batch + recurrence: + kind: invariant + talker_cache_0: + contract: + dtype: float32 + rank: 4 + shape: + - batch + - 1 + - past_sequence_len + - 4 + batch_layout: + kind: request_aligned + axis: 0 + class: semantic + scope: invocation + initializer: talker.setup.present.0.key + recurrence: + kind: bounded + axis: 2 + max: package.talker_context_limit + service_group: talker_cache + management: runtime + release_boundary: invocation + talker_cache_1: + contract: + dtype: float32 + rank: 4 + shape: + - batch + - 1 + - past_sequence_len + - 4 + batch_layout: + kind: request_aligned + axis: 0 + class: semantic + scope: invocation + initializer: talker.setup.present.0.value + recurrence: + kind: bounded + axis: 2 + max: package.talker_context_limit + service_group: talker_cache + management: runtime + release_boundary: invocation + predictor_cache_0: + contract: + dtype: float32 + rank: 4 + shape: + - batch + - 8 + - past_sequence_len + - 128 + batch_layout: + kind: request_aligned + axis: 0 + class: semantic + scope: invocation + initializer: frame.predictor.present.0.key + recurrence: + kind: bounded + axis: 2 + max: package.predictor_context_limit + service_group: predictor_cache + management: runtime + release_boundary: invocation + predictor_cache_1: + contract: + dtype: float32 + rank: 4 + shape: + - batch + - 8 + - past_sequence_len + - 128 + batch_layout: + kind: request_aligned + axis: 0 + class: semantic + scope: invocation + initializer: frame.predictor.present.0.value + recurrence: + kind: bounded + axis: 2 + max: package.predictor_context_limit + service_group: predictor_cache + management: runtime + release_boundary: invocation + predictor_cache_2: + contract: + dtype: float32 + rank: 4 + shape: + - batch + - 8 + - past_sequence_len + - 128 + batch_layout: + kind: request_aligned + axis: 0 + class: semantic + scope: invocation + initializer: frame.predictor.present.1.key + recurrence: + kind: bounded + axis: 2 + max: package.predictor_context_limit + service_group: predictor_cache + management: runtime + release_boundary: invocation + predictor_cache_3: + contract: + dtype: float32 + rank: 4 + shape: + - batch + - 8 + - past_sequence_len + - 128 + batch_layout: + kind: request_aligned + axis: 0 + class: semantic + scope: invocation + initializer: frame.predictor.present.1.value + recurrence: + kind: bounded + axis: 2 + max: package.predictor_context_limit + service_group: predictor_cache + management: runtime + release_boundary: invocation + predictor_cache_4: + contract: + dtype: float32 + rank: 4 + shape: + - batch + - 8 + - past_sequence_len + - 128 + batch_layout: + kind: request_aligned + axis: 0 + class: semantic + scope: invocation + initializer: frame.predictor.present.2.key + recurrence: + kind: bounded + axis: 2 + max: package.predictor_context_limit + service_group: predictor_cache + management: runtime + release_boundary: invocation + predictor_cache_5: + contract: + dtype: float32 + rank: 4 + shape: + - batch + - 8 + - past_sequence_len + - 128 + batch_layout: + kind: request_aligned + axis: 0 + class: semantic + scope: invocation + initializer: frame.predictor.present.2.value + recurrence: + kind: bounded + axis: 2 + max: package.predictor_context_limit + service_group: predictor_cache + management: runtime + release_boundary: invocation + predictor_cache_6: + contract: + dtype: float32 + rank: 4 + shape: + - batch + - 8 + - past_sequence_len + - 128 + batch_layout: + kind: request_aligned + axis: 0 + class: semantic + scope: invocation + initializer: frame.predictor.present.3.key + recurrence: + kind: bounded + axis: 2 + max: package.predictor_context_limit + service_group: predictor_cache + management: runtime + release_boundary: invocation + predictor_cache_7: + contract: + dtype: float32 + rank: 4 + shape: + - batch + - 8 + - past_sequence_len + - 128 + batch_layout: + kind: request_aligned + axis: 0 + class: semantic + scope: invocation + initializer: frame.predictor.present.3.value + recurrence: + kind: bounded + axis: 2 + max: package.predictor_context_limit + service_group: predictor_cache + management: runtime + release_boundary: invocation + predictor_cache_8: + contract: + dtype: float32 + rank: 4 + shape: + - batch + - 8 + - past_sequence_len + - 128 + batch_layout: + kind: request_aligned + axis: 0 + class: semantic + scope: invocation + initializer: frame.predictor.present.4.key + recurrence: + kind: bounded + axis: 2 + max: package.predictor_context_limit + service_group: predictor_cache + management: runtime + release_boundary: invocation + predictor_cache_9: + contract: + dtype: float32 + rank: 4 + shape: + - batch + - 8 + - past_sequence_len + - 128 + batch_layout: + kind: request_aligned + axis: 0 + class: semantic + scope: invocation + initializer: frame.predictor.present.4.value + recurrence: + kind: bounded + axis: 2 + max: package.predictor_context_limit + service_group: predictor_cache + management: runtime + release_boundary: invocation + loop_1_active: + contract: + dtype: bool + rank: 1 + shape: + - 1 + scope: invocation + initializer: package.loop_1_active + recurrence: + kind: invariant + loop_0_active: + contract: + dtype: bool + rank: 1 + shape: + - 1 + scope: invocation + initializer: package.loop_0_active + recurrence: + kind: invariant + serving: + active: active + done: done + accepted_len: accepted_len + state_service: + groups: + talker_cache: + kind: full_attention + sequence_axis: 2 + layout: bnsh + aliasing: forbidden + reuse: + prefix_reusable: true + evictable_prefix: false + capabilities: + snapshot: true + fork: true + ports: + talker: + talker_cache_0: + input: past_key_values.0.key + output: present.0.key + role: key + layer: 0 + talker_cache_1: + input: past_key_values.0.value + output: present.0.value + role: value + layer: 0 + logical_lengths: talker_cache_lengths + predictor_cache: + kind: full_attention + sequence_axis: 2 + layout: bnsh + aliasing: forbidden + reuse: + prefix_reusable: true + evictable_prefix: false + capabilities: + snapshot: true + fork: true + ports: + code_predictor: + predictor_cache_0: + input: past_key_values.0.key + output: present.0.key + role: key + layer: 0 + predictor_cache_1: + input: past_key_values.0.value + output: present.0.value + role: value + layer: 0 + predictor_cache_2: + input: past_key_values.1.key + output: present.1.key + role: key + layer: 1 + predictor_cache_3: + input: past_key_values.1.value + output: present.1.value + role: value + layer: 1 + predictor_cache_4: + input: past_key_values.2.key + output: present.2.key + role: key + layer: 2 + predictor_cache_5: + input: past_key_values.2.value + output: present.2.value + role: value + layer: 2 + predictor_cache_6: + input: past_key_values.3.key + output: present.3.key + role: key + layer: 3 + predictor_cache_7: + input: past_key_values.3.value + output: present.3.value + role: value + layer: 3 + predictor_cache_8: + input: past_key_values.4.key + output: present.4.key + role: key + layer: 4 + predictor_cache_9: + input: past_key_values.4.value + output: present.4.value + role: value + layer: 4 + logical_lengths: predictor_cache_lengths + steps: + - kind: loop + setup: + - kind: invoke + component: tts_state_initializer + inputs: + prompt_tokens: request.prompt_tokens + outputs: + frame_codes: initializer.frame_codes + token_slot: initializer.token_slot + code_history: initializer.code_history + - kind: invoke + component: talker_prefill_embedder + inputs: + text_ids: request.prompt_tokens + outputs: + prefill_embeds: tts.prefill_embeds + trailing_text_embeds: tts.trailing_text_embeds + - kind: invoke + component: talker_state_initializer + inputs: + prefill_embeds: tts.prefill_embeds + outputs: + attention_mask: talker.initializer.attention_mask + position_ids: talker.initializer.position_ids + body_attention_mask: talker.initializer.body_attention_mask + body_position_ids: talker.initializer.body_position_ids + past_key_values.0.key: talker.initializer.past_key_values.0.key + past_key_values.0.value: talker.initializer.past_key_values.0.value + - kind: invoke + component: talker + inputs: + inputs_embeds: tts.prefill_embeds + attention_mask: talker.initializer.attention_mask + position_ids: talker.initializer.position_ids + past_key_values.0.key: talker.initializer.past_key_values.0.key + past_key_values.0.value: talker.initializer.past_key_values.0.value + outputs: + logits: talker.setup.logits + last_hidden_state: talker.setup.hidden + present.0.key: talker.setup.present.0.key + present.0.value: talker.setup.present.0.value + - kind: invoke + component: last_token_logits + inputs: + logits: talker.setup.logits + outputs: + last_logits: setup.group0_logits + - kind: invoke + component: setup_talker_sampler + inputs: + logits: setup.group0_logits + outputs: + token: setup.group0 + - kind: invoke + component: token_to_slot + inputs: + token: setup.group0 + outputs: + slot: setup.group0_slot + - kind: invoke + component: embedding + inputs: + text_ids: request.prompt_tokens + codec_ids: setup.group0_slot + outputs: + text_embeds: setup.unused_text_embeds + codec_embeds: setup.group0_embed + - kind: invoke + component: code_predictor_prefill + inputs: + talker_hidden: talker.setup.hidden + group_0_embed: setup.group0_embed + outputs: + inputs_embeds: setup.predictor_prefill + - kind: invoke + component: predictor_state_initializer + inputs: + prefill_embeds: setup.predictor_prefill + outputs: + attention_mask: setup.predictor.initializer.attention_mask + position_ids: setup.predictor.initializer.position_ids + body_attention_mask: setup.predictor.initializer.body_attention_mask + body_position_ids: setup.predictor.initializer.body_position_ids + past_key_values.0.key: setup.predictor.initializer.past_key_values.0.key + past_key_values.0.value: setup.predictor.initializer.past_key_values.0.value + past_key_values.1.key: setup.predictor.initializer.past_key_values.1.key + past_key_values.1.value: setup.predictor.initializer.past_key_values.1.value + past_key_values.2.key: setup.predictor.initializer.past_key_values.2.key + past_key_values.2.value: setup.predictor.initializer.past_key_values.2.value + past_key_values.3.key: setup.predictor.initializer.past_key_values.3.key + past_key_values.3.value: setup.predictor.initializer.past_key_values.3.value + past_key_values.4.key: setup.predictor.initializer.past_key_values.4.key + past_key_values.4.value: setup.predictor.initializer.past_key_values.4.value + - kind: invoke + component: code_predictor + inputs: + inputs_embeds: setup.predictor_prefill + step_index: package.zero_scalar + attention_mask: setup.predictor.initializer.attention_mask + position_ids: setup.predictor.initializer.position_ids + past_key_values.0.key: setup.predictor.initializer.past_key_values.0.key + past_key_values.0.value: setup.predictor.initializer.past_key_values.0.value + past_key_values.1.key: setup.predictor.initializer.past_key_values.1.key + past_key_values.1.value: setup.predictor.initializer.past_key_values.1.value + past_key_values.2.key: setup.predictor.initializer.past_key_values.2.key + past_key_values.2.value: setup.predictor.initializer.past_key_values.2.value + past_key_values.3.key: setup.predictor.initializer.past_key_values.3.key + past_key_values.3.value: setup.predictor.initializer.past_key_values.3.value + past_key_values.4.key: setup.predictor.initializer.past_key_values.4.key + past_key_values.4.value: setup.predictor.initializer.past_key_values.4.value + outputs: + logits: setup.predictor.logits + codec_embeddings: setup.predictor.codec_embeddings + present.0.key: setup.predictor.present.0.key + present.0.value: setup.predictor.present.0.value + present.1.key: setup.predictor.present.1.key + present.1.value: setup.predictor.present.1.value + present.2.key: setup.predictor.present.2.key + present.2.value: setup.predictor.present.2.value + present.3.key: setup.predictor.present.3.key + present.3.value: setup.predictor.present.3.value + present.4.key: setup.predictor.present.4.key + present.4.value: setup.predictor.present.4.value + - kind: invoke + component: last_token_logits + inputs: + logits: setup.predictor.logits + outputs: + last_logits: setup.group1_logits + - kind: invoke + component: setup_predictor_sampler + inputs: + logits: setup.group1_logits + outputs: + token: setup.group1 + - kind: invoke + component: code_frame_update + inputs: + frame_codes: initializer.frame_codes + token: setup.group0 + index: package.zero_scalar + outputs: + next_frame: setup.frame_group0 + - kind: invoke + component: code_frame_update + inputs: + frame_codes: setup.frame_group0 + token: setup.group1 + index: package.one_scalar + outputs: + next_frame: setup.frame_prefill + - kind: invoke + component: code_predictor_indices + inputs: + iteration: package.setup_predictor_iteration_0 + outputs: + embedding_index: setup.predictor.remaining_0.embedding_index + step_index: setup.predictor.remaining_0.step_index + frame_index: setup.predictor.remaining_0.frame_index + - kind: invoke + component: code_predictor_step_embedder + inputs: + codec_embeddings: setup.predictor.codec_embeddings + token: setup.group1 + embedding_index: setup.predictor.remaining_0.embedding_index + outputs: + inputs_embeds: setup.predictor.remaining_0.inputs_embeds + - kind: invoke + component: code_predictor + inputs: + inputs_embeds: setup.predictor.remaining_0.inputs_embeds + step_index: setup.predictor.remaining_0.step_index + attention_mask: setup.predictor.initializer.body_attention_mask + position_ids: setup.predictor.initializer.body_position_ids + past_key_values.0.key: setup.predictor.present.0.key + past_key_values.0.value: setup.predictor.present.0.value + past_key_values.1.key: setup.predictor.present.1.key + past_key_values.1.value: setup.predictor.present.1.value + past_key_values.2.key: setup.predictor.present.2.key + past_key_values.2.value: setup.predictor.present.2.value + past_key_values.3.key: setup.predictor.present.3.key + past_key_values.3.value: setup.predictor.present.3.value + past_key_values.4.key: setup.predictor.present.4.key + past_key_values.4.value: setup.predictor.present.4.value + outputs: + logits: setup.predictor.remaining_0.logits + codec_embeddings: setup.predictor.remaining_0.codec_embeddings + present.0.key: setup.predictor.remaining_0.present.0.key + present.0.value: setup.predictor.remaining_0.present.0.value + present.1.key: setup.predictor.remaining_0.present.1.key + present.1.value: setup.predictor.remaining_0.present.1.value + present.2.key: setup.predictor.remaining_0.present.2.key + present.2.value: setup.predictor.remaining_0.present.2.value + present.3.key: setup.predictor.remaining_0.present.3.key + present.3.value: setup.predictor.remaining_0.present.3.value + present.4.key: setup.predictor.remaining_0.present.4.key + present.4.value: setup.predictor.remaining_0.present.4.value + - kind: invoke + component: last_token_logits + inputs: + logits: setup.predictor.remaining_0.logits + outputs: + last_logits: setup.predictor.remaining_0.last_logits + - kind: invoke + component: predictor_body_sampler + inputs: + logits: setup.predictor.remaining_0.last_logits + outputs: + token: setup.predictor.remaining_0.token + - kind: invoke + component: code_frame_update + inputs: + frame_codes: setup.frame_prefill + token: setup.predictor.remaining_0.token + index: setup.predictor.remaining_0.frame_index + outputs: + next_frame: setup.predictor.remaining_0.frame + - kind: invoke + component: predictor_step_update + inputs: + attention_mask: setup.predictor.initializer.body_attention_mask + position_ids: setup.predictor.initializer.body_position_ids + outputs: + next_attention_mask: setup.predictor.remaining_0.mask + next_position_ids: setup.predictor.remaining_0.position + - kind: invoke + component: code_predictor_indices + inputs: + iteration: package.setup_predictor_iteration_1 + outputs: + embedding_index: setup.predictor.remaining_1.embedding_index + step_index: setup.predictor.remaining_1.step_index + frame_index: setup.predictor.remaining_1.frame_index + - kind: invoke + component: code_predictor_step_embedder + inputs: + codec_embeddings: setup.predictor.codec_embeddings + token: setup.predictor.remaining_0.token + embedding_index: setup.predictor.remaining_1.embedding_index + outputs: + inputs_embeds: setup.predictor.remaining_1.inputs_embeds + - kind: invoke + component: code_predictor + inputs: + inputs_embeds: setup.predictor.remaining_1.inputs_embeds + step_index: setup.predictor.remaining_1.step_index + attention_mask: setup.predictor.remaining_0.mask + position_ids: setup.predictor.remaining_0.position + past_key_values.0.key: setup.predictor.remaining_0.present.0.key + past_key_values.0.value: setup.predictor.remaining_0.present.0.value + past_key_values.1.key: setup.predictor.remaining_0.present.1.key + past_key_values.1.value: setup.predictor.remaining_0.present.1.value + past_key_values.2.key: setup.predictor.remaining_0.present.2.key + past_key_values.2.value: setup.predictor.remaining_0.present.2.value + past_key_values.3.key: setup.predictor.remaining_0.present.3.key + past_key_values.3.value: setup.predictor.remaining_0.present.3.value + past_key_values.4.key: setup.predictor.remaining_0.present.4.key + past_key_values.4.value: setup.predictor.remaining_0.present.4.value + outputs: + logits: setup.predictor.remaining_1.logits + codec_embeddings: setup.predictor.remaining_1.codec_embeddings + present.0.key: setup.predictor.remaining_1.present.0.key + present.0.value: setup.predictor.remaining_1.present.0.value + present.1.key: setup.predictor.remaining_1.present.1.key + present.1.value: setup.predictor.remaining_1.present.1.value + present.2.key: setup.predictor.remaining_1.present.2.key + present.2.value: setup.predictor.remaining_1.present.2.value + present.3.key: setup.predictor.remaining_1.present.3.key + present.3.value: setup.predictor.remaining_1.present.3.value + present.4.key: setup.predictor.remaining_1.present.4.key + present.4.value: setup.predictor.remaining_1.present.4.value + - kind: invoke + component: last_token_logits + inputs: + logits: setup.predictor.remaining_1.logits + outputs: + last_logits: setup.predictor.remaining_1.last_logits + - kind: invoke + component: predictor_body_sampler + inputs: + logits: setup.predictor.remaining_1.last_logits + outputs: + token: setup.predictor.remaining_1.token + - kind: invoke + component: code_frame_update + inputs: + frame_codes: setup.predictor.remaining_0.frame + token: setup.predictor.remaining_1.token + index: setup.predictor.remaining_1.frame_index + outputs: + next_frame: setup.predictor.remaining_1.frame + - kind: invoke + component: predictor_step_update + inputs: + attention_mask: setup.predictor.remaining_0.mask + position_ids: setup.predictor.remaining_0.position + outputs: + next_attention_mask: setup.predictor.remaining_1.mask + next_position_ids: setup.predictor.remaining_1.position + - kind: invoke + component: code_history_append + inputs: + history: initializer.code_history + frame: setup.predictor.remaining_1.frame + outputs: + next_history: history.setup + - kind: invoke + component: continue_predicate + inputs: + done: package.false + outputs: + continue: talker.setup.continue + steps: + - kind: invoke + component: talker_text_step + inputs: + trailing_text_embeds: tts.trailing_text_embeds + iteration: talker.iteration + outputs: + text_embed: talker.text_embed + - kind: invoke + component: talker_step_embedder + inputs: + frame_codes: last_frame + text_embed: talker.text_embed + outputs: + inputs_embeds: talker.step_embeds + - kind: invoke + component: talker + inputs: + inputs_embeds: talker.step_embeds + attention_mask: talker_mask + position_ids: talker_position + past_key_values.0.key: talker_cache_0 + past_key_values.0.value: talker_cache_1 + outputs: + logits: talker.body.logits + last_hidden_state: talker.body.hidden + present.0.key: talker.body.present.0.key + present.0.value: talker.body.present.0.value + - kind: invoke + component: last_token_logits + inputs: + logits: talker.body.logits + outputs: + last_logits: frame.group0_logits + - kind: invoke + component: talker_sampler + inputs: + logits: frame.group0_logits + outputs: + token: frame.group0 + - kind: invoke + component: token_to_slot + inputs: + token: frame.group0 + outputs: + slot: frame.group0_slot + - kind: invoke + component: embedding + inputs: + text_ids: request.prompt_tokens + codec_ids: frame.group0_slot + outputs: + text_embeds: frame.unused_text_embeds + codec_embeds: frame.group0_embed + - kind: invoke + component: code_predictor_prefill + inputs: + talker_hidden: talker.body.hidden + group_0_embed: frame.group0_embed + outputs: + inputs_embeds: frame.predictor_prefill + - kind: invoke + component: predictor_state_initializer + inputs: + prefill_embeds: frame.predictor_prefill + outputs: + attention_mask: frame.predictor.initializer.attention_mask + position_ids: frame.predictor.initializer.position_ids + body_attention_mask: frame.predictor.initializer.body_attention_mask + body_position_ids: frame.predictor.initializer.body_position_ids + past_key_values.0.key: frame.predictor.initializer.past_key_values.0.key + past_key_values.0.value: frame.predictor.initializer.past_key_values.0.value + past_key_values.1.key: frame.predictor.initializer.past_key_values.1.key + past_key_values.1.value: frame.predictor.initializer.past_key_values.1.value + past_key_values.2.key: frame.predictor.initializer.past_key_values.2.key + past_key_values.2.value: frame.predictor.initializer.past_key_values.2.value + past_key_values.3.key: frame.predictor.initializer.past_key_values.3.key + past_key_values.3.value: frame.predictor.initializer.past_key_values.3.value + past_key_values.4.key: frame.predictor.initializer.past_key_values.4.key + past_key_values.4.value: frame.predictor.initializer.past_key_values.4.value + - kind: invoke + component: code_predictor + inputs: + inputs_embeds: frame.predictor_prefill + step_index: package.zero_scalar + attention_mask: frame.predictor.initializer.attention_mask + position_ids: frame.predictor.initializer.position_ids + past_key_values.0.key: frame.predictor.initializer.past_key_values.0.key + past_key_values.0.value: frame.predictor.initializer.past_key_values.0.value + past_key_values.1.key: frame.predictor.initializer.past_key_values.1.key + past_key_values.1.value: frame.predictor.initializer.past_key_values.1.value + past_key_values.2.key: frame.predictor.initializer.past_key_values.2.key + past_key_values.2.value: frame.predictor.initializer.past_key_values.2.value + past_key_values.3.key: frame.predictor.initializer.past_key_values.3.key + past_key_values.3.value: frame.predictor.initializer.past_key_values.3.value + past_key_values.4.key: frame.predictor.initializer.past_key_values.4.key + past_key_values.4.value: frame.predictor.initializer.past_key_values.4.value + outputs: + logits: frame.predictor.logits + codec_embeddings: frame.predictor.codec_embeddings + present.0.key: frame.predictor.present.0.key + present.0.value: frame.predictor.present.0.value + present.1.key: frame.predictor.present.1.key + present.1.value: frame.predictor.present.1.value + present.2.key: frame.predictor.present.2.key + present.2.value: frame.predictor.present.2.value + present.3.key: frame.predictor.present.3.key + present.3.value: frame.predictor.present.3.value + present.4.key: frame.predictor.present.4.key + present.4.value: frame.predictor.present.4.value + - kind: invoke + component: last_token_logits + inputs: + logits: frame.predictor.logits + outputs: + last_logits: frame.group1_logits + - kind: invoke + component: predictor_prefill_sampler + inputs: + logits: frame.group1_logits + outputs: + token: frame.group1 + - kind: invoke + component: code_frame_update + inputs: + frame_codes: initializer.frame_codes + token: frame.group0 + index: package.zero_scalar + outputs: + next_frame: frame.frame_group0 + - kind: invoke + component: code_frame_update + inputs: + frame_codes: frame.frame_group0 + token: frame.group1 + index: package.one_scalar + outputs: + next_frame: frame.frame_prefill + - kind: loop + setup: + - kind: invoke + component: continue_predicate + inputs: + done: package.false + outputs: + continue: code.setup.continue + steps: + - kind: invoke + component: code_predictor_indices + inputs: + iteration: code.iteration + outputs: + embedding_index: predictor.body.embedding_index + step_index: predictor.body.step_index + frame_index: predictor.body.frame_index + - kind: invoke + component: cache_length_update + inputs: + left: predictor_cache_lengths + right: package.one_batch + outputs: + total: predictor_cache_lengths.next + - kind: invoke + component: code_predictor_step_embedder + inputs: + codec_embeddings: frame.predictor.codec_embeddings + token: code_token + embedding_index: predictor.body.embedding_index + outputs: + inputs_embeds: predictor.body.inputs_embeds + - kind: invoke + component: code_predictor + inputs: + inputs_embeds: predictor.body.inputs_embeds + step_index: predictor.body.step_index + attention_mask: predictor_mask + position_ids: predictor_position + past_key_values.0.key: predictor_cache_0 + past_key_values.0.value: predictor_cache_1 + past_key_values.1.key: predictor_cache_2 + past_key_values.1.value: predictor_cache_3 + past_key_values.2.key: predictor_cache_4 + past_key_values.2.value: predictor_cache_5 + past_key_values.3.key: predictor_cache_6 + past_key_values.3.value: predictor_cache_7 + past_key_values.4.key: predictor_cache_8 + past_key_values.4.value: predictor_cache_9 + outputs: + logits: predictor.body.logits + codec_embeddings: predictor.body.codec_embeddings + present.0.key: predictor.body.present.0.key + present.0.value: predictor.body.present.0.value + present.1.key: predictor.body.present.1.key + present.1.value: predictor.body.present.1.value + present.2.key: predictor.body.present.2.key + present.2.value: predictor.body.present.2.value + present.3.key: predictor.body.present.3.key + present.3.value: predictor.body.present.3.value + present.4.key: predictor.body.present.4.key + present.4.value: predictor.body.present.4.value + - kind: invoke + component: last_token_logits + inputs: + logits: predictor.body.logits + outputs: + last_logits: predictor.body.last_logits + - kind: invoke + component: predictor_body_sampler + inputs: + logits: predictor.body.last_logits + outputs: + token: code.token + - kind: invoke + component: code_frame_update + inputs: + frame_codes: frame + token: code.token + index: predictor.body.frame_index + outputs: + next_frame: frame.inner + - kind: invoke + component: predictor_step_update + inputs: + attention_mask: predictor_mask + position_ids: predictor_position + outputs: + next_attention_mask: predictor.mask.inner + next_position_ids: predictor.position.inner + - kind: invoke + component: continue_predicate + inputs: + done: package.false + outputs: + continue: code.continue + continue_when: loop_1_active + max_iterations: package.remaining_groups + carried: + - cell: frame + next: frame.inner + - cell: code_token + next: code.token + - cell: predictor_mask + next: predictor.mask.inner + - cell: predictor_position + next: predictor.position.inner + - cell: predictor_cache_lengths + next: predictor_cache_lengths.next + - cell: predictor_cache_0 + next: predictor.body.present.0.key + - cell: predictor_cache_1 + next: predictor.body.present.0.value + - cell: predictor_cache_2 + next: predictor.body.present.1.key + - cell: predictor_cache_3 + next: predictor.body.present.1.value + - cell: predictor_cache_4 + next: predictor.body.present.2.key + - cell: predictor_cache_5 + next: predictor.body.present.2.value + - cell: predictor_cache_6 + next: predictor.body.present.3.key + - cell: predictor_cache_7 + next: predictor.body.present.3.value + - cell: predictor_cache_8 + next: predictor.body.present.4.key + - cell: predictor_cache_9 + next: predictor.body.present.4.value + - cell: loop_1_active + next: code.continue + iteration: + value: code.iteration + contract: + dtype: int64 + rank: 0 + shape: [] + - kind: invoke + component: code_history_append + inputs: + history: history + frame: frame + outputs: + next_history: history.outer + - kind: invoke + component: talker_step_update + inputs: + attention_mask: talker_mask + position_ids: talker_position + outputs: + next_attention_mask: talker.mask.body + next_position_ids: talker.position.body + - kind: invoke + component: cache_length_update + inputs: + left: talker_cache_lengths + right: package.one_batch + outputs: + total: talker_cache_lengths.next + - kind: invoke + component: cache_length_update + inputs: + left: package.zero_batch + right: package.one_batch + outputs: + total: accepted_len.next + - kind: invoke + component: continue_predicate + inputs: + done: package.false + outputs: + continue: talker.continue + continue_when: loop_0_active + max_iterations: request.max_iterations + carried: + - cell: active + next: active + - cell: done + next: done + - cell: accepted_len + next: accepted_len.next + - cell: talker_cache_lengths + next: talker_cache_lengths.next + - cell: last_frame + next: frame + initial: setup.frame_prefill + - cell: history + next: history.outer + - cell: talker_mask + next: talker.mask.body + - cell: talker_position + next: talker.position.body + - cell: talker_cache_0 + next: talker.body.present.0.key + - cell: talker_cache_1 + next: talker.body.present.0.value + - cell: loop_0_active + next: talker.continue + iteration: + value: talker.iteration + contract: + dtype: int64 + rank: 1 + shape: + - batch + batch_layout: + kind: request_aligned + axis: 0 + - kind: invoke + component: codec_layout + inputs: + history: history + outputs: + codes: codec.codes + - kind: invoke + component: codec + inputs: + codes: codec.codes + outputs: + waveform: tts.waveform + - kind: emit + value: tts.waveform + output: waveform + mode: replace diff --git a/tests/fixtures/onnx_genai_workflows/video/inference_metadata.yaml b/tests/fixtures/onnx_genai_workflows/video/inference_metadata.yaml new file mode 100644 index 000000000..f31277956 --- /dev/null +++ b/tests/fixtures/onnx_genai_workflows/video/inference_metadata.yaml @@ -0,0 +1,850 @@ +schema_version: v1 +pipeline: + workflow: + manifest: + capabilities: + - workflow_ssa + - linear_effects + - nested_control_flow + - loop_induction_values + - typed_emit + - bounded_state_recurrence + inputs: + request.noise: + contract: + dtype: float32 + rank: 5 + shape: + - batch + - num_frames + - 4 + - height + - width + batch_layout: + kind: request_aligned + axis: 0 + role: + kind: opaque + source: + kind: application + name: noise + required: true + request.max_iterations: + contract: + dtype: int64 + rank: 1 + shape: + - 1 + role: + kind: runtime + version: '1.0' + role: max_iterations + source: + kind: request + required: false + default: 3 + package.false: + contract: + dtype: bool + rank: 1 + shape: + - batch + batch_layout: + kind: request_aligned + axis: 0 + role: + kind: opaque + source: + kind: literal + required: false + default: false + package.one_control: + contract: + dtype: int64 + rank: 1 + shape: + - 1 + role: + kind: opaque + source: + kind: literal + required: false + default: 1 + package.history_limit: + contract: + dtype: int64 + rank: 1 + shape: + - 1 + role: + kind: opaque + source: + kind: literal + required: false + default: 3 + package.cache_frames: + contract: + dtype: int64 + rank: 1 + shape: + - 1 + role: + kind: opaque + source: + kind: literal + required: false + default: 2 + request.encoder_hidden_states: + contract: + dtype: float32 + rank: 3 + shape: + - batch + - prompt_sequence + - 32 + batch_layout: + kind: request_aligned + axis: 0 + role: + kind: opaque + source: + kind: application + name: encoder_hidden_states + required: true + package.loop_0_active: + contract: + dtype: bool + rank: 1 + shape: + - 1 + role: + kind: opaque + source: + kind: literal + required: false + default: true + package.loop_1_active: + contract: + dtype: bool + rank: 1 + shape: + - 1 + role: + kind: opaque + source: + kind: literal + required: false + default: true + outputs: + video: + contract: + dtype: float32 + rank: 5 + shape: + - batch + - 3 + - frames + - 2*latent_height + - 2*latent_width + batch_layout: + kind: request_aligned + axis: 0 + role: video + stage: pre_adapter + components: + transformer: + implementation: + kind: onnx + artifact: transformer/model.onnx + ports: + roles: + encoder_hidden_states: encoder_hidden_states + vae_decoder: + implementation: + kind: onnx + artifact: vae_decoder/model.onnx + model_input: + implementation: + kind: onnx + artifact: policies/model_input.onnx + ports: + inputs: + sample: + dtype: float32 + rank: 5 + shape: + - batch + - frames + - channels + - height + - width + batch_layout: + kind: request_aligned + axis: 0 + step: + dtype: int64 + rank: 1 + shape: + - batch + batch_layout: + kind: request_aligned + axis: 0 + schedule: + dtype: float32 + rank: 1 + shape: + - schedule_length + outputs: + model_input: + dtype: float32 + rank: 5 + shape: + - batch + - frames + - channels + - height + - width + batch_layout: + kind: request_aligned + axis: 0 + solver_step: + implementation: + kind: onnx + artifact: policies/solver_step.onnx + ports: + inputs: + sample: + dtype: float32 + rank: 5 + shape: + - batch + - frames + - channels + - height + - width + batch_layout: + kind: request_aligned + axis: 0 + derivative: + dtype: float32 + rank: 5 + shape: + - batch + - frames + - channels + - height + - width + batch_layout: + kind: request_aligned + axis: 0 + step: + dtype: int64 + rank: 1 + shape: + - batch + batch_layout: + kind: request_aligned + axis: 0 + schedule: + dtype: float32 + rank: 1 + shape: + - schedule_length + outputs: + next_state: + dtype: float32 + rank: 5 + shape: + - batch + - frames + - channels + - height + - width + batch_layout: + kind: request_aligned + axis: 0 + contract: + id: onnx-genai.solver-step + version: '1' + bindings: + state: sample + estimate: derivative + step: step + schedule: schedule + next_state: next_state + continue_predicate: + implementation: + kind: onnx + artifact: policies/continue_predicate.onnx + ports: + inputs: + done: + dtype: bool + rank: 1 + shape: + - batch + batch_layout: + kind: request_aligned + axis: 0 + outputs: + continue: + dtype: bool + rank: 1 + shape: + - 1 + video_latent_init: + implementation: + kind: onnx + artifact: policies/video_latent_init.onnx + ports: + inputs: + noise: + dtype: float32 + rank: 5 + shape: + - batch + - frames + - channels + - height + - width + batch_layout: + kind: request_aligned + axis: 0 + outputs: + latent: + dtype: float32 + rank: 5 + shape: + - batch + - frames + - channels + - height + - width + batch_layout: + kind: request_aligned + axis: 0 + history: + dtype: int64 + rank: 2 + shape: + - batch + - 0 + batch_layout: + kind: request_aligned + axis: 0 + schedule_history_append: + implementation: + kind: onnx + artifact: policies/schedule_history_append.onnx + ports: + inputs: + history: + dtype: int64 + rank: 2 + shape: + - batch + - history + batch_layout: + kind: request_aligned + axis: 0 + timestep: + dtype: int64 + rank: 1 + shape: + - batch + batch_layout: + kind: request_aligned + axis: 0 + outputs: + next: + dtype: int64 + rank: 2 + shape: + - batch + - history + batch_layout: + kind: request_aligned + axis: 0 + video_latent_permute: + implementation: + kind: onnx + artifact: policies/video_latent_permute.onnx + ports: + inputs: + latent: + dtype: float32 + rank: 5 + shape: + - batch + - frames + - channels + - height + - width + batch_layout: + kind: request_aligned + axis: 0 + outputs: + permuted: + dtype: float32 + rank: 5 + shape: + - batch + - channels + - frames + - height + - width + batch_layout: + kind: request_aligned + axis: 0 + video_latent_unscale: + implementation: + kind: onnx + artifact: policies/video_latent_unscale.onnx + ports: + inputs: + latent: + dtype: float32 + rank: 5 + shape: + - batch + - channels + - frames + - height + - width + batch_layout: + kind: request_aligned + axis: 0 + outputs: + unscaled: + dtype: float32 + rank: 5 + shape: + - batch + - channels + - frames + - height + - width + batch_layout: + kind: request_aligned + axis: 0 + video_decode_chunks: + implementation: + kind: onnx + artifact: policies/video_decode_chunks.onnx + ports: + inputs: + latent: + dtype: float32 + rank: 5 + shape: + - batch + - channels + - latent_frames + - height + - width + batch_layout: + kind: request_aligned + axis: 0 + outputs: + count: + dtype: int64 + rank: 1 + shape: + - 1 + video_decode_chunk: + implementation: + kind: onnx + artifact: policies/video_decode_chunk.onnx + ports: + inputs: + latent: + dtype: float32 + rank: 5 + shape: + - batch + - channels + - latent_frames + - height + - width + batch_layout: + kind: request_aligned + axis: 0 + step: + dtype: int64 + rank: 1 + shape: + - batch + batch_layout: + kind: request_aligned + axis: 0 + outputs: + chunk: + dtype: float32 + rank: 5 + shape: + - batch + - channels + - chunk_frames + - height + - width + batch_layout: + kind: request_aligned + axis: 0 + video_conv_cache_init: + implementation: + kind: onnx + artifact: policies/video_conv_cache_init.onnx + ports: + inputs: + latent: + dtype: float32 + rank: 5 + shape: + - batch + - channels + - latent_frames + - height + - width + batch_layout: + kind: request_aligned + axis: 0 + outputs: + conv_cache.conv_in: + dtype: float32 + rank: 5 + shape: + - batch + - 4 + - 0 + - height + - width + batch_layout: + kind: request_aligned + axis: 0 + conv_cache.conv_out: + dtype: float32 + rank: 5 + shape: + - batch + - 3 + - 0 + - 2*height + - 2*width + batch_layout: + kind: request_aligned + axis: 0 + diffusion_schedule: + implementation: + kind: onnx + artifact: policies/diffusion_schedule.onnx + ports: + inputs: {} + outputs: + schedule: + dtype: float32 + rank: 1 + shape: + - 4 + diffusion_timesteps: + implementation: + kind: onnx + artifact: policies/diffusion_timesteps.onnx + ports: + inputs: {} + outputs: + schedule: + dtype: float32 + rank: 1 + shape: + - 3 + schedule_lookup: + implementation: + kind: onnx + artifact: policies/schedule_lookup.onnx + ports: + inputs: + schedule: + dtype: float32 + rank: 1 + shape: + - schedule_length + step: + dtype: int64 + rank: 1 + shape: + - batch + batch_layout: + kind: request_aligned + axis: 0 + outputs: + timestep: + dtype: int64 + rank: 1 + shape: + - batch + batch_layout: + kind: request_aligned + axis: 0 + state: + latent: + contract: + dtype: float32 + rank: 5 + shape: + - batch + - num_frames + - 4 + - height + - width + batch_layout: + kind: request_aligned + axis: 0 + scope: invocation + initializer: latent.initial + recurrence: + kind: invariant + scheduler_history: + contract: + dtype: int64 + rank: 2 + shape: + - batch + - scheduler_history + batch_layout: + kind: request_aligned + axis: 0 + scope: invocation + initializer: scheduler.history.initial + recurrence: + kind: growing + axis: 1 + increment: package.one_control + max: package.history_limit + conv_cache_conv_in: + contract: + dtype: float32 + rank: 5 + shape: + - batch + - 4 + - cache_frames + - latent_height + - latent_width + batch_layout: + kind: request_aligned + axis: 0 + scope: invocation + initializer: conv_cache_conv_in.initial + recurrence: + kind: bounded + axis: 2 + max: package.cache_frames + management: runtime + release_boundary: invocation + conv_cache_conv_out: + contract: + dtype: float32 + rank: 5 + shape: + - batch + - 3 + - cache_frames + - 2*latent_height + - 2*latent_width + batch_layout: + kind: request_aligned + axis: 0 + scope: invocation + initializer: conv_cache_conv_out.initial + recurrence: + kind: bounded + axis: 2 + max: package.cache_frames + management: runtime + release_boundary: invocation + loop_0_active: + contract: + dtype: bool + rank: 1 + shape: + - 1 + scope: invocation + initializer: package.loop_0_active + recurrence: + kind: invariant + loop_1_active: + contract: + dtype: bool + rank: 1 + shape: + - 1 + scope: invocation + initializer: package.loop_1_active + recurrence: + kind: invariant + steps: + - kind: invoke + component: video_latent_init + inputs: + noise: request.noise + outputs: + latent: latent.initial + history: scheduler.history.initial + - kind: loop + setup: + - kind: invoke + component: diffusion_schedule + inputs: {} + outputs: + schedule: diffusion.schedule + - kind: invoke + component: diffusion_timesteps + inputs: {} + outputs: + schedule: diffusion.timesteps + - kind: invoke + component: continue_predicate + inputs: + done: package.false + outputs: + continue: setup.continue + steps: + - kind: invoke + component: schedule_lookup + inputs: + schedule: diffusion.timesteps + step: loop.iteration + outputs: + timestep: diffusion.timestep + - kind: invoke + component: model_input + inputs: + sample: latent + step: loop.iteration + schedule: diffusion.schedule + outputs: + model_input: diffusion.model_input + - kind: invoke + component: transformer + inputs: + sample: diffusion.model_input + timestep: diffusion.timestep + encoder_hidden_states: request.encoder_hidden_states + outputs: + noise_pred: denoiser.estimate + - kind: invoke + component: solver_step + inputs: + sample: latent + derivative: denoiser.estimate + step: loop.iteration + schedule: diffusion.schedule + outputs: + next_state: latent.body + - kind: invoke + component: schedule_history_append + inputs: + history: scheduler_history + timestep: diffusion.timestep + outputs: + next: scheduler_history.body + - kind: invoke + component: continue_predicate + inputs: + done: package.false + outputs: + continue: loop.continue + continue_when: loop_0_active + max_iterations: request.max_iterations + carried: + - cell: latent + next: latent.body + - cell: scheduler_history + next: scheduler_history.body + - cell: loop_0_active + next: loop.continue + iteration: + value: loop.iteration + contract: + dtype: int64 + rank: 1 + shape: + - batch + - kind: invoke + component: video_latent_permute + inputs: + latent: latent + outputs: + permuted: decode.latent_permuted + - kind: invoke + component: video_latent_unscale + inputs: + latent: decode.latent_permuted + outputs: + unscaled: decode.latent + - kind: invoke + component: video_decode_chunks + inputs: + latent: decode.latent + outputs: + count: decode.chunks + - kind: invoke + component: video_conv_cache_init + inputs: + latent: decode.latent + outputs: + conv_cache.conv_in: conv_cache_conv_in.initial + conv_cache.conv_out: conv_cache_conv_out.initial + - kind: loop + setup: + - kind: invoke + component: continue_predicate + inputs: + done: package.false + outputs: + continue: decode.setup.continue + steps: + - kind: invoke + component: video_decode_chunk + inputs: + latent: decode.latent + step: decode.iteration + outputs: + chunk: decode.chunk + - kind: invoke + component: vae_decoder + inputs: + latent_sample: decode.chunk + conv_cache.conv_in: conv_cache_conv_in + conv_cache.conv_out: conv_cache_conv_out + outputs: + sample: decode.frames + conv_cache_out.conv_in: conv_cache_conv_in.body + conv_cache_out.conv_out: conv_cache_conv_out.body + - kind: emit + value: decode.frames + output: video + mode: append + axis: 2 + - kind: invoke + component: continue_predicate + inputs: + done: package.false + outputs: + continue: decode.loop.continue + continue_when: loop_1_active + max_iterations: decode.chunks + carried: + - cell: conv_cache_conv_in + next: conv_cache_conv_in.body + - cell: conv_cache_conv_out + next: conv_cache_conv_out.body + - cell: loop_1_active + next: decode.loop.continue + iteration: + value: decode.iteration + contract: + dtype: int64 + rank: 1 + shape: + - batch diff --git a/tests/fixtures/onnx_genai_workflows/vlm/inference_metadata.yaml b/tests/fixtures/onnx_genai_workflows/vlm/inference_metadata.yaml new file mode 100644 index 000000000..f2dce1a06 --- /dev/null +++ b/tests/fixtures/onnx_genai_workflows/vlm/inference_metadata.yaml @@ -0,0 +1,1536 @@ +schema_version: v1 +preprocessing: + image: + transforms: + - op: decode_rgb + outputs: + - image.transform_0 + - op: resize + mode: pixel_area + interpolation: bicubic + min_pixels: 224 + max_pixels: 224 + size_multiple: 14 + inputs: + - image.transform_0 + outputs: + - image.transform_1 + - op: rescale + scale: 0.00392156862745098 + inputs: + - image.transform_1 + outputs: + - image.transform_2 + - op: normalize + mean: + - 0.5 + - 0.5 + - 0.5 + std: + - 0.5 + - 0.5 + - 0.5 + inputs: + - image.transform_2 + outputs: + - image.transform_3 + - op: patchify + patch_size: 14 + flatten: true + temporal_patch_size: 2 + merge_size: 1 + channel_order: channels_first + inputs: + - image.transform_3 + outputs: + - image.transform_4 + - op: emit_grid_coordinates + inputs: + - image.transform_4 + outputs: + - image.output_grid_dimensions + outputs: + - name: image.pixel_values + content: pixels + dtype: float32 + source: image.transform_4 + contract: + dtype: float32 + rank: 2 + shape: + - 4 + - 1176 + - name: image.grid_thw + content: grid_dimensions + dtype: int64 + source: image.output_grid_dimensions + contract: + dtype: int64 + rank: 2 + shape: + - 1 + - 3 +pipeline: + workflow: + manifest: + adapter_abis: + onnx-genai.image-preprocess: '1' + capabilities: + - workflow_ssa + - linear_effects + - nested_control_flow + - loop_induction_values + - typed_emit + - emit_valid_length + - serving_service_contract + - bounded_state_recurrence + inputs: + request.prompt_tokens: + contract: + dtype: int64 + rank: 2 + shape: + - batch + - sequence + batch_layout: + kind: request_aligned + axis: 0 + role: + kind: runtime + version: '1.0' + role: prompt_tokens + source: + kind: request + required: true + request.image: + contract: + dtype: uint8 + rank: 1 + shape: + - encoded_bytes + role: + kind: runtime + version: '1.0' + role: media + source: + kind: request + required: true + request.max_iterations: + contract: + dtype: int64 + rank: 1 + shape: + - 1 + role: + kind: runtime + version: '1.0' + role: max_output_tokens + source: + kind: request + required: true + request.prompt_lengths: + contract: + dtype: int64 + rank: 1 + shape: + - batch + batch_layout: + kind: request_aligned + axis: 0 + role: + kind: opaque + source: + kind: application + name: prompt_lengths + required: false + default: -1 + request.eos_ids: + contract: + dtype: int64 + rank: 2 + shape: + - batch + - num_eos + batch_layout: + kind: request_aligned + axis: 0 + role: + kind: opaque + source: + kind: application + name: eos_ids + required: false + default: 2 + request.eos_lengths: + contract: + dtype: int64 + rank: 1 + shape: + - batch + batch_layout: + kind: request_aligned + axis: 0 + role: + kind: opaque + source: + kind: application + name: eos_lengths + required: false + default: 1 + request.row_max_iterations: + contract: + dtype: int64 + rank: 1 + shape: + - batch + batch_layout: + kind: request_aligned + axis: 0 + role: + kind: opaque + source: + kind: application + name: row_max_iterations + required: false + default: -1 + package.eos_ids: + contract: + dtype: int64 + rank: 1 + shape: + - 1 + role: + kind: opaque + source: + kind: literal + required: false + default: 2 + package.max_context: + contract: + dtype: int64 + rank: 1 + shape: + - 1 + role: + kind: opaque + source: + kind: literal + required: false + default: 8192 + package.one: + contract: + dtype: int64 + rank: 1 + shape: + - batch + batch_layout: + kind: request_aligned + axis: 0 + role: + kind: opaque + source: + kind: literal + required: false + default: 1 + package.one_step: + contract: + dtype: int64 + rank: 1 + shape: + - 1 + role: + kind: opaque + source: + kind: literal + required: false + default: 1 + package.active: + contract: + dtype: bool + rank: 1 + shape: + - batch + batch_layout: + kind: request_aligned + axis: 0 + role: + kind: opaque + source: + kind: literal + required: false + default: true + package.false: + contract: + dtype: bool + rank: 1 + shape: + - batch + batch_layout: + kind: request_aligned + axis: 0 + role: + kind: opaque + source: + kind: literal + required: false + default: false + package.zero_batch: + contract: + dtype: int64 + rank: 1 + shape: + - batch + batch_layout: + kind: request_aligned + axis: 0 + role: + kind: opaque + source: + kind: literal + required: false + default: 0 + request.temperature: + contract: + dtype: float32 + rank: 1 + shape: + - batch + batch_layout: + kind: request_aligned + axis: 0 + role: + kind: runtime + version: '1.0' + role: sampling_temperature + source: + kind: request + required: false + default: 1.0 + request.top_k: + contract: + dtype: int64 + rank: 1 + shape: + - batch + batch_layout: + kind: request_aligned + axis: 0 + role: + kind: runtime + version: '1.0' + role: sampling_top_k + source: + kind: request + required: false + default: 1 + request.top_p: + contract: + dtype: float32 + rank: 1 + shape: + - batch + batch_layout: + kind: request_aligned + axis: 0 + role: + kind: runtime + version: '1.0' + role: sampling_top_p + source: + kind: request + required: false + default: 1.0 + request.min_p: + contract: + dtype: float32 + rank: 1 + shape: + - batch + batch_layout: + kind: request_aligned + axis: 0 + role: + kind: runtime + version: '1.0' + role: sampling_min_p + source: + kind: request + required: false + default: 0.0 + request.seed: + contract: + dtype: int64 + rank: 1 + shape: + - batch + batch_layout: + kind: request_aligned + axis: 0 + role: + kind: runtime + version: '1.0' + role: seed + source: + kind: request + required: false + default: 0 + request.rng_counter: + contract: + dtype: int64 + rank: 1 + shape: + - batch + batch_layout: + kind: request_aligned + axis: 0 + role: + kind: opaque + source: + kind: application + name: rng_counter + required: false + default: 0 + package.loop_0_active: + contract: + dtype: bool + rank: 1 + shape: + - 1 + role: + kind: opaque + source: + kind: literal + required: false + default: true + outputs: + tokens: + contract: + dtype: int64 + rank: 2 + shape: + - batch + - generated_sequence + batch_layout: + kind: request_aligned + axis: 0 + role: tokens + stage: pre_adapter + components: + vision_encoder: + implementation: + kind: onnx + artifact: vision_encoder/model.onnx + embedding: + implementation: + kind: onnx + artifact: embedding/model.onnx + ports: + roles: + input_ids: token_ids + inputs_embeds: inputs_embeds + decoder: + implementation: + kind: onnx + artifact: decoder/model.onnx + ports: + roles: + inputs_embeds: inputs_embeds + attention_mask: attention_mask + position_ids: position_ids + logits: logits + image_preprocess: + implementation: + kind: adapter + abi: onnx-genai.image-preprocess + version: '1' + ports: + inputs: + encoded: + dtype: uint8 + rank: 1 + shape: + - encoded_bytes + outputs: + pixel_values: + dtype: float32 + rank: 2 + shape: + - 4 + - 1176 + grid_thw: + dtype: int64 + rank: 2 + shape: + - 1 + - 3 + token_sampler: + implementation: + kind: onnx + artifact: policies/token_sampler.onnx + ports: + inputs: + logits: + dtype: float32 + rank: 2 + shape: + - batch + - vocabulary + batch_layout: + kind: request_aligned + axis: 0 + temperature: + dtype: float32 + rank: 1 + shape: + - batch + batch_layout: + kind: request_aligned + axis: 0 + top_k: + dtype: int64 + rank: 1 + shape: + - batch + batch_layout: + kind: request_aligned + axis: 0 + top_p: + dtype: float32 + rank: 1 + shape: + - batch + batch_layout: + kind: request_aligned + axis: 0 + min_p: + dtype: float32 + rank: 1 + shape: + - batch + batch_layout: + kind: request_aligned + axis: 0 + seed: + dtype: int64 + rank: 1 + shape: + - batch + batch_layout: + kind: request_aligned + axis: 0 + counter: + dtype: int64 + rank: 1 + shape: + - batch + batch_layout: + kind: request_aligned + axis: 0 + active: + dtype: bool + rank: 1 + shape: + - batch + batch_layout: + kind: request_aligned + axis: 0 + done: + dtype: bool + rank: 1 + shape: + - batch + batch_layout: + kind: request_aligned + axis: 0 + outputs: + token: + dtype: int64 + rank: 1 + shape: + - batch + batch_layout: + kind: request_aligned + axis: 0 + next_counter: + dtype: int64 + rank: 1 + shape: + - batch + batch_layout: + kind: request_aligned + axis: 0 + contract: + id: onnx-genai.token-sampler + version: '2' + bindings: + logits: logits + token: token + temperature: temperature + top_k: top_k + top_p: top_p + min_p: min_p + active: active + done: done + seed: seed + counter: counter + next_counter: next_counter + parameters: + mode: seeded_stochastic + batching: per_row + inactive_rows: preserve + application_overridable: true + termination: + implementation: + kind: onnx + artifact: policies/termination.onnx + ports: + inputs: + tokens: + dtype: int64 + rank: 1 + shape: + - batch + batch_layout: + kind: request_aligned + axis: 0 + eos_ids: + dtype: int64 + rank: 2 + shape: + - batch + - num_eos + batch_layout: + kind: request_aligned + axis: 0 + eos_lengths: + dtype: int64 + rank: 1 + shape: + - batch + batch_layout: + kind: request_aligned + axis: 0 + iteration: + dtype: int64 + rank: 1 + shape: + - 1 + max_iterations: + dtype: int64 + rank: 1 + shape: + - batch + batch_layout: + kind: request_aligned + axis: 0 + active: + dtype: bool + rank: 1 + shape: + - batch + batch_layout: + kind: request_aligned + axis: 0 + outputs: + done: + dtype: bool + rank: 1 + shape: + - batch + batch_layout: + kind: request_aligned + axis: 0 + next_active: + dtype: bool + rank: 1 + shape: + - batch + batch_layout: + kind: request_aligned + axis: 0 + continue: + dtype: bool + rank: 1 + shape: + - 1 + contract: + id: onnx-genai.termination-predicate + version: '2' + bindings: + tokens: tokens + eos_ids: eos_ids + iteration: iteration + max_iterations: max_iterations + eos_lengths: eos_lengths + active: active + done: done + next_active: next_active + continue: continue + parameters: + batching: per_row + inactive_rows: preserve + token_state_update: + implementation: + kind: onnx + artifact: policies/token_state_update.onnx + ports: + inputs: + current: + dtype: int64 + rank: 2 + shape: + - batch + - 1 + batch_layout: + kind: request_aligned + axis: 0 + update: + dtype: int64 + rank: 2 + shape: + - batch + - 1 + batch_layout: + kind: request_aligned + axis: 0 + active: + dtype: bool + rank: 1 + shape: + - batch + batch_layout: + kind: request_aligned + axis: 0 + done: + dtype: bool + rank: 1 + shape: + - batch + batch_layout: + kind: request_aligned + axis: 0 + outputs: + next: + dtype: int64 + rank: 2 + shape: + - batch + - 1 + batch_layout: + kind: request_aligned + axis: 0 + contract: + id: onnx-genai.state-update + version: '2' + bindings: + current: current + update: update + active: active + done: done + next: next + parameters: + batching: per_row + inactive_rows: preserve + last_token_logits: + implementation: + kind: onnx + artifact: policies/last_token_logits.onnx + ports: + inputs: + logits: + dtype: float32 + rank: 3 + shape: + - batch + - sequence + - vocabulary + batch_layout: + kind: request_aligned + axis: 0 + outputs: + last_logits: + dtype: float32 + rank: 2 + shape: + - batch + - vocabulary + batch_layout: + kind: request_aligned + axis: 0 + decoder_state_initializer: + implementation: + kind: onnx + artifact: policies/decoder_state_initializer.onnx + ports: + inputs: + prompt_tokens: + dtype: int64 + rank: 2 + shape: + - batch + - prompt_sequence + batch_layout: + kind: request_aligned + axis: 0 + prompt_lengths: + dtype: int64 + rank: 1 + shape: + - batch + batch_layout: + kind: request_aligned + axis: 0 + outputs: + attention_mask: + dtype: int64 + rank: 2 + shape: + - batch + - past_sequence + sequence + batch_layout: + kind: request_aligned + axis: 0 + position_ids: + dtype: int64 + rank: 2 + shape: + - batch + - prompt_sequence + batch_layout: + kind: request_aligned + axis: 0 + body_attention_mask: + dtype: int64 + rank: 2 + shape: + - batch + - prompt_sequence + 1 + batch_layout: + kind: request_aligned + axis: 0 + body_position_ids: + dtype: int64 + rank: 2 + shape: + - batch + - 1 + batch_layout: + kind: request_aligned + axis: 0 + token_slot: + dtype: int64 + rank: 2 + shape: + - batch + - 1 + batch_layout: + kind: request_aligned + axis: 0 + generated_lengths: + dtype: int64 + rank: 1 + shape: + - batch + batch_layout: + kind: request_aligned + axis: 0 + past_key_values.0.key: + dtype: float32 + rank: 4 + shape: + - batch + - 2 + - past_sequence + - 8 + batch_layout: + kind: request_aligned + axis: 0 + past_key_values.0.value: + dtype: float32 + rank: 4 + shape: + - batch + - 2 + - past_sequence + - 8 + batch_layout: + kind: request_aligned + axis: 0 + decoder_step_update: + implementation: + kind: onnx + artifact: policies/decoder_step_update.onnx + ports: + inputs: + attention_mask: + dtype: int64 + rank: 2 + shape: + - batch + - context + batch_layout: + kind: request_aligned + axis: 0 + position_ids: + dtype: int64 + rank: 2 + shape: + - batch + - 1 + batch_layout: + kind: request_aligned + axis: 0 + outputs: + next_attention_mask: + dtype: int64 + rank: 2 + shape: + - batch + - context + 1 + batch_layout: + kind: request_aligned + axis: 0 + next_position_ids: + dtype: int64 + rank: 2 + shape: + - batch + - 1 + batch_layout: + kind: request_aligned + axis: 0 + termination_batch_initializer: + implementation: + kind: onnx + artifact: policies/termination_batch_initializer.onnx + ports: + inputs: + input_eos_ids: + dtype: int64 + rank: 2 + shape: + - batch + - num_eos + batch_layout: + kind: request_aligned + axis: 0 + input_eos_lengths: + dtype: int64 + rank: 1 + shape: + - batch + batch_layout: + kind: request_aligned + axis: 0 + input_max_iterations: + dtype: int64 + rank: 1 + shape: + - batch + batch_layout: + kind: request_aligned + axis: 0 + fallback_max_iterations: + dtype: int64 + rank: 1 + shape: + - 1 + active: + dtype: bool + rank: 1 + shape: + - batch + batch_layout: + kind: request_aligned + axis: 0 + outputs: + row_eos_ids: + dtype: int64 + rank: 2 + shape: + - batch + - num_eos + batch_layout: + kind: request_aligned + axis: 0 + eos_lengths: + dtype: int64 + rank: 1 + shape: + - batch + batch_layout: + kind: request_aligned + axis: 0 + max_iterations: + dtype: int64 + rank: 1 + shape: + - batch + batch_layout: + kind: request_aligned + axis: 0 + token_to_slot: + implementation: + kind: onnx + artifact: policies/token_to_slot.onnx + ports: + inputs: + token: + dtype: int64 + rank: 1 + shape: + - batch + batch_layout: + kind: request_aligned + axis: 0 + outputs: + slot: + dtype: int64 + rank: 2 + shape: + - batch + - 1 + batch_layout: + kind: request_aligned + axis: 0 + generated_length_update: + implementation: + kind: onnx + artifact: policies/generated_length_update.onnx + ports: + inputs: + left: + dtype: int64 + rank: 1 + shape: + - batch + batch_layout: + kind: request_aligned + axis: 0 + right: + dtype: int64 + rank: 1 + shape: + - batch + batch_layout: + kind: request_aligned + axis: 0 + active: + dtype: bool + rank: 1 + shape: + - batch + batch_layout: + kind: request_aligned + axis: 0 + done: + dtype: bool + rank: 1 + shape: + - batch + batch_layout: + kind: request_aligned + axis: 0 + outputs: + total: + dtype: int64 + rank: 1 + shape: + - batch + batch_layout: + kind: request_aligned + axis: 0 + cache_length_update: + implementation: + kind: onnx + artifact: policies/cache_length_update.onnx + ports: + inputs: + left: + dtype: int64 + rank: 1 + shape: + - batch + batch_layout: + kind: request_aligned + axis: 0 + right: + dtype: int64 + rank: 1 + shape: + - batch + batch_layout: + kind: request_aligned + axis: 0 + active: + dtype: bool + rank: 1 + shape: + - batch + batch_layout: + kind: request_aligned + axis: 0 + done: + dtype: bool + rank: 1 + shape: + - batch + batch_layout: + kind: request_aligned + axis: 0 + outputs: + total: + dtype: int64 + rank: 1 + shape: + - batch + batch_layout: + kind: request_aligned + axis: 0 + state: + token: + contract: + dtype: int64 + rank: 2 + shape: + - batch + - 1 + batch_layout: + kind: request_aligned + axis: 0 + scope: invocation + initializer: initializer.token_slot + recurrence: + kind: invariant + logits: + contract: + dtype: float32 + rank: 2 + shape: + - batch + - 128 + batch_layout: + kind: request_aligned + axis: 0 + scope: invocation + initializer: decoder.setup.last_logits + recurrence: + kind: invariant + generated_lengths: + contract: + dtype: int64 + rank: 1 + shape: + - batch + batch_layout: + kind: request_aligned + axis: 0 + class: semantic + scope: invocation + initializer: initializer.generated_lengths + recurrence: + kind: invariant + attention_mask: + contract: + dtype: int64 + rank: 2 + shape: + - batch + - context + batch_layout: + kind: request_aligned + axis: 0 + scope: invocation + initializer: initializer.body_attention_mask + recurrence: + kind: growing + axis: 1 + increment: package.one_step + max: package.max_context + active: + contract: + dtype: bool + rank: 1 + shape: + - batch + batch_layout: + kind: request_aligned + axis: 0 + class: semantic + scope: invocation + initializer: package.active + recurrence: + kind: invariant + done: + contract: + dtype: bool + rank: 1 + shape: + - batch + batch_layout: + kind: request_aligned + axis: 0 + class: semantic + scope: invocation + initializer: package.false + recurrence: + kind: invariant + accepted_len: + contract: + dtype: int64 + rank: 1 + shape: + - batch + batch_layout: + kind: request_aligned + axis: 0 + class: semantic + scope: invocation + initializer: package.zero_batch + recurrence: + kind: invariant + cache_lengths: + contract: + dtype: int64 + rank: 1 + shape: + - batch + batch_layout: + kind: request_aligned + axis: 0 + class: semantic + scope: invocation + initializer: package.zero_batch + recurrence: + kind: invariant + rng_counter: + contract: + dtype: int64 + rank: 1 + shape: + - batch + batch_layout: + kind: request_aligned + axis: 0 + class: semantic + scope: invocation + initializer: request.rng_counter + recurrence: + kind: invariant + position_ids: + contract: + dtype: int64 + rank: 2 + shape: + - batch + - 1 + batch_layout: + kind: request_aligned + axis: 0 + scope: invocation + initializer: initializer.body_position_ids + recurrence: + kind: invariant + cache_0: + contract: + dtype: float32 + rank: 4 + shape: + - batch + - 2 + - past_sequence + - 8 + batch_layout: + kind: request_aligned + axis: 0 + scope: invocation + initializer: decoder.setup.present.0.key + recurrence: + kind: bounded + axis: 2 + max: package.max_context + management: runtime + release_boundary: invocation + service_group: decoder_cache + cache_1: + contract: + dtype: float32 + rank: 4 + shape: + - batch + - 2 + - past_sequence + - 8 + batch_layout: + kind: request_aligned + axis: 0 + scope: invocation + initializer: decoder.setup.present.0.value + recurrence: + kind: bounded + axis: 2 + max: package.max_context + management: runtime + release_boundary: invocation + service_group: decoder_cache + loop_0_active: + contract: + dtype: bool + rank: 1 + shape: + - 1 + scope: invocation + initializer: package.loop_0_active + recurrence: + kind: invariant + serving: + active: active + done: done + accepted_len: accepted_len + state_service: + groups: + decoder_cache: + kind: full_attention + sequence_axis: 2 + layout: bnsh + logical_lengths: cache_lengths + aliasing: forbidden + reuse: + prefix_reusable: true + evictable_prefix: false + ports: + decoder: + cache_0: + input: past_key_values.0.key + output: present.0.key + role: key + layer: 0 + cache_1: + input: past_key_values.0.value + output: present.0.value + role: value + layer: 0 + steps: + - kind: loop + setup: + - kind: invoke + component: image_preprocess + inputs: + encoded: request.image + outputs: + pixel_values: image.pixel_values + grid_thw: image.grid_thw + - kind: invoke + component: vision_encoder + inputs: + pixel_values: image.pixel_values + grid_thw: image.grid_thw + outputs: + image_features: vision.image_features + - kind: invoke + component: decoder_state_initializer + inputs: + prompt_tokens: request.prompt_tokens + prompt_lengths: request.prompt_lengths + outputs: + attention_mask: initializer.attention_mask + body_attention_mask: initializer.body_attention_mask + token_slot: initializer.token_slot + generated_lengths: initializer.generated_lengths + position_ids: initializer.position_ids + body_position_ids: initializer.body_position_ids + past_key_values.0.key: initializer.past_key_values.0.key + past_key_values.0.value: initializer.past_key_values.0.value + - kind: invoke + component: termination_batch_initializer + inputs: + input_eos_ids: request.eos_ids + input_eos_lengths: request.eos_lengths + input_max_iterations: request.row_max_iterations + fallback_max_iterations: request.max_iterations + active: package.active + outputs: + row_eos_ids: termination.eos_ids + eos_lengths: termination.eos_lengths + max_iterations: termination.max_iterations + - kind: invoke + component: embedding + inputs: + input_ids: request.prompt_tokens + image_features: vision.image_features + outputs: + inputs_embeds: embedding.setup.embeds + - kind: invoke + component: decoder + inputs: + inputs_embeds: embedding.setup.embeds + attention_mask: initializer.attention_mask + position_ids: initializer.position_ids + past_key_values.0.key: initializer.past_key_values.0.key + past_key_values.0.value: initializer.past_key_values.0.value + outputs: + logits: decoder.setup.logits + present.0.key: decoder.setup.present.0.key + present.0.value: decoder.setup.present.0.value + - kind: invoke + component: last_token_logits + inputs: + logits: decoder.setup.logits + outputs: + last_logits: decoder.setup.last_logits + steps: + - kind: invoke + component: token_sampler + inputs: + logits: logits + temperature: request.temperature + top_k: request.top_k + top_p: request.top_p + min_p: request.min_p + seed: request.seed + counter: rng_counter + active: active + done: done + outputs: + token: sample.body + next_counter: sample.next_counter + - kind: invoke + component: termination + inputs: + tokens: sample.body + eos_ids: termination.eos_ids + eos_lengths: termination.eos_lengths + iteration: loop.iteration + max_iterations: termination.max_iterations + active: active + outputs: + done: loop.done + next_active: loop.next_active + continue: loop.continue + - kind: invoke + component: cache_length_update + inputs: + left: cache_lengths + right: package.one + active: active + done: done + outputs: + total: cache_lengths.next + - kind: invoke + component: cache_length_update + inputs: + left: package.zero_batch + right: package.one + active: active + done: done + outputs: + total: accepted_len.next + - kind: invoke + component: token_to_slot + inputs: + token: sample.body + outputs: + slot: sample.slot + - kind: invoke + component: generated_length_update + inputs: + left: generated_lengths + right: package.one + active: active + done: done + outputs: + total: token.next_lengths + - kind: invoke + component: generated_length_update + inputs: + left: package.zero_batch + right: package.one + active: active + done: done + outputs: + total: token.emitted_length + - kind: invoke + component: token_state_update + inputs: + current: token + update: sample.slot + active: active + done: done + outputs: + next: token.body + - kind: emit + value: token.body + output: tokens + mode: append + valid_length: token.emitted_length + when: active + - kind: invoke + component: embedding + inputs: + input_ids: token.body + image_features: vision.image_features + outputs: + inputs_embeds: embedding.body.embeds + - kind: invoke + component: decoder + inputs: + inputs_embeds: embedding.body.embeds + attention_mask: attention_mask + position_ids: position_ids + past_key_values.0.key: cache_0 + past_key_values.0.value: cache_1 + outputs: + logits: decoder.body.logits + present.0.key: decoder.body.present.0.key + present.0.value: decoder.body.present.0.value + - kind: invoke + component: last_token_logits + inputs: + logits: decoder.body.logits + outputs: + last_logits: decoder.body.last_logits + - kind: invoke + component: decoder_step_update + inputs: + attention_mask: attention_mask + position_ids: position_ids + outputs: + next_attention_mask: decoder_step.body_attention_mask + next_position_ids: decoder_step.body_position_ids + continue_when: loop_0_active + max_iterations: request.max_iterations + carried: + - cell: token + next: token.body + - cell: logits + next: decoder.body.last_logits + - cell: attention_mask + next: decoder_step.body_attention_mask + - cell: generated_lengths + next: token.next_lengths + - cell: rng_counter + next: sample.next_counter + - cell: active + next: loop.next_active + - cell: done + next: loop.done + - cell: accepted_len + next: accepted_len.next + - cell: cache_lengths + next: cache_lengths.next + - cell: position_ids + next: decoder_step.body_position_ids + - cell: cache_0 + next: decoder.body.present.0.key + - cell: cache_1 + next: decoder.body.present.0.value + - cell: loop_0_active + next: loop.continue + termination: generation_eos + iteration: + value: loop.iteration + contract: + dtype: int64 + rank: 1 + shape: + - 1 diff --git a/tests/gemma4_prefill_prefix_test.py b/tests/gemma4_prefill_prefix_test.py index dbccb285d..2fb5cb724 100644 --- a/tests/gemma4_prefill_prefix_test.py +++ b/tests/gemma4_prefill_prefix_test.py @@ -1,135 +1,394 @@ -# Copyright (c) Microsoft Corporation. -# Licensed under the MIT License. - -from __future__ import annotations - -import torch - -from mobius import build_from_module -from mobius._configs import Gemma4Config, VisionConfig -from mobius._registry import registry -from mobius.models.gemma4 import _split_per_layer_projection_weight - - -def _make_config(*, with_vision: bool = False) -> Gemma4Config: - return Gemma4Config( - num_hidden_layers=4, - hidden_size=64, - intermediate_size=128, - num_attention_heads=4, - num_key_value_heads=1, - head_dim=16, - vocab_size=256, - rms_norm_eps=1e-6, - hidden_act="silu", - attn_qk_norm=True, - layer_types=[ - "sliding_attention", - "full_attention", - "sliding_attention", - "full_attention", - ], - sliding_window=8, - global_head_dim=16, - global_rope_theta=10_000.0, - global_partial_rotary_factor=0.25, - final_logit_softcapping=30.0, - hidden_size_per_layer_input=8, - vocab_size_per_layer_input=64, - split_per_layer_embedding=True, - image_token_id=255999 if with_vision else None, - pad_token_id=0, - tie_word_embeddings=False, - num_kv_shared_layers=2, - vision=( - VisionConfig( - hidden_size=32, - intermediate_size=64, - num_hidden_layers=1, - num_attention_heads=2, - patch_size=16, - norm_eps=1e-6, - ) - if with_vision - else None - ), - ) - - -def test_prunes_gemma4_shared_layer_prefix() -> None: - config = _make_config() - module = registry.get("gemma4_text")(config) - model = build_from_module( - module, - config, - task="gemma4-text-generation", - execution_provider="webgpu", - prune_prefill_prefix=True, - )["model"] - - producer = next( - node - for node in model.graph - if node.op_type == "MatMul" and "/per_layer_model_projection/" in node.name - ) - consumer = next( - node - for node in model.graph - if node.op_type == "MatMul" and "/per_layer_model_projection_consumer/" in node.name - ) - assert producer.outputs[0].shape[1] != 1 - assert producer.outputs[0].shape[2] == 16 - assert consumer.inputs[0].shape[1:] == (1, 64) - assert consumer.outputs[0].shape[1:] == (1, 16) - - first_shared_norm = next( - node for node in model.graph if "layers.2/input_layernorm" in node.name - ) - assert first_shared_norm.inputs[0].shape[1:] == (1, 64) - logits = next(value for value in model.graph.outputs if value.name == "logits") - assert logits.shape[1:] == (1, config.vocab_size) - - consumer_embedding_scale = next( - node - for node in model.graph - if node.op_type == "Mul" and "embed_tokens_per_layer_split.2" in node.name - ) - assert any( - node.op_type == "Gather" - and any( - input_value is consumer_embedding_scale.outputs[0] for input_value in node.inputs - ) - for node in model.graph - ) - assert not any(node.op_type == "CastLike" for node in model.graph) - - -def test_multimodal_task_prunes_decoder_prefix() -> None: - config = _make_config(with_vision=True) - module = registry.get("gemma4")(config) - package = build_from_module( - module, - config, - task="gemma4", - execution_provider="webgpu", - prune_prefill_prefix=True, - ) - - logits = next( - value for value in package["decoder"].graph.outputs if value.name == "logits" - ) - assert logits.shape[1:] == (1, config.vocab_size) - - -def test_splits_per_layer_projection_weight() -> None: - config = _make_config() - original = torch.arange(32 * 64, dtype=torch.float32).reshape(32, 64) - state_dict = {"model.per_layer_model_projection.weight": original.clone()} - - _split_per_layer_projection_weight(state_dict, "model.", config) - - assert torch.equal(state_dict["model.per_layer_model_projection.weight"], original[:16]) - assert torch.equal( - state_dict["model.per_layer_model_projection_consumer.weight"], - original[16:], - ) +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT License. + +from __future__ import annotations + +import numpy as np +import onnx_ir as ir +import pytest +import torch + +from mobius import build_from_module +from mobius._configs import Gemma4Config, VisionConfig +from mobius._registry import registry +from mobius._testing.ort_inference import OnnxModelSession +from mobius.models.gemma4 import _split_per_layer_projection_weight + + +def _make_config(*, with_vision: bool = False) -> Gemma4Config: + return Gemma4Config( + num_hidden_layers=4, + hidden_size=64, + intermediate_size=128, + num_attention_heads=4, + num_key_value_heads=1, + head_dim=16, + vocab_size=256, + rms_norm_eps=1e-6, + hidden_act="silu", + attn_qk_norm=True, + layer_types=[ + "sliding_attention", + "full_attention", + "sliding_attention", + "full_attention", + ], + sliding_window=8, + global_head_dim=16, + global_rope_theta=10_000.0, + global_partial_rotary_factor=0.25, + final_logit_softcapping=30.0, + hidden_size_per_layer_input=8, + vocab_size_per_layer_input=64, + split_per_layer_embedding=True, + image_token_id=255999 if with_vision else None, + pad_token_id=0, + tie_word_embeddings=False, + num_kv_shared_layers=2, + vision=( + VisionConfig( + hidden_size=32, + intermediate_size=64, + num_hidden_layers=1, + num_attention_heads=2, + patch_size=16, + norm_eps=1e-6, + ) + if with_vision + else None + ), + ) + + +def test_prunes_gemma4_shared_layer_prefix() -> None: + config = _make_config() + module = registry.get("gemma4_text")(config) + model = build_from_module( + module, + config, + task="gemma4-text-generation", + execution_provider="webgpu", + prune_prefill_prefix=True, + )["model"] + + producer = next( + node + for node in model.graph + if node.op_type == "MatMul" and "/per_layer_model_projection/" in node.name + ) + consumer = next( + node + for node in model.graph + if node.op_type == "MatMul" and "/per_layer_model_projection_consumer/" in node.name + ) + assert producer.outputs[0].shape[1] != 1 + assert producer.outputs[0].shape[2] == 16 + assert consumer.inputs[0].shape[1:] == (1, 64) + assert consumer.outputs[0].shape[1:] == (1, 16) + + first_shared_norm = next( + node for node in model.graph if "layers.2/input_layernorm" in node.name + ) + assert first_shared_norm.inputs[0].shape[1:] == (1, 64) + logits = next(value for value in model.graph.outputs if value.name == "logits") + assert logits.shape[1:] == (1, config.vocab_size) + + consumer_embedding_scale = next( + node + for node in model.graph + if node.op_type == "Mul" and "embed_tokens_per_layer_split.2" in node.name + ) + assert any( + node.op_type == "Gather" + and any( + input_value is consumer_embedding_scale.outputs[0] for input_value in node.inputs + ) + for node in model.graph + ) + assert not any(node.op_type == "CastLike" for node in model.graph) + + +def test_multimodal_task_prunes_decoder_prefix() -> None: + config = _make_config(with_vision=True) + module = registry.get("gemma4")(config) + package = build_from_module( + module, + config, + task="gemma4", + execution_provider="webgpu", + prune_prefill_prefix=True, + ) + + logits = next( + value for value in package["decoder"].graph.outputs if value.name == "logits" + ) + assert logits.shape[1:] == (1, config.vocab_size) + + +def test_splits_per_layer_projection_weight() -> None: + config = _make_config() + original = torch.arange(32 * 64, dtype=torch.float32).reshape(32, 64) + state_dict = {"model.per_layer_model_projection.weight": original.clone()} + + _split_per_layer_projection_weight(state_dict, "model.", config) + + assert torch.equal(state_dict["model.per_layer_model_projection.weight"], original[:16]) + assert torch.equal( + state_dict["model.per_layer_model_projection_consumer.weight"], + original[16:], + ) + + +# --------------------------------------------------------------------------- +# Numerical parity: pruned package must reproduce the unpruned final row +# --------------------------------------------------------------------------- + + +def _parity_config() -> Gemma4Config: + """Tiny hybrid config with a KV-shared tail and two distinct head sizes.""" + return Gemma4Config( + num_hidden_layers=6, + hidden_size=64, + intermediate_size=128, + num_attention_heads=4, + num_key_value_heads=1, + head_dim=16, + vocab_size=256, + rms_norm_eps=1e-6, + hidden_act="silu", + attn_qk_norm=True, + layer_types=[ + "sliding_attention", + "full_attention", + "sliding_attention", + "full_attention", + "sliding_attention", + "full_attention", + ], + sliding_window=8, + # Distinct global head size: full-attention layers cache 32-wide K/V, + # sliding layers cache 16-wide K/V. + global_head_dim=32, + global_rope_theta=10_000.0, + global_partial_rotary_factor=0.25, + final_logit_softcapping=30.0, + hidden_size_per_layer_input=8, + vocab_size_per_layer_input=64, + split_per_layer_embedding=True, + max_position_embeddings=256, + pad_token_id=0, + tie_word_embeddings=False, + num_kv_shared_layers=2, + ) + + +def _build_text_model(config: Gemma4Config, *, execution_provider: str, prune: bool): + module = registry.get("gemma4_text")(config) + return build_from_module( + module, + config, + task="gemma4-text-generation", + execution_provider=execution_provider, + prune_prefill_prefix=prune, + )["model"] + + +def _fill_random_weights(model, seed: int = 0) -> dict[str, np.ndarray]: + rng = np.random.default_rng(seed) + weights: dict[str, np.ndarray] = {} + for initializer in model.graph.initializers.values(): + if initializer.const_value is None: + array = (rng.standard_normal(tuple(initializer.shape)) * 0.05).astype(np.float32) + initializer.const_value = ir.tensor(array, name=initializer.name) + weights[initializer.name] = initializer.const_value.numpy() + return weights + + +def _copy_weights(model, weights: dict[str, np.ndarray]) -> None: + for initializer in model.graph.initializers.values(): + if initializer.name in weights: + initializer.const_value = ir.tensor( + weights[initializer.name], name=initializer.name + ) + elif initializer.const_value is None: + raise AssertionError( + f"pruned graph declares weight {initializer.name!r} that the " + "unpruned graph does not" + ) + + +def _feeds(config: Gemma4Config, seq_len: int) -> dict[str, np.ndarray]: + feeds: dict[str, np.ndarray] = { + "input_ids": (np.arange(1, seq_len + 1, dtype=np.int64) % config.vocab_size)[None], + "attention_mask": np.ones((1, seq_len), dtype=np.int64), + "position_ids": np.arange(seq_len, dtype=np.int64)[None], + } + # Cache-owning layers are the contiguous prefix before the KV-shared tail; + # each carries its own head size (global layers are double-wide). + kv_layers = config.num_hidden_layers - (config.num_kv_shared_layers or 0) + for index in range(kv_layers): + head_dim = 32 if config.layer_types[index] == "full_attention" else 16 + empty = np.zeros((1, config.num_key_value_heads, 0, head_dim), dtype=np.float32) + feeds[f"past_key_values.{index}.key"] = empty + feeds[f"past_key_values.{index}.value"] = empty.copy() + return feeds + + +@pytest.mark.parametrize( + "execution_provider", + # "default" exercises the opset-24 Attention + RotaryEmbedding path; + # "cpu" exercises the fused GroupQueryAttention path. + ["default", "cpu"], +) +@pytest.mark.parametrize("seq_len", [5, 12]) +def test_pruned_prefill_matches_unpruned_final_row( + execution_provider: str, seq_len: int +) -> None: + """Prefill-prefix pruning must be a pure graph-surface optimisation. + + Regression guard for the Gemma 4 mid-stack truncation: at the first + KV-shared layer the hidden states narrow to a single query position, so the + per-layer RoPE ``(cos, sin)`` caches and the additive attention bias must + narrow with them. Without that, the ``RotaryEmbedding``/``Attention`` path + fails outright at load/run time; ``seq_len=12`` additionally reaches past + the 8-token sliding window so global (full-attention) layers exercise a + different key extent from the sliding ones. + """ + config = _parity_config() + base = _build_text_model(config, execution_provider=execution_provider, prune=False) + pruned = _build_text_model(config, execution_provider=execution_provider, prune=True) + + weights = _fill_random_weights(base) + _copy_weights(pruned, weights) + + feeds = _feeds(config, seq_len) + base_out = OnnxModelSession(base).run(feeds) + pruned_out = OnnxModelSession(pruned).run(feeds) + + assert set(base_out) == set(pruned_out), "pruning changed the model's output surface" + + # Logits: the pruned package emits only the final row. + expected_logits = base_out["logits"][:, -1:, :] + assert pruned_out["logits"].shape == expected_logits.shape + np.testing.assert_allclose(pruned_out["logits"], expected_logits, atol=1e-4, rtol=0) + assert np.argmax(pruned_out["logits"][0, 0]) == np.argmax(expected_logits[0, 0]) + + # KV cache: pruning must not touch cache-owning layers at all. + kv_layers = config.num_hidden_layers - (config.num_kv_shared_layers or 0) + present_names = sorted(name for name in base_out if name.startswith("present.")) + assert len(present_names) == 2 * kv_layers, ( + f"expected {kv_layers} cache-owning layers, got {present_names}" + ) + for name in present_names: + np.testing.assert_allclose(pruned_out[name], base_out[name], atol=1e-5, rtol=0) + + # Double head size survives pruning: global layers stay twice as wide. + assert base_out["present.0.key"].shape[-1] == config.head_dim + assert base_out["present.1.key"].shape[-1] == config.global_head_dim + + +def _dual_head_dim_config() -> Gemma4Config: + """A Gemma 4 text config whose global layers use a wider head than sliding ones.""" + config = _make_config() + config.global_head_dim = 32 + return config + + +def test_metadata_splits_cache_groups_by_attention_kind() -> None: + """Gemma 4's two cache geometries must surface as two declared state groups. + + Local/sliding layers store ``head_dim``-wide entries and are prefix-evictable; + global/full-attention layers store ``global_head_dim``-wide entries and keep + the entire history. A single undifferentiated group would let a runtime apply + sliding-window eviction to the global layers and corrupt them. + """ + from mobius.integrations.onnx_genai.workflow_metadata import ( + build_decoder_workflow_metadata, + ) + + config = _dual_head_dim_config() + module = registry.get("gemma4_text")(config) + pkg = build_from_module(module, config, task="gemma4-text-generation") + + metadata = build_decoder_workflow_metadata(pkg, config) + workflow = metadata["pipeline"]["workflow"] + groups = workflow["serving"]["state_service"]["groups"] + + assert set(groups) == { + "decoder_cache_sliding_attention", + "decoder_cache_full_attention", + } + assert groups["decoder_cache_sliding_attention"]["kind"] == "sliding_attention" + assert groups["decoder_cache_full_attention"]["kind"] == "full_attention" + # Only the sliding layers may drop their prefix. + assert groups["decoder_cache_sliding_attention"]["reuse"]["evictable_prefix"] is True + assert groups["decoder_cache_full_attention"]["reuse"]["evictable_prefix"] is False + + head_dims: dict[str, set[int]] = {} + for name, group in groups.items(): + cells = {cell for ports in group["ports"].values() for cell in ports} + head_dims[name] = {workflow["state"][cell]["contract"]["shape"][-1] for cell in cells} + assert head_dims["decoder_cache_sliding_attention"] == {config.head_dim} + assert head_dims["decoder_cache_full_attention"] == {config.global_head_dim} + + +def test_metadata_declares_no_cache_for_kv_shared_layers() -> None: + """KV-shared layers borrow K/V and must not own phantom cache slots.""" + from mobius.integrations.onnx_genai.workflow_metadata import ( + build_decoder_workflow_metadata, + ) + + config = _dual_head_dim_config() + module = registry.get("gemma4_text")(config) + pkg = build_from_module(module, config, task="gemma4-text-generation") + + metadata = build_decoder_workflow_metadata(pkg, config) + workflow = metadata["pipeline"]["workflow"] + groups = workflow["serving"]["state_service"]["groups"] + cells = { + cell + for group in groups.values() + for ports in group["ports"].values() + for cell in ports + } + + cache_owning_layers = config.num_hidden_layers - config.num_kv_shared_layers + assert len(cells) == 2 * cache_owning_layers + + present_outputs = [ + value.name for value in pkg["model"].graph.outputs if value.name.startswith("present.") + ] + assert len(present_outputs) == 2 * cache_owning_layers + assert max(int(name.split(".")[1]) for name in present_outputs) == cache_owning_layers - 1 + + +def test_metadata_cache_cells_are_runtime_managed_and_request_aligned() -> None: + """Cache cells the runtime allocates must declare ownership and a row axis.""" + from mobius.integrations.onnx_genai.workflow_metadata import ( + build_decoder_workflow_metadata, + ) + + config = _dual_head_dim_config() + module = registry.get("gemma4_text")(config) + pkg = build_from_module(module, config, task="gemma4-text-generation") + + workflow = build_decoder_workflow_metadata(pkg, config)["pipeline"]["workflow"] + groups = workflow["serving"]["state_service"]["groups"] + for group in groups.values(): + # Runtime-private storage: no allocator/paging/slot policy is serialized. + assert "storage" not in group + assert "paging" not in group + # This package is exported on the opset-24 ``Attention`` path, whose + # cache grows by concatenation, so ``present`` is a fresh tensor rather + # than a view onto ``past``. + assert group["aliasing"] == "forbidden" + for ports in group["ports"].values(): + for cell in ports: + state = workflow["state"][cell] + assert state["management"] == "runtime" + assert state["release_boundary"] == "invocation" + assert state["contract"]["batch_layout"] == { + "kind": "request_aligned", + "axis": 0, + } + assert "slot_ids" not in workflow["serving"] diff --git a/tests/generate_onnx_genai_validation_packages.py b/tests/generate_onnx_genai_validation_packages.py new file mode 100644 index 000000000..05d341829 --- /dev/null +++ b/tests/generate_onnx_genai_validation_packages.py @@ -0,0 +1,1123 @@ +from __future__ import annotations + +import argparse +import json +import shutil +from pathlib import Path + +import numpy as np +import onnx_ir as ir +import yaml +from onnxscript import GraphBuilder +from safetensors.numpy import save_file + +from mobius._model_package import ModelPackage +from mobius._passes import RemoveDeadGraphInputsPass +from mobius.adapter_io import load_peft_adapter +from mobius.adapters import ( + AdapterArtifact, + AdapterServiceOptions, + AdapterTarget, + AdapterTargetDescriptor, + AdapterTargetManifest, + AdapterWeights, + fingerprint_model_weights, +) +from mobius.integrations.onnx_genai import write_onnx_genai_config +from mobius.integrations.onnx_genai.auto_export_test import ( + _Cfg, + _VlmCfg, +) +from mobius.integrations.onnx_genai.inference_metadata import ( + add_adapter_service_to_metadata, +) +from mobius.integrations.onnx_genai.workflow_metadata import ( + write_audio_codec_workflow_metadata, + write_diffusion_workflow_metadata, + write_encoder_embedding_workflow_metadata, + write_language_diffusion_workflow_metadata, + write_speculative_workflow_metadata, + write_video_diffusion_workflow_metadata, +) +from mobius.models.bert import BertModel +from mobius.models.bert_test import PROTBERT_TINY_CONFIG +from mobius.models.esm import EsmModel +from mobius.models.esm_test import TINY_CONFIG as ESM2_TINY_CONFIG +from mobius.models.qwen3_tts import Qwen3TTSForConditionalGeneration +from mobius.models.qwen3_tts_test import _TINY_CONFIG +from mobius.tasks import FeatureExtractionTask, TTSTask + + +def _materialize_deterministic_initializers(package: ModelPackage) -> None: + """Give real tiny producer graphs environment-independent synthetic weights.""" + for model in package.values(): + for value in model.graph.initializers.values(): + if value.const_value is not None and not value.dtype.is_floating_point(): + continue + shape = [int(dimension) for dimension in value.shape] + value.const_value = ir.tensor(np.zeros(shape, dtype=value.dtype.numpy())) + + +def _tts_package() -> ModelPackage: + package = TTSTask().build( + Qwen3TTSForConditionalGeneration(_TINY_CONFIG), + _TINY_CONFIG, + ) + _materialize_deterministic_initializers(package) + graph, builder = _graph("codec") + codes = builder.input("codes", ir.DataType.INT64, ["batch", 4, "frames"]) + waveform = builder.op.Cast( + builder.op.Slice( + codes, + builder.op.Constant(value_ints=[0]), + builder.op.Constant(value_ints=[1]), + builder.op.Constant(value_ints=[1]), + ), + to=ir.DataType.FLOAT, + ) + builder.add_output( + _typed(waveform, ir.DataType.FLOAT, ["batch", 1, "frames"]), + "waveform", + ) + package["codec"] = ir.Model(graph, ir_version=11) + return package + + +def _graph(name: str) -> tuple[ir.Graph, GraphBuilder]: + graph = ir.Graph([], [], nodes=[], name=name, opset_imports={"": 24}) + return graph, GraphBuilder(graph) + + +def _typed(value: ir.Value, dtype: ir.DataType, shape: list[int | str]) -> ir.Value: + value.type = ir.TensorType(dtype) + value.shape = ir.Shape(shape) + return value + + +def _executable_decoder_package() -> ModelPackage: + graph, builder = _graph("decoder") + input_ids = builder.input("input_ids", ir.DataType.INT64, ["batch", "sequence"]) + builder.input("attention_mask", ir.DataType.INT64, ["batch", "total_sequence"]) + builder.input("position_ids", ir.DataType.INT64, ["batch", "sequence"]) + past = builder.input( + "past_key_values.0.key", + ir.DataType.FLOAT, + ["batch", 2, "past_sequence", 8], + ) + shape = builder.op.Shape(input_ids) + logits = builder.op.ConstantOfShape( + builder.op.Concat(shape, builder.op.Constant(value_ints=[128]), axis=0), + value=ir.tensor([0.0]), + ) + batch = builder.op.Shape(input_ids, start=0, end=1) + sequence = builder.op.Shape(input_ids, start=1, end=2) + cache_shape = builder.op.Concat( + batch, + builder.op.Constant(value_ints=[2]), + sequence, + builder.op.Constant(value_ints=[8]), + axis=0, + ) + update = builder.op.ConstantOfShape(cache_shape, value=ir.tensor([0.0])) + present = builder.op.Concat(past, update, axis=2) + builder.add_output( + _typed(logits, ir.DataType.FLOAT, ["batch", "sequence", 128]), + "logits", + ) + builder.add_output( + _typed(present, ir.DataType.FLOAT, ["batch", 2, "present_sequence", 8]), + "present.0.key", + ) + config = _Cfg() + config.eos_token_id = 127 + return ModelPackage({"model": ir.Model(graph, ir_version=11)}, config=config) + + +def _executable_static_cache_package() -> ModelPackage: + """A runnable decoder that scatters into fixed-capacity KV buffers. + + The appending decoder fixture grows its cache with ``Concat``; this one + keeps a preallocated ``[batch, capacity, kv_hidden]`` buffer and writes each + step at a per-row cursor with ``TensorScatter``. That is the shape the + workflow's ``indexed_scatter`` state discipline describes, so the engine + conformance run exercises the write cursor and the fixed-capacity carry + rather than only the growing-tensor path. + """ + capacity = 16 + graph, builder = _graph("decoder") + input_ids = builder.input("input_ids", ir.DataType.INT64, ["batch", "sequence"]) + builder.input("position_ids", ir.DataType.INT64, ["batch", "sequence"]) + key_cache = builder.input("key_cache.0", ir.DataType.FLOAT, ["batch", capacity, 8]) + value_cache = builder.input("value_cache.0", ir.DataType.FLOAT, ["batch", capacity, 8]) + write_indices = builder.input("write_indices", ir.DataType.INT64, ["batch"]) + builder.input("nonpad_kv_seqlen", ir.DataType.INT64, ["batch"]) + + shape = builder.op.Shape(input_ids) + logits = builder.op.ConstantOfShape( + builder.op.Concat(shape, builder.op.Constant(value_ints=[128]), axis=0), + value=ir.tensor([0.0]), + ) + # This step's keys/values: (batch, sequence, 8), scattered into the buffer + # at row `write_indices[b]` along the capacity axis. + update_shape = builder.op.Concat(shape, builder.op.Constant(value_ints=[8]), axis=0) + update = builder.op.ConstantOfShape(update_shape, value=ir.tensor([0.0])) + updated_key = builder.op.TensorScatter(key_cache, update, write_indices, axis=1) + updated_value = builder.op.TensorScatter(value_cache, update, write_indices, axis=1) + + builder.add_output( + _typed(logits, ir.DataType.FLOAT, ["batch", "sequence", 128]), + "logits", + ) + builder.add_output( + _typed(updated_key, ir.DataType.FLOAT, ["batch", capacity, 8]), + "updated_key_cache.0", + ) + builder.add_output( + _typed(updated_value, ir.DataType.FLOAT, ["batch", capacity, 8]), + "updated_value_cache.0", + ) + config = _Cfg() + config.eos_token_id = 127 + return ModelPackage({"model": ir.Model(graph, ir_version=11)}, config=config) + + +def _executable_vlm_package() -> ModelPackage: + vision_graph, vision_builder = _graph("vision_encoder") + pixel_values = vision_builder.input("pixel_values", ir.DataType.FLOAT, [4, 1176]) + vision_builder.input("grid_thw", ir.DataType.INT64, [1, 3]) + image_scalar = vision_builder.op.ReduceMean(pixel_values) + image_features = vision_builder.op.Expand( + image_scalar, + vision_builder.op.Constant(value_ints=[1, 4, 32]), + ) + vision_builder.add_output( + _typed(image_features, ir.DataType.FLOAT, [1, 4, 32]), + "image_features", + ) + + embedding_graph, embedding_builder = _graph("embedding") + input_ids = embedding_builder.input("input_ids", ir.DataType.INT64, ["batch", "sequence"]) + image_features = embedding_builder.input( + "image_features", ir.DataType.FLOAT, ["batch", 4, 32] + ) + token_values = embedding_builder.op.Cast( + embedding_builder.op.Unsqueeze(input_ids, [2]), + to=ir.DataType.FLOAT, + ) + token_shape = embedding_builder.op.Concat( + embedding_builder.op.Shape(input_ids), + embedding_builder.op.Constant(value_ints=[32]), + axis=0, + ) + image_bias = embedding_builder.op.ReduceMean(image_features) + inputs_embeds = embedding_builder.op.Add( + embedding_builder.op.Expand(token_values, token_shape), + image_bias, + ) + embedding_builder.add_output( + _typed(inputs_embeds, ir.DataType.FLOAT, ["batch", "sequence", 32]), + "inputs_embeds", + ) + + decoder_graph, decoder_builder = _graph("decoder") + inputs_embeds = decoder_builder.input( + "inputs_embeds", ir.DataType.FLOAT, ["batch", "sequence", 32] + ) + decoder_builder.input( + "attention_mask", + ir.DataType.INT64, + ["batch", "past_sequence + sequence"], + ) + decoder_builder.input("position_ids", ir.DataType.INT64, ["batch", "sequence"]) + past_key = decoder_builder.input( + "past_key_values.0.key", + ir.DataType.FLOAT, + ["batch", 2, "past_sequence", 8], + ) + past_value = decoder_builder.input( + "past_key_values.0.value", + ir.DataType.FLOAT, + ["batch", 2, "past_sequence", 8], + ) + batch = decoder_builder.op.Shape(inputs_embeds, start=0, end=1) + sequence = decoder_builder.op.Shape(inputs_embeds, start=1, end=2) + logits_shape = decoder_builder.op.Concat( + batch, + sequence, + decoder_builder.op.Constant(value_ints=[128]), + axis=0, + ) + logits = decoder_builder.op.Expand( + decoder_builder.op.ReduceMean(inputs_embeds, axes=[2], keepdims=1), + logits_shape, + ) + cache_shape = decoder_builder.op.Concat( + batch, + decoder_builder.op.Constant(value_ints=[2]), + sequence, + decoder_builder.op.Constant(value_ints=[8]), + axis=0, + ) + cache_update = decoder_builder.op.Add( + decoder_builder.op.ConstantOfShape( + cache_shape, + value=ir.tensor([0.0]), + ), + decoder_builder.op.ReduceMean(inputs_embeds), + ) + present_key = decoder_builder.op.Concat(past_key, cache_update, axis=2) + present_value = decoder_builder.op.Concat(past_value, cache_update, axis=2) + decoder_builder.add_output( + _typed(logits, ir.DataType.FLOAT, ["batch", "sequence", 128]), + "logits", + ) + decoder_builder.add_output( + _typed( + present_key, + ir.DataType.FLOAT, + ["batch", 2, "total_sequence", 8], + ), + "present.0.key", + ) + decoder_builder.add_output( + _typed( + present_value, + ir.DataType.FLOAT, + ["batch", 2, "total_sequence", 8], + ), + "present.0.value", + ) + return ModelPackage( + { + "vision_encoder": ir.Model(vision_graph, ir_version=11), + "embedding": ir.Model(embedding_graph, ir_version=11), + "decoder": ir.Model(decoder_graph, ir_version=11), + }, + config=_VlmCfg(), + ) + + +def _executable_diffusion_package() -> ModelPackage: + text_graph, text_builder = _graph("text_encoder") + input_ids = text_builder.input( + "input_ids", ir.DataType.INT64, ["batch", "prompt_sequence"] + ) + hidden_shape = text_builder.op.Concat( + text_builder.op.Shape(input_ids), + text_builder.op.Constant(value_ints=[32]), + axis=0, + ) + hidden = text_builder.op.Expand( + text_builder.op.Cast( + text_builder.op.Unsqueeze(input_ids, [2]), + to=ir.DataType.FLOAT, + ), + hidden_shape, + ) + text_builder.add_output( + _typed( + hidden, + ir.DataType.FLOAT, + ["batch", "prompt_sequence", 32], + ), + "encoder_hidden_states", + ) + + denoiser_graph, denoiser_builder = _graph("denoiser") + sample = denoiser_builder.input( + "sample", ir.DataType.FLOAT, ["batch", 4, "height", "width"] + ) + timestep = denoiser_builder.input("timestep", ir.DataType.FLOAT, ["batch"]) + conditioning = denoiser_builder.input( + "encoder_hidden_states", + ir.DataType.FLOAT, + ["batch", "prompt_sequence", 32], + ) + batch = denoiser_builder.op.Shape(sample, start=0, end=1) + scalar_shape = denoiser_builder.op.Concat( + batch, + denoiser_builder.op.Constant(value_ints=[1, 1, 1]), + axis=0, + ) + timestep_bias = denoiser_builder.op.Reshape(timestep, scalar_shape) + conditioning_bias = denoiser_builder.op.Reshape( + denoiser_builder.op.ReduceMean(conditioning, axes=[1, 2]), + scalar_shape, + ) + estimate = denoiser_builder.op.Add( + sample, + denoiser_builder.op.Add(timestep_bias, conditioning_bias), + ) + denoiser_builder.add_output( + _typed( + estimate, + ir.DataType.FLOAT, + ["batch", 4, "height", "width"], + ), + "noise_pred", + ) + + vae_graph, vae_builder = _graph("vae_decoder") + latent = vae_builder.input("latent", ir.DataType.FLOAT, ["batch", 4, "height", "width"]) + image = vae_builder.op.Slice( + latent, + vae_builder.op.Constant(value_ints=[0]), + vae_builder.op.Constant(value_ints=[3]), + vae_builder.op.Constant(value_ints=[1]), + ) + vae_builder.add_output( + _typed( + image, + ir.DataType.FLOAT, + ["batch", 3, "height", "width"], + ), + "image", + ) + return ModelPackage( + { + "text_encoder": ir.Model(text_graph, ir_version=11), + "denoiser": ir.Model(denoiser_graph, ir_version=11), + "vae_decoder": ir.Model(vae_graph, ir_version=11), + } + ) + + +def _causal_temporal_step( + builder, + value: ir.Value, + cache: ir.Value, + taps: int, +) -> tuple[ir.Value, ir.Value]: + """One causal temporal convolution over frames, threaded through a cache. + + Reproduces the structure a real causal video decoder relies on: the frames + that precede the current chunk come from the cache, the first chunk of a + clip replicates its own first frame instead, and the tail of the padded + input becomes the next chunk's cache. Returns the filtered frames and the + cache to carry. + """ + op = builder.op + padded = op.Concat(cache, value, axis=2) + # A zero-length cache means this is the clip's first chunk, so the missing + # history is the first frame repeated -- the same branch-free trick the real + # decoder uses to avoid a first-chunk special case. + have = op.Min( + op.Shape(cache, start=2, end=3), + op.Constant(value_ints=[taps]), + ) + front = op.Sub(op.Constant(value_ints=[taps]), have) + padded = op.Pad( + padded, + op.Concat(front, op.Constant(value_ints=[0]), axis=0), + None, + op.Constant(value_ints=[2]), + mode="edge", + ) + length = op.Shape(padded, start=2, end=3) + current = op.Slice( + padded, + op.Constant(value_ints=[taps]), + length, + op.Constant(value_ints=[2]), + ) + history = op.Slice( + padded, + op.Constant(value_ints=[0]), + op.Sub(length, op.Constant(value_ints=[taps])), + op.Constant(value_ints=[2]), + ) + filtered = op.Mul( + op.Add(current, history), + op.CastLike(op.Constant(value_float=0.5), current), + ) + next_cache = op.Slice( + padded, + op.Sub(length, op.Constant(value_ints=[taps])), + length, + op.Constant(value_ints=[2]), + ) + return filtered, next_cache + + +def _executable_video_package() -> ModelPackage: + """A rank-5 video denoiser plus a causal, chunked video decoder. + + Small enough to run anywhere, but structurally a video pipeline: the latent + carries a temporal axis through the denoise loop, and the decoder expands + frames, works at two spatial resolutions, and keeps per-resolution + convolution caches so a clip can be decoded a chunk at a time. + """ + denoiser_graph, denoiser_builder = _graph("transformer") + op = denoiser_builder.op + sample = denoiser_builder.input( + "sample", ir.DataType.FLOAT, ["batch", "num_frames", 4, "height", "width"] + ) + timestep = denoiser_builder.input("timestep", ir.DataType.INT64, ["batch"]) + conditioning = denoiser_builder.input( + "encoder_hidden_states", ir.DataType.FLOAT, ["batch", "prompt_sequence", 32] + ) + scalar_shape = op.Concat( + op.Shape(sample, start=0, end=1), + op.Constant(value_ints=[1, 1, 1, 1]), + axis=0, + ) + timestep_bias = op.Reshape( + op.Div(op.Cast(timestep, to=ir.DataType.FLOAT), op.Constant(value_float=1000.0)), + scalar_shape, + ) + conditioning_bias = op.Reshape(op.ReduceMean(conditioning, axes=[1, 2]), scalar_shape) + # A frame-dependent term, so a stage that collapsed or reordered the temporal + # axis would change the result rather than silently pass. + frame_index = op.Cast( + op.Range( + op.Squeeze(op.Constant(value_ints=[0])), + op.Squeeze(op.Shape(sample, start=1, end=2)), + op.Squeeze(op.Constant(value_ints=[1])), + ), + to=ir.DataType.FLOAT, + ) + frame_bias = op.Div( + op.Reshape(frame_index, op.Constant(value_ints=[1, -1, 1, 1, 1])), + op.Constant(value_float=100.0), + ) + estimate = op.Add( + op.Mul(sample, op.Constant(value_float=0.5)), + op.Add(op.Add(timestep_bias, conditioning_bias), frame_bias), + ) + denoiser_builder.add_output( + _typed(estimate, ir.DataType.FLOAT, ["batch", "num_frames", 4, "height", "width"]), + "noise_pred", + ) + + vae_graph, vae_builder = _graph("vae_decoder") + op = vae_builder.op + latent = vae_builder.input( + "latent_sample", + ir.DataType.FLOAT, + ["batch", 4, "latent_frames", "latent_height", "latent_width"], + ) + cache_in = vae_builder.input( + "conv_cache.conv_in", + ir.DataType.FLOAT, + ["batch", 4, "cache_frames", "latent_height", "latent_width"], + ) + cache_out_port = vae_builder.input( + "conv_cache.conv_out", + ir.DataType.FLOAT, + ["batch", 3, "cache_frames", "2*latent_height", "2*latent_width"], + ) + filtered, next_cache_in = _causal_temporal_step(vae_builder, latent, cache_in, 1) + # Nearest-neighbour expansion in time and space: [B,C,T,H,W] -> [B,C,2T,2H,2W]. + shape = op.Shape(filtered) + batch = op.Slice(shape, [0], [1], [0]) + channels = op.Slice(shape, [1], [2], [0]) + frames = op.Slice(shape, [2], [3], [0]) + height = op.Slice(shape, [3], [4], [0]) + width = op.Slice(shape, [4], [5], [0]) + one = op.Constant(value_ints=[1]) + two = op.Constant(value_ints=[2]) + expanded = op.Reshape( + filtered, + op.Concat(batch, channels, frames, one, height, one, width, one, axis=0), + ) + expanded = op.Expand( + expanded, + op.Concat(batch, channels, frames, two, height, two, width, two, axis=0), + ) + expanded = op.Reshape( + expanded, + op.Concat( + batch, + channels, + op.Mul(frames, two), + op.Mul(height, two), + op.Mul(width, two), + axis=0, + ), + ) + rgb = op.Slice(expanded, [0], [3], [1]) + decoded, next_cache_out = _causal_temporal_step(vae_builder, rgb, cache_out_port, 1) + vae_builder.add_output( + _typed( + decoded, + ir.DataType.FLOAT, + ["batch", 3, "frames", "2*latent_height", "2*latent_width"], + ), + "sample", + ) + vae_builder.add_output( + _typed( + next_cache_in, + ir.DataType.FLOAT, + ["batch", 4, "cache_frames", "latent_height", "latent_width"], + ), + "conv_cache_out.conv_in", + ) + vae_builder.add_output( + _typed( + next_cache_out, + ir.DataType.FLOAT, + ["batch", 3, "cache_frames", "2*latent_height", "2*latent_width"], + ), + "conv_cache_out.conv_out", + ) + vae_model = ir.Model(vae_graph, ir_version=11) + vae_model.metadata_props["mobius.conv_cache.spatial_scale.conv_cache.conv_in"] = "1" + vae_model.metadata_props["mobius.conv_cache.spatial_scale.conv_cache.conv_out"] = "2" + return ModelPackage( + { + "transformer": ir.Model(denoiser_graph, ir_version=11), + "vae_decoder": vae_model, + } + ) + + +def _executable_speculative_package() -> ModelPackage: + proposer_graph, proposer_builder = _graph("proposer") + tokens = proposer_builder.input("tokens", ir.DataType.INT64, ["batch", 4]) + proposer_builder.input("proposal_budget", ir.DataType.INT64, ["batch"]) + proposal_scores = proposer_builder.op.ConstantOfShape( + proposer_builder.op.Concat( + proposer_builder.op.Shape(tokens), + proposer_builder.op.Constant(value_ints=[32]), + axis=0, + ), + value=ir.tensor([0.0]), + ) + proposer_builder.add_output( + _typed( + proposer_builder.op.Identity(tokens), + ir.DataType.INT64, + ["batch", 4], + ), + "proposed_tokens", + ) + proposer_builder.add_output( + _typed(proposal_scores, ir.DataType.FLOAT, ["batch", 4, 32]), + "proposal_scores", + ) + + verifier_graph, verifier_builder = _graph("verifier") + proposed = verifier_builder.input("proposed_tokens", ir.DataType.INT64, ["batch", 4]) + past = verifier_builder.input( + "past_key_values.0.key", + ir.DataType.FLOAT, + ["batch", 2, "past_sequence", 8], + ) + batch = verifier_builder.op.Shape(proposed, start=0, end=1) + row = verifier_builder.op.Range( + verifier_builder.op.Constant(value_int=0), + verifier_builder.op.Squeeze(batch, [0]), + verifier_builder.op.Constant(value_int=1), + ) + reject_at = verifier_builder.op.Min( + verifier_builder.op.Add(row, verifier_builder.op.Constant(value_int=1)), + verifier_builder.op.Constant(value_int=3), + ) + indices = verifier_builder.op.Unsqueeze(reject_at, [1]) + corrections = verifier_builder.op.Expand( + verifier_builder.op.Constant(value_int=31), + batch, + ) + target_tokens = verifier_builder.op.ScatterElements( + proposed, + indices, + verifier_builder.op.Unsqueeze(corrections, [1]), + axis=1, + ) + target_scores = verifier_builder.op.OneHot( + target_tokens, + verifier_builder.op.Constant(value_int=32), + verifier_builder.op.Constant(value_floats=[0.0, 1.0]), + axis=-1, + ) + cache_update = verifier_builder.op.ConstantOfShape( + verifier_builder.op.Concat( + batch, + verifier_builder.op.Constant(value_ints=[2, 4, 8]), + axis=0, + ), + value=ir.tensor([0.0]), + ) + present = verifier_builder.op.Concat(past, cache_update, axis=2) + verifier_builder.add_output( + _typed(target_scores, ir.DataType.FLOAT, ["batch", 4, 32]), + "target_scores", + ) + verifier_builder.add_output( + _typed( + present, + ir.DataType.FLOAT, + ["batch", 2, "past_sequence + 4", 8], + ), + "present.0.key", + ) + return ModelPackage( + { + "proposer": ir.Model(proposer_graph, ir_version=11), + "verifier": ir.Model(verifier_graph, ir_version=11), + } + ) + + +def _executable_masked_package() -> ModelPackage: + graph, builder = _graph("masked_denoiser") + input_ids = builder.input("input_ids", ir.DataType.INT64, ["batch", "sequence"]) + logits = builder.op.ConstantOfShape( + builder.op.Concat( + builder.op.Shape(input_ids), + builder.op.Constant(value_ints=[128]), + axis=0, + ), + value=ir.tensor([0.0]), + ) + builder.add_output( + _typed(logits, ir.DataType.FLOAT, ["batch", "sequence", 128]), + "logits", + ) + builder.add_output( + _typed( + builder.op.Identity(input_ids), + ir.DataType.INT64, + ["batch", "sequence"], + ), + "proposed_tokens", + ) + return ModelPackage({"model": ir.Model(graph, ir_version=11)}) + + +def _esm2_embedding_package() -> ModelPackage: + """A tiny but real ESM-2 protein encoder. + + Built from the same producer path as ``facebook/esm2_t6_8M_UR50D`` so the + fixture exercises rotary positions, pre-norm blocks and the token-dropout + rescale rather than a stand-in graph. ESM-2 has no token-type embedding, so + the saved artifact carries only ``input_ids`` and ``attention_mask`` -- the + asymmetry against ProtBert below is the point of shipping both. + """ + config = ESM2_TINY_CONFIG + package = FeatureExtractionTask().build(EsmModel(config), config) + # The feature-extraction task offers ``token_type_ids`` to every encoder, + # but ESM-2 has no token-type embedding and never reads it. The real export + # path drops the dead input during optimization, so drop it here too -- + # otherwise the committed metadata would declare a port the shipped + # artifact does not expose. + RemoveDeadGraphInputsPass()(package["model"]) + _materialize_deterministic_initializers(package) + return package + + +def _protbert_embedding_package() -> ModelPackage: + """A tiny but real ProtBert-shaped encoder (BERT with an amino-acid vocab).""" + config = PROTBERT_TINY_CONFIG + package = FeatureExtractionTask().build(BertModel(config), config) + _materialize_deterministic_initializers(package) + return package + + +def _executable_codec_package() -> ModelPackage: + encoder_graph, encoder_builder = _graph("encoder") + waveform = encoder_builder.input( + "waveform", + ir.DataType.FLOAT, + ["batch", 1, "audio_samples"], + ) + encoder_builder.add_output( + _typed( + encoder_builder.op.Identity(waveform), + ir.DataType.FLOAT, + ["batch", 1, "audio_samples"], + ), + "codes", + ) + + decoder_graph, decoder_builder = _graph("decoder") + codes = decoder_builder.input( + "codes", + ir.DataType.FLOAT, + ["batch", 1, "audio_samples"], + ) + decoder_builder.add_output( + _typed( + decoder_builder.op.Identity(codes), + ir.DataType.FLOAT, + ["batch", 1, "audio_samples"], + ), + "waveform", + ) + return ModelPackage( + { + "encoder": ir.Model(encoder_graph, ir_version=11), + "decoder": ir.Model(decoder_graph, ir_version=11), + } + ) + + +def _adapter_package(source_root: Path) -> ModelPackage: + graph, builder = _graph("decoder") + activations = builder.input("activations", ir.DataType.FLOAT, ["batch", 2]) + weight = ir.Value( + name="projection", + const_value=ir.tensor(np.eye(2, dtype=np.float32)), + type=ir.TensorType(ir.DataType.FLOAT), + shape=ir.Shape([2, 2]), + ) + graph.initializers.add(weight) + projection = builder.op.MatMul(activations, weight) + projection.producer().name = "projection" + projection.name = "projection.output" + builder.add_output( + _typed(projection, ir.DataType.FLOAT, ["batch", 2]), + "projection.output", + ) + model = ir.Model(graph, ir_version=11) + target = AdapterTarget("decoder", "projection") + descriptor = AdapterTargetDescriptor( + target, + semantic_name="projection", + node_name="projection", + output_name="projection.output", + input_size=2, + output_size=2, + layer_index=0, + rank=1, + alpha=1.0, + activation_dtype=ir.DataType.FLOAT, + graph_input_a="lora.projection.a", + graph_input_b="lora.projection.b", + graph_input_scale="lora.projection.scale", + ) + fingerprint = fingerprint_model_weights({"decoder": model}, (descriptor,)) + manifest = AdapterTargetManifest(fingerprint, (descriptor,)) + package = ModelPackage( + {"decoder": model}, + adapter_target_manifest=manifest, + adapter_service_options=AdapterServiceOptions( + active="request.active", + max_adapters=2, + cache_max_entries=2, + preserve_source_format=True, + ), + ) + for name, a, b in ( + ( + "red", + np.array([[1.0, 0.0]], dtype=np.float32), + np.array([[1.0], [2.0]], dtype=np.float32), + ), + ( + "blue", + np.array([[0.0, 1.0]], dtype=np.float32), + np.array([[3.0], [4.0]], dtype=np.float32), + ), + ( + "green", + np.array([[1.0, 1.0]], dtype=np.float32), + np.array([[1.0], [1.0]], dtype=np.float32), + ), + ): + package.add_adapter_artifact( + AdapterArtifact( + name, + fingerprint, + (AdapterWeights(target, ir.tensor(a), ir.tensor(b), 1.0),), + identity=name, + version="1", + ) + ) + peft_source = source_root / "peft" + peft_source.mkdir(parents=True, exist_ok=True) + (peft_source / "adapter_config.json").write_text( + json.dumps( + { + "base_model_name_or_path": "synthetic/adapter-base", + "r": 1, + "lora_alpha": 1.0, + "target_modules": ["projection"], + "revision": "synthetic-revision", + }, + sort_keys=True, + separators=(",", ":"), + ) + ) + save_file( + { + "base_model.projection.lora_A.weight": np.array([[1.0, -1.0]], dtype=np.float32), + "base_model.projection.lora_B.weight": np.array([[0.5], [0.25]], dtype=np.float32), + }, + peft_source / "adapter_model.safetensors", + ) + package.add_adapter_artifact( + load_peft_adapter( + peft_source, + name="peft", + base_fingerprint=fingerprint, + target_bindings={"projection": target}, + ) + ) + return package + + +def _write_adapter_metadata(package: ModelPackage, directory: Path) -> None: + metadata = { + "schema_version": "v1", + "pipeline": { + "workflow": { + "manifest": { + "adapter_abis": {"onnx-genai.parameter-overlay": "1"}, + "capabilities": [ + "workflow_ssa", + "typed_emit", + "parameter_adapters", + "heterogeneous_adapter_batching", + ], + }, + "inputs": { + "request.active": { + "contract": { + "dtype": "bool", + "rank": 1, + "shape": ["batch"], + }, + "role": { + "kind": "runtime", + "version": "1.0", + "role": "adapter_active", + }, + "source": {"kind": "request"}, + }, + "activations": { + "contract": { + "dtype": "float32", + "rank": 2, + "shape": ["batch", 2], + }, + "role": {"kind": "opaque"}, + "source": {"kind": "application", "name": "activations"}, + }, + }, + "outputs": { + "result": { + "contract": { + "dtype": "float32", + "rank": 2, + "shape": ["batch", 2], + }, + "role": "tensor", + "stage": "pre_adapter", + } + }, + "components": { + "decoder": { + "implementation": { + "kind": "onnx", + "artifact": "model.onnx", + }, + # An ONNX component declares no port contracts: the + # artifact shipped beside this metadata is authoritative + # for its own ports, and the runtime resolves them + # against the live session rather than against a copy + # that can drift. Only an adapter, which ships no graph, + # has to state its ports here. + }, + "overlay": { + "implementation": { + "kind": "adapter", + "abi": "onnx-genai.parameter-overlay", + "version": "1", + }, + "ports": { + "inputs": { + "input": { + "dtype": "float32", + "rank": 2, + "shape": ["batch", 2], + } + }, + "outputs": { + "output": { + "dtype": "float32", + "rank": 2, + "shape": ["batch", 2], + } + }, + }, + "contract": { + "id": "onnx-genai.parameter-overlay", + "version": "1", + "bindings": {"input": "input", "output": "output"}, + "parameters": { + "action": "apply", + "component": "decoder", + "parameter": "projection", + }, + }, + }, + }, + "state": {}, + "steps": [ + { + "kind": "invoke", + "component": "overlay", + "inputs": {"input": "activations"}, + "outputs": {"output": "adapted"}, + }, + { + "kind": "emit", + "value": "adapted", + "output": "result", + "mode": "replace", + }, + ], + } + }, + } + add_adapter_service_to_metadata(metadata, package, str(directory)) + with open(directory / "inference_metadata.yaml", "w", encoding="utf-8") as handle: + yaml.safe_dump(metadata, handle, sort_keys=False) + + +def generate_packages(output: Path) -> Path: + """Write every conformance package under ``output``. + + The graphs are deterministic functions of this file, so they are generated + on demand rather than committed. Callers that need the artifacts — the + validator, the runtime conformance suite, the tests that resolve a port + against the graph that really exposes it — materialize them here. + """ + args = argparse.Namespace(output=output) + decoder = _executable_decoder_package() + static_cache = _executable_static_cache_package() + packages = { + "decoder": (decoder, {"config": decoder.config}), + "static_cache": (static_cache, {"config": static_cache.config}), + "vlm": (_executable_vlm_package(), {}), + "diffusion": ( + _executable_diffusion_package(), + {"guidance_scale": 1.0}, + ), + "tts": (_tts_package(), {}), + } + for name, (package, options) in packages.items(): + directory = args.output / name + package.save(str(directory), progress_bar=False, check_weights=False) + write_onnx_genai_config(package, str(directory), **options) + + # A second image-diffusion fixture that exercises every optional part of the + # workflow at once: classifier-free guidance from a negative prompt, a + # multistep solver with a carried history cell, per-step trajectory emits, + # a seeded latent drawn inside the workflow, and a scaled VAE input. + guided = _executable_diffusion_package() + directory = args.output / "diffusion_guided" + guided.save(str(directory), progress_bar=False, check_weights=False) + write_diffusion_workflow_metadata( + guided, + str(directory), + num_inference_steps=3, + schedule=[8.0, 4.0, 1.0, 0.0], + timesteps=[900.0, 600.0, 300.0], + solver="multistep", + scale_model_input=False, + decoder_input_scale=1.0 / 0.18215, + guidance_scale=7.5, + latent_source="seed", + latent_row_shape=[4, 4, 4], + ) + + speculative = _executable_speculative_package() + directory = args.output / "speculative" + speculative.save(str(directory), progress_bar=False, check_weights=False) + write_speculative_workflow_metadata( + speculative, + str(directory), + grammar_guidance=True, + adaptive_k_max=4, + ) + + masked = _executable_masked_package() + directory = args.output / "masked" + masked.save(str(directory), progress_bar=False, check_weights=False) + write_language_diffusion_workflow_metadata( + masked, + str(directory), + num_inference_steps=8, + ) + + video = _executable_video_package() + directory = args.output / "video" + video.save(str(directory), progress_bar=False, check_weights=False) + write_video_diffusion_workflow_metadata( + video, + str(directory), + num_inference_steps=3, + schedule=[0.9, 0.6, 0.3, 1.0], + timesteps=[600.0, 300.0, 0.0], + solver="ddim", + clip_sample_range=1.0, + scaling_factor=1.15258426, + ) + + codec = _executable_codec_package() + directory = args.output / "codec" + codec.save(str(directory), progress_bar=False, check_weights=False) + write_audio_codec_workflow_metadata(codec, str(directory)) + + # Two encoder-embedding packages rather than one: ESM-2 has no token-type + # embedding and ProtBert does, so together they pin down that the producer + # declares the ports the artifact actually exposes instead of the ports the + # feature-extraction task signature offers. + esm2 = _esm2_embedding_package() + directory = args.output / "esm2_protein_embeddings" + esm2.save(str(directory), progress_bar=False, check_weights=False) + write_encoder_embedding_workflow_metadata(esm2, str(directory), ESM2_TINY_CONFIG) + + protbert = _protbert_embedding_package() + directory = args.output / "protbert_protein_embeddings" + protbert.save(str(directory), progress_bar=False, check_weights=False) + write_encoder_embedding_workflow_metadata(protbert, str(directory), PROTBERT_TINY_CONFIG) + + directory = args.output / "adapter" + source_root = directory / ".sources" + adapter = _adapter_package(source_root) + adapter.save(str(directory), progress_bar=False) + _write_adapter_metadata(adapter, directory) + shutil.rmtree(source_root) + + (args.output / "README.md").write_text( + """# ONNX GenAI workflow conformance fixtures + +Generated by `tests/generate_onnx_genai_validation_packages.py` for semantic +validation and runtime conformance. The ONNX GenAI revision these are checked +against is pinned in `.github/workflows/main.yml`; it is not restated here, +because a second copy of a SHA only ever drifts from the first. + +Only the metadata is committed under `tests/fixtures/onnx_genai_workflows`. +The graphs and adapter weights are a deterministic function of this script, so +committing them would store megabytes of bytes no reviewer can read. CI +regenerates them, compares the metadata against the committed copy, and runs +validation and conformance against the regenerated tree. Tests that need the +graphs use the `materialized_workflow_packages` fixture. + +The decoder, VLM, diffusion, masked diffusion, speculative, and codec packages +contain executable synthetic models. The TTS fixture uses the real tiny +Qwen3-TTS producer graphs with deterministic synthetic weights. The two +protein-encoder fixtures use the real tiny ESM-2 and ProtBert producer graphs, +also with deterministic synthetic weights; they cover the non-generative +embedding shape, which has one invocation, no carried state and no sampler. +No downloaded model weights are included. The adapter fixture covers +authoritative target metadata, portable artifacts, ordered heterogeneous +composition, inactive rows, compaction, and request-epoch slot reuse. +""", + encoding="utf-8", + ) + return args.output + + +def main() -> None: + parser = argparse.ArgumentParser() + parser.add_argument("output", type=Path) + generate_packages(parser.parse_args().output) + + +if __name__ == "__main__": + main() diff --git a/tests/integration_test.py b/tests/integration_test.py index d6e8127c6..0c535dd28 100644 --- a/tests/integration_test.py +++ b/tests/integration_test.py @@ -1599,6 +1599,7 @@ def test_vision_features_parity(self, model_id: str): _ENCODER_MODELS = [ pytest.param("google-bert/bert-base-uncased", False, id="bert-base"), pytest.param("distilbert/distilbert-base-uncased", False, id="distilbert-base"), + pytest.param("facebook/esm2_t6_8M_UR50D", False, id="esm2-8m"), pytest.param( "FacebookAI/roberta-base", False, @@ -1617,6 +1618,23 @@ def test_vision_features_parity(self, model_id: str): ), ] +# Two inputs of deliberately unequal length, so batching them forces padding. +# Protein models have an amino-acid vocabulary, so they get real sequences +# (hemoglobin alpha and a ubiquitin fragment) rather than English. +_UNEQUAL_LENGTH_PROMPTS = { + "protein": [ + "MVLSPADKTNVKAAWGKVGAHAGEYGAEALERMFLSFPTTKTYFPHFDLSH", + "MQIFVKTLTGKTITLEVEPSDTIENVKAKIQDKEGIPPDQQRLIFAGKQLEDGRTLSDYNIQKESTLHLVLRLRGG", + ], + "text": [ + "The capital of France is Paris.", + ( + "Encoder models read the whole sequence at once, so padding must not " + "change a row's contextual embeddings." + ), + ], +} + @pytest.mark.integration @pytest.mark.integration_fast @@ -1667,6 +1685,83 @@ def test_hidden_states_match(self, model_id: str, trust_remote_code: bool): atol=1e-3, ) + def test_padded_batch_matches_huggingface_and_unpadded_rows( + self, model_id: str, trust_remote_code: bool + ): + """Batch two unequal-length inputs and check every row three ways. + + This is the assertion the attention-mask path lives or dies on. An + encoder that hands the raw rank-2 mask to ``op.Attention`` cannot even + run at batch > 1 -- the mask fails to broadcast to + ``(batch, heads, q, kv)`` -- and were it to run it would add a 0/1 bias + where a large negative one is required, so padded positions would leak + into every valid residue. Checking each row against both HuggingFace + and its own unpadded run separates "the export is wrong" from "the + export is right but padding-sensitive". + """ + from mobius._testing.torch_reference import ( + load_torch_encoder_model, + torch_encoder_forward, + ) + + onnx_model = build(model_id, dtype="f32", load_weights=True) + torch_model, tokenizer = load_torch_encoder_model(model_id) + prompts = _UNEQUAL_LENGTH_PROMPTS["protein" if "esm" in model_id.lower() else "text"] + session = _make_session(onnx_model) + + def feeds_for(texts: list[str]) -> dict[str, np.ndarray]: + batch = tokenizer(texts, return_tensors="np", padding=True) + feeds: dict[str, np.ndarray] = { + "input_ids": batch["input_ids"].astype(np.int64), + "attention_mask": batch["attention_mask"].astype(np.int64), + } + if "token_type_ids" in session.input_names: + token_type_ids = batch.get("token_type_ids") + feeds["token_type_ids"] = ( + np.zeros_like(feeds["input_ids"]) + if token_type_ids is None + else token_type_ids.astype(np.int64) + ) + return feeds + + try: + batched = feeds_for(prompts) + mask = batched["attention_mask"] + assert mask.shape[0] == 2 + assert mask.min() == 0, "prompts must differ in length so the batch pads" + + batched_out = session.run(batched)["last_hidden_state"] + torch_hidden = torch_encoder_forward( + torch_model, + batched["input_ids"], + mask, + batched.get("token_type_ids"), + ) + + for row, prompt in enumerate(prompts): + length = int(mask[row].sum()) + got, ref = batched_out[row, :length], torch_hidden[row, :length] + + # Scale-free metrics rather than an elementwise tolerance. A + # contextual embedding is used as a direction: what matters is + # that the vector matches, not that every one of its small + # components survives fp32 cancellation through six layers. + rel_l2 = np.linalg.norm(got - ref) / np.linalg.norm(ref) + cosine = (got * ref).sum(-1) / ( + np.linalg.norm(got, axis=-1) * np.linalg.norm(ref, axis=-1) + ) + assert rel_l2 < 1e-3, f"row {row}: relative L2 error {rel_l2:.2e}" + assert cosine.min() > 1 - 1e-5, ( + f"row {row}: worst per-token cosine {cosine.min():.8f}" + ) + + solo = session.run(feeds_for([prompt]))["last_hidden_state"] + # Padding is bit-exact, not merely close: a padded position that + # reached a valid one would show up here as a real difference. + np.testing.assert_array_equal(solo[0, :length], got) + finally: + session.close() + # --------------------------------------------------------------------------- # Seq2seq models (BART, T5) diff --git a/tests/onnx_genai_workflow_conformance.rs b/tests/onnx_genai_workflow_conformance.rs new file mode 100644 index 000000000..cf9930895 --- /dev/null +++ b/tests/onnx_genai_workflow_conformance.rs @@ -0,0 +1,706 @@ +//! Runtime conformance for the checked Mobius workflow packages. +//! +//! This file is copied into `onnx-genai-engine/tests` by Mobius CI so every +//! package is executed by the authoritative ONNX GenAI workflow runtime. + +#![allow(clippy::field_reassign_with_default)] + +use onnx_genai_engine::{ + AdapterActivation, AdapterSelection, Engine, EngineConfig, GenerateOptions, + GeneratePrompt, GenerateRequest, PipelineGenerateRequest, + pipeline::{PipelineEngine, WorkflowOutputRole}, +}; +use onnx_genai_ort::{DataType, Value}; +use std::path::PathBuf; + +fn root(name: &str) -> anyhow::Result { + // No default: the packages are generated rather than committed, so there + // is no fixed path to fall back to. A relative guess would resolve inside + // the ONNX GenAI checkout this file is copied into and fail with a missing + // package rather than a missing directory. + let root = std::env::var_os("MOBIUS_WORKFLOW_CONFORMANCE_DIR") + .map(PathBuf::from) + .ok_or_else(|| { + anyhow::anyhow!( + "set MOBIUS_WORKFLOW_CONFORMANCE_DIR to a directory of generated \ + packages, e.g. `python tests/generate_onnx_genai_validation_packages.py \ + validation/generated` from a Mobius checkout" + ) + })?; + Ok(root.join(name)) +} + +fn options(max_new_tokens: usize) -> GenerateOptions { + let mut options = GenerateOptions::default(); + options.max_new_tokens = max_new_tokens; + options.seed = Some(7); + options +} + +fn adapter_request( + active: &[bool], + values: &[f32], + selection: &AdapterSelection, +) -> anyhow::Result { + let batch = i64::try_from(selection.rows.len())?; + let mut segments = vec![-1i64; selection.rows.len() * 2]; + let mut adapter_counts = vec![0i64; selection.rows.len()]; + let mut adapter_scales = vec![0.0f32; selection.rows.len() * 2]; + for (row, activations) in selection.rows.iter().enumerate() { + adapter_counts[row] = i64::try_from(activations.len())?; + for (slot, activation) in activations.iter().enumerate() { + segments[row * 2 + slot] = match activation.adapter.as_str() { + "blue" => 0, + "green" => 1, + "red" => 3, + other => anyhow::bail!("unknown test adapter {other}"), + }; + adapter_scales[row * 2 + slot] = activation.scale; + } + } + Ok(PipelineGenerateRequest::new(GenerateRequest { + prompt: GeneratePrompt::TokenIds(vec![]), + options: Default::default(), + }) + .with_input( + "request.adapter_segments", + Value::from_slice_i64(&segments, &[batch, 2])?, + ) + .with_input( + "request.adapter_counts", + Value::from_slice_i64(&adapter_counts, &[batch])?, + ) + .with_input( + "request.adapter_scales", + Value::from_slice_f32(&adapter_scales, &[batch, 2])?, + ) + .with_input( + "request.active", + Value::from_raw_bytes( + active.iter().map(|value| u8::from(*value)).collect(), + &[batch], + DataType::Bool, + )?, + ) + .with_input( + "activations", + Value::from_slice_f32(values, &[batch, 2])?, + )) +} + +/// Adapter composition is positional: the runtime keys its cache by the +/// adapter set a batch row asks for, not by any serialized slot identity. The +/// rows below therefore describe order, per-row composition, and compaction, +/// while the request table that maps a row back to a caller stays private to +/// the runtime. +#[test] +fn mobius_parameter_adapters_preserve_order_rows_and_compaction() -> anyhow::Result<()> { + let mut engine = Engine::from_pipeline_dir(&root("adapter")?, EngineConfig::default())?; + let selection = AdapterSelection::default() + .with_row([AdapterActivation::new("red", 1.0)]) + .with_row([AdapterActivation::new("blue", 1.0)]) + .with_row([ + AdapterActivation::new("red", 0.5), + AdapterActivation::new("blue", 1.0), + ]); + let output = engine.run_pipeline(adapter_request( + &[true, false, true], + &[1.0, 2.0, 3.0, 4.0, 5.0, 6.0], + &selection, + )?)?; + assert_eq!( + output["result"].to_vec_f32()?, + vec![2.0, 4.0, 3.0, 4.0, 25.5, 35.0] + ); + // Compaction reorders the batch: the surviving rows keep their composition + // while moving to new positions. + let compacted_selection = AdapterSelection::default() + .with_row([ + AdapterActivation::new("red", 0.5), + AdapterActivation::new("blue", 1.0), + ]) + .with_row([AdapterActivation::new("red", 1.0)]); + let compacted = engine.run_pipeline(adapter_request( + &[true, true], + &[5.0, 6.0, 1.0, 2.0], + &compacted_selection, + )?)?; + assert_eq!( + compacted["result"].to_vec_f32()?, + vec![25.5, 35.0, 2.0, 4.0] + ); + // A row that asks for no adapter is passed through unmodified. + let unadapted = AdapterSelection::default().with_row([]); + let bare = engine.run_pipeline(adapter_request(&[true], &[1.0, 2.0], &unadapted)?)?; + assert_eq!(bare["result"].to_vec_f32()?, vec![1.0, 2.0]); + let reused = AdapterSelection::default().with_row([AdapterActivation::new("blue", 1.0)]); + for _ in 0..2 { + let output = engine.run_pipeline(adapter_request(&[true], &[1.0, 2.0], &reused)?)?; + assert_eq!(output["result"].to_vec_f32()?, vec![7.0, 10.0]); + } + let green = AdapterSelection::default().with_row([AdapterActivation::new("green", 1.0)]); + let output = engine.run_pipeline(adapter_request(&[true], &[1.0, 2.0], &green)?)?; + assert_eq!(output["result"].to_vec_f32()?, vec![4.0, 5.0]); + let red = AdapterSelection::default().with_row([AdapterActivation::new("red", 1.0)]); + for _ in 0..2 { + let output = engine.run_pipeline(adapter_request(&[true], &[1.0, 2.0], &red)?)?; + assert_eq!(output["result"].to_vec_f32()?, vec![2.0, 4.0]); + } + let diagnostic = engine.adapter_lifecycle_diagnostic(); + assert_eq!(diagnostic.loads, 4); + assert!(diagnostic.cache_hits > 0); + assert_eq!(diagnostic.evictions, 2); + assert_eq!(diagnostic.reloads, 1); + assert_eq!(diagnostic.capture_invalidations, 2); + assert!(diagnostic.replayed_plans > 0); + Ok(()) +} + +fn assert_batched_policy_super_island(engine: &PipelineEngine) { + let diagnostics = engine.execution_island_diagnostics(); + let island = diagnostics + .iter() + .find(|island| { + ["token_sampler", "termination", "token_state_update"] + .iter() + .all(|component| island.components.iter().any(|item| item == component)) + }) + .unwrap_or_else(|| { + panic!( + "sampler, termination, and state update must share one execution island: \ + {diagnostics:#?}" + ) + }); + assert!(island.runs > 0, "batched policy island must execute"); + assert_eq!(island.session_runs, island.runs); + assert!( + island.component_boundaries_elided >= 2, + "the three policy components must execute as a fused island" + ); + if island.device.starts_with("cuda:") { + assert_eq!(island.fallback_reason, None); + } else { + assert_eq!( + island.fallback_reason.as_deref(), + Some("island is not placed on CUDA") + ); + } +} + +fn decoder_batch_request( + input_ids: &[i64], + batch: i64, + sequence: i64, + prompt_lengths: &[i64], + active: &[bool], + max_new_tokens: usize, +) -> anyhow::Result { + let seeds = (0..batch).collect::>(); + decoder_batch_request_with_seeds( + input_ids, + batch, + sequence, + prompt_lengths, + active, + &seeds, + max_new_tokens, + ) +} + +fn decoder_batch_request_with_seeds( + input_ids: &[i64], + batch: i64, + sequence: i64, + prompt_lengths: &[i64], + active: &[bool], + seeds: &[i64], + max_new_tokens: usize, +) -> anyhow::Result { + let bool_bytes = active.iter().map(|value| u8::from(*value)).collect(); + let zeros = vec![0_i64; usize::try_from(batch)?]; + let ones = vec![1_i64; usize::try_from(batch)?]; + let negative_ones = vec![-1_i64; usize::try_from(batch)?]; + let floats_zero = vec![0.0_f32; usize::try_from(batch)?]; + let floats_one = vec![1.0_f32; usize::try_from(batch)?]; + assert_eq!(seeds.len(), usize::try_from(batch)?); + Ok(PipelineGenerateRequest::new(GenerateRequest { + prompt: GeneratePrompt::TokenIds(vec![0]), + options: options(max_new_tokens), + }) + .with_input( + "request.input_ids", + Value::from_slice_i64(input_ids, &[batch, sequence])?, + ) + .with_input( + "request.prompt_lengths", + Value::from_slice_i64(prompt_lengths, &[batch])?, + ) + .with_input( + "package.active", + Value::from_raw_bytes(bool_bytes, &[batch], DataType::Bool)?, + ) + .with_input( + "package.not_done", + Value::from_raw_bytes(vec![0; usize::try_from(batch)?], &[batch], DataType::Bool)?, + ) + .with_input("package.one_token", Value::from_slice_i64(&ones, &[batch])?) + .with_input( + "request.eos_ids", + Value::from_slice_i64(&vec![2_i64; usize::try_from(batch)?], &[batch, 1])?, + ) + .with_input( + "request.eos_lengths", + Value::from_slice_i64(&ones, &[batch])?, + ) + .with_input( + "request.row_max_iterations", + Value::from_slice_i64(&negative_ones, &[batch])?, + ) + .with_input( + "request.temperature", + Value::from_slice_f32(&floats_one, &[batch])?, + ) + .with_input("request.top_k", Value::from_slice_i64(&ones, &[batch])?) + .with_input( + "request.top_p", + Value::from_slice_f32(&floats_one, &[batch])?, + ) + .with_input( + "request.min_p", + Value::from_slice_f32(&floats_zero, &[batch])?, + ) + .with_input("request.seed", Value::from_slice_i64(seeds, &[batch])?) + .with_input( + "request.rng_counter", + Value::from_slice_i64(&zeros, &[batch])?, + ) + .with_input( + "package.cache_lengths", + Value::from_slice_i64(&zeros, &[batch])?, + ) + .with_input( + "package.zero_batch", + Value::from_slice_i64(&zeros, &[batch])?, + )) +} + +#[test] +fn mobius_decoder_workflow_executes() -> anyhow::Result<()> { + let mut engine = Engine::from_pipeline_dir(&root("decoder")?, EngineConfig::default())?; + let output = engine.run_pipeline_outputs( + PipelineGenerateRequest::new(GenerateRequest { + prompt: GeneratePrompt::TokenIds(vec![4, 5]), + options: options(3), + }), + )?; + assert_eq!( + engine + .structured_output_for_role(&output, WorkflowOutputRole::Tokens) + .expect("decoder must emit tokens") + .to_vec_i64()? + .len(), + 3 + ); + assert_batched_policy_super_island(&engine); + Ok(()) +} + +/// The fixed-capacity decoder must execute the same way the appending one does. +/// +/// Its cache never grows: every step scatters into a preallocated buffer at a +/// per-row cursor. That exercises the `indexed_scatter` state discipline, the +/// invariant loop cells and the `package.cache_capacity` input, none of which +/// the appending `decoder` fixture reaches. +#[test] +fn mobius_static_cache_workflow_executes() -> anyhow::Result<()> { + let mut engine = + Engine::from_pipeline_dir(&root("static_cache")?, EngineConfig::default())?; + let output = engine.run_pipeline_outputs(PipelineGenerateRequest::new(GenerateRequest { + prompt: GeneratePrompt::TokenIds(vec![4, 5]), + options: options(3), + }))?; + assert_eq!( + engine + .structured_output_for_role(&output, WorkflowOutputRole::Tokens) + .expect("static-cache decoder must emit tokens") + .to_vec_i64()? + .len(), + 3 + ); + Ok(()) +} + +#[test] +fn mobius_decoder_rows_match_independent_runs_and_dynamic_batch_replay() -> anyhow::Result<()> { + let mut engine = Engine::from_pipeline_dir(&root("decoder")?, EngineConfig::default())?; + let generated = engine.generate_with_pipeline_request(decoder_batch_request( + &[4, 5], + 1, + 2, + &[2], + &[true], + 3, + )?)?; + assert_eq!(generated.token_ids.len(), 3); + + let first = decoder_batch_request(&[4, 5], 1, 2, &[2], &[true], 3)?; + let first_output = engine.run_pipeline_outputs(first)?; + let first_tokens = engine + .structured_output_for_role(&first_output, WorkflowOutputRole::Tokens) + .expect("batch-one decoder must emit tokens") + .to_vec_i64()?; + + let multi_row_error = engine + .generate_with_pipeline_request(decoder_batch_request( + &[4, 5, 6, 0], + 2, + 2, + &[2, 1], + &[true, true], + 3, + )?) + .expect_err("generate must not flatten multiple semantic rows"); + assert!( + multi_row_error + .to_string() + .contains("multi-row ragged output"), + "{multi_row_error:#}" + ); + + let batched = decoder_batch_request(&[4, 5, 6, 0], 2, 2, &[2, 1], &[true, true], 3)?; + let batched_output = engine.run_pipeline_outputs(batched)?; + let rows = engine.output_rows_for_role(&batched_output, WorkflowOutputRole::Tokens); + assert_eq!(rows.len(), 2); + assert_eq!(rows[0].0, 0); + assert_eq!(rows[0].1.to_vec_i64()?, first_tokens); + + let mut independent = Engine::from_pipeline_dir(&root("decoder")?, EngineConfig::default())?; + let second = decoder_batch_request(&[6, 0], 1, 2, &[1], &[true], 3)?; + let second_output = independent.run_pipeline_outputs(second)?; + let second_tokens = independent + .structured_output_for_role(&second_output, WorkflowOutputRole::Tokens) + .expect("independent second row must emit tokens") + .to_vec_i64()?; + assert_eq!(rows[1].0, 1); + assert_eq!(rows[1].1.to_vec_i64()?, second_tokens); + + let stable_before = engine + .execution_island_diagnostics() + .iter() + .map(|island| island.stable_binding_runs) + .sum::(); + let compacted = decoder_batch_request_with_seeds( + &[6, 0, 4, 5], + 2, + 2, + &[1, 2], + &[true, true], + &[1, 0], + 3, + )?; + let compacted_output = engine.run_pipeline_outputs(compacted)?; + let compacted_rows = + engine.output_rows_for_role(&compacted_output, WorkflowOutputRole::Tokens); + let compacted_row = |semantic_id| { + compacted_rows + .iter() + .find(|(row_id, _)| *row_id == semantic_id) + .expect("compacted semantic row must be present") + .1 + .to_vec_i64() + }; + // Semantic row ids are positional: the runtime maps a batch row back to a + // caller through its own private request table, so the reordered batch + // reports the sequence it actually placed in each row. + assert_eq!(compacted_row(0)?, second_tokens); + assert_eq!(compacted_row(1)?, first_tokens); + let stable_after = engine + .execution_island_diagnostics() + .iter() + .map(|island| island.stable_binding_runs) + .sum::(); + assert!( + stable_after > stable_before, + "same-shape row compaction must reuse stable island bindings" + ); + + // This decoder concatenates `present` onto `past`, so its cache is dynamic + // and its attention mask is a dense carry that grows one column per step + // for every row at once. A partially active batch is therefore not + // expressible: preserving an inactive row would require its mask to keep + // the narrower width the rest of the batch has already outgrown. The + // runtime says so instead of silently corrupting the held row, and this + // asserts that contract rather than leaving it to chance. + for active in [[true, false], [false, true]] { + let mixed = decoder_batch_request(&[4, 5, 6, 0], 2, 2, &[2, 1], &active, 3)?; + let Err(error) = engine.run_pipeline_outputs(mixed) else { + panic!("a growing dense carry cannot hold an inactive row"); + }; + assert!( + format!("{error:#}").contains("cannot preserve inactive rows"), + "{error:#}" + ); + } + + let replay = decoder_batch_request(&[6, 0], 1, 2, &[1], &[true], 3)?; + let replay_output = engine.run_pipeline_outputs(replay)?; + assert_eq!( + engine + .structured_output_for_role(&replay_output, WorkflowOutputRole::Tokens) + .expect("batch-one replay must emit tokens") + .to_vec_i64()?, + second_tokens + ); + Ok(()) +} + +#[test] +fn mobius_vlm_workflow_executes_complete_image_path() -> anyhow::Result<()> { + let mut engine = Engine::from_pipeline_dir(&root("vlm")?, EngineConfig::default())?; + let png = vec![ + 137, 80, 78, 71, 13, 10, 26, 10, 0, 0, 0, 13, 73, 72, 68, 82, 0, 0, 0, 1, 0, 0, 0, 1, 8, 2, + 0, 0, 0, 144, 119, 83, 222, 0, 0, 0, 12, 73, 68, 65, 84, 120, 156, 99, 248, 207, 192, 0, 0, + 3, 1, 1, 0, 201, 254, 146, 239, 0, 0, 0, 0, 73, 69, 78, 68, 174, 66, 96, 130, + ]; + let png_len = i64::try_from(png.len())?; + let request = PipelineGenerateRequest::new(GenerateRequest { + prompt: GeneratePrompt::TokenIds(vec![4, 5]), + options: options(2), + }) + .with_input( + "request.image", + Value::from_raw_bytes(png, &[png_len], DataType::Uint8)?, + ); + let output = engine.run_pipeline_outputs(request)?; + assert_eq!( + engine + .structured_output_for_role(&output, WorkflowOutputRole::Tokens) + .expect("VLM must emit tokens") + .shape(), + [1, 2] + ); + assert_batched_policy_super_island(&engine); + Ok(()) +} + +#[test] +fn mobius_euler_diffusion_workflow_executes_complete_path() -> anyhow::Result<()> { + let mut engine = Engine::from_pipeline_dir(&root("diffusion")?, EngineConfig::default())?; + let request = PipelineGenerateRequest::new(GenerateRequest { + prompt: GeneratePrompt::TokenIds(vec![1, 2]), + options: options(2), + }) + .with_input( + "request.noise", + Value::from_slice_f32(&[1.0; 64], &[1, 4, 4, 4])?, + ); + let output = engine.run_pipeline_outputs(request)?; + assert_eq!(output["image"].shape(), [1, 3, 4, 4]); + assert!( + output["image"] + .to_vec_f32()? + .iter() + .all(|value| value.is_finite()) + ); + Ok(()) +} + +#[test] +fn mobius_masked_diffusion_workflow_executes() -> anyhow::Result<()> { + let mut engine = Engine::from_pipeline_dir(&root("masked")?, EngineConfig::default())?; + let request = PipelineGenerateRequest::new(GenerateRequest { + prompt: GeneratePrompt::TokenIds(vec![0, 0]), + options: options(3), + }) + .with_input( + "masked_positions", + Value::from_raw_bytes(vec![1, 0], &[1, 2], DataType::Bool)?, + ) + .with_input("rng_offset", Value::from_slice_i64(&[0], &[1])?); + let output = engine.run_pipeline_outputs(request)?; + assert_eq!(output["tokens"].shape(), [1, 2]); + Ok(()) +} + +#[test] +fn mobius_codec_workflow_executes() -> anyhow::Result<()> { + let mut engine = Engine::from_pipeline_dir(&root("codec")?, EngineConfig::default())?; + let request = + PipelineGenerateRequest::new(GenerateRequest::new(GeneratePrompt::TokenIds(vec![]))) + .with_input( + "request.waveform", + Value::from_slice_f32(&[0.25, -0.5], &[1, 1, 2])?, + ); + let output = engine.run_pipeline_outputs(request)?; + assert_eq!(output["waveform"].to_vec_f32()?, [0.25, -0.5]); + Ok(()) +} + +fn tts_request(prompt_tokens: &[i64], batch: i64) -> anyhow::Result { + let rows = usize::try_from(batch)?; + assert_eq!(prompt_tokens.len(), rows * 2); + Ok( + PipelineGenerateRequest::new(GenerateRequest { + prompt: GeneratePrompt::TokenIds(vec![0]), + options: options(1), + }) + .with_input( + "request.prompt_tokens", + Value::from_slice_i64(prompt_tokens, &[batch, 2])?, + ) + .with_input( + "package.false", + Value::from_raw_bytes(vec![0; rows], &[batch], DataType::Bool)?, + ) + .with_input( + "package.zero_batch", + Value::from_slice_i64(&vec![0; rows], &[batch])?, + ) + .with_input( + "package.one_batch", + Value::from_slice_i64(&vec![1; rows], &[batch])?, + ) + .with_input( + "package.true", + Value::from_raw_bytes(vec![1; rows], &[batch], DataType::Bool)?, + ), + ) +} + +#[test] +fn mobius_tts_workflow_executes_real_producer_graphs() -> anyhow::Result<()> { + let mut engine = Engine::from_pipeline_dir(&root("tts")?, EngineConfig::default())?; + let output = engine.run_pipeline_outputs(tts_request(&[1, 2], 1)?)?; + assert_eq!(output["waveform"].shape()[..2], [1, 1]); + let first = output["waveform"].to_vec_f32()?; + assert!(!first.is_empty()); + + let mut independent = Engine::from_pipeline_dir(&root("tts")?, EngineConfig::default())?; + let second_output = independent.run_pipeline_outputs(tts_request(&[3, 4], 1)?)?; + let second = second_output["waveform"].to_vec_f32()?; + + let batched = engine.run_pipeline_outputs(tts_request(&[1, 2, 3, 4], 2)?)?; + let frames = first.len(); + assert_eq!(batched["waveform"].shape(), [2, 1, i64::try_from(frames)?]); + let batched_waveform = batched["waveform"].to_vec_f32()?; + assert_eq!(&batched_waveform[..frames], first); + assert_eq!(&batched_waveform[frames..], second); + + let stable_before = engine + .execution_island_diagnostics() + .iter() + .map(|island| island.stable_binding_runs) + .sum::(); + let compacted = engine.run_pipeline_outputs(tts_request(&[3, 4, 1, 2], 2)?)?; + let compacted_waveform = compacted["waveform"].to_vec_f32()?; + assert_eq!(&compacted_waveform[..frames], second); + assert_eq!(&compacted_waveform[frames..], first); + let stable_after = engine + .execution_island_diagnostics() + .iter() + .map(|island| island.stable_binding_runs) + .sum::(); + assert!( + stable_after > stable_before, + "same-shape nested TTS compaction must preserve stable bindings" + ); + + let reused = engine.run_pipeline_outputs(tts_request(&[3, 4], 1)?)?; + assert_eq!(reused["waveform"].to_vec_f32()?, second); + Ok(()) +} + +#[test] +fn mobius_speculative_workflow_executes_rejection_and_correction() -> anyhow::Result<()> { + let mut engine = Engine::from_pipeline_dir(&root("speculative")?, EngineConfig::default())?; + let request = PipelineGenerateRequest::new(GenerateRequest { + prompt: GeneratePrompt::TokenIds(vec![1, 2, 3, 4]), + options: options(1), + }) + .with_input( + "verifier.past_key_values.0.key", + Value::from_slice_f32(&[], &[1, 2, 0, 8])?, + ) + .with_input("grammar.initial_state", Value::from_slice_i64(&[0], &[1])?) + .with_input( + "grammar.transition_table", + Value::from_slice_i64(&[0; 32], &[1, 32])?, + ) + .with_input("adaptive.current_k", Value::from_slice_i64(&[4], &[1])?) + .with_input( + "adaptive.estimates", + Value::from_slice_f32(&[0.0; 24], &[1, 24])?, + ) + .with_input("telemetry.draft_ms", Value::from_slice_f32(&[1.0], &[1])?) + .with_input("telemetry.target_ms", Value::from_slice_f32(&[1.0], &[1])?); + let output = engine.run_pipeline_outputs(request)?; + assert_eq!(output["tokens.row.0"].to_vec_i64()?, [1, 31]); + Ok(()) +} + +fn video_request(latent_frames: i64, batch: i64) -> anyhow::Result { + let rows = usize::try_from(batch)?; + // [batch, latent_frames, channels, height, width]. Generating from the flat + // index keeps row 0 and the leading frames identical across shapes, so the + // comparisons below isolate the runtime's handling of the temporal axis. + let elements = batch * latent_frames * 4 * 2 * 2; + let noise: Vec = (0..elements) + .map(|index| (index % 11) as f32 / 11.0 - 0.5) + .collect(); + Ok(PipelineGenerateRequest::new(GenerateRequest { + prompt: GeneratePrompt::TokenIds(vec![]), + options: options(3), + }) + .with_input( + "request.noise", + Value::from_slice_f32(&noise, &[batch, latent_frames, 4, 2, 2])?, + ) + .with_input( + "request.encoder_hidden_states", + Value::from_slice_f32(&vec![0.25; rows * 2 * 32], &[batch, 2, 32])?, + ) + .with_input( + "package.false", + Value::from_raw_bytes(vec![0; rows], &[batch], DataType::Bool)?, + )) +} + +#[test] +fn mobius_video_diffusion_workflow_publishes_causal_temporal_chunks() -> anyhow::Result<()> { + let mut engine = Engine::from_pipeline_dir(&root("video")?, EngineConfig::default())?; + + // Three latent frames decode as a single chunk and expand 2x in time. + let short = engine.run_pipeline_outputs(video_request(3, 1)?)?; + assert_eq!(short["video"].shape(), [1, 3, 6, 4, 4]); + let short_frames = short["video"].to_vec_f32()?; + assert!(short_frames.iter().all(|value| value.is_finite())); + let frame = |frames: &[f32], total: usize, channel: usize, time: usize| { + let start = (channel * total + time) * 16; + frames[start..start + 16].to_vec() + }; + + // Five latent frames decode as two causal chunks (three frames, then two). + // The clip is the concatenation along time, and what the first chunk already + // published must not change once the second one runs: that is what the + // decoder's carried convolution caches are for. + let long = engine.run_pipeline_outputs(video_request(5, 1)?)?; + assert_eq!(long["video"].shape(), [1, 3, 10, 4, 4]); + let long_frames = long["video"].to_vec_f32()?; + for channel in 0..3 { + for time in 0..6 { + assert_eq!( + frame(&short_frames, 6, channel, time), + frame(&long_frames, 10, channel, time), + "chunk boundary changed already-published frame {time}" + ); + } + } + + // A batched request decodes independent clips, and the caches from the + // previous invocations are gone: row 0 reproduces the single-row clip. + let batched = engine.run_pipeline_outputs(video_request(3, 2)?)?; + assert_eq!(batched["video"].shape(), [2, 3, 6, 4, 4]); + let batched_frames = batched["video"].to_vec_f32()?; + assert_eq!(&batched_frames[..short_frames.len()], &short_frames[..]); + Ok(()) +} diff --git a/tests/static_cache_metadata_test.py b/tests/static_cache_metadata_test.py new file mode 100644 index 000000000..eef86097f --- /dev/null +++ b/tests/static_cache_metadata_test.py @@ -0,0 +1,479 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT License. + +"""Workflow metadata for static (fixed-capacity, indexed) KV caches. + +A static-cache export does not append to a growing tensor: it scatters each +step's keys and values into a preallocated buffer at a per-row cursor. The +published metadata therefore has to describe three things a dynamic cache +never needs — where the write lands (``write_indices``), how much of the +buffer is valid afterwards (``nonpad_kv_seqlen``), and how large the buffer is +(``package.cache_capacity``) — and it has to say that the cache tensors are +loop *invariant* rather than growing. + +These tests pin that contract against real exported packages, not synthetic +graphs, so a change to the exporter's port names or scatter axis fails here +rather than at runtime. +""" + +from __future__ import annotations + +from typing import Any + +import onnx_ir as ir +import pytest + +from mobius import registry +from mobius._configs import ArchitectureConfig +from mobius._constants import ( + STATIC_CACHE_KV_SEQUENCE_LENGTH, + STATIC_CACHE_SEQUENCE_AXIS, + STATIC_CACHE_WRITE_INDICES, +) +from mobius.integrations.onnx_genai.workflow_metadata import ( + build_decoder_workflow_metadata, +) +from mobius.tasks import CausalLMTask + +CAPACITY = 128 + + +def _text_config(**overrides) -> ArchitectureConfig: + params = { + "num_hidden_layers": 2, + "hidden_size": 64, + "intermediate_size": 128, + "num_attention_heads": 4, + "num_key_value_heads": 2, + "head_dim": 16, + "vocab_size": 256, + "rms_norm_eps": 1e-6, + "hidden_act": "silu", + "max_position_embeddings": 512, + } + params.update(overrides) + return ArchitectureConfig(**params) + + +def _static_package(**overrides): + config = _text_config(**overrides) + module = registry.get("qwen2")(config) + task = CausalLMTask(static_cache=True, max_seq_len=CAPACITY) + return task.build(module, config), config + + +def _cache_cells(workflow) -> list[str]: + """Names of the loop cells the state service publishes as cache buffers.""" + group = next(iter(workflow["serving"]["state_service"]["groups"].values())) + return list(group["ports"]["model"]) + + +def _model_invoke(steps) -> dict[str, str]: + """Input bindings of the neural component's invoke step in *steps*.""" + return next(step for step in steps if step.get("component") == "model")["inputs"] + + +def _scatter_group(metadata) -> tuple[str, dict]: + """The state-service group whose buffers are written by an indexed scatter.""" + groups = metadata["pipeline"]["workflow"]["serving"]["state_service"]["groups"] + return next( + (name, group) + for name, group in groups.items() + if group.get("update", {}).get("kind") == "indexed_scatter" + ) + + +def _static_cache_abi(metadata) -> dict: + """Recover the whole scatter ABI from the workflow and nothing else. + + This is deliberately written the way a runtime lowers a one-component + workflow onto a direct decode path. If it can reconstruct every port a + driver needs, then the workflow is a complete description and republishing + the same facts under a second top-level key would only create two truths + that can disagree. + + Note what it does not read: the component declares no port contracts, so + every name here comes from the scatter discipline and the state pairs. The + ports themselves are checked against the graph, which is the only thing + entitled to say what a port's dtype and shape are. + """ + workflow = metadata["pipeline"]["workflow"] + _, group = _scatter_group(metadata) + update = group["update"] + component = next(iter(update["write_indices_ports"])) + aliases = group["ports"][component] + buffers = [ + aliases[cell] for cell in sorted(aliases, key=lambda cell: int(cell.rsplit("_")[-1])) + ] + return { + "component": component, + "write_indices_input": update["write_indices_ports"][component], + "kv_sequence_length_input": update["kv_length_ports"][component], + "capacity": workflow["inputs"][update["capacity"]]["default"], + "cache_inputs": [alias["input"] for alias in buffers], + "cache_outputs": [alias["output"] for alias in buffers], + } + + +def _graph_ports(pkg) -> dict[str, Any]: + """Every port of the exported decoder, keyed by name. + + The artifact is authoritative for dtype, rank and shape, so assertions + about those resolve here rather than against a transcription in the + metadata — a transcription could agree with itself while disagreeing with + the graph a runtime actually binds. + """ + model = pkg["model"] + return {str(value.name): value for value in (*model.graph.inputs, *model.graph.outputs)} + + +@pytest.fixture(scope="module") +def static_built(): + pkg, config = _static_package() + return pkg, build_decoder_workflow_metadata(pkg, config) + + +@pytest.fixture(scope="module") +def static_workflow(static_built): + return static_built[1] + + +@pytest.fixture(scope="module") +def static_ports(static_built): + return _graph_ports(static_built[0]) + + +@pytest.fixture(scope="module") +def mixed(): + """Gemma 4 text with `--features static-cache`: one static + one dynamic geometry.""" + from gemma4_prefill_prefix_test import _make_config + + from mobius.tasks._gemma4 import Gemma4TextCausalLMTask + + config = _make_config() + # Widen the global head so a collapsed cache group would be observably wrong. + config.global_head_dim = 32 + module = registry.get("gemma4_text")(config) + pkg = Gemma4TextCausalLMTask(static_cache=True, max_seq_len=CAPACITY).build(module, config) + return build_decoder_workflow_metadata(pkg, config) + + +class TestStaticCacheAbiLivesOnlyInTheWorkflow: + """The workflow is the package's single description of its scatter ABI.""" + + def test_no_second_top_level_port_declaration(self, static_workflow): + # `model` carries package-wide geometry and capabilities, never a copy + # of the port ABI: two declarations of one fact can drift apart, and + # nothing in the format says which one wins. + assert "io" not in static_workflow.get("model", {}) + + def test_control_ports_are_declared_by_the_scatter_discipline(self, static_workflow): + abi = _static_cache_abi(static_workflow) + assert abi["write_indices_input"] == STATIC_CACHE_WRITE_INDICES + assert abi["kv_sequence_length_input"] == STATIC_CACHE_KV_SEQUENCE_LENGTH + + def test_control_ports_are_real_ports_of_the_component( + self, static_workflow, static_ports + ): + # Naming a port the component does not expose would bind nothing, and + # the graph is the only thing that can settle whether it does. + abi = _static_cache_abi(static_workflow) + for role in ("write_indices_input", "kv_sequence_length_input"): + value = static_ports[abi[role]] + # Both are per-row integer vectors, which is exactly why they are + # declared rather than recognized by shape. + assert value.dtype == ir.DataType.INT64 + assert len(value.shape) == 1 + + def test_every_buffer_pair_is_declared(self, static_workflow, static_ports): + abi = _static_cache_abi(static_workflow) + assert abi["cache_inputs"] == [ + "key_cache.0", + "value_cache.0", + "key_cache.1", + "value_cache.1", + ] + assert abi["cache_outputs"] == [f"updated_{name}" for name in abi["cache_inputs"]] + for name in (*abi["cache_inputs"], *abi["cache_outputs"]): + assert static_ports[name].shape[STATIC_CACHE_SEQUENCE_AXIS] == CAPACITY + + def test_capacity_is_recoverable(self, static_workflow): + assert _static_cache_abi(static_workflow)["capacity"] == CAPACITY + + +class TestStaticCacheWorkflow: + """The loop body has to carry the buffers and drive the write cursor.""" + + def test_capacity_is_a_declared_workflow_input(self, static_workflow): + workflow = static_workflow["pipeline"]["workflow"] + capacity = workflow["inputs"]["package.cache_capacity"] + assert capacity["source"] == {"kind": "literal"} + assert capacity["default"] == CAPACITY + assert capacity["required"] is False + assert capacity["contract"]["dtype"] == "int64" + assert capacity["contract"]["rank"] == 1 + + def test_cache_cells_are_invariant_not_growing(self, static_workflow): + workflow = static_workflow["pipeline"]["workflow"] + for cell_name in _cache_cells(workflow): + cell = workflow["state"][cell_name] + assert cell["recurrence"] == {"kind": "invariant"} + # The buffer keeps its full capacity every step. + assert cell["contract"]["shape"][STATIC_CACHE_SEQUENCE_AXIS] == CAPACITY + + def test_write_cursor_and_valid_length_are_bound_each_phase(self, static_workflow): + loop = static_workflow["pipeline"]["workflow"]["steps"][0] + + setup_inputs = _model_invoke(loop["setup"]) + # Prefill starts every row at slot 0 and ends with the prompt length. + assert setup_inputs[STATIC_CACHE_WRITE_INDICES] == "initializer.write_indices" + assert setup_inputs[STATIC_CACHE_KV_SEQUENCE_LENGTH] == "initializer.cache_lengths" + + body_inputs = _model_invoke(loop["steps"]) + # Decode writes at the length carried in from the previous step and + # reports the length that step produced. + assert body_inputs[STATIC_CACHE_WRITE_INDICES] == "cache_lengths" + assert body_inputs[STATIC_CACHE_KV_SEQUENCE_LENGTH] == "cache_lengths.next" + + def test_buffers_are_carried_by_the_loop_not_regrown(self, static_workflow): + loop = static_workflow["pipeline"]["workflow"]["steps"][0] + carried = {entry["cell"]: entry["next"] for entry in loop["carried"]} + for cell in _cache_cells(static_workflow["pipeline"]["workflow"]): + assert carried[cell].startswith("decoder.body.updated_") + + def test_state_service_publishes_an_indexed_scatter_discipline(self, static_workflow): + groups = static_workflow["pipeline"]["workflow"]["serving"]["state_service"]["groups"] + assert len(groups) == 1 + group = next(iter(groups.values())) + assert group["update"] == { + "kind": "indexed_scatter", + "write_indices": "cache_lengths", + "capacity": "package.cache_capacity", + "write_indices_ports": {"model": STATIC_CACHE_WRITE_INDICES}, + "kv_length_ports": {"model": STATIC_CACHE_KV_SEQUENCE_LENGTH}, + } + assert group["logical_lengths"] == "cache_lengths" + assert group["sequence_axis"] == STATIC_CACHE_SEQUENCE_AXIS + assert group["layout"] == "bsh" + # Scattering writes in place, so the runtime may alias the buffers. + assert group["aliasing"] == "permitted" + # Every buffer port pair is published so a runtime can bind them. + ports = group["ports"]["model"] + assert len(ports) == 4 + assert ports["cache_0"] == { + "input": "key_cache.0", + "output": "updated_key_cache.0", + "role": "key", + "layer": 0, + } + + def test_control_ports_are_not_advertised_as_request_inputs(self, static_workflow): + # They are derived from loop state, so a caller must not be asked + # to supply them. + inputs = static_workflow["pipeline"]["workflow"]["inputs"] + assert STATIC_CACHE_WRITE_INDICES not in inputs + assert STATIC_CACHE_KV_SEQUENCE_LENGTH not in inputs + + +class TestStaticCachePortDerivation: + """The ABI is read from the graph, never assumed.""" + + def test_capacity_follows_the_requested_max_sequence_length(self): + config = _text_config() + module = registry.get("qwen2")(config) + pkg = CausalLMTask(static_cache=True, max_seq_len=64).build(module, config) + metadata = build_decoder_workflow_metadata(pkg, config) + capacity = metadata["pipeline"]["workflow"]["inputs"]["package.cache_capacity"] + assert capacity["default"] == 64 + + def test_grouped_query_layouts_are_declared_flat(self, static_workflow): + # The exporter stores the static cache as (batch, capacity, kv_hidden); + # publishing a 4-D BNSH shape would misdescribe the buffer a runtime + # has to allocate. + workflow = static_workflow["pipeline"]["workflow"] + contract = workflow["state"][_cache_cells(workflow)[0]]["contract"] + assert contract["rank"] == 3 + # 2 kv heads x 16 head_dim + assert contract["shape"] == ["batch", CAPACITY, 32] + + def test_layer_count_follows_the_config(self): + pkg, config = _static_package(num_hidden_layers=3) + abi = _static_cache_abi(build_decoder_workflow_metadata(pkg, config)) + assert abi["cache_inputs"] == [ + "key_cache.0", + "value_cache.0", + "key_cache.1", + "value_cache.1", + "key_cache.2", + "value_cache.2", + ] + + +class TestHeterogeneousStaticCache: + """Gemma 4 mixes a static full-attention cache with a dynamic sliding one. + + Its two cache geometries have different ranks, layouts, sequence axes and + head dimensions, and its KV-shared suffix owns no cache at all. Publishing + one undifferentiated group — or listing the borrowing layers as if they + owned buffers — would have a runtime allocate caches that do not exist and + scatter into a sliding cache that is appended to. + """ + + def test_only_cache_owning_layers_are_declared(self, mixed): + # layer_types = [sliding, full, sliding, full] with the last two layers + # sharing KV: exactly one layer owns a static buffer. + abi = _static_cache_abi(mixed) + assert abi["cache_inputs"] == ["key_cache.1", "value_cache.1"] + assert abi["cache_outputs"] == ["updated_key_cache.1", "updated_value_cache.1"] + + def test_each_geometry_gets_its_own_update_discipline(self, mixed): + groups = mixed["pipeline"]["workflow"]["serving"]["state_service"]["groups"] + assert set(groups) == { + "decoder_cache_full_attention", + "decoder_cache_sliding_attention", + } + full = groups["decoder_cache_full_attention"] + sliding = groups["decoder_cache_sliding_attention"] + + assert full["update"]["kind"] == "indexed_scatter" + assert full["sequence_axis"] == STATIC_CACHE_SEQUENCE_AXIS + assert full["layout"] == "bsh" + assert full["aliasing"] == "permitted" + + # The sliding layers still append into a growing BNSH tensor. + assert "update" not in sliding or sliding["update"]["kind"] == "append" + assert sliding["sequence_axis"] == 2 + assert sliding["layout"] == "bnsh" + assert sliding["aliasing"] == "forbidden" + assert sliding["reuse"]["evictable_prefix"] is True + + def test_dual_head_dims_survive_the_split(self, mixed): + workflow = mixed["pipeline"]["workflow"] + groups = workflow["serving"]["state_service"]["groups"] + state = workflow["state"] + + full_cell = next(iter(groups["decoder_cache_full_attention"]["ports"]["model"])) + sliding_cell = next(iter(groups["decoder_cache_sliding_attention"]["ports"]["model"])) + # 1 kv head x global_head_dim 32, flattened into a fixed buffer. + assert state[full_cell]["contract"]["shape"] == ["batch", CAPACITY, 32] + # 1 kv head x head_dim 16, still growing. + assert state[sliding_cell]["contract"]["shape"] == [ + "batch", + 1, + "past_sequence_len", + 16, + ] + + def test_a_hybrid_decoder_keeps_its_padding_mask(self, mixed): + # The dynamic sliding layers still build their bias from attention_mask; + # dropping it because *some* layers are static loses padding. + loop = mixed["pipeline"]["workflow"]["steps"][0] + assert "attention_mask" in _model_invoke(loop["setup"]) + + +class TestFp8KvCacheMetadata: + """FP8 storage is a graph-visible fact, so the metadata must repeat it. + + A runtime allocates the KV buffers from these contracts. Publishing + ``float16`` for a cache the graph declares as ``float8_e4m3fn`` would size + every buffer at twice the bytes the model reads, so the declared dtype has + to be whatever the graph actually says — never the model's compute dtype. + """ + + @staticmethod + @pytest.fixture(scope="class") + def fp8_workflow(): + import onnx_ir as ir + + from mobius._optimizations import optimize_model + + config = _text_config() + module = registry.get("qwen2")(config) + pkg = CausalLMTask().build(module, config) + optimize_model( + pkg["model"], + ep="cuda", + dtype=ir.DataType.FLOAT16, + model_role="decoder", + fp8_kv_cache=True, + ) + return pkg, build_decoder_workflow_metadata(pkg, config) + + def test_graph_ports_are_fp8(self, fp8_workflow): + pkg, _ = fp8_workflow + import onnx_ir as ir + + caches = [ + value + for value in [*pkg["model"].graph.inputs, *pkg["model"].graph.outputs] + if value.name.startswith(("past_key_values.", "present.")) + ] + assert caches + assert {value.dtype for value in caches} == {ir.DataType.FLOAT8E4M3FN} + + def test_state_contracts_declare_the_graph_dtype(self, fp8_workflow): + _, metadata = fp8_workflow + workflow = metadata["pipeline"]["workflow"] + cells = _cache_cells(workflow) + assert cells + for cell in cells: + assert workflow["state"][cell]["contract"]["dtype"] == "float8_e4m3fn" + + def test_carried_cache_is_still_an_appending_cache(self, fp8_workflow): + # Quantizing the cells changes their dtype, not their update discipline. + _, metadata = fp8_workflow + group = next( + iter( + metadata["pipeline"]["workflow"]["serving"]["state_service"]["groups"].values() + ) + ) + assert group["sequence_axis"] == 2 + assert group["layout"] == "bnsh" + assert group.get("update", {}).get("kind") != "indexed_scatter" + + +class TestFeatureCombinations: + """Which feature pairs are representable, and which are refused and why.""" + + def test_fp8_requires_an_operator_that_can_dequantize_the_cache(self): + # A static-cache graph scatters into buffers read by ai.onnx Attention, + # which has no k_scale/v_scale inputs. Retyping those buffers would + # declare FP8 over bytes that are read as float16, so the build must + # refuse rather than emit either a wrong graph or a silently fp16 one. + import onnx_ir as ir + + from mobius._optimizations import optimize_model + + config = _text_config() + module = registry.get("qwen2")(config) + pkg = CausalLMTask(static_cache=True, max_seq_len=CAPACITY).build(module, config) + with pytest.raises(ValueError, match="no GroupQueryAttention KV cache"): + optimize_model( + pkg["model"], + ep="cuda", + dtype=ir.DataType.FLOAT16, + model_role="decoder", + fp8_kv_cache=True, + ) + + def test_static_cache_survives_cuda_optimization(self): + # Optimizing must not rewrite the scatter into an appending cache. + import onnx_ir as ir + + from mobius._optimizations import optimize_model + + config = _text_config() + module = registry.get("qwen2")(config) + pkg = CausalLMTask(static_cache=True, max_seq_len=CAPACITY).build(module, config) + optimize_model( + pkg["model"], ep="cuda", dtype=ir.DataType.FLOAT16, model_role="decoder" + ) + metadata = build_decoder_workflow_metadata(pkg, config) + _, group = _scatter_group(metadata) + assert group["update"]["kind"] == "indexed_scatter" + assert _static_cache_abi(metadata)["cache_inputs"] == [ + "key_cache.0", + "value_cache.0", + "key_cache.1", + "value_cache.1", + ] diff --git a/tests/wav2vec2_ctc_metadata_integration_test.py b/tests/wav2vec2_ctc_metadata_integration_test.py new file mode 100644 index 000000000..aedff1c0d --- /dev/null +++ b/tests/wav2vec2_ctc_metadata_integration_test.py @@ -0,0 +1,132 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT License. + +"""End-to-end metadata-driven parity for a real Wav2Vec2 CTC checkpoint. + +Exports ``facebook/wav2vec2-base-960h`` with Mobius, publishes the one-file +inference metadata, then drives the whole pipeline — preprocessing, encoder, +frame argmax, CTC collapse, transcript — using nothing but the emitted +document, and compares every stage against HuggingFace on real audio. + +This is a frame-synchronous package: the encoder runs exactly once and there is +no generation loop. +""" + +from __future__ import annotations + +import tempfile +from pathlib import Path + +import numpy as np +import onnx_ir as ir +import pytest +import soundfile as sf +import torch +import transformers + +from mobius import build +from mobius._configs import MMSConfig +from mobius.integrations.onnx_genai import ctc_runtime +from mobius.integrations.onnx_genai.auto_export import write_onnx_genai_config + +_MODEL_ID = "facebook/wav2vec2-base-960h" +_REVISION = "22aad52d435eb6dbaf354bdad9b0da84ce7d6156" +_AUDIO = Path("testdata") / "652-129742-0006.flac" +_EXPECTED = ( + "CAULIFLOWER MAYONAISE TAKE COLD BOILED CULIFLOWER BREAK INTO BRANCHES " + "ADDING SALT PEPPER AND VINEGAR TO SEASON" +) + + +def _export(directory: Path) -> None: + package = build(_MODEL_ID) + ir.save( + package["model"], + str(directory / "model.onnx"), + external_data="model.onnx.data", + ) + write_onnx_genai_config(package, str(directory), config=package.config, source=_MODEL_ID) + + +def _hf_logits(waveform: np.ndarray, mask: np.ndarray | None = None) -> np.ndarray: + model = transformers.Wav2Vec2ForCTC.from_pretrained( + _MODEL_ID, revision=_REVISION, dtype=torch.float32 + ).eval() + inputs = {"input_values": torch.tensor(waveform)} + if mask is not None: + inputs["attention_mask"] = torch.tensor(mask) + with torch.no_grad(): + return model(**inputs).logits.numpy() + + +@pytest.mark.integration +def test_wav2vec2_ctc_metadata_pipeline_matches_huggingface(): + audio, sample_rate = sf.read(str(_AUDIO)) + audio = audio.astype(np.float32) + assert np.any(audio != 0) + + with tempfile.TemporaryDirectory() as raw: + directory = Path(raw) + _export(directory) + result = ctc_runtime.transcribe(str(directory), [audio], sample_rate) + + processor = transformers.Wav2Vec2Processor.from_pretrained(_MODEL_ID, revision=_REVISION) + reference = _hf_logits( + processor(audio, sampling_rate=sample_rate, return_tensors="pt").input_values.numpy() + ) + + # Logits agree to float32 accumulation noise. + assert result["logits"].shape == reference.shape + np.testing.assert_allclose(result["logits"], reference, atol=5e-3) + + # Every downstream decode stage agrees exactly. + assert result["argmax_ids"][0] == reference.argmax(-1)[0].tolist() + assert result["collapsed_ids"][0] == ctc_runtime.collapse_ctc( + reference.argmax(-1)[0].tolist(), blank_id=0, collapse_repeats=True + ) + assert ( + result["transcripts"][0] + == processor.batch_decode(torch.tensor(reference).argmax(-1))[0] + ) + assert result["transcripts"][0] == _EXPECTED + + +@pytest.mark.integration +def test_wav2vec2_ctc_padded_batch_is_segmented_by_declared_frame_lengths(): + audio, sample_rate = sf.read(str(_AUDIO)) + audio = audio.astype(np.float32) + rows = [audio, audio[: 4 * sample_rate]] # unequal lengths + + with tempfile.TemporaryDirectory() as raw: + directory = Path(raw) + _export(directory) + metadata = ctc_runtime.load_metadata(str(directory)) + preprocessed = ctc_runtime.run_audio_preprocessing( + metadata["preprocessing"]["audio"], rows, sample_rate + ) + result = ctc_runtime.transcribe(str(directory), rows, sample_rate) + + _, profile = ctc_runtime.select_profile(metadata, "transcription") + # Group-normalizing over the padded time axis makes rows interdependent, + # which the package must declare rather than leave for a caller to discover. + assert profile["batch_invariance"] == "padding_sensitive" + assert profile["decoding"]["lengths"] == "frame_lengths" + + config = MMSConfig.from_transformers( + transformers.AutoConfig.from_pretrained(_MODEL_ID, revision=_REVISION) + ) + expected_frames = [config.feature_extract_output_length(row.shape[0]) for row in rows] + assert [len(ids) for ids in result["argmax_ids"]] == expected_frames + + # HuggingFace fed the identical padded batch is the reference for a padded + # run; a solo run is a different computation for this checkpoint. + reference = _hf_logits(preprocessed["input_values"], preprocessed["attention_mask"]) + processor = transformers.Wav2Vec2Processor.from_pretrained(_MODEL_ID, revision=_REVISION) + for row, frames in enumerate(expected_frames): + valid = reference[row, :frames] + np.testing.assert_allclose(result["logits"][row, :frames], valid, atol=5e-3) + assert result["argmax_ids"][row] == valid.argmax(-1).tolist() + assert ( + result["transcripts"][row] + == processor.batch_decode(torch.tensor(valid[None]).argmax(-1))[0] + )